wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
apache--tvm/python/tvm/contrib/cutlass/build.py
T
wehub-resource-sync 26446540fa
Lint / lint (push) Waiting to run
Details
CI / MacOS (push) Waiting to run
Details
CI / Windows (push) Waiting to run
Details
chore: import upstream snapshot with attribution
2026-07-13 13:36:25 +08:00

908 lines
30 KiB
Python
Raw Permalink Blame History

# 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 View Git Blame Copy Permalink