chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,907 @@
|
||||
# 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, dangerous-default-value, arguments-differ
|
||||
# ruff: noqa: F821
|
||||
"""Driver for partitioning and building a Relax module for CUTLASS offload."""
|
||||
|
||||
import itertools
|
||||
import logging
|
||||
import multiprocessing
|
||||
import operator
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
from functools import reduce
|
||||
|
||||
from tvm_ffi import register_global_func
|
||||
|
||||
import tvm
|
||||
from tvm import relax, runtime
|
||||
from tvm.support.nvcc import get_cuda_version
|
||||
from tvm.topi.utils import get_const_tuple
|
||||
|
||||
from .gen_conv2d import CutlassConv2DProfiler
|
||||
from .gen_gemm import CutlassGemmProfiler
|
||||
from .library import ConvKind, LayoutType
|
||||
|
||||
logger = logging.getLogger("cutlass")
|
||||
|
||||
|
||||
def has_cutlass():
|
||||
"""Returns true if the CUTLASS custom codegen is available"""
|
||||
return tvm.get_global_func("relax.ext.cutlass", True) is not None
|
||||
|
||||
|
||||
def _get_cutlass_path():
|
||||
invalid_paths = []
|
||||
for rel in ["../../../../", "../../../", "../../"]:
|
||||
tvm_root = os.path.join(os.path.dirname(os.path.realpath(__file__)), rel)
|
||||
cutlass_path = os.path.join(tvm_root, "3rdparty/cutlass")
|
||||
if os.path.exists(cutlass_path):
|
||||
return cutlass_path
|
||||
invalid_paths.append(cutlass_path)
|
||||
raise AssertionError(f"The CUTLASS root directory not found in: {invalid_paths}")
|
||||
|
||||
|
||||
def _get_cutlass_compile_options(sm, threads, use_fast_math=False):
|
||||
cutlass_root = _get_cutlass_path()
|
||||
cutlass_include = os.path.join(cutlass_root, "include")
|
||||
cutlass_util_include = os.path.join(cutlass_root, "tools/util/include")
|
||||
cutlass_attention_include = os.path.join(cutlass_root, "examples/41_fused_multi_head_attention")
|
||||
cutlass_fpA_intB_gemm_include = os.path.join(cutlass_root, "../cutlass_fpA_intB_gemm")
|
||||
flash_attn_include = os.path.join(cutlass_root, "../libflash_attn/include")
|
||||
|
||||
kwargs = {}
|
||||
kwargs["cc"] = "nvcc"
|
||||
kwargs["options"] = [
|
||||
"-c",
|
||||
"-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1",
|
||||
f"-gencode=arch=compute_{sm},code=[sm_{sm},compute_{sm}]",
|
||||
"-DNDEBUG",
|
||||
"-Xcompiler=-fPIC",
|
||||
"-Xcompiler=-Wconversion",
|
||||
"-Xcompiler=-fno-strict-aliasing",
|
||||
"-Xcompiler=-fvisibility=hidden",
|
||||
"-O3",
|
||||
"-std=c++17",
|
||||
f"-I{cutlass_include}",
|
||||
f"-I{cutlass_util_include}",
|
||||
f"-I{cutlass_attention_include}",
|
||||
f"-I{cutlass_fpA_intB_gemm_include}",
|
||||
f"-I{flash_attn_include}",
|
||||
]
|
||||
if use_fast_math:
|
||||
kwargs["options"].append("-DCUTLASS_USE_TANH_FOR_SIGMOID")
|
||||
cuda_ver = get_cuda_version()
|
||||
if cuda_ver >= (11, 2):
|
||||
ncpu = multiprocessing.cpu_count() if threads < 0 else threads
|
||||
kwargs["options"].append(f"-t {ncpu}")
|
||||
return kwargs
|
||||
|
||||
|
||||
def select_gemm_kernel(
|
||||
cutlass_profiler,
|
||||
op_type,
|
||||
MM,
|
||||
KK,
|
||||
NN,
|
||||
out_dtype,
|
||||
arg0_dtype,
|
||||
arg1_dtype,
|
||||
use_3xtf32,
|
||||
batched,
|
||||
find_first_valid,
|
||||
use_multiprocessing,
|
||||
):
|
||||
"""Run CUTLASS profiler to select the best kernel, or return the default one for dynamic
|
||||
workloads."""
|
||||
if any(isinstance(s, tvm.tirx.Any) for s in [MM, KK, NN]):
|
||||
out = cutlass_profiler.get_default(
|
||||
op_type, out_dtype, arg0_dtype, arg1_dtype, use_3xtf32, batched=batched
|
||||
)
|
||||
name, cutlass_op_def = out["name"], out["opdef"]
|
||||
logger.info("Picked the default kernel %s", name)
|
||||
else:
|
||||
name, cutlass_op_def, _ = cutlass_profiler.profile(
|
||||
op_type,
|
||||
MM,
|
||||
NN,
|
||||
KK,
|
||||
out_dtype,
|
||||
arg0_dtype,
|
||||
arg1_dtype,
|
||||
use_3xtf32,
|
||||
batched=batched,
|
||||
find_first_valid=find_first_valid,
|
||||
use_multiprocessing=use_multiprocessing,
|
||||
)
|
||||
if not find_first_valid:
|
||||
logger.info("The best kernel is %s", name)
|
||||
else:
|
||||
logger.info("Picked the first kernel found %s", name)
|
||||
|
||||
return name, cutlass_op_def
|
||||
|
||||
|
||||
def handle_batch_matmul(
|
||||
cutlass_profiler,
|
||||
op_type,
|
||||
arg0_shape,
|
||||
arg1_shape,
|
||||
out_dtype,
|
||||
arg0_dtype,
|
||||
arg1_dtype,
|
||||
use_3xtf32,
|
||||
find_first_valid,
|
||||
use_multiprocessing,
|
||||
):
|
||||
"""Profile and select a kernel for batch_matmul op workload."""
|
||||
MM = arg0_shape[1]
|
||||
KK = arg0_shape[2]
|
||||
NN = arg1_shape[1]
|
||||
|
||||
name, cutlass_op_def = select_gemm_kernel(
|
||||
cutlass_profiler,
|
||||
op_type,
|
||||
MM,
|
||||
KK,
|
||||
NN,
|
||||
out_dtype,
|
||||
arg0_dtype,
|
||||
arg1_dtype,
|
||||
use_3xtf32,
|
||||
True,
|
||||
find_first_valid,
|
||||
use_multiprocessing,
|
||||
)
|
||||
|
||||
return {
|
||||
"batch": arg0_shape[0],
|
||||
"batch_stride_A": arg0_shape[1] * arg0_shape[2],
|
||||
"batch_stride_B": arg1_shape[1] * arg1_shape[2],
|
||||
"batch_stride_C": arg0_shape[1] * arg1_shape[1],
|
||||
"cutlass_op_def": cutlass_op_def,
|
||||
"cutlass_op_name": name,
|
||||
"lda": "K",
|
||||
"ldb": "K",
|
||||
"ldc": "N",
|
||||
}
|
||||
|
||||
|
||||
def handle_dense(
|
||||
cutlass_profiler,
|
||||
op_type,
|
||||
arg0_shape,
|
||||
arg1_shape,
|
||||
out_dtype,
|
||||
arg0_dtype,
|
||||
arg1_dtype,
|
||||
use_3xtf32,
|
||||
find_first_valid,
|
||||
use_multiprocessing,
|
||||
):
|
||||
"""Profile and select a kernel for dense op workload."""
|
||||
MM = arg0_shape[0]
|
||||
KK = arg0_shape[1]
|
||||
NN = arg1_shape[0]
|
||||
|
||||
name, cutlass_op_def = select_gemm_kernel(
|
||||
cutlass_profiler,
|
||||
op_type,
|
||||
MM,
|
||||
KK,
|
||||
NN,
|
||||
out_dtype,
|
||||
arg0_dtype,
|
||||
arg1_dtype,
|
||||
use_3xtf32,
|
||||
False,
|
||||
find_first_valid,
|
||||
use_multiprocessing,
|
||||
)
|
||||
|
||||
assert "tn_align" in name, "Only supports (row_major, col_major) input layout for now."
|
||||
|
||||
return {
|
||||
"cutlass_op_def": cutlass_op_def,
|
||||
"cutlass_op_name": name,
|
||||
"lda": "K",
|
||||
"ldb": "K",
|
||||
"ldc": "N",
|
||||
}
|
||||
|
||||
|
||||
def handle_conv2d(
|
||||
cutlass_profiler,
|
||||
op_type,
|
||||
d_shape,
|
||||
w_shape,
|
||||
padding,
|
||||
strides,
|
||||
dilation,
|
||||
out_dtype,
|
||||
data_dtype,
|
||||
weight_dtype,
|
||||
use_3xtf32,
|
||||
split_k_slices,
|
||||
profile_all_alignments,
|
||||
find_first_valid,
|
||||
use_multiprocessing,
|
||||
):
|
||||
"""Profile and select a kernel for conv2d op workload."""
|
||||
if "conv2d_transpose" in op_type:
|
||||
conv_kind = ConvKind.Dgrad
|
||||
elif "backward_weight" in op_type:
|
||||
conv_kind = ConvKind.Wgrad
|
||||
else:
|
||||
conv_kind = ConvKind.Fprop
|
||||
|
||||
if any(isinstance(s, tvm.tirx.Any) for s in d_shape):
|
||||
out = cutlass_profiler.get_default(
|
||||
op_type, out_dtype, data_dtype, weight_dtype, use_3xtf32, conv_kind, strides
|
||||
)
|
||||
name, cutlass_op_def = out["name"], out["opdef"]
|
||||
logger.info("Picked the default kernel %s", name)
|
||||
else:
|
||||
name, cutlass_op_def, _ = cutlass_profiler.profile(
|
||||
op_type,
|
||||
d_shape,
|
||||
w_shape,
|
||||
padding,
|
||||
strides,
|
||||
dilation,
|
||||
out_dtype,
|
||||
data_dtype,
|
||||
weight_dtype,
|
||||
use_3xtf32,
|
||||
conv_kind,
|
||||
split_k_slices,
|
||||
profile_all_alignments,
|
||||
find_first_valid=find_first_valid,
|
||||
use_multiprocessing=use_multiprocessing,
|
||||
)
|
||||
if not find_first_valid:
|
||||
logger.info("The best kernel is %s", name)
|
||||
else:
|
||||
logger.info("Picked the first kernel found %s", name)
|
||||
|
||||
return {"cutlass_op_def": cutlass_op_def, "cutlass_op_name": name}
|
||||
|
||||
|
||||
def num_cutlass_partitions(mod):
|
||||
return sum([(1 if "cutlass" in var.name_hint else 0) for var in mod.get_global_vars()])
|
||||
|
||||
|
||||
def tune_cutlass_kernels(
|
||||
mod,
|
||||
sm,
|
||||
use_3xtf32=True,
|
||||
split_k_slices=[1],
|
||||
profile_all_alignments=False,
|
||||
find_first_valid=False,
|
||||
use_multiprocessing=False,
|
||||
tmp_dir="./tmp",
|
||||
):
|
||||
"""Given a module partitioned for CUTLASS offloading, profile each workload to select which
|
||||
kernels to emit.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mod : IRModule
|
||||
The IRModule with cutlass partitions.
|
||||
|
||||
sm : int
|
||||
An integer specifying the compute capability. For example, 75 for Turing and
|
||||
80 or 86 for Ampere.
|
||||
|
||||
use_3xtf32 : bool
|
||||
Wheter or not use slower but very accurate (compared to tf32) 3xtf32 mode for
|
||||
fp32 inputs on tensorcore.
|
||||
|
||||
split_k_slices : list of int
|
||||
Split factor candidates for split-K GEMM. If split-K > 1, the GEMM K-loop is computed in
|
||||
parallel across split-K blocks, and a separate global reduction kernel is launched to
|
||||
accumulate partial reductions. The profiler will pick the best split-k factor from the
|
||||
given candidate list. Note that the larger split-K factor requires a larger workspace.
|
||||
Currently, parallel split-k has been tested only for wgrad. For GEMM and other conv2d
|
||||
kinds, split_k_slices is ignored.
|
||||
|
||||
profile_all_alignments : bool
|
||||
When True, profile all kernal variants with smaller alignments than the largest possible.
|
||||
|
||||
find_first_valid : bool
|
||||
Whether or not profile all candidate kernels, or stop profiling after
|
||||
the first applicable kernel is found.
|
||||
|
||||
use_multiprocessing : bool
|
||||
Whether or not compile profiler executables for different kernels in parallel.
|
||||
|
||||
tmp_dir : string, optional
|
||||
A temporary directory where intermediate compiled artifacts will be stored.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mod : IRModule
|
||||
The updated module annotated with cutlass profiling information.
|
||||
|
||||
num_cutlass_partition : int
|
||||
The number of partitioned functions created for CUTLASS.
|
||||
"""
|
||||
gemm_profiler = CutlassGemmProfiler(sm, _get_cutlass_path(), tmp_dir)
|
||||
conv2d_profiler = CutlassConv2DProfiler(sm, _get_cutlass_path(), tmp_dir)
|
||||
num_cutlass_partition = 0
|
||||
for var in mod.get_global_vars():
|
||||
fun_name = var.name_hint
|
||||
func = mod[fun_name]
|
||||
if "cutlass" in fun_name:
|
||||
num_cutlass_partition += 1
|
||||
new_func = tune_cutlass_function(
|
||||
func,
|
||||
use_3xtf32,
|
||||
split_k_slices,
|
||||
profile_all_alignments,
|
||||
find_first_valid,
|
||||
use_multiprocessing,
|
||||
gemm_profiler,
|
||||
conv2d_profiler,
|
||||
)
|
||||
mod.update_func(var, new_func)
|
||||
|
||||
return mod, num_cutlass_partition
|
||||
|
||||
|
||||
def _get_call_node(expr: relax.Expr, op_name: str) -> relax.Call | None:
|
||||
node = None
|
||||
|
||||
def fvisit(e):
|
||||
nonlocal node
|
||||
if isinstance(e, relax.Call) and e.op.name == op_name:
|
||||
node = e
|
||||
|
||||
relax.analysis.post_order_visit(expr, fvisit)
|
||||
return node
|
||||
|
||||
|
||||
def _extract_relax_function_signature(f):
|
||||
signature = {}
|
||||
|
||||
for i, arg in enumerate(f.params):
|
||||
ty = arg.ty
|
||||
if isinstance(ty, relax.TensorType):
|
||||
signature[f"arg{i}_shape"] = get_const_tuple(ty.shape)
|
||||
signature[f"arg{i}_dtype"] = ty.dtype
|
||||
elif isinstance(ty, relax.ShapeType):
|
||||
signature[f"arg{i}_shape"] = get_const_tuple(ty.values)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
ret_ty = f.ret_ty
|
||||
if ret_ty.shape is not None:
|
||||
signature["ret_shape"] = get_const_tuple(ret_ty.shape)
|
||||
else:
|
||||
signature["ret_shape"] = None
|
||||
signature["ret_dtype"] = ret_ty.dtype
|
||||
|
||||
return signature
|
||||
|
||||
|
||||
def _extract_arg_idx(pattern_name, f):
|
||||
extract_func = tvm.get_global_func("relax.contrib.extract_arg_idx")
|
||||
arg_indices = extract_func(pattern_name, f)
|
||||
return {k: int(v) for k, v in arg_indices.items()}
|
||||
|
||||
|
||||
def is_shape_valid_for_cutlass_matmul(
|
||||
lhs_shape: Sequence[tvm.ir.Expr],
|
||||
rhs_shape: Sequence[tvm.ir.Expr],
|
||||
) -> bool:
|
||||
"""
|
||||
Check whether the shape of inputs can be handled by CUTLASS GEMM.
|
||||
|
||||
The stride-based batch matmul in CUTLASS cannot handle cases that some of
|
||||
the batch dimensions need to be stretched while others don't. This means
|
||||
it can only handle ND x ND whose batch dimensions match exactly on both side,
|
||||
as well as ND x 2D and 2D x ND. For example, it cannot handle matmul with shape
|
||||
(2, 1, 4, 8) x (2, 3, 8, 16), because the batch stride of lhs is not constant.
|
||||
"""
|
||||
if not isinstance(lhs_shape[-1], tvm.tirx.expr.IntImm | int):
|
||||
# Reduction axis must be constant
|
||||
return False
|
||||
|
||||
lhs_batches = reduce(operator.mul, lhs_shape[:-2], 1)
|
||||
rhs_batches = reduce(operator.mul, rhs_shape[:-2], 1)
|
||||
if lhs_batches == 1 or rhs_batches == 1:
|
||||
# This could be regular matmul or batch matmul with shape ND x 2D or 2D x ND
|
||||
return True
|
||||
|
||||
analyzer = tvm.arith.Analyzer()
|
||||
# If one side has less dimensions, use 1 to fill the gap
|
||||
batch_dim_pairs = list(
|
||||
itertools.zip_longest(
|
||||
list(lhs_shape)[-3::-1], # Remove the last two dimensions and reverse
|
||||
list(rhs_shape)[-3::-1],
|
||||
fillvalue=1,
|
||||
)
|
||||
)
|
||||
return all(analyzer.can_prove_equal(p[0], p[1]) for p in batch_dim_pairs)
|
||||
|
||||
|
||||
@relax.expr_functor.mutator
|
||||
class CutlassRelaxFunctionAnnotator(relax.PyExprMutator):
|
||||
"""A Relax function mutator that tunes and annotates CUTLASS composite functions
|
||||
with shape, dtype and generated templates.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mod,
|
||||
conv2d_profiler: CutlassConv2DProfiler,
|
||||
gemm_profiler: CutlassGemmProfiler,
|
||||
options,
|
||||
):
|
||||
super().__init__(mod)
|
||||
self.options = options
|
||||
self.conv2d_profiler = conv2d_profiler
|
||||
self.gemm_profiler = gemm_profiler
|
||||
|
||||
def handle_conv2d(self, f, op_type):
|
||||
"""Tune and annotate a conv2d op."""
|
||||
signature = _extract_relax_function_signature(f)
|
||||
arg_idx = _extract_arg_idx(op_type, f)
|
||||
op_attrs = _get_call_node(f.body, "relax.nn.conv2d").attrs
|
||||
|
||||
data_arg = f"arg{arg_idx['lhs']}"
|
||||
weight_arg = f"arg{arg_idx['rhs']}"
|
||||
|
||||
d_shape = signature[f"{data_arg}_shape"]
|
||||
w_shape = signature[f"{weight_arg}_shape"]
|
||||
out_shape = signature["ret_shape"]
|
||||
data_dtype = signature[f"{data_arg}_dtype"]
|
||||
weight_dtype = signature[f"{weight_arg}_dtype"]
|
||||
out_dtype = signature["ret_dtype"]
|
||||
padding = op_attrs["padding"]
|
||||
strides = op_attrs["strides"]
|
||||
dilation = op_attrs["dilation"]
|
||||
conv_kind = ConvKind.Fprop
|
||||
|
||||
use_3xtf32 = self.options.get("use_3xtf32", False)
|
||||
profile_all_alignments = self.options.get("profile_all_alignments", False)
|
||||
find_first_valid = self.options.get("find_first_valid", True)
|
||||
use_multiprocessing = self.options.get("use_multiprocessing", True)
|
||||
split_k_slices = self.options.get("split_k_slices", [1])
|
||||
|
||||
op_name, op_def, _ = self.conv2d_profiler.profile(
|
||||
op_type,
|
||||
d_shape,
|
||||
w_shape,
|
||||
padding,
|
||||
strides,
|
||||
dilation,
|
||||
out_dtype,
|
||||
data_dtype,
|
||||
weight_dtype,
|
||||
use_3xtf32,
|
||||
conv_kind,
|
||||
split_k_slices,
|
||||
profile_all_alignments,
|
||||
find_first_valid=find_first_valid,
|
||||
use_multiprocessing=use_multiprocessing,
|
||||
)
|
||||
|
||||
attrs = {
|
||||
"op_type": op_type,
|
||||
"data_arg_idx": arg_idx["lhs"],
|
||||
"weight_arg_idx": arg_idx["rhs"],
|
||||
"bias_arg_idx": arg_idx.get("bias"),
|
||||
"residual_arg_idx": arg_idx.get("residual"),
|
||||
"arg0_dtype": data_dtype,
|
||||
"arg1_dtype": weight_dtype,
|
||||
"ret_dtype": out_dtype,
|
||||
"arg0_shape": d_shape,
|
||||
"arg1_shape": w_shape,
|
||||
"ret_shape": out_shape,
|
||||
"strides": strides,
|
||||
"padding": padding,
|
||||
"dilation": dilation,
|
||||
"cutlass_op_name": op_name,
|
||||
"cutlass_op_def": op_def,
|
||||
}
|
||||
|
||||
residual_arg = arg_idx.get("residual")
|
||||
|
||||
if residual_arg:
|
||||
residual_shape = signature[f"arg{residual_arg}_shape"]
|
||||
attrs["residual_shape"] = residual_shape
|
||||
elif "residual" in op_type:
|
||||
attrs["residual_shape"] = d_shape
|
||||
|
||||
return f.with_attrs(attrs)
|
||||
|
||||
def handle_decode_matmul(self, f, op_type):
|
||||
"""Annotate a decode -> matmul op."""
|
||||
arg_idx = _extract_arg_idx(op_type, f)
|
||||
signature = _extract_relax_function_signature(f)
|
||||
lhs_arg = f"arg{arg_idx['lhs']}"
|
||||
rhs_arg = f"arg{arg_idx['w_encoded']}"
|
||||
lhs_shape = signature[f"{lhs_arg}_shape"]
|
||||
rhs_shape = signature[f"{rhs_arg}_shape"]
|
||||
ret_shape = signature["ret_shape"]
|
||||
scale_arg = f"arg{arg_idx['scales']}"
|
||||
scale_shape = signature[f"{scale_arg}_shape"]
|
||||
N = ret_shape[-1]
|
||||
|
||||
attrs = {
|
||||
"op_type": op_type,
|
||||
"lhs_arg_idx": arg_idx["lhs"],
|
||||
"rhs_arg_idx": arg_idx["w_encoded"],
|
||||
"scales_arg_idx": arg_idx["scales"],
|
||||
"bias_arg_idx": arg_idx.get("bias"),
|
||||
"activation": "identity",
|
||||
}
|
||||
# TODO(wuwei): find a better way to get group size
|
||||
attrs["group_size"] = 64 if len(scale_shape) == 2 and scale_shape[0] != 1 else -1
|
||||
|
||||
attrs["batch_rank"] = len(lhs_shape[:-1])
|
||||
attrs["M"] = reduce(operator.mul, lhs_shape[:-1], 1)
|
||||
|
||||
attrs["bias_stride"] = 0
|
||||
|
||||
if "bias" in arg_idx:
|
||||
bias_shape = signature[f"arg{arg_idx['bias']}_shape"]
|
||||
bias_shape_1d = reduce(operator.mul, bias_shape, 1)
|
||||
if bias_shape_1d != bias_shape[-1]:
|
||||
attrs["bias_stride"] = bias_shape[-1]
|
||||
|
||||
if N == rhs_shape[1]:
|
||||
attrs["weight_nbit"] = 8
|
||||
else:
|
||||
assert N == rhs_shape[1] * 2
|
||||
attrs["weight_nbit"] = 4
|
||||
|
||||
if "residual" in op_type:
|
||||
residual_pos = op_type.find("residual_")
|
||||
postfix = op_type[residual_pos + len("residual_") :]
|
||||
|
||||
if postfix.startswith("multiply"):
|
||||
binary_op = "multiply"
|
||||
else:
|
||||
binary_op = "plus"
|
||||
|
||||
if "relu" in postfix:
|
||||
unary_op = "relu"
|
||||
else:
|
||||
unary_op = "identity"
|
||||
|
||||
activation = "identity"
|
||||
|
||||
for act in ["relu", "silu", "gelu"]:
|
||||
if act in op_type[op_type.find("matmul_") + len("matmul_") : residual_pos]:
|
||||
activation = act
|
||||
break
|
||||
|
||||
attrs.update(
|
||||
{
|
||||
"unary_op": unary_op,
|
||||
"binary_op": binary_op,
|
||||
"activation": activation,
|
||||
"residual_arg_idx": arg_idx["residual"],
|
||||
}
|
||||
)
|
||||
else:
|
||||
for act in ["relu", "silu", "gelu"]:
|
||||
if act in op_type:
|
||||
attrs["activation"] = act
|
||||
break
|
||||
|
||||
return f.with_attrs(attrs)
|
||||
|
||||
def handle_matmul(self, f, op_type):
|
||||
"""Tune and annotate a matmul op."""
|
||||
signature = _extract_relax_function_signature(f)
|
||||
arg_idx = _extract_arg_idx(op_type, f)
|
||||
|
||||
lhs_arg = f"arg{arg_idx['lhs']}"
|
||||
rhs_arg = f"arg{arg_idx['rhs']}"
|
||||
|
||||
lhs_shape = signature[f"{lhs_arg}_shape"]
|
||||
rhs_shape = signature[f"{rhs_arg}_shape"]
|
||||
out_shape = signature["ret_shape"]
|
||||
lhs_dtype = signature[f"{lhs_arg}_dtype"]
|
||||
rhs_dtype = signature[f"{rhs_arg}_dtype"]
|
||||
out_dtype = signature["ret_dtype"]
|
||||
|
||||
if not is_shape_valid_for_cutlass_matmul(lhs_shape, rhs_shape):
|
||||
raise ValueError(f"Cannot handle the input shapes, lhs: {lhs_shape}, rhs: {rhs_shape}")
|
||||
|
||||
MM = lhs_shape[-2]
|
||||
KK = lhs_shape[-1]
|
||||
if "transposed" in op_type:
|
||||
NN = rhs_shape[-2]
|
||||
ldb = "K"
|
||||
layout_b = LayoutType.ColumnMajor
|
||||
else:
|
||||
NN = rhs_shape[-1]
|
||||
ldb = "N"
|
||||
layout_b = LayoutType.RowMajor
|
||||
|
||||
lhs_batches = reduce(operator.mul, lhs_shape[:-2], 1)
|
||||
rhs_batches = reduce(operator.mul, rhs_shape[:-2], 1)
|
||||
if lhs_batches == 1 and rhs_batches == 1:
|
||||
# Regular matmul
|
||||
is_batched = False
|
||||
batch_attrs = {}
|
||||
else:
|
||||
is_batched = True
|
||||
batch_attrs = {
|
||||
# If both lhs_batches and rhs_batches are greater than 1,
|
||||
# they must be equal. This is checked by is_shape_valid_for_cutlass_matmul.
|
||||
"batch": lhs_batches if rhs_batches == 1 else rhs_batches,
|
||||
"batch_stride_A": 0 if lhs_batches == 1 else MM * KK,
|
||||
"batch_stride_B": 0 if rhs_batches == 1 else KK * NN,
|
||||
"batch_stride_C": MM * NN,
|
||||
}
|
||||
|
||||
use_3xtf32 = self.options.get("use_3xtf32", False)
|
||||
find_first_valid = self.options.get("find_first_valid", True)
|
||||
use_multiprocessing = self.options.get("use_multiprocessing", True)
|
||||
|
||||
op_name, op_def, _ = self.gemm_profiler.profile(
|
||||
op_type,
|
||||
MM,
|
||||
NN,
|
||||
KK,
|
||||
out_dtype,
|
||||
lhs_dtype,
|
||||
rhs_dtype,
|
||||
use_3xtf32,
|
||||
batched=is_batched,
|
||||
find_first_valid=find_first_valid,
|
||||
use_multiprocessing=use_multiprocessing,
|
||||
layout_b=layout_b,
|
||||
)
|
||||
|
||||
return f.with_attrs(
|
||||
{
|
||||
"op_type": op_type,
|
||||
"lhs_arg_idx": arg_idx["lhs"],
|
||||
"rhs_arg_idx": arg_idx["rhs"],
|
||||
"residual_arg_idx": arg_idx.get("residual"),
|
||||
"bias_arg_idx": arg_idx.get("bias"),
|
||||
"arg0_dtype": signature["arg0_dtype"],
|
||||
"arg1_dtype": signature["arg1_dtype"],
|
||||
"ret_dtype": out_dtype,
|
||||
"arg0_shape": signature["arg0_shape"],
|
||||
"arg1_shape": signature["arg1_shape"],
|
||||
"ret_shape": out_shape,
|
||||
"lda": "K",
|
||||
"ldb": ldb,
|
||||
"ldc": "N",
|
||||
"cutlass_op_name": op_name,
|
||||
"cutlass_op_def": op_def,
|
||||
**batch_attrs,
|
||||
}
|
||||
)
|
||||
|
||||
def handle_attention(self, f, op_type):
|
||||
"""Annotate an attention op."""
|
||||
signature = _extract_relax_function_signature(f)
|
||||
|
||||
if _get_call_node(f.body, "relax.nn.attention") is not None:
|
||||
attention_node = _get_call_node(f.body, "relax.nn.attention")
|
||||
op_attrs = attention_node.attrs
|
||||
elif _get_call_node(f.body, "relax.nn.attention_bias") is not None:
|
||||
attention_node = _get_call_node(f.body, "relax.nn.attention_bias")
|
||||
op_attrs = attention_node.attrs
|
||||
elif _get_call_node(f.body, "relax.nn.attention_var_len") is not None:
|
||||
attention_node = _get_call_node(f.body, "relax.nn.attention_var_len")
|
||||
op_attrs = attention_node.attrs
|
||||
else:
|
||||
raise ValueError("Cannot find call node for attention")
|
||||
arg = {}
|
||||
|
||||
if "stacked_attention" in op_type:
|
||||
arg["arg0_dtype"] = signature["arg0_dtype"]
|
||||
q_shape = get_const_tuple(attention_node.args[0].ty.shape)
|
||||
k_shape = get_const_tuple(attention_node.args[1].ty.shape)
|
||||
v_shape = get_const_tuple(attention_node.args[2].ty.shape)
|
||||
if len(attention_node.args) == 4:
|
||||
arg["bias_shape"] = get_const_tuple(attention_node.args[3].ty.shape)
|
||||
arg["bias_dtype"] = attention_node.args[3].ty.dtype
|
||||
|
||||
qkv_layout = "qkv_stacked"
|
||||
else:
|
||||
# arg0: q, arg1: k, arg2: v, arg3: bias, arg4: workspace
|
||||
arg["arg0_shape"] = q_shape = signature["arg0_shape"]
|
||||
arg["arg1_shape"] = k_shape = signature["arg1_shape"]
|
||||
arg["arg2_shape"] = v_shape = signature["arg2_shape"]
|
||||
arg["arg0_dtype"] = signature["arg0_dtype"]
|
||||
arg["arg1_dtype"] = signature["arg1_dtype"]
|
||||
arg["arg2_dtype"] = signature["arg2_dtype"]
|
||||
|
||||
if "arg4_dtype" in signature:
|
||||
arg["bias_dtype"] = signature["arg3_dtype"]
|
||||
if "arg4_shape" in signature:
|
||||
arg["bias_shape"] = signature["arg3_shape"]
|
||||
|
||||
qkv_layout = "default"
|
||||
|
||||
out_shape = signature["ret_shape"]
|
||||
out_dtype = signature["ret_dtype"]
|
||||
num_batches, num_queries, num_q_heads, head_dim = q_shape
|
||||
_, num_keys, num_kv_heads, _ = k_shape
|
||||
_, _, _, head_dim_value = v_shape
|
||||
scale = op_attrs.scale
|
||||
|
||||
if op_attrs.causal_mask is None:
|
||||
custom_mask_type = 0
|
||||
elif op_attrs.causal_mask == "TopLeft":
|
||||
custom_mask_type = 1
|
||||
elif op_attrs.causal_mask == "BottomRight":
|
||||
custom_mask_type = 2
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
attrs = {
|
||||
"op_type": op_type,
|
||||
"ret_dtype": out_dtype,
|
||||
"ret_shape": out_shape,
|
||||
"num_batches": num_batches,
|
||||
"num_queries": num_queries,
|
||||
"num_keys": num_keys,
|
||||
"num_q_heads": num_q_heads,
|
||||
"num_kv_heads": num_kv_heads,
|
||||
"head_dim": head_dim,
|
||||
"head_dim_value": head_dim_value,
|
||||
"scale": scale,
|
||||
"arch": self.options["sm"],
|
||||
"qkv_layout": qkv_layout,
|
||||
"custom_mask_type": custom_mask_type,
|
||||
**arg,
|
||||
}
|
||||
|
||||
if "var_len" in op_type:
|
||||
arg_idx = _extract_arg_idx(op_type, f)
|
||||
for arg in ["seqstart_q", "seqstart_k", "max_seqlen_q", "max_seqlen_k"]:
|
||||
if arg in arg_idx:
|
||||
attrs[arg + "_idx"] = arg_idx[arg]
|
||||
|
||||
if op_attrs.window_size:
|
||||
attrs["window_size"] = op_attrs.window_size
|
||||
|
||||
return f.with_attrs(attrs)
|
||||
|
||||
def handle_norm(self, f, op_type):
|
||||
"""Annotate a layer or rms norm op."""
|
||||
signature = _extract_relax_function_signature(f)
|
||||
attrs = {}
|
||||
attrs["batch_rank"] = len(signature["arg0_shape"][:-1])
|
||||
attrs["M"] = reduce(operator.mul, signature["arg0_shape"][:-1], 1)
|
||||
attrs["N"] = signature["arg0_shape"][-1]
|
||||
dtype = signature["arg0_dtype"]
|
||||
attrs["data_type"] = {"float32": "float", "float16": "cutlass::half_t"}[str(dtype)]
|
||||
|
||||
if "rms" in op_type:
|
||||
attrs["rms_eps"] = self.options.get("rms_eps", 1e-5)
|
||||
else:
|
||||
attrs["layer_norm_eps"] = self.options.get("layer_nrom_eps", 1e-5)
|
||||
|
||||
return f.with_attrs(attrs)
|
||||
|
||||
def visit_function_(self, f):
|
||||
if "Composite" not in f.attrs:
|
||||
body = super().visit_expr(f.body)
|
||||
return relax.Function(f.params, body, f.ret_ty, f.is_pure, f.attrs, f.span)
|
||||
|
||||
op_type = f.attrs["Composite"]
|
||||
|
||||
if "conv2d" in op_type:
|
||||
return self.handle_conv2d(f, op_type)
|
||||
elif "decode" in op_type:
|
||||
return self.handle_decode_matmul(f, op_type)
|
||||
elif "matmul" in op_type:
|
||||
return self.handle_matmul(f, op_type)
|
||||
elif "attention" in op_type:
|
||||
return self.handle_attention(f, op_type)
|
||||
elif "layer_norm" in op_type or "rms_norm" in op_type:
|
||||
return self.handle_norm(f, op_type)
|
||||
|
||||
raise ValueError(f"Unsupported composite {op_type}")
|
||||
|
||||
def visit_span(self, span):
|
||||
return span
|
||||
|
||||
|
||||
@register_global_func("contrib.cutlass.tune_relax_function")
|
||||
def profile_relax_function(functions, options):
|
||||
"""Tune and annotate CUTLASS composite functions with shape, dtype and generated templates."""
|
||||
tmp_dir = options.get("tmp_dir", "./tmp")
|
||||
sm = options.get("sm", 80)
|
||||
conv2d_profiler = CutlassConv2DProfiler(sm, _get_cutlass_path(), tmp_dir)
|
||||
gemm_profiler = CutlassGemmProfiler(sm, _get_cutlass_path(), tmp_dir)
|
||||
|
||||
annotated_functions = []
|
||||
|
||||
for f in functions:
|
||||
annotator = CutlassRelaxFunctionAnnotator(
|
||||
tvm.IRModule.from_expr(f), conv2d_profiler, gemm_profiler, options
|
||||
)
|
||||
annotated_functions.append(annotator.visit_expr(f))
|
||||
|
||||
return annotated_functions
|
||||
|
||||
|
||||
@register_global_func("contrib.cutlass.compile")
|
||||
def compile_cutlass_module(c_source_module, options):
|
||||
"""Compile all CUTLASS kernels in the given C-source module.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
c_source_module: runtime.Module
|
||||
A C-source module containing CUTLASS kernels.
|
||||
|
||||
options: dict
|
||||
Compilation options. Currently recognizes
|
||||
"sm": The target architecture (compute capability), for example 75 or 80 (default: 80)
|
||||
"threads": The number of threads to use in NVCC parallel compilation (default:
|
||||
use all logical cores)
|
||||
"use_fast_math": Whether or not to use faster but approximate arithmetic in some
|
||||
CUTLASS epilogues (default: False)
|
||||
|
||||
Returns
|
||||
-------
|
||||
rt_mod : runtime.Module
|
||||
A runtime module where all cutlass kernels have been compiled.
|
||||
"""
|
||||
tmp_dir = options.get("tmp_dir", "./tmp")
|
||||
defaults = {"sm": 80, "threads": -1, "use_fast_math": False}
|
||||
compile_config = {key: options.get(key, val) for key, val in defaults.items()}
|
||||
|
||||
function_names = c_source_module.get_function("get_func_names")()
|
||||
compile_options = _get_cutlass_compile_options(**compile_config)
|
||||
lib_path = os.path.join(tmp_dir, "cutlass.o")
|
||||
logger.info("Compiling generated CUTLASS code")
|
||||
c_source_module.export_library(lib_path, workspace_dir=tmp_dir, **compile_options)
|
||||
|
||||
# Recover static library
|
||||
return tvm.runtime.load_static_library(lib_path, function_names)
|
||||
|
||||
|
||||
def finalize_modules(lib, lib_path="compile.so", tmp_dir="./tmp"):
|
||||
"""Returns lib with any C source, LLVM and static library modules complied and linked in ready
|
||||
for use by the graph or AOT executors. This method is not specific to CUTLASS, however it does
|
||||
assume nvcc will be used for final compilation and linking. It is provided here for
|
||||
convenience.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
lib : runtime.Module
|
||||
The output from build.
|
||||
|
||||
lib_path : string
|
||||
The path to a shared library which will be generated as the result of the build process.
|
||||
|
||||
tmp_dir : string
|
||||
A temporary directory where intermediate compiled artifacts will be stored.
|
||||
|
||||
Returns
|
||||
-------
|
||||
updated_lib : runtime.Module
|
||||
The updated library with all compilation and linking completed.
|
||||
|
||||
"""
|
||||
lib_path = os.path.join(tmp_dir, lib_path)
|
||||
lib.export_library(lib_path, workspace_dir=tmp_dir, cc="nvcc")
|
||||
return runtime.load_module(lib_path)
|
||||
Reference in New Issue
Block a user