wehub-resource-sync
353 lines
11 KiB
Python
353 lines
11 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=invalid-name
|
|
"""GEMM kernel generator and profiler for CUTLASS."""
|
|
|
|
import os
|
|
import pickle
|
|
from functools import partial
|
|
|
|
from .gemm_operation import EmitGemmInstance, GemmOperation
|
|
from .gemm_profiler import GemmProfilerEmitter
|
|
from .gen_tensor_op import EPILOGUE_MAP, GENERATOR_FUNC_TABLE, ProfilerEngine
|
|
from .library import (
|
|
DataType,
|
|
DataTypeTag,
|
|
EpilogueFunctor,
|
|
LayoutType,
|
|
SwizzlingFunctor,
|
|
TensorDescription,
|
|
)
|
|
|
|
|
|
def create_gemm_operator_with_epilogue(
|
|
op_type,
|
|
tile_description,
|
|
data_type,
|
|
alignment,
|
|
swizzling_functor,
|
|
batched=False,
|
|
layout_b=LayoutType.ColumnMajor,
|
|
):
|
|
"""
|
|
Instantiate a cutlass kernel from the given configuration,
|
|
along with the epilouge functor
|
|
"""
|
|
element_a, element_b, element_c, element_epilogue = data_type
|
|
|
|
A = TensorDescription(element_a, LayoutType.RowMajor, alignment)
|
|
B = TensorDescription(element_b, layout_b, alignment)
|
|
C = TensorDescription(element_c, LayoutType.RowMajor, alignment)
|
|
|
|
if batched:
|
|
swizzling_functor = SwizzlingFunctor.Batched
|
|
|
|
if "residual" in op_type:
|
|
if "hardswish" in op_type:
|
|
activation = "cutlass::epilogue::thread::HardSwish"
|
|
elif "silu" in op_type:
|
|
activation = "cutlass::epilogue::thread::SiLu"
|
|
elif "sigmoid" in op_type:
|
|
activation = "cutlass::epilogue::thread::Sigmoid"
|
|
elif "gelu" in op_type:
|
|
activation = "cutlass::epilogue::thread::GELU"
|
|
elif "relu" in op_type:
|
|
activation = "cutlass::epilogue::thread::ReLu"
|
|
else:
|
|
activation = "cutlass::epilogue::thread::Identity"
|
|
|
|
binary_op = "cutlass::multiplies" if "residual_multiply" in op_type else "cutlass::plus"
|
|
unary_op = (
|
|
"cutlass::epilogue::thread::ReLu"
|
|
if op_type.endswith("relu")
|
|
else "cutlass::epilogue::thread::Identity"
|
|
)
|
|
residual_block_info = {
|
|
"activation": activation,
|
|
"binary_op": binary_op,
|
|
"unary_op": unary_op,
|
|
}
|
|
epilogue = EpilogueFunctor.LinearCombinationResidualBlock
|
|
no_beta_scaling = False
|
|
else:
|
|
residual_block_info = None
|
|
epilogue, no_beta_scaling = EPILOGUE_MAP[op_type]
|
|
|
|
op = GemmOperation(
|
|
tile_description.minimum_compute_capability,
|
|
tile_description,
|
|
A,
|
|
B,
|
|
C,
|
|
element_epilogue,
|
|
epilogue,
|
|
swizzling_functor,
|
|
)
|
|
|
|
return (
|
|
op.procedural_name(),
|
|
EmitGemmInstance().emit(
|
|
op,
|
|
no_beta_scaling=no_beta_scaling,
|
|
batched=batched,
|
|
residual_block_info=residual_block_info,
|
|
),
|
|
)
|
|
|
|
|
|
def enumerate_gemm_operators(
|
|
tile_descriptions,
|
|
data_type,
|
|
alignment_constraints,
|
|
swizzling_functor=SwizzlingFunctor.Identity8,
|
|
layout_b=LayoutType.ColumnMajor,
|
|
):
|
|
"""Exhaustively instantiate all kernels from a given configuration."""
|
|
ret = []
|
|
kernel_emitter = EmitGemmInstance()
|
|
profiler_emitter = GemmProfilerEmitter()
|
|
|
|
element_a, element_b, element_c, element_epilogue = data_type
|
|
|
|
for tile_description in tile_descriptions:
|
|
for alignment in alignment_constraints:
|
|
A = TensorDescription(element_a, LayoutType.RowMajor, alignment)
|
|
B = TensorDescription(element_b, layout_b, alignment)
|
|
C = TensorDescription(element_c, LayoutType.RowMajor, alignment)
|
|
|
|
if element_c == DataType.s32 and A.alignment == 1:
|
|
tile_description.threadblock_shape[0] = min(
|
|
tile_description.threadblock_shape[0], 128
|
|
)
|
|
tile_description.threadblock_shape[1] = min(
|
|
tile_description.threadblock_shape[1], 128
|
|
)
|
|
|
|
op = GemmOperation(
|
|
tile_description.minimum_compute_capability,
|
|
tile_description,
|
|
A,
|
|
B,
|
|
C,
|
|
element_epilogue,
|
|
EpilogueFunctor.LinearCombination,
|
|
swizzling_functor,
|
|
)
|
|
|
|
src = profiler_emitter.emit(
|
|
op.procedural_name(),
|
|
kernel_emitter.emit(op, batched=False),
|
|
DataTypeTag[element_a],
|
|
DataTypeTag[element_b],
|
|
DataTypeTag[element_c],
|
|
op.leading_dim(),
|
|
)
|
|
|
|
ret.append(
|
|
{
|
|
"src": src,
|
|
"op": op,
|
|
"name": op.procedural_name(),
|
|
"tile_description": tile_description,
|
|
"alignment": alignment,
|
|
"data_type": data_type,
|
|
"swizzle_functor": swizzling_functor,
|
|
}
|
|
)
|
|
|
|
return ret
|
|
|
|
|
|
# TODO(masahi): A sensible way to pick reasonable default kernels
|
|
DEFAULT_KERNELS = {
|
|
75: {
|
|
("float16", "float16"): "cutlass_tensorop_h1688gemm_128x64_32x2_tn_align1",
|
|
("float16", "float32"): "cutlass_tensorop_s1688gemm_f16_64x64_32x2_tn_align1",
|
|
},
|
|
# align1 variants do not seem to be available for sm80
|
|
80: {
|
|
("float16", "float16"): "cutlass_tensorop_h1688gemm_128x64_32x2_tn_align1",
|
|
("float16", "float32"): "cutlass_tensorop_s1688gemm_f16_64x64_32x2_tn_align1",
|
|
# two kernels for tf32 and 3xtf32
|
|
("float32", "float32"): (
|
|
"cutlass_tensorop_s1688gemm_128x64_32x3_tn_align1",
|
|
"cutlass_tensorop_s1688gemm_64x64_16x3_tn_align1",
|
|
),
|
|
},
|
|
}
|
|
|
|
|
|
class CutlassGemmProfiler:
|
|
"""Profile all candidate kernels and select the best one."""
|
|
|
|
def __init__(self, sm, cutlass_path, binary_path):
|
|
assert sm in GENERATOR_FUNC_TABLE and sm in DEFAULT_KERNELS, f"sm{sm} not supported yet."
|
|
self.engine = ProfilerEngine(sm, cutlass_path, binary_path)
|
|
self.sm = sm
|
|
self.cache_path = os.path.join(binary_path, "cutlass_gemm_cache.pickle")
|
|
if os.path.exists(self.cache_path):
|
|
self.cache = pickle.load(open(self.cache_path, "rb"))
|
|
else:
|
|
self.cache = {}
|
|
|
|
def get_default(
|
|
self,
|
|
op_type,
|
|
out_dtype,
|
|
arg0_dtype,
|
|
arg1_dtype,
|
|
use_3xtf32=True,
|
|
batched=False,
|
|
layout_b=LayoutType.ColumnMajor,
|
|
):
|
|
"""Return the default kernel for the requested architecture.
|
|
For now, the default kernel was picked arbitrary.
|
|
"""
|
|
ops = GENERATOR_FUNC_TABLE[self.sm](
|
|
out_dtype,
|
|
arg0_dtype,
|
|
arg1_dtype,
|
|
partial(enumerate_gemm_operators, layout_b=layout_b),
|
|
lambda align: align == 1, # Only request align1 kernels
|
|
use_3xtf32,
|
|
profile_all_alignments=True, # To include all align1 kernels
|
|
# TODO(masahi): Invesitigate when fp32 accumulation is needed for gemm
|
|
accumlator_dtype=out_dtype,
|
|
)
|
|
|
|
default_kernel_name = DEFAULT_KERNELS[self.sm][(arg0_dtype, out_dtype)]
|
|
|
|
if arg0_dtype == "float32":
|
|
default_kernel_name = (
|
|
default_kernel_name[0] if not use_3xtf32 else default_kernel_name[1]
|
|
)
|
|
|
|
filtered = list(filter(lambda op: op["name"] == default_kernel_name, ops))
|
|
assert len(filtered) == 1
|
|
op = filtered[0]
|
|
name, opdef = create_gemm_operator_with_epilogue(
|
|
op_type,
|
|
op["tile_description"],
|
|
op["data_type"],
|
|
op["alignment"],
|
|
op["swizzle_functor"],
|
|
batched=batched,
|
|
layout_b=layout_b,
|
|
)
|
|
op.update({"name": name, "opdef": opdef})
|
|
return op
|
|
|
|
def select_op(
|
|
self,
|
|
M,
|
|
N,
|
|
K,
|
|
out_dtype,
|
|
arg0_dtype,
|
|
arg1_dtype,
|
|
use_3xtf32,
|
|
profile_all_alignments=False,
|
|
find_first_valid=False,
|
|
use_multiprocessing=False,
|
|
layout_b=LayoutType.ColumnMajor,
|
|
):
|
|
"""
|
|
Profile and select the best kernel from candidate kernels.
|
|
See the documentation for the profile method below.
|
|
"""
|
|
if (M, N, K) in self.cache:
|
|
op = self.cache[(M, N, K)]
|
|
return op
|
|
|
|
# TODO(masahi): CUTLASS alignment check on gemm kernels is too restrictive.
|
|
# See https://github.com/NVIDIA/cutlass/issues/362.
|
|
# When the above issue is resolved, we can remove the alignment check on M below.
|
|
|
|
ops = GENERATOR_FUNC_TABLE[self.sm](
|
|
out_dtype,
|
|
arg0_dtype,
|
|
arg1_dtype,
|
|
partial(enumerate_gemm_operators, layout_b=layout_b),
|
|
lambda align: all([dim % align == 0 for dim in [M, N, K]]),
|
|
use_3xtf32,
|
|
profile_all_alignments=profile_all_alignments,
|
|
# TODO(masahi): Invesitigate when fp32 accumulation is needed for gemm
|
|
accumlator_dtype=out_dtype,
|
|
)
|
|
|
|
if not find_first_valid:
|
|
self.engine.compile_all(ops, use_multiprocessing)
|
|
|
|
for op in ops:
|
|
out = self.engine.evaluate(op, [M, N, K])
|
|
op["runtime"] = out
|
|
if out < float("inf") and find_first_valid:
|
|
self.cache[(M, N, K)] = op
|
|
return op
|
|
|
|
op = min(ops, key=lambda i: i["runtime"])
|
|
self.cache[(M, N, K)] = op
|
|
with open(self.cache_path, "wb") as f:
|
|
pickle.dump(self.cache, f)
|
|
return op
|
|
|
|
def profile(
|
|
self,
|
|
op_type,
|
|
M,
|
|
N,
|
|
K,
|
|
out_dtype,
|
|
arg0_dtype,
|
|
arg1_dtype,
|
|
use_3xtf32=True,
|
|
profile_all_alignments=False,
|
|
find_first_valid=False,
|
|
use_multiprocessing=False,
|
|
batched=False,
|
|
layout_b=LayoutType.ColumnMajor,
|
|
):
|
|
"""Profile and select the best kernel from candidate kernels.
|
|
If find_first_valid is True, return immediately after the first applicable kernel is found.
|
|
If use_multiprocessing is True, compile all profiler executables in parallel.
|
|
"""
|
|
op = self.select_op(
|
|
M,
|
|
N,
|
|
K,
|
|
out_dtype,
|
|
arg0_dtype,
|
|
arg1_dtype,
|
|
use_3xtf32,
|
|
profile_all_alignments=profile_all_alignments,
|
|
find_first_valid=find_first_valid,
|
|
use_multiprocessing=use_multiprocessing,
|
|
layout_b=layout_b,
|
|
)
|
|
|
|
name, opdef = create_gemm_operator_with_epilogue(
|
|
op_type,
|
|
op["tile_description"],
|
|
op["data_type"],
|
|
op["alignment"],
|
|
op["swizzle_functor"],
|
|
batched=batched,
|
|
layout_b=layout_b,
|
|
)
|
|
|
|
return name, opdef, op["runtime"]
|