chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,20 @@
|
||||
# isort: skip_file
|
||||
# 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.
|
||||
"""BYOC support for CUTLASS."""
|
||||
|
||||
from .build import has_cutlass, num_cutlass_partitions, finalize_modules
|
||||
@@ -0,0 +1,21 @@
|
||||
# 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.
|
||||
"""FFI API for CUTLASS BYOC."""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("contrib.cutlass", __name__)
|
||||
@@ -0,0 +1,327 @@
|
||||
# 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
|
||||
# ruff: noqa: E501
|
||||
"""Generator for CUTLASS attention kernels."""
|
||||
|
||||
from .library import substitute_template
|
||||
|
||||
|
||||
def instantiate_attention_template(attrs):
|
||||
"""Return CUTLASS host code for fused multi head attention
|
||||
based on a template and the provided attribute map."""
|
||||
|
||||
bias_template = """
|
||||
TVM_FFI_ICHECK(${bias}->ndim == 4); // B, N, S, S'
|
||||
|
||||
p.attn_bias_ptr = reinterpret_cast<T *>(${bias}->data);
|
||||
p.bias_strideM = ${bias_strideM};
|
||||
p.bias_strideH = ${bias_strideH};
|
||||
p.bias_strideB = ${bias_strideB};
|
||||
"""
|
||||
|
||||
var_len_template = """
|
||||
p.seqstart_q_ptr = (int32_t*)${seqstart_q}->data;
|
||||
p.seqstart_k_ptr = (int32_t*)${seqstart_k}->data;
|
||||
p.num_queries = ((int32_t*)${max_seqlen_q}->data)[0];
|
||||
p.num_batches = ${seqstart_q}->shape[0] - 1;
|
||||
"""
|
||||
|
||||
qkv_template = {
|
||||
"default": """
|
||||
p.query_ptr = reinterpret_cast<T *>(${query}->data);
|
||||
p.key_ptr = reinterpret_cast<T *>(${key}->data);
|
||||
p.value_ptr = reinterpret_cast<T *>(${value}->data);
|
||||
TVM_FFI_ICHECK(${query}->ndim == 4); // B, S, N, H
|
||||
TVM_FFI_ICHECK(${key}->ndim == 4); // B, S', N, H
|
||||
TVM_FFI_ICHECK(${value}->ndim == 4); // B, S', N, H'
|
||||
|
||||
// stride for N
|
||||
p.q_strideH = p.head_dim; // H
|
||||
p.k_strideH = p.head_dim; // H
|
||||
p.v_strideH = p.head_dim_value; // H'
|
||||
|
||||
// stride for S
|
||||
p.q_strideM = p.q_strideH * p.num_heads; // H * N
|
||||
p.k_strideM = p.k_strideH * p.num_heads; // H * N
|
||||
p.v_strideM = p.v_strideH * p.num_heads; // H' * N
|
||||
|
||||
// stride for B
|
||||
p.q_strideB = p.q_strideM * p.num_queries; // H * N * S
|
||||
p.k_strideB = p.k_strideM * p.num_keys; // H * N * S'
|
||||
p.v_strideB = p.v_strideM * p.num_keys; // H'* N * S'
|
||||
""",
|
||||
"qkv_stacked": """
|
||||
p.query_ptr = reinterpret_cast<T *>(${qkv}->data);
|
||||
p.key_ptr = reinterpret_cast<T *>(${qkv}->data) + p.head_dim * p.num_heads;
|
||||
p.value_ptr = reinterpret_cast<T *>(${qkv}->data) + p.head_dim * p.num_heads * 2;
|
||||
TVM_FFI_ICHECK(${qkv}->ndim == 3); // B, S, NH + NH + NH'
|
||||
|
||||
// stride for N
|
||||
p.q_strideH = p.head_dim; // H
|
||||
p.k_strideH = p.head_dim; // H
|
||||
p.v_strideH = p.head_dim_value; // H'
|
||||
|
||||
// stride for S
|
||||
p.q_strideM = p.k_strideM = p.v_strideM =
|
||||
p.q_strideH * p.num_heads +
|
||||
p.k_strideH * p.num_heads +
|
||||
p.v_strideH * p.num_heads; // H * N + H * N + H * N'
|
||||
|
||||
// stride for B
|
||||
p.q_strideB = p.k_strideB = p.v_strideB =
|
||||
p.q_strideM * p.num_queries; // (H * N + H * N + H * N') * S
|
||||
""",
|
||||
}
|
||||
|
||||
template = """
|
||||
using T = ${data_type};
|
||||
|
||||
using Attention =
|
||||
AttentionKernel<T,
|
||||
/*ArchTag=*/${arch},
|
||||
/*is_aligned=*/${kIsAligned},
|
||||
/*queries_per_block=*/${kQueriesPerBlock},
|
||||
/*keys_per_block=*/${kKeysPerBlock},
|
||||
/*kMaxK=*/${kMaxK},
|
||||
/*supports_dropout=*/${kSupportsDropout},
|
||||
/*supports_bias=*/${kSupportsBias}
|
||||
>;
|
||||
|
||||
typename Attention::Params p;
|
||||
p.logsumexp_ptr = nullptr;
|
||||
p.output_ptr = reinterpret_cast<T *>(out0->data);
|
||||
|
||||
p.output_accum_ptr = nullptr;
|
||||
uint64_t accumulator_buf_size = ${output_size} * sizeof(Attention::output_accum_t);
|
||||
bool accumulator_buf_allocated = false;
|
||||
if (Attention::kNeedsOutputAccumulatorBuffer) {
|
||||
if (accumulator_buf_size <= ${workspace}->shape[0]) {
|
||||
p.output_accum_ptr = static_cast<float*>(${workspace}->data);
|
||||
} else {
|
||||
accumulator_buf_allocated = true;
|
||||
cudaMalloc(
|
||||
&p.output_accum_ptr,
|
||||
accumulator_buf_size
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
p.num_heads = ${num_heads}; // N
|
||||
p.num_batches = ${num_batches}; // B
|
||||
p.head_dim = ${head_dim}; // H
|
||||
p.head_dim_value = ${head_dim_value}; // H'
|
||||
p.num_queries = ${num_queries}; // S
|
||||
p.num_keys = ${num_keys}; // S'
|
||||
p.scale = ${scale};
|
||||
p.custom_mask_type = ${custom_mask_type};
|
||||
|
||||
|
||||
p.o_strideM = p.head_dim_value * p.num_heads; // H' * N
|
||||
TVM_FFI_ICHECK(out0->ndim == 4); // B, S, N, H'
|
||||
|
||||
${qkv_template}
|
||||
${bias_template}
|
||||
${var_len_template}
|
||||
|
||||
constexpr auto kernel_fn = attention_kernel_batched_impl<Attention>;
|
||||
int smem_bytes = sizeof(typename Attention::SharedStorage);
|
||||
if (smem_bytes > 0xc000) {
|
||||
static bool once = [&]() {
|
||||
cudaFuncSetAttribute(
|
||||
kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
|
||||
return true;
|
||||
}();
|
||||
}
|
||||
|
||||
TVM_FFI_ICHECK(Attention::check_supported(p));
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
|
||||
|
||||
kernel_fn<<<p.getBlocksGrid(), p.getThreadsGrid(), smem_bytes, stream>>>(p);
|
||||
|
||||
if (accumulator_buf_allocated) {
|
||||
cudaFree(p.output_accum_ptr);
|
||||
}
|
||||
"""
|
||||
|
||||
template = substitute_template(
|
||||
template,
|
||||
{
|
||||
"qkv_template": qkv_template[attrs["qkv_layout"]],
|
||||
"bias_template": bias_template if "bias" in attrs else "",
|
||||
"var_len_template": var_len_template if "seqstart_q" in attrs else "",
|
||||
},
|
||||
)
|
||||
|
||||
return substitute_template(template, attrs)
|
||||
|
||||
|
||||
def instantiate_flash_attention_template(attrs):
|
||||
"""Return host code for flash attention."""
|
||||
|
||||
template = """
|
||||
int q_head_stride = ${head_dim};
|
||||
int k_head_stride = ${head_dim};
|
||||
int v_head_stride = ${head_dim};
|
||||
int o_head_stride = ${head_dim};
|
||||
int q_row_stride = q_head_stride * ${num_q_heads};
|
||||
int k_row_stride = k_head_stride * ${num_kv_heads};
|
||||
int v_row_stride = v_head_stride * ${num_kv_heads};
|
||||
int o_row_stride = o_head_stride * ${num_q_heads};
|
||||
int q_batch_stride = q_row_stride * ${num_queries};
|
||||
int k_batch_stride = k_row_stride * ${num_keys};
|
||||
int v_batch_stride = v_row_stride * ${num_keys};
|
||||
int o_batch_stride = o_row_stride * ${num_queries};
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
|
||||
|
||||
flash_attn::flash_attention_forward(
|
||||
static_cast<const cutlass::half_t*>(${query}->data),
|
||||
static_cast<const cutlass::half_t*>(${key}->data),
|
||||
static_cast<const cutlass::half_t*>(${value}->data),
|
||||
static_cast<cutlass::half_t*>(out0->data),
|
||||
${num_batches},
|
||||
${num_queries},
|
||||
${num_keys},
|
||||
${num_q_heads},
|
||||
${num_kv_heads},
|
||||
${head_dim},
|
||||
q_batch_stride,
|
||||
k_batch_stride,
|
||||
v_batch_stride,
|
||||
o_batch_stride,
|
||||
q_head_stride,
|
||||
k_head_stride,
|
||||
v_head_stride,
|
||||
o_head_stride,
|
||||
q_row_stride,
|
||||
k_row_stride,
|
||||
v_row_stride,
|
||||
o_row_stride,
|
||||
${scale},
|
||||
${is_causal},
|
||||
${window_size_left},
|
||||
${window_size_right},
|
||||
stream);
|
||||
"""
|
||||
|
||||
template_stacked = """
|
||||
int q_head_stride = ${head_dim};
|
||||
int k_head_stride = ${head_dim};
|
||||
int v_head_stride = ${head_dim};
|
||||
int o_head_stride = ${head_dim};
|
||||
int row_stride = q_head_stride * ${num_q_heads} +
|
||||
k_head_stride * ${num_kv_heads} +
|
||||
v_head_stride * ${num_kv_heads};
|
||||
int q_row_stride = row_stride;
|
||||
int k_row_stride = row_stride;
|
||||
int v_row_stride = row_stride;
|
||||
int o_row_stride = o_head_stride * ${num_q_heads};
|
||||
|
||||
int q_batch_stride = q_row_stride * ${num_queries};
|
||||
int k_batch_stride = k_row_stride * ${num_keys};
|
||||
int v_batch_stride = v_row_stride * ${num_keys};
|
||||
int o_batch_stride = o_row_stride * ${num_queries};
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
|
||||
|
||||
flash_attn::flash_attention_forward(
|
||||
static_cast<const cutlass::half_t*>(${qkv}->data),
|
||||
static_cast<const cutlass::half_t*>(${qkv}->data) + ${head_dim} * ${num_q_heads},
|
||||
static_cast<const cutlass::half_t*>(${qkv}->data) + ${head_dim} * (${num_q_heads} + ${num_kv_heads}),
|
||||
static_cast<cutlass::half_t*>(out0->data),
|
||||
${num_batches},
|
||||
${num_queries},
|
||||
${num_keys},
|
||||
${num_q_heads},
|
||||
${num_kv_heads},
|
||||
${head_dim},
|
||||
q_batch_stride,
|
||||
k_batch_stride,
|
||||
v_batch_stride,
|
||||
o_batch_stride,
|
||||
q_head_stride,
|
||||
k_head_stride,
|
||||
v_head_stride,
|
||||
o_head_stride,
|
||||
q_row_stride,
|
||||
k_row_stride,
|
||||
v_row_stride,
|
||||
o_row_stride,
|
||||
${scale},
|
||||
${is_causal},
|
||||
${window_size_left},
|
||||
${window_size_right},
|
||||
stream);
|
||||
"""
|
||||
|
||||
if "qkv" in attrs:
|
||||
return substitute_template(template_stacked, attrs)
|
||||
|
||||
return substitute_template(template, attrs)
|
||||
|
||||
|
||||
def instantiate_flash_attention_var_len_template(attrs):
|
||||
"""Return host code for flash attention with variable sequence lengths."""
|
||||
|
||||
template = """
|
||||
int _max_seqlen_q = ((int32_t*)${max_seqlen_q}->data)[0];
|
||||
int _max_seqlen_k = ((int32_t*)${max_seqlen_k}->data)[0];
|
||||
|
||||
int batch_size = ${seqstart_q}->shape[0] - 1;
|
||||
|
||||
int q_head_stride = ${head_dim};
|
||||
int k_head_stride = ${head_dim};
|
||||
int v_head_stride = ${head_dim};
|
||||
int o_head_stride = ${head_dim};
|
||||
int q_row_stride = q_head_stride * ${num_q_heads};
|
||||
int k_row_stride = k_head_stride * ${num_kv_heads};
|
||||
int v_row_stride = v_head_stride * ${num_kv_heads};
|
||||
int o_row_stride = o_head_stride * ${num_q_heads};
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
|
||||
|
||||
flash_attn::flash_attention_var_len_forward(
|
||||
static_cast<const cutlass::half_t*>(${query}->data),
|
||||
static_cast<const cutlass::half_t*>(${key}->data),
|
||||
static_cast<const cutlass::half_t*>(${value}->data),
|
||||
static_cast<const int*>(${seqstart_q}->data),
|
||||
static_cast<const int*>(${seqstart_k}->data),
|
||||
static_cast<cutlass::half_t*>(out0->data),
|
||||
batch_size,
|
||||
_max_seqlen_q,
|
||||
_max_seqlen_k,
|
||||
${num_q_heads},
|
||||
${num_kv_heads},
|
||||
${head_dim},
|
||||
q_head_stride,
|
||||
k_head_stride,
|
||||
v_head_stride,
|
||||
o_head_stride,
|
||||
q_row_stride,
|
||||
k_row_stride,
|
||||
v_row_stride,
|
||||
o_row_stride,
|
||||
${scale},
|
||||
${is_causal},
|
||||
// For SWA, is_causal must be false.
|
||||
${is_causal} ? _max_seqlen_k : ${window_size_left},
|
||||
${window_size_right},
|
||||
stream);
|
||||
"""
|
||||
|
||||
return substitute_template(template, attrs)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,552 @@
|
||||
# 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, unused-wildcard-import, wildcard-import
|
||||
# ruff: noqa: E501, F403, F405
|
||||
"""Generator for CUTLASS Conv2D kernels."""
|
||||
|
||||
from .library import *
|
||||
|
||||
|
||||
class Conv2dOperation:
|
||||
"""Describes various attributes for instantiating Conv2d kernels."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conv_kind,
|
||||
iterator_algorithm,
|
||||
arch,
|
||||
tile_description,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
element_epilogue,
|
||||
stride_support,
|
||||
epilogue_functor=EpilogueFunctor.LinearCombination,
|
||||
swizzling_functor=SwizzlingFunctor.Identity1,
|
||||
split_k_slices=1,
|
||||
):
|
||||
self.operation_kind = OperationKind.Conv2d
|
||||
self.arch = arch
|
||||
self.tile_description = tile_description
|
||||
self.conv_kind = conv_kind
|
||||
self.A = A
|
||||
self.B = B
|
||||
self.C = C
|
||||
self.element_epilogue = element_epilogue
|
||||
self.epilogue_functor = epilogue_functor
|
||||
self.iterator_algorithm = iterator_algorithm
|
||||
self.stride_support = stride_support
|
||||
self.swizzling_functor = swizzling_functor
|
||||
self.split_k_slices = split_k_slices
|
||||
|
||||
def accumulator_type(self):
|
||||
return self.tile_description.math_instruction.element_accumulator
|
||||
|
||||
def core_name(self):
|
||||
"""The basic operation kind is prefixed with a letter indicating the accumulation type."""
|
||||
intermediate_type = ""
|
||||
|
||||
if self.tile_description.math_instruction.opcode_class == OpcodeClass.TensorOp:
|
||||
inst_shape = "{}{}{}".format(*self.tile_description.math_instruction.instruction_shape)
|
||||
if (
|
||||
self.tile_description.math_instruction.element_a != self.A.element
|
||||
and self.tile_description.math_instruction.element_a != self.accumulator_type()
|
||||
):
|
||||
intermediate_type = DataTypeNames[self.tile_description.math_instruction.element_a]
|
||||
else:
|
||||
inst_shape = ""
|
||||
|
||||
return f"{ShortDataTypeNames[self.accumulator_type()]}{inst_shape}{intermediate_type}{ConvKindNames[self.conv_kind]}_{IteratorAlgorithmNames[self.iterator_algorithm]}"
|
||||
|
||||
def extended_name(self):
|
||||
"""Append data types if they differ from compute type."""
|
||||
if (
|
||||
self.C.element != self.tile_description.math_instruction.element_accumulator
|
||||
and self.A.element != self.tile_description.math_instruction.element_accumulator
|
||||
):
|
||||
extended_name = "${element_c}_${core_name}_${element_a}"
|
||||
elif (
|
||||
self.C.element == self.tile_description.math_instruction.element_accumulator
|
||||
and self.A.element != self.tile_description.math_instruction.element_accumulator
|
||||
):
|
||||
extended_name = "${core_name}_${element_a}"
|
||||
else:
|
||||
extended_name = "${core_name}"
|
||||
|
||||
extended_name = substitute_template(
|
||||
extended_name,
|
||||
{
|
||||
"element_a": DataTypeNames[self.A.element],
|
||||
"element_c": DataTypeNames[self.C.element],
|
||||
"core_name": self.core_name(),
|
||||
},
|
||||
)
|
||||
|
||||
return extended_name
|
||||
|
||||
def layout_name(self):
|
||||
return f"{ShortLayoutTypeNames[self.A.layout]}"
|
||||
|
||||
def procedural_name(self):
|
||||
"""
|
||||
The full procedural name indicates architecture, extended name, tile size, and layout.
|
||||
"""
|
||||
opcode_class_name = OpcodeClassNames[self.tile_description.math_instruction.opcode_class]
|
||||
|
||||
threadblock = f"{self.tile_description.threadblock_shape[0]}x{self.tile_description.threadblock_shape[1]}_{self.tile_description.threadblock_shape[2]}x{self.tile_description.stages}"
|
||||
|
||||
if self.stride_support == StrideSupport.Unity:
|
||||
configuration_name = (
|
||||
"cutlass_${opcode_class}_${extended_name}_${threadblock}"
|
||||
"_${layout}_align${alignment}_unity_stride"
|
||||
)
|
||||
else:
|
||||
configuration_name = (
|
||||
"cutlass_${opcode_class}_${extended_name}_${threadblock}"
|
||||
"_${layout}_align${alignment}"
|
||||
)
|
||||
|
||||
if self.split_k_slices > 1:
|
||||
configuration_name += f"_splitk{self.split_k_slices}"
|
||||
|
||||
return substitute_template(
|
||||
configuration_name,
|
||||
{
|
||||
"opcode_class": opcode_class_name,
|
||||
"extended_name": self.extended_name(),
|
||||
"threadblock": threadblock,
|
||||
"layout": self.layout_name(),
|
||||
"alignment": f"{self.A.alignment}",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class EmitConv2dInstance:
|
||||
"""Responsible for emitting a CUTLASS template definition."""
|
||||
|
||||
def __init__(self):
|
||||
self.epilogue_default = """
|
||||
${epilogue_functor}<
|
||||
${element_c},
|
||||
${epilogue_vector_length},
|
||||
${element_accumulator},
|
||||
${element_epilogue}
|
||||
>"""
|
||||
|
||||
self.epilogue_no_beta_scaling = """
|
||||
${epilogue_functor}<
|
||||
${element_c},
|
||||
${epilogue_vector_length},
|
||||
${element_accumulator},
|
||||
${element_epilogue},
|
||||
cutlass::epilogue::thread::ScaleType::NoBetaScaling
|
||||
>"""
|
||||
|
||||
self.epilogue_residual_block = """
|
||||
${epilogue_functor}<
|
||||
${element_c},
|
||||
${element_accumulator},
|
||||
${element_epilogue},
|
||||
${element_c},
|
||||
${epilogue_vector_length},
|
||||
${activation},
|
||||
${binary_op},
|
||||
${unary_op}
|
||||
>"""
|
||||
|
||||
self.epilogue_wgrad = """
|
||||
${epilogue_functor}<
|
||||
${element_c},
|
||||
4,
|
||||
float,
|
||||
float
|
||||
>"""
|
||||
|
||||
self.template = """
|
||||
// Conv2d${conv_kind_name} ${iterator_algorithm_name} kernel instance "${operation_name}"
|
||||
using ${operation_name} =
|
||||
typename cutlass::conv::kernel::DefaultConv2d${conv_kind_name}${conv_kernel_postfix}<
|
||||
${element_a},
|
||||
${layout_a},
|
||||
${element_b},
|
||||
${layout_b},
|
||||
${element_c},
|
||||
${layout_c},
|
||||
${element_accumulator},
|
||||
${opcode_class},
|
||||
${arch},
|
||||
cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,
|
||||
cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k} >,
|
||||
cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,
|
||||
${epilogue},
|
||||
${swizzling_functor}, // cutlass::gemm::threadblock::GemmSplitKIdentityThreadblockSwizzle<>,
|
||||
${stages},
|
||||
${math_operator},
|
||||
${iterator_algorithm},
|
||||
${stride_support},
|
||||
${align_a},
|
||||
${align_b}
|
||||
>::Kernel;
|
||||
|
||||
${reduction}
|
||||
"""
|
||||
|
||||
self.reduction_template = """
|
||||
using EpilogueOutputOp = ${epilogue};
|
||||
using ReductionOp = cutlass::reduction::thread::ReduceAdd<
|
||||
${element_accumulator},
|
||||
${element_accumulator},
|
||||
EpilogueOutputOp::kCount
|
||||
>;
|
||||
|
||||
using ReductionKernel = cutlass::reduction::kernel::ReduceSplitK<
|
||||
cutlass::MatrixShape<4, 32 * EpilogueOutputOp::kCount>,
|
||||
EpilogueOutputOp,
|
||||
ReductionOp
|
||||
>;
|
||||
|
||||
using ReductionDevice = cutlass::reduction::device::ReduceSplitK<ReductionKernel>;
|
||||
using ReductionStrideIndex = typename ReductionDevice::StrideIndex;
|
||||
"""
|
||||
|
||||
def emit(
|
||||
self, operation, no_beta_scaling=False, residual_block_info=False, emit_reduction=False
|
||||
):
|
||||
"""Instantiate a Conv2d kernel from given `operation`."""
|
||||
warp_shape = [
|
||||
int(
|
||||
operation.tile_description.threadblock_shape[idx]
|
||||
/ operation.tile_description.warp_count[idx]
|
||||
)
|
||||
for idx in range(3)
|
||||
]
|
||||
|
||||
epilogue_vector_length = int(
|
||||
min(operation.C.alignment * DataTypeSize[operation.C.element], 128)
|
||||
/ DataTypeSize[operation.C.element]
|
||||
)
|
||||
|
||||
element_c = operation.C.element
|
||||
use_split_k_wgrad = operation.conv_kind == ConvKind.Wgrad and operation.split_k_slices > 1
|
||||
# Gemm output always fp32 in wgrad with split k
|
||||
element_c_gemm = DataType.f32 if use_split_k_wgrad else element_c
|
||||
|
||||
if emit_reduction:
|
||||
epilogue_reduction = substitute_template(
|
||||
self.epilogue_wgrad,
|
||||
{
|
||||
"epilogue_functor": EpilogueFunctorTag[operation.epilogue_functor],
|
||||
"element_c": DataTypeTag[element_c],
|
||||
},
|
||||
)
|
||||
reduction = substitute_template(
|
||||
self.reduction_template,
|
||||
{
|
||||
"epilogue": epilogue_reduction,
|
||||
"operation_name": operation.procedural_name(),
|
||||
"element_accumulator": DataTypeTag[operation.accumulator_type()],
|
||||
},
|
||||
)
|
||||
gemm_template = substitute_template(self.template, {"reduction": reduction})
|
||||
else:
|
||||
gemm_template = substitute_template(self.template, {"reduction": ""})
|
||||
|
||||
values = {
|
||||
"operation_name": operation.procedural_name(),
|
||||
"conv_kind": ConvKindTag[operation.conv_kind],
|
||||
"conv_kind_name": ConvKindNames[operation.conv_kind].capitalize(),
|
||||
"element_a": DataTypeTag[operation.A.element],
|
||||
"layout_a": LayoutTag[operation.A.layout],
|
||||
"element_b": DataTypeTag[operation.B.element],
|
||||
"layout_b": LayoutTag[operation.B.layout],
|
||||
"element_c": DataTypeTag[element_c_gemm],
|
||||
"layout_c": LayoutTag[operation.C.layout],
|
||||
"element_accumulator": DataTypeTag[operation.accumulator_type()],
|
||||
"opcode_class": OpcodeClassTag[
|
||||
operation.tile_description.math_instruction.opcode_class
|
||||
],
|
||||
"arch": f"cutlass::arch::Sm{operation.arch}",
|
||||
"threadblock_shape_m": str(operation.tile_description.threadblock_shape[0]),
|
||||
"threadblock_shape_n": str(operation.tile_description.threadblock_shape[1]),
|
||||
"threadblock_shape_k": str(operation.tile_description.threadblock_shape[2]),
|
||||
"warp_shape_m": str(warp_shape[0]),
|
||||
"warp_shape_n": str(warp_shape[1]),
|
||||
"warp_shape_k": str(warp_shape[2]),
|
||||
"instruction_shape_m": str(
|
||||
operation.tile_description.math_instruction.instruction_shape[0]
|
||||
),
|
||||
"instruction_shape_n": str(
|
||||
operation.tile_description.math_instruction.instruction_shape[1]
|
||||
),
|
||||
"instruction_shape_k": str(
|
||||
operation.tile_description.math_instruction.instruction_shape[2]
|
||||
),
|
||||
"epilogue_vector_length": str(epilogue_vector_length),
|
||||
"epilogue_functor": EpilogueFunctorTag[operation.epilogue_functor],
|
||||
"element_epilogue": str(DataTypeTag[operation.element_epilogue]),
|
||||
"swizzling_functor": SwizzlingFunctorTag[operation.swizzling_functor],
|
||||
"stages": str(operation.tile_description.stages),
|
||||
"iterator_algorithm": IteratorAlgorithmTag[operation.iterator_algorithm],
|
||||
"iterator_algorithm_name": IteratorAlgorithmNames[
|
||||
operation.iterator_algorithm
|
||||
].capitalize(),
|
||||
"stride_support": StrideSupportTag[operation.stride_support],
|
||||
"math_operator": MathOperationTag[
|
||||
operation.tile_description.math_instruction.math_operation
|
||||
],
|
||||
"align_a": str(operation.A.alignment),
|
||||
"align_b": str(operation.B.alignment),
|
||||
"conv_kernel_postfix": "",
|
||||
}
|
||||
|
||||
if use_split_k_wgrad:
|
||||
# Even if the output is fp16, gemm output is always fp32 for split k wgrad.
|
||||
epilogue_gemm = substitute_template(
|
||||
self.epilogue_wgrad,
|
||||
{
|
||||
"epilogue_functor": EpilogueFunctorTag[operation.epilogue_functor],
|
||||
"element_c": "float",
|
||||
},
|
||||
)
|
||||
template = substitute_template(gemm_template, {"epilogue": epilogue_gemm})
|
||||
elif residual_block_info:
|
||||
template = substitute_template(
|
||||
gemm_template, {"epilogue": self.epilogue_residual_block}
|
||||
)
|
||||
values.update(
|
||||
{
|
||||
"unary_op": residual_block_info["unary_op"],
|
||||
"binary_op": residual_block_info["binary_op"],
|
||||
"activation": residual_block_info["activation"],
|
||||
"conv_kernel_postfix": "WithBroadcast",
|
||||
}
|
||||
)
|
||||
elif no_beta_scaling:
|
||||
template = substitute_template(
|
||||
gemm_template, {"epilogue": self.epilogue_no_beta_scaling}
|
||||
)
|
||||
else:
|
||||
template = substitute_template(gemm_template, {"epilogue": self.epilogue_default})
|
||||
|
||||
return substitute_template(template, values)
|
||||
|
||||
|
||||
def instantiate_conv2d_template(attrs):
|
||||
"""Return CUTLASS host code for conv2d based on a template and the provided attribute map."""
|
||||
template = """
|
||||
${cutlass_op_def}
|
||||
|
||||
using Conv2d = cutlass::conv::device::ImplicitGemmConvolution<${cutlass_op_name}>;
|
||||
using ElementInputA = Conv2d::ElementA;
|
||||
using ElementInputB = Conv2d::ElementB;
|
||||
using ElementComputeEpilogue = Conv2d::ElementAccumulator;
|
||||
int N = ${N};
|
||||
int H = ${H};
|
||||
int W = ${W};
|
||||
int C = ${C};
|
||||
int K = ${K};
|
||||
int R = ${R};
|
||||
int S = ${S};
|
||||
int P = ${P};
|
||||
int Q = ${Q};
|
||||
int pad_h = ${pad_h};
|
||||
int pad_w = ${pad_w};
|
||||
int stride_h = ${stride_h};
|
||||
int stride_w = ${stride_w};
|
||||
int dilation_h = ${dilation_h};
|
||||
int dilation_w = ${dilation_w};
|
||||
int split_k_slices = ${split_k_slices};
|
||||
cutlass::conv::Conv2dProblemSize problem_size(N, H, W, C, K, R, S, P, Q, pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w, cutlass::conv::Mode::kCrossCorrelation, split_k_slices);
|
||||
const cutlass::conv::SplitKMode split_k_mode = cutlass::conv::SplitKMode::${split_k_mode};
|
||||
|
||||
void* ptr_a = (void*)(${data_arg}->data);
|
||||
void* ptr_b = (void*)(${weight_arg}->data);
|
||||
${bias_decl}
|
||||
${residual_decl}
|
||||
void* ptr_out = (void*)(out0->data);
|
||||
|
||||
ElementComputeEpilogue alpha = ElementComputeEpilogue(1);
|
||||
ElementComputeEpilogue beta = ElementComputeEpilogue(${beta});
|
||||
using cutlass::layout::TensorNHWC;
|
||||
auto activation_shape = TensorNHWC::packed(cutlass::make_Coord(N, H, W, C));
|
||||
auto weight_shape = TensorNHWC::packed(cutlass::make_Coord(K, R, S, C));
|
||||
auto output_shape = TensorNHWC::packed(cutlass::make_Coord(N, P, Q, K));
|
||||
${residual_shape_decl}
|
||||
|
||||
TensorNHWC layout_A(${A_shape});
|
||||
TensorNHWC layout_B(${B_shape});
|
||||
TensorNHWC layout_C(${C_shape});
|
||||
TensorNHWC layout_D(${D_shape});
|
||||
|
||||
using ElementOutput = ${ElementOutput};
|
||||
cutlass::TensorRef<ElementOutput, TensorNHWC> tensor_c{static_cast<ElementOutput*>(${tensor_c}), ${tensor_c_layout}};
|
||||
cutlass::TensorRef<ElementOutput, TensorNHWC> tensor_d{static_cast<ElementOutput*>(ptr_out), layout_D};
|
||||
typename Conv2d::Arguments arguments{
|
||||
problem_size,
|
||||
{static_cast<ElementInputA*>(ptr_a), layout_A},
|
||||
{static_cast<ElementInputB*>(ptr_b), layout_B},
|
||||
${tensor_c_arg},
|
||||
${tensor_d_arg},
|
||||
{${alpha_beta}},
|
||||
split_k_mode
|
||||
${additional_args}
|
||||
};
|
||||
Conv2d conv2d_op;
|
||||
size_t workspace_size = conv2d_op.get_workspace_size(arguments);
|
||||
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
|
||||
cutlass::Status status = conv2d_op.can_implement(arguments);
|
||||
TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
|
||||
${split_k_reset}
|
||||
status = conv2d_op.initialize(arguments, workspace.get());
|
||||
TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
|
||||
${split_k_update}
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${data_arg}->device.device_id));
|
||||
|
||||
status = conv2d_op(stream);
|
||||
TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
|
||||
${split_k_reduction}
|
||||
"""
|
||||
|
||||
split_k_reset = """
|
||||
arguments.ref_D.reset(reinterpret_cast<ElementComputeEpilogue*>(workspace.get()), layout_D);
|
||||
"""
|
||||
|
||||
split_k_update = """
|
||||
arguments.output_op = {ElementComputeEpilogue(1), ElementComputeEpilogue(0)};
|
||||
status = conv2d_op.update(arguments, workspace.get());
|
||||
TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
|
||||
"""
|
||||
|
||||
split_k_reduction = """
|
||||
ReductionDevice reduction_op;
|
||||
const static cutlass::conv::Operator kConvolutionalOperator = Conv2d::kConvolutionalOperator;
|
||||
typename ReductionDevice::Arguments reduction_args(
|
||||
cutlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, problem_size).mn(),
|
||||
problem_size.split_k_slices,
|
||||
cutlass::conv::implicit_gemm_tensor_c_size(kConvolutionalOperator, problem_size),
|
||||
{
|
||||
reinterpret_cast<Conv2d::ElementAccumulator*> (workspace.get()),
|
||||
ReductionStrideIndex(tensor_c.stride()[Conv2d::UnderlyingKernel::kTensorCStrideIdx])
|
||||
},
|
||||
{
|
||||
tensor_d.data(),
|
||||
ReductionStrideIndex(tensor_d.stride()[Conv2d::UnderlyingKernel::kTensorCStrideIdx])
|
||||
},
|
||||
{
|
||||
tensor_c.data(),
|
||||
ReductionStrideIndex(tensor_c.stride()[Conv2d::UnderlyingKernel::kTensorCStrideIdx])
|
||||
},
|
||||
{alpha, beta}
|
||||
);
|
||||
status = reduction_op.initialize(reduction_args, nullptr);
|
||||
status = reduction_op();
|
||||
"""
|
||||
op_type = attrs["op_type"]
|
||||
has_bias = "bias" in op_type
|
||||
use_split_k = "splitk" in attrs["cutlass_op_name"]
|
||||
is_wgrad = "backward_weight" in op_type
|
||||
is_dgrad = "conv2d_transpose" in op_type
|
||||
has_residual_block = "residual" in op_type
|
||||
no_bias_scaling = op_type not in [
|
||||
"cutlass.conv2d_bias_sigmoid",
|
||||
"cutlass.conv2d_bias_silu",
|
||||
"cutlass.conv2d_bias_hardswish",
|
||||
]
|
||||
|
||||
aux_map = {}
|
||||
|
||||
if (not has_bias or no_bias_scaling) and not has_residual_block:
|
||||
aux_map["beta"] = 0
|
||||
else:
|
||||
aux_map["beta"] = 1
|
||||
|
||||
if has_residual_block:
|
||||
aux_map["bias_decl"] = "void* ptr_bias = (void*)(${bias_arg}->data);\n"
|
||||
aux_map["residual_decl"] = "void* ptr_residual = (void*)(${residual_arg}->data);"
|
||||
aux_map["tensor_c"] = "ptr_residual"
|
||||
aux_map["tensor_c_layout"] = "layout_C"
|
||||
elif has_bias:
|
||||
aux_map["bias_decl"] = "void* ptr_c_bias = (void*)(${bias_arg}->data);\n"
|
||||
aux_map["residual_decl"] = ""
|
||||
aux_map["tensor_c"] = "ptr_c_bias"
|
||||
aux_map["tensor_c_layout"] = "cutlass::layout::TensorNHWC::Stride(0)"
|
||||
else:
|
||||
aux_map["bias_decl"] = ""
|
||||
aux_map["residual_decl"] = ""
|
||||
aux_map["tensor_c"] = "ptr_out"
|
||||
aux_map["tensor_c_layout"] = "layout_C"
|
||||
|
||||
if has_bias and no_bias_scaling and not has_residual_block:
|
||||
aux_map["alpha_beta"] = "alpha"
|
||||
else:
|
||||
aux_map["alpha_beta"] = "alpha, beta"
|
||||
|
||||
if has_residual_block:
|
||||
aux_map["additional_args"] = ", static_cast<ElementOutput*>(ptr_bias), nullptr, 0, K"
|
||||
else:
|
||||
aux_map["additional_args"] = ""
|
||||
|
||||
aux_map["residual_shape_decl"] = ""
|
||||
|
||||
if is_wgrad:
|
||||
aux_map["A_shape"] = "output_shape"
|
||||
aux_map["B_shape"] = "activation_shape"
|
||||
aux_map["C_shape"] = "weight_shape"
|
||||
aux_map["D_shape"] = "weight_shape"
|
||||
elif is_dgrad:
|
||||
aux_map["A_shape"] = "output_shape"
|
||||
aux_map["B_shape"] = "weight_shape"
|
||||
aux_map["C_shape"] = "activation_shape"
|
||||
aux_map["D_shape"] = "activation_shape"
|
||||
else:
|
||||
aux_map["A_shape"] = "activation_shape"
|
||||
aux_map["B_shape"] = "weight_shape"
|
||||
aux_map["D_shape"] = "output_shape"
|
||||
|
||||
if has_residual_block:
|
||||
res_shape = list(attrs.pop("residual_shape"))
|
||||
shape_str = f"cutlass::make_Coord({res_shape[0]}, {res_shape[1]}, {res_shape[2]}, K)"
|
||||
aux_map["residual_shape_decl"] = (
|
||||
f"auto residual_shape = TensorNHWC::packed({shape_str});"
|
||||
)
|
||||
aux_map["C_shape"] = "residual_shape"
|
||||
|
||||
if res_shape == [int(attrs[c]) for c in ["N", "H", "W", "K"]]:
|
||||
aux_map["tensor_c_layout"] = "layout_C"
|
||||
else:
|
||||
# bias-like residual input
|
||||
aux_map["tensor_c_layout"] = "cutlass::layout::TensorNHWC::Stride(0)"
|
||||
else:
|
||||
aux_map["C_shape"] = "output_shape"
|
||||
|
||||
if use_split_k:
|
||||
aux_map["ElementOutput"] = "EpilogueOutputOp::ElementOutput"
|
||||
aux_map["tensor_c_arg"] = "{nullptr, TensorNHWC()}"
|
||||
aux_map["tensor_d_arg"] = "{nullptr, TensorNHWC()}"
|
||||
aux_map["split_k_reset"] = split_k_reset
|
||||
aux_map["split_k_update"] = split_k_update
|
||||
aux_map["split_k_reduction"] = split_k_reduction
|
||||
else:
|
||||
aux_map["ElementOutput"] = "Conv2d::ElementC"
|
||||
aux_map["tensor_c_arg"] = "tensor_c"
|
||||
aux_map["tensor_d_arg"] = "tensor_d"
|
||||
aux_map["split_k_reset"] = aux_map["split_k_update"] = aux_map["split_k_reduction"] = ""
|
||||
|
||||
template = substitute_template(template, aux_map)
|
||||
|
||||
return substitute_template(template, attrs)
|
||||
@@ -0,0 +1,216 @@
|
||||
# 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=import-outside-toplevel, invalid-name
|
||||
# ruff: noqa: E501
|
||||
"""Instantiate a C++ source for profiling CUTLASS kernels."""
|
||||
|
||||
from .library import DataTypeTag
|
||||
|
||||
|
||||
class Conv2dProfilerEmitter:
|
||||
"""Emit a C++ source for profiling CUTLASS kernels."""
|
||||
|
||||
def __init__(self):
|
||||
from jinja2 import Template
|
||||
|
||||
self.reduction = """
|
||||
ReductionDevice reduction_op;
|
||||
static cutlass::conv::Operator const kConvolutionalOperator = ImplicitGemm::kConvolutionalOperator;
|
||||
typename ReductionDevice::Arguments reduction_args(
|
||||
cutlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, problem_size).mn(),
|
||||
problem_size.split_k_slices,
|
||||
cutlass::conv::implicit_gemm_tensor_c_size(kConvolutionalOperator, problem_size),
|
||||
{
|
||||
reinterpret_cast<ImplicitGemm::ElementC*> (workspace.get()),
|
||||
ReductionStrideIndex(tensor_c.stride()[ImplicitGemm::UnderlyingKernel::kTensorCStrideIdx])
|
||||
},
|
||||
{
|
||||
tensor_d.device_data(),
|
||||
ReductionStrideIndex(tensor_d.stride()[ImplicitGemm::UnderlyingKernel::kTensorCStrideIdx])
|
||||
},
|
||||
{
|
||||
tensor_c.device_data(),
|
||||
ReductionStrideIndex(tensor_c.stride()[ImplicitGemm::UnderlyingKernel::kTensorCStrideIdx])
|
||||
},
|
||||
{ElementComputeEpilogue(1), ElementComputeEpilogue(0)}
|
||||
);
|
||||
|
||||
reduction_op.initialize(reduction_args, nullptr);
|
||||
reduction_op();
|
||||
"""
|
||||
|
||||
self.template = Template(
|
||||
"""
|
||||
#include <iostream>
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/conv/kernel/default_conv2d_fprop.h"
|
||||
#include "cutlass/conv/kernel/default_conv2d_wgrad.h"
|
||||
#include "cutlass/conv/kernel/default_conv2d_dgrad.h"
|
||||
#include "cutlass/conv/device/implicit_gemm_convolution.h"
|
||||
#include "cutlass/util/command_line.h"
|
||||
#include "cutlass/util/host_tensor.h"
|
||||
#include "cutlass/util/reference/host/tensor_fill.h"
|
||||
#include "cutlass/reduction/device/reduce_split_k.h"
|
||||
#include "cutlass/reduction/thread/reduction_operators.h"
|
||||
|
||||
#define CUTLASS_CHECK(status) \
|
||||
{ \
|
||||
cutlass::Status error = status; \
|
||||
if (error != cutlass::Status::kSuccess) { \
|
||||
std::cerr << "Got cutlass error: " << cutlassGetStatusString(error) << " at: " << __LINE__ \
|
||||
<< std::endl; \
|
||||
exit(EXIT_FAILURE); \
|
||||
} \
|
||||
}
|
||||
|
||||
{{OperatorDef}}
|
||||
using ImplicitGemm = cutlass::conv::device::ImplicitGemmConvolution<{{OperatorName}}>;
|
||||
|
||||
struct Options {
|
||||
cutlass::Tensor4DCoord input_size;
|
||||
cutlass::Tensor4DCoord filter_size;
|
||||
cutlass::Tensor4DCoord padding;
|
||||
cutlass::MatrixCoord conv_stride;
|
||||
cutlass::MatrixCoord dilation;
|
||||
|
||||
void parse(int argc, char const **args) {
|
||||
cutlass::CommandLine cmd(argc, args);
|
||||
cmd.get_cmd_line_argument("n", input_size.n());
|
||||
cmd.get_cmd_line_argument("h", input_size.h());
|
||||
cmd.get_cmd_line_argument("w", input_size.w());
|
||||
cmd.get_cmd_line_argument("c", input_size.c());
|
||||
cmd.get_cmd_line_argument("k", filter_size.n());
|
||||
cmd.get_cmd_line_argument("r", filter_size.h());
|
||||
cmd.get_cmd_line_argument("s", filter_size.w());
|
||||
int pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w;
|
||||
cmd.get_cmd_line_argument("pad_h", pad_h);
|
||||
cmd.get_cmd_line_argument("pad_w", pad_w);
|
||||
cmd.get_cmd_line_argument("stride_h", stride_h);
|
||||
cmd.get_cmd_line_argument("stride_w", stride_w);
|
||||
cmd.get_cmd_line_argument("dilation_h", dilation_h);
|
||||
cmd.get_cmd_line_argument("dilation_w", dilation_w);
|
||||
filter_size.c() = input_size.c();
|
||||
padding = {pad_h, pad_h, pad_w, pad_w};
|
||||
conv_stride = {stride_h, stride_w};
|
||||
dilation = {dilation_h, dilation_w};
|
||||
}
|
||||
|
||||
cutlass::Tensor4DCoord output_size() const {
|
||||
auto dilated_h = (filter_size.h() - 1) * dilation.row() + 1;
|
||||
auto dilated_w = (filter_size.w() - 1) * dilation.column() + 1;
|
||||
auto h = (input_size.h() + padding.n() + padding.h() - dilated_h) / conv_stride.row() + 1;
|
||||
auto w = (input_size.w() + padding.w() + padding.c() - dilated_w) / conv_stride.column() + 1;
|
||||
return cutlass::Tensor4DCoord(input_size.n(), h, w, filter_size.n());
|
||||
}
|
||||
};
|
||||
|
||||
double profile_convolution(Options const &options) {
|
||||
using ElementOutput = {{ElementOutput}};
|
||||
using ElementInputA = typename ImplicitGemm::ElementA;
|
||||
using ElementInputB = typename ImplicitGemm::ElementB;
|
||||
|
||||
int split_k_slices = {{SplitK}};
|
||||
cutlass::conv::Conv2dProblemSize problem_size(
|
||||
options.input_size,
|
||||
options.filter_size,
|
||||
options.padding,
|
||||
options.conv_stride,
|
||||
options.dilation,
|
||||
options.output_size(),
|
||||
cutlass::conv::Mode::kCrossCorrelation,
|
||||
split_k_slices
|
||||
);
|
||||
|
||||
auto conv_kind = ImplicitGemm::kConvolutionalOperator;
|
||||
auto a_extent = implicit_gemm_tensor_a_extent(conv_kind, problem_size);
|
||||
auto b_extent = implicit_gemm_tensor_b_extent(conv_kind, problem_size);
|
||||
auto c_extent = implicit_gemm_tensor_c_extent(conv_kind, problem_size);
|
||||
|
||||
using LayoutC = typename ImplicitGemm::LayoutC;
|
||||
cutlass::HostTensor<ElementInputA, typename ImplicitGemm::LayoutA> tensor_a(a_extent);
|
||||
cutlass::HostTensor<ElementInputB, typename ImplicitGemm::LayoutB> tensor_b(b_extent);
|
||||
cutlass::HostTensor<ElementOutput, typename ImplicitGemm::LayoutC> tensor_c(c_extent);
|
||||
cutlass::HostTensor<ElementOutput, LayoutC> tensor_d(c_extent);
|
||||
cutlass::HostTensor<ImplicitGemm::ElementC, LayoutC> tensor_c_gemm(c_extent);
|
||||
|
||||
using ElementComputeEpilogue = typename ImplicitGemm::ElementCompute;
|
||||
|
||||
cutlass::conv::SplitKMode const split_k_mode = split_k_slices > 1 ?
|
||||
cutlass::conv::SplitKMode::kParallel : cutlass::conv::SplitKMode::kSerial;
|
||||
|
||||
typename ImplicitGemm::Arguments arguments{
|
||||
problem_size,
|
||||
tensor_a.device_ref(),
|
||||
tensor_b.device_ref(),
|
||||
tensor_c_gemm.device_ref(),
|
||||
tensor_c_gemm.device_ref(),
|
||||
{ElementComputeEpilogue(1), ElementComputeEpilogue(0)},
|
||||
split_k_mode,
|
||||
};
|
||||
|
||||
ImplicitGemm implicit_gemm_op;
|
||||
size_t workspace_size = implicit_gemm_op.get_workspace_size(arguments);
|
||||
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
|
||||
auto status = implicit_gemm_op.can_implement(arguments);
|
||||
CUTLASS_CHECK(status);
|
||||
|
||||
status = implicit_gemm_op.initialize(arguments, workspace.get());
|
||||
CUTLASS_CHECK(status);
|
||||
status = implicit_gemm_op();
|
||||
CUTLASS_CHECK(status);
|
||||
|
||||
cudaEvent_t events[2];
|
||||
for (auto & event : events) {
|
||||
cudaEventCreate(&event);
|
||||
}
|
||||
cudaEventRecord(events[0]);
|
||||
|
||||
for (int iteration = 0; iteration < 100; ++iteration) {
|
||||
auto status = implicit_gemm_op();
|
||||
CUTLASS_CHECK(status);
|
||||
{{Reduction}}
|
||||
}
|
||||
|
||||
cudaEventRecord(events[1]);
|
||||
cudaEventSynchronize(events[1]);
|
||||
float runtime_ms = 0;
|
||||
cudaEventElapsedTime(&runtime_ms, events[0], events[1]);
|
||||
|
||||
for (auto event : events) {
|
||||
(void)cudaEventDestroy(event);
|
||||
}
|
||||
return double(runtime_ms) / 100.0;
|
||||
}
|
||||
|
||||
int main(int argc, char const **args) {
|
||||
Options options;
|
||||
options.parse(argc, args);
|
||||
std::cout << profile_convolution(options) << std::endl;
|
||||
return 0;
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
def emit(self, op_def, op_name, element_output, split_k_slices=1):
|
||||
src = self.template.render(
|
||||
OperatorDef=op_def,
|
||||
OperatorName=op_name,
|
||||
ElementOutput=DataTypeTag[element_output],
|
||||
SplitK=split_k_slices,
|
||||
Reduction=self.reduction if split_k_slices > 1 else "",
|
||||
)
|
||||
return src
|
||||
@@ -0,0 +1,478 @@
|
||||
# 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, unused-wildcard-import, wildcard-import, pointless-exception-statement
|
||||
# ruff: noqa: E501, F403, F405
|
||||
"""Generator for CUTLASS GEMM kernels."""
|
||||
|
||||
from .library import *
|
||||
|
||||
|
||||
class GemmOperation:
|
||||
"""Describes various attributes for instantiating GEMM kernels."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
arch,
|
||||
tile_description,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
element_epilogue,
|
||||
epilogue_functor=EpilogueFunctor.LinearCombination,
|
||||
swizzling_functor=SwizzlingFunctor.Identity8,
|
||||
):
|
||||
self.operation_kind = OperationKind.Gemm
|
||||
self.arch = arch
|
||||
self.tile_description = tile_description
|
||||
self.A = A
|
||||
self.B = B
|
||||
self.C = C
|
||||
self.element_epilogue = element_epilogue
|
||||
self.epilogue_functor = epilogue_functor
|
||||
self.swizzling_functor = swizzling_functor
|
||||
|
||||
def accumulator_type(self):
|
||||
return self.tile_description.math_instruction.element_accumulator
|
||||
|
||||
def short_math_name(self):
|
||||
return ShortDataTypeNames[self.accumulator_type()]
|
||||
|
||||
def core_name(self):
|
||||
"""The basic operation kind is prefixed with a letter indicating the accumulation type."""
|
||||
inst_shape = ""
|
||||
intermediate_type = ""
|
||||
|
||||
if (
|
||||
self.tile_description.math_instruction.opcode_class == OpcodeClass.TensorOp
|
||||
or self.tile_description.math_instruction.opcode_class == OpcodeClass.WmmaTensorOp
|
||||
):
|
||||
inst_shape = "{}{}{}".format(*self.tile_description.math_instruction.instruction_shape)
|
||||
if (
|
||||
self.tile_description.math_instruction.element_a != self.A.element
|
||||
and self.tile_description.math_instruction.element_a
|
||||
!= self.tile_description.math_instruction.element_accumulator
|
||||
):
|
||||
intermediate_type = DataTypeNames[self.tile_description.math_instruction.element_a]
|
||||
|
||||
return f"{self.short_math_name()}{inst_shape}{intermediate_type}gemm"
|
||||
|
||||
def extended_name(self):
|
||||
"""Append data types if they differ from compute type."""
|
||||
if (
|
||||
self.C.element != self.tile_description.math_instruction.element_accumulator
|
||||
and self.A.element != self.tile_description.math_instruction.element_accumulator
|
||||
):
|
||||
extended_name = "${element_c}_${core_name}_${element_a}"
|
||||
elif (
|
||||
self.C.element == self.tile_description.math_instruction.element_accumulator
|
||||
and self.A.element != self.tile_description.math_instruction.element_accumulator
|
||||
):
|
||||
extended_name = "${core_name}_${element_a}"
|
||||
else:
|
||||
extended_name = "${core_name}"
|
||||
|
||||
extended_name = substitute_template(
|
||||
extended_name,
|
||||
{
|
||||
"element_a": DataTypeNames[self.A.element],
|
||||
"element_c": DataTypeNames[self.C.element],
|
||||
"core_name": self.core_name(),
|
||||
},
|
||||
)
|
||||
|
||||
return extended_name
|
||||
|
||||
def layout_name(self):
|
||||
return f"{ShortLayoutTypeNames[self.A.layout]}{ShortLayoutTypeNames[self.B.layout]}"
|
||||
|
||||
def procedural_name(self):
|
||||
"""The full procedural name indicates architecture, extended name, tile size,
|
||||
and layout.
|
||||
"""
|
||||
threadblock = self.tile_description.procedural_name()
|
||||
opcode_class_name = OpcodeClassNames[self.tile_description.math_instruction.opcode_class]
|
||||
|
||||
return substitute_template(
|
||||
"cutlass_${opcode_class}_${extended_name}_${threadblock}_${layout}_align${alignment}",
|
||||
{
|
||||
"opcode_class": opcode_class_name,
|
||||
"extended_name": self.extended_name(),
|
||||
"threadblock": threadblock,
|
||||
"layout": self.layout_name(),
|
||||
"alignment": f"{self.A.alignment}",
|
||||
},
|
||||
)
|
||||
|
||||
def leading_dim(self):
|
||||
"""lda, ldb, ldc, according to the leading dimension."""
|
||||
if self.A.layout == LayoutType.RowMajor:
|
||||
lda = "K"
|
||||
elif self.A.layout == LayoutType.ColumnMajor:
|
||||
lda = "M"
|
||||
else:
|
||||
ValueError("The layout of A is not implemented.")
|
||||
|
||||
if self.B.layout == LayoutType.RowMajor:
|
||||
ldb = "N"
|
||||
elif self.B.layout == LayoutType.ColumnMajor:
|
||||
ldb = "K"
|
||||
else:
|
||||
ValueError("The layout of B is not implemented.")
|
||||
|
||||
if self.C.layout == LayoutType.RowMajor:
|
||||
ldc = "N"
|
||||
elif self.C.layout == LayoutType.ColumnMajor:
|
||||
ldc = "M"
|
||||
else:
|
||||
ValueError("The layout of B is not implemented.")
|
||||
|
||||
return substitute_template(
|
||||
"int lda = ${lda_val};\n\tint ldb = ${ldb_val};\n\tint ldc = ${ldc_val};\n",
|
||||
{"lda_val": lda, "ldb_val": ldb, "ldc_val": ldc},
|
||||
)
|
||||
|
||||
|
||||
class EmitGemmInstance:
|
||||
"""Responsible for emitting a CUTLASS template definition."""
|
||||
|
||||
def __init__(self):
|
||||
self.epilogue_default = """
|
||||
${epilogue_functor}<
|
||||
${element_c},
|
||||
${epilogue_vector_length},
|
||||
${element_accumulator},
|
||||
${element_epilogue}
|
||||
>"""
|
||||
|
||||
self.epilogue_no_beta_scaling = """
|
||||
${epilogue_functor}<
|
||||
${element_c},
|
||||
${epilogue_vector_length},
|
||||
${element_accumulator},
|
||||
${element_epilogue},
|
||||
cutlass::epilogue::thread::ScaleType::NoBetaScaling
|
||||
>"""
|
||||
|
||||
self.epilogue_residual_block = """
|
||||
${epilogue_functor}<
|
||||
${element_c},
|
||||
${element_accumulator},
|
||||
${element_epilogue},
|
||||
${element_c},
|
||||
${epilogue_vector_length},
|
||||
${activation},
|
||||
${binary_op},
|
||||
${unary_op}
|
||||
>"""
|
||||
|
||||
self.gemm_template = """
|
||||
// Gemm operator ${operation_name}
|
||||
using Operation_${operation_name} = cutlass::gemm::device::${kernel_name}<
|
||||
${element_a}, ${layout_a},
|
||||
${element_b}, ${layout_b},
|
||||
${element_c}, ${layout_c},
|
||||
${element_accumulator},
|
||||
${opcode_class},
|
||||
${arch},
|
||||
cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,
|
||||
cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k}>,
|
||||
cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,
|
||||
${epilogue},
|
||||
${swizzling_functor},
|
||||
${stages},
|
||||
${align_a},
|
||||
${align_b}
|
||||
>;
|
||||
"""
|
||||
|
||||
def emit(self, operation, no_beta_scaling=False, batched=False, residual_block_info=False):
|
||||
"""Instantiate a GEMM kernel from given `operation`."""
|
||||
warp_shape = [
|
||||
operation.tile_description.threadblock_shape[idx]
|
||||
// operation.tile_description.warp_count[idx]
|
||||
for idx in range(3)
|
||||
]
|
||||
epilogue_vector_length = (
|
||||
min(operation.C.alignment * DataTypeSize[operation.C.element], 128)
|
||||
// DataTypeSize[operation.C.element]
|
||||
)
|
||||
values = {
|
||||
"operation_name": operation.procedural_name(),
|
||||
"element_a": DataTypeTag[operation.A.element],
|
||||
"layout_a": LayoutTag[operation.A.layout],
|
||||
"element_b": DataTypeTag[operation.B.element],
|
||||
"layout_b": LayoutTag[operation.B.layout],
|
||||
"element_c": DataTypeTag[operation.C.element],
|
||||
"layout_c": LayoutTag[operation.C.layout],
|
||||
"element_accumulator": DataTypeTag[operation.accumulator_type()],
|
||||
"opcode_class": OpcodeClassTag[
|
||||
operation.tile_description.math_instruction.opcode_class
|
||||
],
|
||||
"arch": f"cutlass::arch::Sm{operation.arch}",
|
||||
"threadblock_shape_m": str(operation.tile_description.threadblock_shape[0]),
|
||||
"threadblock_shape_n": str(operation.tile_description.threadblock_shape[1]),
|
||||
"threadblock_shape_k": str(operation.tile_description.threadblock_shape[2]),
|
||||
"warp_shape_m": str(warp_shape[0]),
|
||||
"warp_shape_n": str(warp_shape[1]),
|
||||
"warp_shape_k": str(warp_shape[2]),
|
||||
"instruction_shape_m": str(
|
||||
operation.tile_description.math_instruction.instruction_shape[0]
|
||||
),
|
||||
"instruction_shape_n": str(
|
||||
operation.tile_description.math_instruction.instruction_shape[1]
|
||||
),
|
||||
"instruction_shape_k": str(
|
||||
operation.tile_description.math_instruction.instruction_shape[2]
|
||||
),
|
||||
"epilogue_vector_length": str(epilogue_vector_length),
|
||||
"element_epilogue": str(DataTypeTag[operation.element_epilogue]),
|
||||
"epilogue_functor": EpilogueFunctorTag[operation.epilogue_functor],
|
||||
"swizzling_functor": SwizzlingFunctorTag[operation.swizzling_functor],
|
||||
"stages": str(operation.tile_description.stages),
|
||||
"align_a": str(operation.A.alignment),
|
||||
"align_b": str(operation.B.alignment),
|
||||
"math_operation": MathOperationTag[
|
||||
operation.tile_description.math_instruction.math_operation
|
||||
],
|
||||
}
|
||||
|
||||
values["kernel_name"] = "GemmBatched" if batched else "Gemm"
|
||||
|
||||
if residual_block_info:
|
||||
values["kernel_name"] = "GemmUniversalWithBroadcast"
|
||||
template = substitute_template(
|
||||
self.gemm_template, {"epilogue": self.epilogue_residual_block}
|
||||
)
|
||||
values.update(
|
||||
{
|
||||
"unary_op": residual_block_info["unary_op"],
|
||||
"binary_op": residual_block_info["binary_op"],
|
||||
"activation": residual_block_info["activation"],
|
||||
}
|
||||
)
|
||||
elif no_beta_scaling:
|
||||
template = substitute_template(
|
||||
self.gemm_template, {"epilogue": self.epilogue_no_beta_scaling}
|
||||
)
|
||||
else:
|
||||
template = substitute_template(self.gemm_template, {"epilogue": self.epilogue_default})
|
||||
|
||||
return substitute_template(template, values)
|
||||
|
||||
|
||||
def instantiate_gemm_template(attrs):
|
||||
"""Return CUTLASS host code for GEMM based on a template and the provided attribute map."""
|
||||
|
||||
argument_template_default = """
|
||||
typename ${kernel}::Arguments arguments{
|
||||
problem_size,
|
||||
{static_cast<ElementInputA*>(ptr_a), ${lda}}, ${batch_stride_A}
|
||||
{static_cast<ElementInputB*>(ptr_b), ${ldb}}, ${batch_stride_B}
|
||||
{static_cast<ElementOutput*>(${ptr_c}), ${c_stride}}, ${batch_stride_C}
|
||||
{static_cast<ElementOutput*>(ptr_out), ${ldc}}, ${batch_stride_D}
|
||||
{${alpha_beta}},
|
||||
${split_k_slices_or_batch}
|
||||
};
|
||||
"""
|
||||
|
||||
# See cutlass/gemm/kernel/gemm_with_fused_epilogue.h
|
||||
argument_template_residual = """
|
||||
typename ${kernel}::Arguments arguments{
|
||||
cutlass::gemm::GemmUniversalMode::${gemm_universal_mode},
|
||||
problem_size,
|
||||
${split_k_slices_or_batch}, // batch_count
|
||||
{${alpha_beta}},
|
||||
static_cast<ElementInputA*>(ptr_a),
|
||||
static_cast<ElementInputB*>(ptr_b),
|
||||
static_cast<ElementOutput*>(ptr_residual),
|
||||
static_cast<ElementOutput*>(ptr_out),
|
||||
static_cast<ElementOutput*>(ptr_bias),
|
||||
nullptr, // ptr_Tensor
|
||||
${batch_stride_A}
|
||||
${batch_stride_B}
|
||||
${batch_stride_C}
|
||||
${batch_stride_D}
|
||||
0, // batch_stride_Vector,
|
||||
0, // batch_stride_Tensor,
|
||||
${lda},
|
||||
${ldb},
|
||||
${ldc},
|
||||
${ldc},
|
||||
0, // ldv, the stride for bias
|
||||
0, // ldt
|
||||
};
|
||||
"""
|
||||
|
||||
template = """
|
||||
using ElementInputA = ${ElementInputA};
|
||||
using ElementInputB = ${ElementInputB};
|
||||
using ElementOutput = ${ElementOutput};
|
||||
using ElementComputeEpilogue = ${ElementOutput};
|
||||
|
||||
${cutlass_op_def}
|
||||
|
||||
using ${kernel} = Operation_${cutlass_op_name};
|
||||
int M = ${M};
|
||||
int N = ${N};
|
||||
int K = ${K};
|
||||
cutlass::gemm::GemmCoord problem_size(M, N, K);
|
||||
ElementComputeEpilogue alpha = ElementComputeEpilogue(1);
|
||||
ElementComputeEpilogue beta = ElementComputeEpilogue(${beta});
|
||||
void* ptr_a = (void*)(${lhs_arg}->data);
|
||||
void* ptr_b = (void*)(${rhs_arg}->data);
|
||||
${bias_decl}
|
||||
${residual_decl}
|
||||
void* ptr_out = (void*)(out0->data);
|
||||
|
||||
${argument}
|
||||
size_t workspace_size = ${kernel}::get_workspace_size(arguments);
|
||||
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
|
||||
${kernel} gemm_op;
|
||||
cutlass::Status status = gemm_op.can_implement(arguments);
|
||||
TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
|
||||
status = gemm_op.initialize(arguments, workspace.get());
|
||||
TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${A_arg}->device.device_id));
|
||||
|
||||
status = gemm_op(stream);
|
||||
TVM_FFI_ICHECK(status == cutlass::Status::kSuccess);
|
||||
"""
|
||||
op_type = attrs["op_type"]
|
||||
has_bias = "bias" in op_type
|
||||
is_gelu = "gelu" in op_type
|
||||
batched = "batch" in attrs
|
||||
has_residual_block = "residual" in op_type
|
||||
aux_map = {"kernel": "Gemm"}
|
||||
|
||||
if has_bias:
|
||||
aux_map.update(
|
||||
{
|
||||
"bias_decl": "void* ptr_bias = (void*)(${bias_arg}->data);\n",
|
||||
"ptr_c": "ptr_bias",
|
||||
"c_stride": (
|
||||
"(${bias_arg}->ndim == 1 ||"
|
||||
" ${bias_arg}->shape[${bias_arg}->ndim - 2] == 1) ? 0 : " + attrs["ldc"]
|
||||
),
|
||||
}
|
||||
)
|
||||
else:
|
||||
aux_map.update({"bias_decl": "", "ptr_c": "ptr_out", "c_stride": attrs["ldc"]})
|
||||
|
||||
if is_gelu or has_residual_block:
|
||||
# GeLU epilogue does not compile with NoBetaScaling, so we explicitly specify the scale.
|
||||
aux_map["beta"] = 1
|
||||
else:
|
||||
aux_map["beta"] = 0
|
||||
|
||||
if has_bias and not is_gelu and not has_residual_block:
|
||||
aux_map["alpha_beta"] = "alpha"
|
||||
else:
|
||||
aux_map["alpha_beta"] = "alpha, beta"
|
||||
|
||||
for key in ["batch_stride_A", "batch_stride_B", "batch_stride_C"]:
|
||||
if not batched and not has_residual_block:
|
||||
aux_map[key] = ""
|
||||
else:
|
||||
aux_map[key] = attrs.get(key, "0") + ","
|
||||
|
||||
aux_map["batch_stride_D"] = aux_map["batch_stride_C"]
|
||||
if has_bias and batched and not has_residual_block:
|
||||
aux_map["batch_stride_C"] = "0,"
|
||||
|
||||
if batched:
|
||||
attrs["split_k_slices_or_batch"] = attrs["batch"]
|
||||
else:
|
||||
attrs["split_k_slices_or_batch"] = 1
|
||||
|
||||
if has_residual_block:
|
||||
template = substitute_template(template, {"argument": argument_template_residual})
|
||||
aux_map["residual_decl"] = "void* ptr_residual = (void*)(${residual_arg}->data);\n"
|
||||
aux_map["gemm_universal_mode"] = "kBatched" if batched else "kGemm"
|
||||
else:
|
||||
template = substitute_template(template, {"argument": argument_template_default})
|
||||
aux_map["residual_decl"] = ""
|
||||
|
||||
template = substitute_template(template, aux_map)
|
||||
|
||||
return substitute_template(template, attrs)
|
||||
|
||||
|
||||
def emit_fp16A_intB_matmul(attrs):
|
||||
"""Return CUTLASS host code for fp16 A and int4 or int8 B GEMM."""
|
||||
if attrs["group_size"] > 0:
|
||||
attrs["quant_op"] = "cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY"
|
||||
else:
|
||||
attrs["quant_op"] = "cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY"
|
||||
attrs["group_size"] = "k"
|
||||
|
||||
attrs["template_common"] = substitute_template(
|
||||
"""
|
||||
using namespace fastertransformer;
|
||||
constexpr auto QuantOp = ${quant_op};
|
||||
|
||||
int m = ${M};
|
||||
int n = ${B_arg}->shape[1] * ${float_per_int};
|
||||
int k = ${B_arg}->shape[0];
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(
|
||||
TVMFFIEnvGetStream(kDLCUDA, ${A_arg}->device.device_id));
|
||||
""",
|
||||
attrs,
|
||||
)
|
||||
|
||||
template = """
|
||||
${template_common}
|
||||
gemm_fp16_int_bias_act<${weight_dtype}, QuantOp>(static_cast<cutlass::half_t*>(${A_arg}->data),
|
||||
static_cast<${weight_dtype}*>(${B_arg}->data),
|
||||
static_cast<cutlass::half_t*>(${scales_arg}->data),
|
||||
${bias},
|
||||
static_cast<cutlass::half_t*>(out0->data),
|
||||
"${activation}",
|
||||
m, n, k, ${group_size}, ${bias_stride}, nullptr, 0, stream);
|
||||
"""
|
||||
|
||||
template_residual = """
|
||||
${template_common}
|
||||
gemm_fp16_int_bias_act_residual<${weight_dtype}, QuantOp>(
|
||||
static_cast<cutlass::half_t*>(${A_arg}->data),
|
||||
static_cast<${weight_dtype}*>(${B_arg}->data),
|
||||
static_cast<cutlass::half_t*>(${scales_arg}->data),
|
||||
${bias},
|
||||
static_cast<cutlass::half_t*>(${residual_arg}->data),
|
||||
static_cast<cutlass::half_t*>(out0->data),
|
||||
"${activation}", "${binary_op}", "${unary_op}",
|
||||
m, n, k, ${group_size}, nullptr, 0, stream);
|
||||
"""
|
||||
|
||||
if "residual_arg" in attrs:
|
||||
if "bias_arg" in attrs:
|
||||
bias = "static_cast<cutlass::half_t*>(${bias_arg}->data)"
|
||||
else:
|
||||
bias = "nullptr"
|
||||
|
||||
template_residual = substitute_template(template_residual, {"bias": bias})
|
||||
return substitute_template(template_residual, attrs)
|
||||
|
||||
if "bias_arg" in attrs:
|
||||
template = substitute_template(
|
||||
template, {"bias": "static_cast<cutlass::half_t*>(${bias_arg}->data)"}
|
||||
)
|
||||
else:
|
||||
template = substitute_template(template, {"bias": "nullptr"})
|
||||
|
||||
return substitute_template(template, attrs)
|
||||
@@ -0,0 +1,196 @@
|
||||
# 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=import-outside-toplevel, invalid-name
|
||||
"""Instantiate a C++ source for profiling CUTLASS kernels."""
|
||||
|
||||
|
||||
class GemmProfilerEmitter:
|
||||
"""Emit a C++ source for profiling CUTLASS kernels."""
|
||||
|
||||
def __init__(self):
|
||||
from jinja2 import Template
|
||||
|
||||
self.template = Template(
|
||||
"""
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <vector>
|
||||
#include <chrono>
|
||||
|
||||
#include "cuda_runtime.h"
|
||||
#include "cutlass/gemm/device/gemm.h"
|
||||
|
||||
#define CUTLASS_CHECK(status) \\
|
||||
{ \\
|
||||
cutlass::Status error = status; \\
|
||||
if (error != cutlass::Status::kSuccess) { \\
|
||||
std::cerr << "Got cutlass error: " << cutlassGetStatusString(error) << " at: " << __LINE__ \\
|
||||
<< std::endl; \\
|
||||
exit(EXIT_FAILURE); \\
|
||||
} \\
|
||||
}
|
||||
|
||||
#define CUDA_CHECK(status) \\
|
||||
{ \\
|
||||
cudaError_t error = status; \\
|
||||
if (error != cudaSuccess) { \\
|
||||
std::cerr << "Got bad CUDA status: " << cudaGetErrorString(error) \\
|
||||
<< " at line: " << __LINE__ << std::endl; \\
|
||||
exit(EXIT_FAILURE); \\
|
||||
} \\
|
||||
}
|
||||
|
||||
template<typename DTypeA, typename DTypeB, typename DTypeC>
|
||||
cudaError_t CutlassGemm(
|
||||
int M,
|
||||
int N,
|
||||
int K,
|
||||
DTypeC alpha,
|
||||
DTypeA const *A,
|
||||
int lda,
|
||||
DTypeB const *B,
|
||||
int ldb,
|
||||
DTypeC beta,
|
||||
DTypeC *C,
|
||||
int ldc) {
|
||||
using namespace std::chrono;
|
||||
{{OperatorDef}}
|
||||
Operation_{{OperatorName}} gemm_operator;
|
||||
Operation_{{OperatorName}}::Arguments args({M, N, K},
|
||||
{A, lda},
|
||||
{B, ldb},
|
||||
{C, ldc},
|
||||
{C, ldc},
|
||||
{alpha, beta});
|
||||
cutlass::Status status = gemm_operator(args);
|
||||
CUTLASS_CHECK(status)
|
||||
|
||||
high_resolution_clock::time_point t1 = high_resolution_clock::now();
|
||||
for (int i = 0; i < 100; ++i) {
|
||||
status = gemm_operator(args);
|
||||
}
|
||||
cudaDeviceSynchronize();
|
||||
high_resolution_clock::time_point t2 = high_resolution_clock::now();
|
||||
duration<double> time_span = duration_cast<duration<double>>(t2 - t1);
|
||||
std::cout << time_span.count() << std::endl;
|
||||
return cudaSuccess;
|
||||
}
|
||||
|
||||
|
||||
template<typename DType>
|
||||
cudaError_t AllocateMatrix(DType **matrix, int ldm, int rows, int columns, int seed = 0) {
|
||||
cudaError_t result;
|
||||
|
||||
size_t sizeof_matrix = sizeof(DType) * rows * columns;
|
||||
|
||||
// Allocate device memory.
|
||||
result = cudaMalloc(reinterpret_cast<void **>(matrix), sizeof_matrix);
|
||||
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "Failed to allocate matrix: "
|
||||
<< cudaGetErrorString(result) << std::endl;
|
||||
return result;
|
||||
}
|
||||
|
||||
// Clear the allocation.
|
||||
result = cudaMemset(*matrix, 0, sizeof_matrix);
|
||||
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "Failed to clear matrix device memory: "
|
||||
<< cudaGetErrorString(result) << std::endl;
|
||||
return result;
|
||||
}
|
||||
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "Failed to initialize matrix: "
|
||||
<< cudaGetErrorString(result) << std::endl;
|
||||
return result;
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
template<typename DTypeA, typename DTypeB, typename DTypeC>
|
||||
cudaError_t TestCutlassGemm(int M, int N, int K, DTypeC alpha, DTypeC beta) {
|
||||
cudaError_t result;
|
||||
|
||||
{{LeadingDim}}
|
||||
// size_t sizeof_C = sizeof(DTypeC) * ldc * N;
|
||||
DTypeA *A;
|
||||
DTypeB *B;
|
||||
DTypeC *C_cutlass;
|
||||
result = AllocateMatrix<DTypeA>(&A, lda, M, K, 0);
|
||||
if (result != cudaSuccess) {
|
||||
return result;
|
||||
}
|
||||
result = AllocateMatrix<DTypeB>(&B, ldb, K, N, 17);
|
||||
if (result != cudaSuccess) {
|
||||
cudaFree(A);
|
||||
return result;
|
||||
}
|
||||
result = AllocateMatrix<DTypeC>(&C_cutlass, ldc, M, N, 101);
|
||||
if (result != cudaSuccess) {
|
||||
cudaFree(A);
|
||||
cudaFree(B);
|
||||
return result;
|
||||
}
|
||||
result = CutlassGemm<DTypeA, DTypeB, DTypeC>(M, N, K, alpha, A, lda, B, ldb,
|
||||
beta, C_cutlass, ldc);
|
||||
if (result != cudaSuccess) {
|
||||
std::cerr << "CUTLASS GEMM kernel failed: "
|
||||
<< cudaGetErrorString(result) << std::endl;
|
||||
cudaFree(C_cutlass);
|
||||
cudaFree(B);
|
||||
cudaFree(A);
|
||||
|
||||
return result;
|
||||
}
|
||||
cudaFree(C_cutlass);
|
||||
cudaFree(B);
|
||||
cudaFree(A);
|
||||
return cudaSuccess;
|
||||
}
|
||||
|
||||
int main(int argc, const char *arg[]) {
|
||||
int problem[3] = { 4096, 4096, 4096 };
|
||||
for (int i = 1; i < argc && i < 4; ++i) {
|
||||
std::stringstream ss(arg[i]);
|
||||
ss >> problem[i - 1];
|
||||
}
|
||||
float scalars[2] = { 1, 0 };
|
||||
cudaError_t result = TestCutlassGemm< {{DTypeA}}, {{DTypeB}}, {{DTypeC}}>(
|
||||
problem[0], // GEMM M dimension
|
||||
problem[1], // GEMM N dimension
|
||||
problem[2], // GEMM K dimension
|
||||
static_cast<{{DTypeC}}>(scalars[0]), // alpha
|
||||
static_cast<{{DTypeC}}>(scalars[1]) // beta
|
||||
);
|
||||
return result == cudaSuccess ? 0 : -1;
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
def emit(self, op_name, op_def, dtype_a, dtype_b, dtype_c, ld):
|
||||
src = self.template.render(
|
||||
OperatorName=op_name,
|
||||
OperatorDef=op_def,
|
||||
DTypeA=dtype_a,
|
||||
DTypeB=dtype_b,
|
||||
DTypeC=dtype_c,
|
||||
LeadingDim=ld,
|
||||
)
|
||||
return src
|
||||
@@ -0,0 +1,392 @@
|
||||
# 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
|
||||
# ruff: noqa: E501
|
||||
"""Conv2d kernel generator and profiler for CUTLASS."""
|
||||
|
||||
import os
|
||||
import pickle
|
||||
from functools import partial
|
||||
|
||||
from .conv2d_operation import Conv2dOperation, EmitConv2dInstance
|
||||
from .conv2d_profiler import Conv2dProfilerEmitter
|
||||
from .gen_gemm import CutlassGemmProfiler
|
||||
from .gen_tensor_op import EPILOGUE_MAP, GENERATOR_FUNC_TABLE, ProfilerEngine
|
||||
from .library import (
|
||||
ConvKind,
|
||||
DataType,
|
||||
EpilogueFunctor,
|
||||
IteratorAlgorithm,
|
||||
LayoutType,
|
||||
StrideSupport,
|
||||
SwizzlingFunctor,
|
||||
TensorDescription,
|
||||
)
|
||||
|
||||
|
||||
def create_conv2d_operator_with_epilogue(
|
||||
conv_kind,
|
||||
stride_support,
|
||||
op_type,
|
||||
tile_description,
|
||||
data_type,
|
||||
alignment,
|
||||
alignment_epilogue,
|
||||
swizzling_functor,
|
||||
split_k_slices,
|
||||
):
|
||||
"""
|
||||
Instantiate a cutlass kernel from the given configuration,
|
||||
along with the epilouge functor
|
||||
"""
|
||||
if "residual" in op_type:
|
||||
activation_map = {
|
||||
"cutlass.conv2d_bias_hardswish": "cutlass::epilogue::thread::HardSwish",
|
||||
"cutlass.conv2d_bias_silu": "cutlass::epilogue::thread::SiLu",
|
||||
"cutlass.conv2d_bias_sigmoid": "cutlass::epilogue::thread::Sigmoid",
|
||||
"cutlass.conv2d_bias_relu": "cutlass::epilogue::thread::ReLu",
|
||||
"cutlass.conv2d_bias": "cutlass::epilogue::thread::Identity",
|
||||
}
|
||||
prefix = op_type[: op_type.find("_residual")]
|
||||
activation = activation_map[prefix]
|
||||
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]
|
||||
|
||||
element_a, element_b, element_c, element_epilogue = data_type
|
||||
|
||||
A = TensorDescription(element_a, LayoutType.TensorNHWC, alignment)
|
||||
B = TensorDescription(element_b, LayoutType.TensorNHWC, alignment)
|
||||
C = TensorDescription(element_c, LayoutType.TensorNHWC, alignment_epilogue)
|
||||
|
||||
op = Conv2dOperation(
|
||||
conv_kind,
|
||||
IteratorAlgorithm.Optimized,
|
||||
tile_description.minimum_compute_capability,
|
||||
tile_description,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
element_epilogue,
|
||||
stride_support,
|
||||
epilogue,
|
||||
swizzling_functor,
|
||||
split_k_slices,
|
||||
)
|
||||
|
||||
name = op.procedural_name()
|
||||
opdef = EmitConv2dInstance().emit(
|
||||
op,
|
||||
no_beta_scaling=no_beta_scaling,
|
||||
residual_block_info=residual_block_info,
|
||||
emit_reduction=split_k_slices > 1,
|
||||
)
|
||||
|
||||
return name, opdef
|
||||
|
||||
|
||||
def enumerate_conv2d_operators(
|
||||
conv_kind,
|
||||
stride_support,
|
||||
split_k_slices,
|
||||
alignment_c,
|
||||
tile_descriptions,
|
||||
data_type,
|
||||
alignment_constraints,
|
||||
swizzling_functor=SwizzlingFunctor.Identity4,
|
||||
):
|
||||
"""Exhaustively instantiate all kernels from a given configuration."""
|
||||
ret = []
|
||||
|
||||
kernel_emitter = EmitConv2dInstance()
|
||||
profiler_emitter = Conv2dProfilerEmitter()
|
||||
|
||||
element_a, element_b, element_c, element_epilogue = data_type
|
||||
|
||||
if conv_kind == ConvKind.Dgrad and stride_support == StrideSupport.Strided:
|
||||
swizzling_functor = SwizzlingFunctor.StridedDgradIdentity1
|
||||
|
||||
for split_k_slice in split_k_slices:
|
||||
for tile in tile_descriptions:
|
||||
for alignmentAB in alignment_constraints:
|
||||
for alignmentC in alignment_c:
|
||||
A = TensorDescription(element_a, LayoutType.TensorNHWC, alignmentAB)
|
||||
B = TensorDescription(element_b, LayoutType.TensorNHWC, alignmentAB)
|
||||
C = TensorDescription(element_c, LayoutType.TensorNHWC, alignmentC)
|
||||
|
||||
if element_c == DataType.s32 and A.alignment == 1:
|
||||
tile.threadblock_shape[0] = min(tile.threadblock_shape[0], 128)
|
||||
tile.threadblock_shape[1] = min(tile.threadblock_shape[1], 128)
|
||||
|
||||
op = Conv2dOperation(
|
||||
conv_kind,
|
||||
IteratorAlgorithm.Optimized,
|
||||
tile.minimum_compute_capability,
|
||||
tile,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
element_epilogue,
|
||||
stride_support,
|
||||
EpilogueFunctor.LinearCombination,
|
||||
swizzling_functor,
|
||||
split_k_slice,
|
||||
)
|
||||
|
||||
ret.append(
|
||||
{
|
||||
"src": profiler_emitter.emit(
|
||||
kernel_emitter.emit(op, emit_reduction=split_k_slice > 1),
|
||||
op.procedural_name(),
|
||||
element_output=element_c,
|
||||
split_k_slices=split_k_slice,
|
||||
),
|
||||
"name": op.procedural_name(),
|
||||
"tile_description": tile,
|
||||
"alignment": alignmentAB,
|
||||
"alignment_epilogue": alignmentC,
|
||||
"data_type": data_type,
|
||||
"swizzle_functor": swizzling_functor,
|
||||
"split_k_slices": split_k_slice,
|
||||
}
|
||||
)
|
||||
|
||||
return ret
|
||||
|
||||
|
||||
class CutlassConv2DProfiler:
|
||||
"""Profile all candidate kernels and select the best one."""
|
||||
|
||||
def __init__(self, sm, cutlass_path, binary_path):
|
||||
self.gemm_profiler = CutlassGemmProfiler(sm, cutlass_path, binary_path)
|
||||
self.sm = sm
|
||||
assert sm in GENERATOR_FUNC_TABLE, f"sm{sm} not supported yet."
|
||||
self.engine = ProfilerEngine(sm, cutlass_path, binary_path)
|
||||
self.cache_path = os.path.join(binary_path, "cutlass_conv2d_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,
|
||||
conv_kind=ConvKind.Fprop,
|
||||
stride=(1, 1),
|
||||
):
|
||||
"""Return the default kernel for the requested architecture.
|
||||
For now, the default kernel was picked arbitrary.
|
||||
"""
|
||||
gemm_profile_result = self.gemm_profiler.get_default(
|
||||
op_type, out_dtype, arg0_dtype, arg1_dtype, use_3xtf32
|
||||
)
|
||||
tile_description = gemm_profile_result["tile_description"]
|
||||
alignment = gemm_profile_result["alignment"]
|
||||
data_type = gemm_profile_result["data_type"]
|
||||
stride_support = StrideSupport.Strided if stride[0] > 1 else StrideSupport.Unity
|
||||
|
||||
if conv_kind == ConvKind.Dgrad and stride_support == StrideSupport.Strided:
|
||||
swizzling_functor = SwizzlingFunctor.StridedDgradIdentity1
|
||||
else:
|
||||
swizzling_functor = SwizzlingFunctor.Identity4
|
||||
|
||||
name, opdef = create_conv2d_operator_with_epilogue(
|
||||
conv_kind,
|
||||
stride_support,
|
||||
op_type,
|
||||
tile_description,
|
||||
data_type,
|
||||
alignment,
|
||||
alignment,
|
||||
swizzling_functor,
|
||||
split_k_slices=1,
|
||||
)
|
||||
return {"name": name, "opdef": opdef}
|
||||
|
||||
def select_op(
|
||||
self,
|
||||
d_shape,
|
||||
w_shape,
|
||||
padding,
|
||||
stride,
|
||||
dilation,
|
||||
out_dtype,
|
||||
data_dtype,
|
||||
weight_dtype,
|
||||
use_3xtf32,
|
||||
conv_kind,
|
||||
stride_support,
|
||||
split_k_slices,
|
||||
profile_all_alignments=False,
|
||||
find_first_valid=False,
|
||||
use_multiprocessing=False,
|
||||
):
|
||||
"""
|
||||
Profile and select the best kernel from candidate kernels.
|
||||
See the documentation for the profile method below.
|
||||
"""
|
||||
N, H, W, IC = d_shape
|
||||
OC, R, S, _ = w_shape
|
||||
|
||||
workload = (
|
||||
N,
|
||||
H,
|
||||
W,
|
||||
IC,
|
||||
OC,
|
||||
R,
|
||||
S,
|
||||
padding[0],
|
||||
padding[1],
|
||||
stride[0],
|
||||
stride[1],
|
||||
dilation[0],
|
||||
dilation[1],
|
||||
)
|
||||
|
||||
if workload in self.cache:
|
||||
return self.cache[workload]
|
||||
|
||||
def alignments(dtype):
|
||||
if dtype in ["float16"]:
|
||||
alignments = [8, 4, 2, 1]
|
||||
elif dtype in ["float", "float32"]:
|
||||
alignments = [4, 2, 1]
|
||||
else:
|
||||
raise ValueError(f"Unsupported data type: {dtype}")
|
||||
return alignments
|
||||
|
||||
alignments_c = [align for align in alignments(out_dtype) if OC % align == 0]
|
||||
|
||||
if not profile_all_alignments:
|
||||
alignments_c = [alignments_c[0]]
|
||||
|
||||
ops = GENERATOR_FUNC_TABLE[self.sm](
|
||||
out_dtype,
|
||||
data_dtype,
|
||||
weight_dtype,
|
||||
partial(
|
||||
enumerate_conv2d_operators,
|
||||
conv_kind,
|
||||
stride_support,
|
||||
split_k_slices,
|
||||
alignments_c,
|
||||
),
|
||||
lambda align: all([dim % align == 0 for dim in [IC]]),
|
||||
use_3xtf32,
|
||||
profile_all_alignments,
|
||||
# Use fp32 accumulation for wgrad to align with cuDNN
|
||||
accumlator_dtype="float32" if conv_kind == ConvKind.Wgrad else out_dtype,
|
||||
)
|
||||
|
||||
if not find_first_valid:
|
||||
self.engine.compile_all(ops, use_multiprocessing)
|
||||
|
||||
args = "--n={} --h={} --w={} --c={} --k={} --r={} --s={} --pad_h={} --pad_w={} --stride_h={} --stride_w={} --dilation_h={} --dilation_w={}".format(
|
||||
*workload
|
||||
)
|
||||
|
||||
for op in ops:
|
||||
out = self.engine.evaluate(op, args.split(" "))
|
||||
op["runtime"] = out
|
||||
if out < float("inf") and find_first_valid:
|
||||
self.cache[workload] = op
|
||||
return op
|
||||
|
||||
op = min(ops, key=lambda i: i["runtime"])
|
||||
self.cache[workload] = op
|
||||
with open(self.cache_path, "wb") as f:
|
||||
pickle.dump(self.cache, f)
|
||||
return op
|
||||
|
||||
def profile(
|
||||
self,
|
||||
op_type,
|
||||
d_shape,
|
||||
w_shape,
|
||||
padding,
|
||||
stride,
|
||||
dilation,
|
||||
out_dtype,
|
||||
data_dtype,
|
||||
weight_dtype,
|
||||
use_3xtf32=True,
|
||||
conv_kind=ConvKind.Fprop,
|
||||
split_k_slices=[1],
|
||||
profile_all_alignments=False,
|
||||
find_first_valid=False,
|
||||
use_multiprocessing=False,
|
||||
):
|
||||
"""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.
|
||||
"""
|
||||
# Dgrad requires Unity stride when stride == (1, 1)
|
||||
stride_support = (
|
||||
StrideSupport.Unity
|
||||
if stride[0] == 1 and stride[1] == 1 and conv_kind == ConvKind.Dgrad
|
||||
else StrideSupport.Strided
|
||||
)
|
||||
|
||||
op = self.select_op(
|
||||
d_shape,
|
||||
w_shape,
|
||||
padding,
|
||||
stride,
|
||||
dilation,
|
||||
out_dtype,
|
||||
data_dtype,
|
||||
weight_dtype,
|
||||
use_3xtf32,
|
||||
conv_kind,
|
||||
stride_support,
|
||||
split_k_slices,
|
||||
profile_all_alignments,
|
||||
find_first_valid,
|
||||
use_multiprocessing,
|
||||
)
|
||||
|
||||
name, opdef = create_conv2d_operator_with_epilogue(
|
||||
conv_kind,
|
||||
stride_support,
|
||||
op_type,
|
||||
op["tile_description"],
|
||||
op["data_type"],
|
||||
op["alignment"],
|
||||
op["alignment_epilogue"],
|
||||
op["swizzle_functor"],
|
||||
op["split_k_slices"],
|
||||
)
|
||||
|
||||
return name, opdef, op["runtime"]
|
||||
@@ -0,0 +1,352 @@
|
||||
# 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"]
|
||||
@@ -0,0 +1,908 @@
|
||||
# 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
|
||||
# ruff: noqa: F821
|
||||
"""Common functions and classes for CUTLASS GEMM and Conv2d geneator."""
|
||||
|
||||
import logging
|
||||
import math
|
||||
import multiprocessing
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import tempfile
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
from tvm.runtime import Object
|
||||
from tvm.tirx import IntImm
|
||||
|
||||
from . import _ffi_api as ffi
|
||||
from .attention_operation import (
|
||||
instantiate_attention_template,
|
||||
instantiate_flash_attention_template,
|
||||
instantiate_flash_attention_var_len_template,
|
||||
)
|
||||
from .conv2d_operation import instantiate_conv2d_template
|
||||
from .gemm_operation import emit_fp16A_intB_matmul, instantiate_gemm_template
|
||||
from .layer_norm_operation import instantiate_layer_norm_template
|
||||
from .library import (
|
||||
DataType,
|
||||
DataTypeSize,
|
||||
DataTypeTag,
|
||||
EpilogueFunctor,
|
||||
MathInstruction,
|
||||
MathOperation,
|
||||
OpcodeClass,
|
||||
TileDescription,
|
||||
)
|
||||
from .rms_norm_operation import instantiate_rms_norm_template
|
||||
|
||||
logger = logging.getLogger("cutlass")
|
||||
|
||||
|
||||
dtype_map = {
|
||||
"int8": DataType.s8,
|
||||
"uint8": DataType.u8,
|
||||
"int32": DataType.s32,
|
||||
"float32": DataType.f32,
|
||||
"float16": DataType.f16,
|
||||
}
|
||||
|
||||
|
||||
def generate_tensor_op_common(
|
||||
math_instructions, alignment_constraints, get_tile_descriptions, op_creator
|
||||
):
|
||||
"""Common kernel generator to be used by archtecture specific generators."""
|
||||
ops = []
|
||||
for math_inst in math_instructions:
|
||||
tile_descriptions = get_tile_descriptions(math_inst)
|
||||
data_type = [
|
||||
math_inst.element_a,
|
||||
math_inst.element_b,
|
||||
math_inst.element_c,
|
||||
math_inst.element_accumulator,
|
||||
]
|
||||
|
||||
out = op_creator(tile_descriptions, data_type, alignment_constraints)
|
||||
|
||||
ops.extend(out)
|
||||
|
||||
return ops
|
||||
|
||||
|
||||
def generate_sm50_simt(out_dtype, arg0_dtype, arg1_dtype, op_creator, accumulator_dtype="float32"):
|
||||
"""Gemerate GEMM or Conv2D SIMT kernels"""
|
||||
# pylint: disable=unused-argument
|
||||
min_cc = 50
|
||||
max_cc = 1024
|
||||
if arg0_dtype == "float32" and arg1_dtype == "float32":
|
||||
assert out_dtype == "float32" and accumulator_dtype == "float32"
|
||||
math_instructions = [
|
||||
MathInstruction(
|
||||
[1, 1, 1],
|
||||
DataType.f32,
|
||||
DataType.f32,
|
||||
DataType.f32,
|
||||
DataType.f32,
|
||||
OpcodeClass.Simt,
|
||||
MathOperation.multiply_add,
|
||||
)
|
||||
]
|
||||
alignment_constraints = [1]
|
||||
tile_descriptions = [
|
||||
([128, 128, 8], 2, [4, 2, 1], min_cc, max_cc),
|
||||
([128, 64, 8], 2, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 128, 8], 2, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 64, 8], 2, [2, 1, 1], min_cc, max_cc),
|
||||
([128, 32, 8], 2, [2, 1, 1], min_cc, max_cc),
|
||||
([32, 128, 8], 2, [1, 2, 1], min_cc, max_cc),
|
||||
]
|
||||
|
||||
def get_tile_descriptions(math_inst):
|
||||
return [
|
||||
TileDescription(threadblock_shape, stages, warp_count, math_inst, min_cc, max_cc)
|
||||
for threadblock_shape, stages, warp_count, min_cc, max_cc in tile_descriptions
|
||||
]
|
||||
|
||||
return generate_tensor_op_common(
|
||||
math_instructions, alignment_constraints, get_tile_descriptions, op_creator
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def generate_sm75_tensor_op_1688(
|
||||
out_dtype,
|
||||
arg0_dtype,
|
||||
arg1_dtype,
|
||||
op_creator,
|
||||
check_align,
|
||||
_,
|
||||
profile_all_alignments=False,
|
||||
accumlator_dtype="float32",
|
||||
):
|
||||
"""Generate GEMM or Conv2D kernels for Turing."""
|
||||
assert out_dtype in ["float32", "float16", "int32"]
|
||||
min_cc = 75
|
||||
max_cc = 1024
|
||||
|
||||
if arg0_dtype == "float16" and arg1_dtype == "float16":
|
||||
math_instructions = [
|
||||
MathInstruction(
|
||||
[16, 8, 8],
|
||||
DataType.f16,
|
||||
DataType.f16,
|
||||
dtype_map[out_dtype],
|
||||
dtype_map[accumlator_dtype],
|
||||
OpcodeClass.TensorOp,
|
||||
MathOperation.multiply_add,
|
||||
)
|
||||
]
|
||||
alignment_constraints = [8, 4, 2, 1]
|
||||
tile_descriptions = [
|
||||
([256, 128, 32], 2, [4, 2, 1], min_cc, max_cc),
|
||||
([128, 256, 32], 2, [2, 4, 1], min_cc, max_cc),
|
||||
([128, 128, 32], 2, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 128, 32], 2, [2, 2, 1], min_cc, max_cc),
|
||||
([128, 64, 32], 2, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 64, 32], 2, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 128, 64], 2, [1, 2, 2], min_cc, max_cc),
|
||||
]
|
||||
|
||||
elif "int8" in arg0_dtype and "int8" in arg1_dtype:
|
||||
assert out_dtype == "int32"
|
||||
math_instructions = [
|
||||
MathInstruction(
|
||||
[8, 8, 16],
|
||||
dtype_map[arg0_dtype],
|
||||
dtype_map[arg1_dtype],
|
||||
DataType.s32,
|
||||
DataType.s32,
|
||||
OpcodeClass.TensorOp,
|
||||
MathOperation.multiply_add_saturate,
|
||||
)
|
||||
]
|
||||
alignment_constraints = [16, 8, 4, 2, 1]
|
||||
tile_descriptions = [
|
||||
([256, 128, 64], 2, [4, 2, 1], min_cc, max_cc),
|
||||
([128, 256, 64], 2, [2, 4, 1], min_cc, max_cc),
|
||||
([128, 128, 64], 2, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 256, 64], 2, [1, 4, 1], min_cc, max_cc),
|
||||
([256, 64, 64], 2, [4, 1, 1], min_cc, max_cc),
|
||||
([64, 128, 64], 2, [2, 2, 1], min_cc, max_cc),
|
||||
([128, 64, 64], 2, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 64, 64], 2, [2, 2, 1], min_cc, max_cc),
|
||||
]
|
||||
elif arg0_dtype == "float32" and arg1_dtype == "float32" and out_dtype == "float32":
|
||||
return generate_sm50_simt(out_dtype, arg0_dtype, arg1_dtype, op_creator, accumlator_dtype)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
alignment_constraints = [align for align in alignment_constraints if check_align(align)]
|
||||
assert len(alignment_constraints) > 0
|
||||
|
||||
if not profile_all_alignments:
|
||||
alignment_constraints = [alignment_constraints[0]]
|
||||
|
||||
def get_tile_descriptions(math_inst):
|
||||
return [
|
||||
TileDescription(threadblock_shape, stages, warp_count, math_inst, min_cc, max_cc)
|
||||
for threadblock_shape, stages, warp_count, min_cc, max_cc in tile_descriptions
|
||||
]
|
||||
|
||||
return generate_tensor_op_common(
|
||||
math_instructions, alignment_constraints, get_tile_descriptions, op_creator
|
||||
)
|
||||
|
||||
|
||||
def generate_sm80_tensor_op_16816(
|
||||
out_dtype,
|
||||
arg0_dtype,
|
||||
arg1_dtype,
|
||||
op_creator,
|
||||
check_align,
|
||||
use_3xtf32=True,
|
||||
profile_all_alignments=False,
|
||||
accumlator_dtype="float32",
|
||||
):
|
||||
"""Generate GEMM or Conv2D kernels for Ampere."""
|
||||
min_cc = 80
|
||||
max_cc = 1024
|
||||
max_cc_smem_limited = 80
|
||||
|
||||
def get_default_tile_descriptions(block_k_factor):
|
||||
return [
|
||||
([128, 256, int(32 * block_k_factor)], 3, [2, 4, 1], min_cc, max_cc),
|
||||
([256, 128, int(32 * block_k_factor)], 3, [4, 2, 1], min_cc, max_cc),
|
||||
([256, 64, int(32 * block_k_factor)], 3, [4, 1, 1], min_cc, max_cc),
|
||||
([256, 64, int(32 * block_k_factor)], 4, [4, 1, 1], min_cc, max_cc),
|
||||
([64, 256, int(32 * block_k_factor)], 4, [1, 4, 1], min_cc, max_cc),
|
||||
([128, 128, int(32 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
|
||||
([128, 128, int(32 * block_k_factor)], 4, [2, 2, 1], min_cc, max_cc),
|
||||
([128, 128, int(32 * block_k_factor)], 5, [2, 2, 1], min_cc, max_cc),
|
||||
([128, 64, int(32 * block_k_factor)], 6, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 128, int(32 * block_k_factor)], 6, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 64, int(32 * block_k_factor)], 10, [2, 2, 1], min_cc, max_cc),
|
||||
([256, 128, int(64 * block_k_factor)], 3, [4, 2, 1], min_cc, max_cc_smem_limited),
|
||||
([128, 256, int(64 * block_k_factor)], 3, [2, 4, 1], min_cc, max_cc_smem_limited),
|
||||
([256, 64, int(64 * block_k_factor)], 4, [4, 1, 1], min_cc, max_cc_smem_limited),
|
||||
([64, 256, int(64 * block_k_factor)], 4, [1, 4, 1], min_cc, max_cc_smem_limited),
|
||||
([128, 128, int(64 * block_k_factor)], 4, [2, 2, 1], min_cc, max_cc),
|
||||
([256, 64, int(64 * block_k_factor)], 3, [4, 1, 1], min_cc, max_cc),
|
||||
([64, 256, int(64 * block_k_factor)], 3, [1, 4, 1], min_cc, max_cc),
|
||||
([128, 128, int(64 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
|
||||
([128, 64, int(64 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 128, int(64 * block_k_factor)], 3, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 64, int(64 * block_k_factor)], 5, [2, 2, 1], min_cc, max_cc),
|
||||
]
|
||||
|
||||
if arg0_dtype == "float16" and arg1_dtype == "float16":
|
||||
math_instructions = [
|
||||
MathInstruction(
|
||||
[16, 8, 16],
|
||||
DataType.f16,
|
||||
DataType.f16,
|
||||
dtype_map[out_dtype],
|
||||
dtype_map[accumlator_dtype],
|
||||
OpcodeClass.TensorOp,
|
||||
MathOperation.multiply_add,
|
||||
)
|
||||
]
|
||||
alignment_constraints = [8, 4, 2]
|
||||
tile_descriptions = get_default_tile_descriptions(1)
|
||||
elif arg0_dtype == "float32" and arg1_dtype == "float32":
|
||||
math_instructions = [
|
||||
MathInstruction(
|
||||
[16, 8, 8],
|
||||
DataType.f32,
|
||||
DataType.f32,
|
||||
DataType.f32,
|
||||
DataType.f32,
|
||||
OpcodeClass.TensorOp,
|
||||
MathOperation.multiply_add_fast_f32 if use_3xtf32 else MathOperation.multiply_add,
|
||||
)
|
||||
]
|
||||
alignment_constraints = [4, 2, 1]
|
||||
|
||||
if use_3xtf32:
|
||||
# tf32
|
||||
tile_descriptions = [
|
||||
([128, 128, 16], 4, [4, 2, 1], min_cc, max_cc),
|
||||
([128, 128, 16], 3, [4, 2, 1], min_cc, max_cc),
|
||||
([256, 64, 16], 3, [4, 2, 1], min_cc, max_cc),
|
||||
([64, 256, 16], 3, [2, 4, 1], min_cc, max_cc),
|
||||
([128, 64, 16], 4, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 128, 16], 4, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 64, 16], 3, [2, 2, 1], min_cc, max_cc),
|
||||
([128, 128, 32], 3, [4, 2, 1], min_cc, max_cc),
|
||||
([256, 64, 32], 3, [4, 2, 1], min_cc, max_cc_smem_limited),
|
||||
([64, 256, 32], 3, [2, 4, 1], min_cc, max_cc_smem_limited),
|
||||
([128, 64, 32], 3, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 128, 32], 3, [2, 2, 1], min_cc, max_cc),
|
||||
([64, 64, 32], 3, [2, 2, 1], min_cc, max_cc),
|
||||
]
|
||||
else:
|
||||
tile_descriptions = get_default_tile_descriptions(0.5)
|
||||
else:
|
||||
assert out_dtype == "int32"
|
||||
math_instructions = [
|
||||
MathInstruction(
|
||||
[16, 8, 32],
|
||||
dtype_map[arg0_dtype],
|
||||
dtype_map[arg1_dtype],
|
||||
DataType.s32,
|
||||
DataType.s32,
|
||||
OpcodeClass.TensorOp,
|
||||
MathOperation.multiply_add_saturate,
|
||||
)
|
||||
]
|
||||
alignment_constraints = [16, 8, 4]
|
||||
tile_descriptions = get_default_tile_descriptions(2)
|
||||
|
||||
def get_tile_descriptions(math_inst):
|
||||
return [
|
||||
TileDescription(threadblock_shape, stages, warp_count, math_inst, min_cc, max_cc)
|
||||
for threadblock_shape, stages, warp_count, min_cc, max_cc in tile_descriptions
|
||||
]
|
||||
|
||||
alignment_constraints = [align for align in alignment_constraints if check_align(align)]
|
||||
|
||||
if len(alignment_constraints) > 0 and not profile_all_alignments:
|
||||
alignment_constraints = [alignment_constraints[0]]
|
||||
|
||||
if arg0_dtype != "float32" and arg1_dtype != "float32":
|
||||
sm75_kernels = generate_sm75_tensor_op_1688(
|
||||
out_dtype,
|
||||
arg0_dtype,
|
||||
arg1_dtype,
|
||||
op_creator,
|
||||
check_align,
|
||||
False,
|
||||
profile_all_alignments,
|
||||
accumlator_dtype=accumlator_dtype,
|
||||
)
|
||||
else:
|
||||
# TF32 (float32 + float32 case) is only supported on sm80
|
||||
sm75_kernels = []
|
||||
|
||||
if len(alignment_constraints) > 0:
|
||||
sm80_kernels = generate_tensor_op_common(
|
||||
math_instructions, alignment_constraints, get_tile_descriptions, op_creator
|
||||
)
|
||||
else:
|
||||
sm80_kernels = []
|
||||
|
||||
# TODO(masahi): For int8 kernels, The CUTLASS generator modifies the output tensor alignment
|
||||
# after ops are created. Revisit how important this modification is.
|
||||
# for op in operations:
|
||||
# if op.tile_description.threadblock_shape[1] >= 128:
|
||||
# op.C.alignment = 16
|
||||
# else:
|
||||
# op.C.alignment = 8
|
||||
|
||||
return sm75_kernels + sm80_kernels
|
||||
|
||||
|
||||
GENERATOR_FUNC_TABLE = {75: generate_sm75_tensor_op_1688, 80: generate_sm80_tensor_op_16816}
|
||||
|
||||
|
||||
# (Epilogue functor name, no_beta_scaling)
|
||||
EPILOGUE_MAP = {
|
||||
"cutlass.dense": (EpilogueFunctor.LinearCombination, False),
|
||||
"cutlass.dense_bias": (EpilogueFunctor.LinearCombinationBias, True),
|
||||
"cutlass.dense_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
|
||||
"cutlass.dense_bias_gelu_fp16": (EpilogueFunctor.LinearCombinationGelu, False),
|
||||
"cutlass.dense_bias_gelu_fp32": (EpilogueFunctor.LinearCombinationGelu, False),
|
||||
"cutlass.matmul": (EpilogueFunctor.LinearCombination, False),
|
||||
"cutlass.matmul_bias": (EpilogueFunctor.LinearCombinationBias, True),
|
||||
"cutlass.matmul_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
|
||||
"cutlass.matmul_bias_gelu": (EpilogueFunctor.LinearCombinationGelu, False),
|
||||
"cutlass.matmul_transposed": (EpilogueFunctor.LinearCombination, False),
|
||||
"cutlass.matmul_transposed_bias": (EpilogueFunctor.LinearCombinationBias, True),
|
||||
"cutlass.matmul_transposed_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
|
||||
"cutlass.matmul_transposed_bias_gelu": (EpilogueFunctor.LinearCombinationGelu, False),
|
||||
"cutlass.batch_matmul": (EpilogueFunctor.LinearCombination, False),
|
||||
"cutlass.conv2d_bias_hardswish": (EpilogueFunctor.LinearCombinationHardSwish, False),
|
||||
"cutlass.conv2d_bias_silu": (EpilogueFunctor.LinearCombinationSilu, False),
|
||||
"cutlass.conv2d_bias_sigmoid": (EpilogueFunctor.LinearCombinationSigmoid, False),
|
||||
"cutlass.conv2d_bias_relu": (EpilogueFunctor.LinearCombinationRelu, True),
|
||||
"cutlass.conv2d_bias": (EpilogueFunctor.LinearCombinationBias, True),
|
||||
"cutlass.conv2d": (EpilogueFunctor.LinearCombination, False),
|
||||
"cutlass.conv2d_transpose": (EpilogueFunctor.LinearCombination, False),
|
||||
"cutlass.conv2d_backward_weight": (EpilogueFunctor.LinearCombination, False),
|
||||
}
|
||||
|
||||
|
||||
class ProfilerEngine:
|
||||
"""Compile and run a given profiler executable."""
|
||||
|
||||
def __init__(self, cuda_arch, cutlass_path, binary_prefix):
|
||||
self.cuda_arch = cuda_arch
|
||||
self.binary_prefix = binary_prefix
|
||||
self.cutlass = cutlass_path
|
||||
self.cflags = f"-I{cutlass_path}/include -I{cutlass_path}/tools/util/include -O3 -std=c++17"
|
||||
self.cflags += " -DCUTLASS_ENABLE_TENSOR_CORE_MMA=1"
|
||||
self.cflags += (
|
||||
f" -gencode=arch=compute_{cuda_arch},code=[sm_{cuda_arch},compute_{cuda_arch}]"
|
||||
)
|
||||
self.cflags += " -Xcompiler=-Wconversion -Xcompiler=-fno-strict-aliasing"
|
||||
self.cmd = "nvcc {cflags} {src} -o {output}"
|
||||
|
||||
def _compile(self, op):
|
||||
os.makedirs(self.binary_prefix, exist_ok=True)
|
||||
opath = os.path.join(self.binary_prefix, op["name"])
|
||||
if os.path.exists(opath):
|
||||
return
|
||||
fi = tempfile.NamedTemporaryFile("w", delete=False, prefix=self.binary_prefix, suffix=".cu")
|
||||
fi.write(op["src"])
|
||||
fi.close()
|
||||
cmd = self.cmd.format(cflags=self.cflags, src=fi.name, output=opath)
|
||||
logger.info("invoking compilation %s", cmd)
|
||||
os.system(cmd)
|
||||
os.unlink(fi.name)
|
||||
|
||||
def compile_all(self, ops, use_multiprocessing=False):
|
||||
"""Compile all profiler executables."""
|
||||
if use_multiprocessing:
|
||||
pool = multiprocessing.Pool(multiprocessing.cpu_count())
|
||||
pool.map(self._compile, ops)
|
||||
else:
|
||||
for op in ops:
|
||||
self._compile(op)
|
||||
|
||||
def evaluate(self, op, args):
|
||||
"""Run the profiler executable corresponding to op_name with args."""
|
||||
op_name = op["name"]
|
||||
opath = os.path.join(self.binary_prefix, op_name)
|
||||
if not os.path.exists(opath):
|
||||
self._compile(op)
|
||||
if not os.path.exists(opath):
|
||||
# Bail out if compilation fails for a whatever reason (e.g. static assert failure)
|
||||
return float("inf")
|
||||
cmd = [opath]
|
||||
for arg in args:
|
||||
cmd.append(str(arg))
|
||||
try:
|
||||
logger.info("invoking evaluation %s", cmd)
|
||||
sp = subprocess.run(cmd, capture_output=True, check=True)
|
||||
rt = float(sp.stdout)
|
||||
if rt == 0.0:
|
||||
# This seems to happen with split-k using invalid split-k-slices
|
||||
rt = float("inf")
|
||||
logger.info("%s, %f", op_name, rt)
|
||||
except subprocess.CalledProcessError:
|
||||
rt = float("inf")
|
||||
return rt
|
||||
|
||||
|
||||
class CodegenResult(Object):
|
||||
"""The holder for the generated code and required headers."""
|
||||
|
||||
def __init__(self, code, headers):
|
||||
self.__init_handle_by_constructor__(ffi.CodegenResult, code, headers)
|
||||
|
||||
|
||||
def _get_optional_int_annotation(annotations, key, default=None):
|
||||
value = annotations.get(key, default)
|
||||
if value is None:
|
||||
return default
|
||||
return int(value)
|
||||
|
||||
|
||||
@tvm_ffi.register_global_func("contrib.cutlass.instantiate_template")
|
||||
def instantiate_template(func_name, annotations, func_args):
|
||||
"""Return CUTLASS host code based on a template and the provided annotations.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func_name: str
|
||||
A string to identify the type of the kernel (dense/matmul, batched_matmul, or conv2d).
|
||||
|
||||
annotations: tvm_ffi.Map
|
||||
Key and value pairs annotated during kernel selection.
|
||||
|
||||
func_args: list
|
||||
Names of the function arguments.
|
||||
|
||||
Returns
|
||||
-------
|
||||
codegen_result : CodegenResult
|
||||
Generated CUTLASS host code and required header-file names.
|
||||
"""
|
||||
attrs = {}
|
||||
|
||||
for k in ["lda", "ldb", "ldc", "cutlass_op_def", "cutlass_op_name", "op_type"]:
|
||||
if k in annotations:
|
||||
attrs[k] = annotations[k]
|
||||
|
||||
headers = ["tvm/ffi/function.h", "tvm/ffi/extra/c_env_api.h"]
|
||||
|
||||
if "relu" in func_name:
|
||||
headers.append("cutlass/epilogue/thread/linear_combination_bias_relu.h")
|
||||
elif "gelu" in func_name:
|
||||
headers.append("cutlass/epilogue/thread/linear_combination_gelu.h")
|
||||
elif "sigmoid" in func_name:
|
||||
headers.append("cutlass/epilogue/thread/linear_combination_sigmoid.h")
|
||||
elif "silu" in func_name:
|
||||
headers.append("cutlass/epilogue/thread/linear_combination_silu.h")
|
||||
elif "hardswish" in func_name:
|
||||
headers.append("cutlass/epilogue/thread/linear_combination_hardswish.h")
|
||||
else:
|
||||
headers.append("cutlass/epilogue/thread/linear_combination.h")
|
||||
|
||||
if "residual" in func_name:
|
||||
headers.append("cutlass/epilogue/thread/linear_combination_residual_block.h")
|
||||
|
||||
def get_dim(shape_annot, var_name, axis_idx, batched_offset=0):
|
||||
if isinstance(shape_annot, IntImm):
|
||||
return str(int(shape_annot))
|
||||
return f"{var_name}->shape[{batched_offset + axis_idx}]"
|
||||
|
||||
def get_batch_stride(stride_annot, arg0_idx, arg1_idx, arg0_axis_idx, arg1_axis_idx):
|
||||
if isinstance(stride_annot, IntImm):
|
||||
return str(int(stride_annot))
|
||||
dim1 = func_args[arg0_idx] + f"->shape[{arg0_axis_idx}]"
|
||||
dim2 = func_args[arg1_idx] + f"->shape[{arg1_axis_idx}]"
|
||||
return dim1 + " * " + dim2
|
||||
|
||||
def get_flattened_batch_dim(arg_name, batch_rank):
|
||||
return " * ".join([f"{arg_name}->shape[{i}]" for i in range(batch_rank)])
|
||||
|
||||
if "decode_matmul" in func_name:
|
||||
headers.append("cutlass_kernels/fpA_intB_gemm.h")
|
||||
lhs_arg_idx = _get_optional_int_annotation(annotations, "lhs_arg_idx", 0)
|
||||
rhs_arg_idx = _get_optional_int_annotation(annotations, "rhs_arg_idx", 1)
|
||||
scales_arg_idx = _get_optional_int_annotation(annotations, "scales_arg_idx", 2)
|
||||
bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
|
||||
residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)
|
||||
|
||||
attrs["A_arg"] = func_args[lhs_arg_idx]
|
||||
attrs["B_arg"] = func_args[rhs_arg_idx]
|
||||
attrs["scales_arg"] = func_args[scales_arg_idx]
|
||||
attrs["activation"] = annotations.get("activation", "identity")
|
||||
attrs["bias_stride"] = annotations["bias_stride"]
|
||||
attrs["M"] = annotations["M"]
|
||||
attrs["group_size"] = annotations["group_size"]
|
||||
|
||||
if not isinstance(attrs["M"], tvm.tirx.IntImm):
|
||||
attrs["M"] = get_flattened_batch_dim(
|
||||
func_args[lhs_arg_idx], int(annotations["batch_rank"])
|
||||
)
|
||||
|
||||
if bias_arg_idx is not None:
|
||||
attrs["bias_arg"] = func_args[bias_arg_idx]
|
||||
|
||||
if residual_arg_idx is not None:
|
||||
attrs["residual_arg"] = func_args[residual_arg_idx]
|
||||
attrs["binary_op"] = annotations["binary_op"]
|
||||
attrs["unary_op"] = annotations["unary_op"]
|
||||
|
||||
if annotations["weight_nbit"] == 4:
|
||||
attrs["weight_dtype"] = "cutlass::uint4b_t"
|
||||
attrs["float_per_int"] = 2
|
||||
else:
|
||||
assert annotations["weight_nbit"] == 8
|
||||
attrs["weight_dtype"] = "uint8_t"
|
||||
attrs["float_per_int"] = 1
|
||||
|
||||
code = emit_fp16A_intB_matmul(attrs)
|
||||
return CodegenResult(code, headers)
|
||||
|
||||
elif "dense" in func_name or "matmul" in func_name:
|
||||
batched = "batch" in annotations
|
||||
# dense is equal to transposed_matmul
|
||||
transposed = "transposed" in func_name or "dense" in func_name
|
||||
lhs_arg_idx = _get_optional_int_annotation(annotations, "lhs_arg_idx", 0)
|
||||
rhs_arg_idx = _get_optional_int_annotation(annotations, "rhs_arg_idx", 1)
|
||||
if "bias" in func_name:
|
||||
bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", 2)
|
||||
else:
|
||||
bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
|
||||
residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)
|
||||
|
||||
lhs_arg = func_args[lhs_arg_idx]
|
||||
rhs_arg = func_args[rhs_arg_idx]
|
||||
lhs_shape = annotations[f"arg{lhs_arg_idx}_shape"]
|
||||
rhs_shape = annotations[f"arg{rhs_arg_idx}_shape"]
|
||||
lhs_batched_offset = len(lhs_shape) - 2
|
||||
rhs_batched_offset = len(rhs_shape) - 2
|
||||
|
||||
attrs["lhs_arg"] = lhs_arg
|
||||
attrs["rhs_arg"] = rhs_arg
|
||||
|
||||
if bias_arg_idx is not None:
|
||||
attrs["bias_arg"] = func_args[bias_arg_idx]
|
||||
if residual_arg_idx is not None:
|
||||
attrs["residual_arg"] = func_args[residual_arg_idx]
|
||||
|
||||
attrs["ElementInputA"] = DataTypeTag[dtype_map[annotations[f"arg{lhs_arg_idx}_dtype"]]]
|
||||
attrs["ElementInputB"] = DataTypeTag[dtype_map[annotations[f"arg{rhs_arg_idx}_dtype"]]]
|
||||
attrs["ElementOutput"] = DataTypeTag[dtype_map[annotations["ret_dtype"]]]
|
||||
|
||||
attrs["K"] = lhs_shape[lhs_batched_offset + 1]
|
||||
attrs["M"] = get_dim(lhs_shape[lhs_batched_offset], lhs_arg, 0, lhs_batched_offset)
|
||||
|
||||
if transposed:
|
||||
attrs["N"] = get_dim(rhs_shape[rhs_batched_offset], rhs_arg, 0, rhs_batched_offset)
|
||||
else:
|
||||
attrs["N"] = get_dim(rhs_shape[rhs_batched_offset + 1], rhs_arg, 1, rhs_batched_offset)
|
||||
|
||||
if batched:
|
||||
headers.append("cutlass/gemm/device/gemm_batched.h")
|
||||
|
||||
def get_batch_on_arg(arg_name, arg_shape):
|
||||
return " * ".join(f"{arg_name}->shape[{i}]" for i in range(len(arg_shape) - 2))
|
||||
|
||||
if isinstance(annotations["batch"], IntImm):
|
||||
attrs["batch"] = str(int(annotations["batch"]))
|
||||
elif annotations["batch_stride_A"] == 0:
|
||||
# 2D x ND
|
||||
attrs["batch"] = get_batch_on_arg(rhs_arg, rhs_shape)
|
||||
else:
|
||||
# ND x 2D or ND x ND
|
||||
attrs["batch"] = get_batch_on_arg(lhs_arg, lhs_shape)
|
||||
|
||||
attrs["batch_stride_A"] = get_batch_stride(
|
||||
annotations["batch_stride_A"],
|
||||
lhs_arg_idx,
|
||||
lhs_arg_idx,
|
||||
lhs_batched_offset,
|
||||
lhs_batched_offset + 1,
|
||||
)
|
||||
attrs["batch_stride_B"] = get_batch_stride(
|
||||
annotations["batch_stride_B"],
|
||||
rhs_arg_idx,
|
||||
rhs_arg_idx,
|
||||
rhs_batched_offset,
|
||||
rhs_batched_offset + 1,
|
||||
)
|
||||
|
||||
if transposed:
|
||||
attrs["batch_stride_C"] = get_batch_stride(
|
||||
annotations["batch_stride_C"],
|
||||
lhs_arg_idx,
|
||||
rhs_arg_idx,
|
||||
lhs_batched_offset,
|
||||
rhs_batched_offset,
|
||||
)
|
||||
else:
|
||||
attrs["batch_stride_C"] = get_batch_stride(
|
||||
annotations["batch_stride_C"],
|
||||
lhs_arg_idx,
|
||||
rhs_arg_idx,
|
||||
lhs_batched_offset,
|
||||
rhs_batched_offset + 1,
|
||||
)
|
||||
else:
|
||||
headers.append("cutlass/gemm/device/gemm.h")
|
||||
|
||||
if "residual" in func_name:
|
||||
headers.append("cutlass/gemm/device/gemm_universal_with_broadcast.h")
|
||||
|
||||
code = instantiate_gemm_template(attrs)
|
||||
return CodegenResult(code, headers)
|
||||
|
||||
elif "conv2d" in func_name:
|
||||
data_arg_idx = _get_optional_int_annotation(annotations, "data_arg_idx", 0)
|
||||
weight_arg_idx = _get_optional_int_annotation(annotations, "weight_arg_idx", 1)
|
||||
bias_arg_idx = _get_optional_int_annotation(annotations, "bias_arg_idx", None)
|
||||
residual_arg_idx = _get_optional_int_annotation(annotations, "residual_arg_idx", None)
|
||||
|
||||
attrs["data_arg"] = func_args[data_arg_idx]
|
||||
attrs["weight_arg"] = func_args[weight_arg_idx]
|
||||
|
||||
if bias_arg_idx is not None:
|
||||
attrs["bias_arg"] = func_args[bias_arg_idx]
|
||||
if residual_arg_idx is not None:
|
||||
attrs["residual_arg"] = func_args[residual_arg_idx]
|
||||
|
||||
activation_shape = annotations[f"arg{data_arg_idx}_shape"]
|
||||
weight_shape = annotations[f"arg{weight_arg_idx}_shape"]
|
||||
output_shape = annotations["ret_shape"]
|
||||
|
||||
if "conv2d_transpose" in func_name:
|
||||
headers.append("cutlass/conv/kernel/default_conv2d_dgrad.h")
|
||||
activation_shape = output_shape
|
||||
output_shape = annotations["arg0_shape"]
|
||||
elif "backward" in func_name:
|
||||
headers.append("cutlass/conv/kernel/default_conv2d_wgrad.h")
|
||||
activation_shape = annotations["arg1_shape"]
|
||||
weight_shape = output_shape
|
||||
output_shape = annotations["arg0_shape"]
|
||||
elif "residual" in func_name:
|
||||
headers.append("cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h")
|
||||
else:
|
||||
headers.append("cutlass/conv/kernel/default_conv2d_fprop.h")
|
||||
|
||||
headers.append("cutlass/conv/device/implicit_gemm_convolution.h")
|
||||
|
||||
op_name = attrs["cutlass_op_name"]
|
||||
|
||||
if "splitk" in op_name:
|
||||
headers += [
|
||||
"cutlass/reduction/device/reduce_split_k.h",
|
||||
"cutlass/reduction/thread/reduction_operators.h",
|
||||
]
|
||||
|
||||
data_arg = attrs["data_arg"]
|
||||
attrs["N"] = get_dim(activation_shape[0], data_arg, 0)
|
||||
attrs["H"] = get_dim(activation_shape[1], data_arg, 1)
|
||||
attrs["W"] = get_dim(activation_shape[2], data_arg, 2)
|
||||
attrs["C"] = activation_shape[3]
|
||||
attrs["P"] = get_dim(output_shape[1], "out0", 1)
|
||||
attrs["Q"] = get_dim(output_shape[2], "out0", 2)
|
||||
attrs["K"] = output_shape[3]
|
||||
attrs["R"] = weight_shape[1]
|
||||
attrs["S"] = weight_shape[2]
|
||||
attrs["pad_h"] = annotations["padding"][0]
|
||||
attrs["pad_w"] = annotations["padding"][1]
|
||||
attrs["stride_h"] = annotations["strides"][0]
|
||||
attrs["stride_w"] = annotations["strides"][1]
|
||||
attrs["dilation_h"] = annotations["dilation"][0]
|
||||
attrs["dilation_w"] = annotations["dilation"][1]
|
||||
|
||||
if "splitk" in op_name:
|
||||
attrs["split_k_mode"] = "kParallel"
|
||||
attrs["split_k_slices"] = str(re.search(r"splitk(\d+)", op_name).group(1))
|
||||
else:
|
||||
attrs["split_k_mode"] = "kSerial"
|
||||
attrs["split_k_slices"] = 1
|
||||
|
||||
if "residual_shape" in annotations:
|
||||
attrs["residual_shape"] = annotations["residual_shape"]
|
||||
|
||||
code = instantiate_conv2d_template(attrs)
|
||||
return CodegenResult(code, headers)
|
||||
|
||||
elif "attention" in func_name:
|
||||
is_var_len = "var_len" in func_name
|
||||
data_type = dtype_map[annotations["arg0_dtype"]]
|
||||
|
||||
attrs["qkv_layout"] = annotations["qkv_layout"]
|
||||
if attrs["qkv_layout"] == "default":
|
||||
attrs["query"] = func_args[0]
|
||||
attrs["key"] = func_args[1]
|
||||
attrs["value"] = func_args[2]
|
||||
attrs["num_queries"] = s = get_dim(annotations["num_queries"], func_args[0], 1)
|
||||
attrs["num_keys"] = get_dim(annotations["num_keys"], func_args[1], 1)
|
||||
if len(func_args) > 4 and not is_var_len: # +1 for workspace, the last arg
|
||||
attrs["bias"] = func_args[3]
|
||||
elif attrs["qkv_layout"] == "qkv_stacked":
|
||||
attrs["qkv"] = func_args[0]
|
||||
attrs["num_queries"] = s = annotations["num_queries"]
|
||||
attrs["num_keys"] = annotations["num_keys"]
|
||||
if len(func_args) > 2 and not is_var_len: # +1 for workspace, the last arg
|
||||
attrs["bias"] = func_args[1]
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
attrs["data_type"] = DataTypeTag[data_type]
|
||||
attrs["num_batches"] = b = annotations["num_batches"]
|
||||
attrs["head_dim"] = h = annotations["head_dim"]
|
||||
attrs["head_dim_value"] = h_v = annotations["head_dim_value"]
|
||||
attrs["kMaxK"] = max(int(attrs["head_dim"]), int(attrs["head_dim_value"]))
|
||||
attrs["scale"] = (
|
||||
float(1 / math.sqrt(h.value)) if annotations["scale"] is None else annotations["scale"]
|
||||
)
|
||||
|
||||
if is_var_len:
|
||||
attrs["seqstart_q"] = func_args[int(annotations["seqstart_q_idx"])]
|
||||
attrs["seqstart_k"] = func_args[int(annotations["seqstart_k_idx"])]
|
||||
attrs["max_seqlen_q"] = func_args[int(annotations["max_seqlen_q_idx"])]
|
||||
attrs["max_seqlen_k"] = func_args[int(annotations["max_seqlen_k_idx"])]
|
||||
|
||||
is_mqa = annotations["num_q_heads"] != annotations["num_kv_heads"]
|
||||
|
||||
use_flash = (
|
||||
annotations["ret_dtype"] == "float16"
|
||||
and "bias" not in attrs
|
||||
and int(attrs["head_dim"]) <= 256
|
||||
and int(attrs["head_dim"]) % 8 == 0
|
||||
and int(attrs["head_dim"]) == int(attrs["head_dim_value"])
|
||||
# For the causal case (custom mask = "BottomRight"), only use flash for multi-query
|
||||
# attention workloads. Otherwise, CUTLASS fMHA seems faster for causal attention
|
||||
# with a single query.
|
||||
# In addition, sliding-window attention is only supported by flash.
|
||||
and (
|
||||
int(annotations["custom_mask_type"]) == 0
|
||||
or (int(annotations["custom_mask_type"]) == 2 and is_mqa)
|
||||
or (int(annotations["custom_mask_type"]) == 2 and "window_size" in annotations)
|
||||
)
|
||||
# Flash v2 is currently not supported for sm < 80
|
||||
and int(annotations["arch"]) >= 80
|
||||
)
|
||||
|
||||
# See https://github.com/Dao-AILab/flash-attention/blob/
|
||||
# 92dd5703ecdb99aa4a4aee9817f28557907403a2/csrc/flash_attn/flash_api.cpp#L111-L116
|
||||
if "window_size" in annotations:
|
||||
assert use_flash, "Sliding-window attention is supported only by Flash Attention."
|
||||
assert int(annotations["custom_mask_type"]) == 2, (
|
||||
"Sliding-window attention is only supported for causal with bottom right mask."
|
||||
)
|
||||
attrs["window_size_left"] = int(annotations["window_size"]) - 1
|
||||
attrs["window_size_right"] = 0
|
||||
attrs["is_causal"] = False
|
||||
else:
|
||||
if int(annotations["custom_mask_type"]) == 2:
|
||||
attrs["window_size_left"] = attrs["num_keys"]
|
||||
attrs["window_size_right"] = 0
|
||||
attrs["is_causal"] = True
|
||||
else:
|
||||
attrs["window_size_left"] = -1
|
||||
attrs["window_size_right"] = -1
|
||||
attrs["is_causal"] = False
|
||||
|
||||
if use_flash:
|
||||
headers.append("flash.h")
|
||||
attrs["num_q_heads"] = annotations["num_q_heads"]
|
||||
attrs["num_kv_heads"] = annotations["num_kv_heads"]
|
||||
|
||||
if is_var_len:
|
||||
code = instantiate_flash_attention_var_len_template(attrs)
|
||||
else:
|
||||
code = instantiate_flash_attention_template(attrs)
|
||||
else:
|
||||
headers.append("kernel_forward.h")
|
||||
|
||||
assert not is_mqa, (
|
||||
"The number of query and KV heads need to be the same for CUTLASS fMHA."
|
||||
)
|
||||
|
||||
attrs["num_heads"] = n = annotations["num_q_heads"]
|
||||
|
||||
data_type_size = DataTypeSize[data_type]
|
||||
if (data_type_size * h // 8) % 16 == 0 and (data_type_size * h_v // 8) % 16 == 0:
|
||||
attrs["kIsAligned"] = True
|
||||
elif (h % 4 == 0) and (h_v % 4 == 0):
|
||||
attrs["kIsAligned"] = False
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
if h_v > 64:
|
||||
attrs["kQueriesPerBlock"] = 32
|
||||
attrs["kKeysPerBlock"] = 128
|
||||
attrs["kSingleValueIteration"] = h_v <= 128
|
||||
else:
|
||||
attrs["kQueriesPerBlock"] = 64
|
||||
attrs["kKeysPerBlock"] = 64
|
||||
attrs["kSingleValueIteration"] = True
|
||||
|
||||
assert attrs["scale"] > 0 or attrs["scale"] < 0, (
|
||||
"Cutlass may generate nan occasionally when scale == 0.0"
|
||||
)
|
||||
attrs["arch"] = "cutlass::arch::Sm{}".format(annotations["arch"])
|
||||
attrs["kSupportsDropout"] = False
|
||||
|
||||
attrs["output_size"] = f"{b} * {s} * {n} * {h_v}"
|
||||
|
||||
attrs["custom_mask_type"] = annotations["custom_mask_type"]
|
||||
|
||||
for arg in func_args:
|
||||
if "workspace" in arg:
|
||||
attrs["workspace"] = arg
|
||||
if "bias" in attrs:
|
||||
attrs["kSupportsBias"] = True
|
||||
if len(annotations["bias_shape"]) == 4:
|
||||
strides = "p.num_keys"
|
||||
if annotations["bias_shape"][2] == 1:
|
||||
attrs["bias_strideM"] = 0
|
||||
else:
|
||||
attrs["bias_strideM"] = strides
|
||||
strides = f"p.num_queries * {strides}"
|
||||
if annotations["bias_shape"][1] == 1:
|
||||
attrs["bias_strideH"] = 0
|
||||
else:
|
||||
attrs["bias_strideH"] = strides
|
||||
strides = f"p.num_heads * {strides}"
|
||||
if annotations["bias_shape"][0] == 1:
|
||||
attrs["bias_strideB"] = 0
|
||||
else:
|
||||
attrs["bias_strideB"] = strides
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
else:
|
||||
# To support negative scale in current Cutlass implementation,
|
||||
# kSupportsBias should be set true, or there are nan's as result.
|
||||
attrs["kSupportsBias"] = attrs["scale"] < 0
|
||||
|
||||
code = instantiate_attention_template(attrs)
|
||||
|
||||
return CodegenResult(code, headers)
|
||||
elif "layer_norm" in func_name:
|
||||
headers.append("cutlass/util/device_layernorm.h")
|
||||
headers.append("cutlass/layout/matrix.h")
|
||||
attrs = {"input": func_args[0], "gamma": func_args[1], "beta": func_args[2]}
|
||||
attrs.update(dict(annotations))
|
||||
|
||||
if not isinstance(attrs["M"], tvm.tirx.IntImm):
|
||||
attrs["M"] = get_flattened_batch_dim(func_args[0], int(attrs["batch_rank"]))
|
||||
|
||||
code = instantiate_layer_norm_template(attrs)
|
||||
return CodegenResult(code, headers)
|
||||
elif "rms_norm" in func_name:
|
||||
headers.append("cutlass/util/device_rmsnorm.h")
|
||||
headers.append("cutlass/layout/matrix.h")
|
||||
attrs = {"input": func_args[0], "weight": func_args[1]}
|
||||
attrs.update(dict(annotations))
|
||||
|
||||
if not isinstance(attrs["M"], tvm.tirx.IntImm):
|
||||
attrs["M"] = get_flattened_batch_dim(func_args[0], int(attrs["batch_rank"]))
|
||||
|
||||
code = instantiate_rms_norm_template(attrs)
|
||||
return CodegenResult(code, headers)
|
||||
|
||||
raise ValueError(f"Do not have a template for {func_name}")
|
||||
@@ -0,0 +1,48 @@
|
||||
# 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
|
||||
# ruff: noqa: E501
|
||||
"""Generator for CUTLASS layer norm kernels."""
|
||||
|
||||
from .library import substitute_template
|
||||
|
||||
|
||||
def instantiate_layer_norm_template(attrs):
|
||||
"""
|
||||
Return CUTLASS host code for layer norm based on
|
||||
a template and the provided attribute map.
|
||||
"""
|
||||
template = """
|
||||
using data_type = ${data_type};
|
||||
using namespace cutlass::layout;
|
||||
|
||||
int M = ${M};
|
||||
int N = ${N};
|
||||
cutlass::MatrixCoord size(M, N);
|
||||
auto layout_2D = RowMajor::packed(size);
|
||||
auto layout_channels = RowMajor::packed({1, N});
|
||||
|
||||
cutlass::TensorRef<data_type, RowMajor> _input((data_type*)${input}->data, layout_2D);
|
||||
cutlass::TensorRef<data_type, RowMajor> _gamma((data_type*)${gamma}->data, layout_channels);
|
||||
cutlass::TensorRef<data_type, RowMajor> _beta((data_type*)${beta}->data, layout_channels);
|
||||
cutlass::TensorRef<data_type, RowMajor> _output((data_type*)out0->data, layout_2D);
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${input}->device.device_id));
|
||||
|
||||
cutlass::layernorm(size, _output, _input, _gamma, _beta, stream);
|
||||
"""
|
||||
return substitute_template(template, attrs)
|
||||
@@ -0,0 +1,301 @@
|
||||
# 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,line-too-long
|
||||
# ruff: noqa: E501
|
||||
"""Various type definitions to help instantiate CUTLASS kernels."""
|
||||
|
||||
import enum
|
||||
import re
|
||||
from enum import auto as enum_auto
|
||||
|
||||
from tvm.tirx.expr import FloatImm, IntImm
|
||||
|
||||
|
||||
class GeneratorTarget(enum.Enum):
|
||||
Library = enum_auto()
|
||||
|
||||
|
||||
class DataType(enum.Enum):
|
||||
f16 = enum_auto()
|
||||
f32 = enum_auto()
|
||||
s8 = enum_auto()
|
||||
u8 = enum_auto()
|
||||
s32 = enum_auto()
|
||||
|
||||
|
||||
ShortDataTypeNames = {DataType.f16: "h", DataType.f32: "s", DataType.s32: "i"}
|
||||
|
||||
|
||||
DataTypeNames = {
|
||||
DataType.f16: "f16",
|
||||
DataType.f32: "f32",
|
||||
DataType.s8: "s8",
|
||||
DataType.u8: "u8",
|
||||
DataType.s32: "s32",
|
||||
}
|
||||
|
||||
DataTypeTag = {
|
||||
DataType.f16: "cutlass::half_t",
|
||||
DataType.f32: "float",
|
||||
DataType.s8: "int8_t",
|
||||
DataType.s32: "int32_t",
|
||||
DataType.u8: "uint8_t",
|
||||
}
|
||||
|
||||
DataTypeSize = {
|
||||
DataType.f16: 16,
|
||||
DataType.f32: 32,
|
||||
DataType.u8: 8,
|
||||
DataType.s8: 8,
|
||||
DataType.s32: 32,
|
||||
}
|
||||
|
||||
|
||||
class MathOperation(enum.Enum):
|
||||
multiply_add = enum_auto()
|
||||
multiply_add_saturate = enum_auto()
|
||||
multiply_add_fast_f32 = enum_auto()
|
||||
|
||||
|
||||
MathOperationTag = {
|
||||
MathOperation.multiply_add: "cutlass::arch::OpMultiplyAdd",
|
||||
MathOperation.multiply_add_saturate: "cutlass::arch::OpMultiplyAddSaturate",
|
||||
MathOperation.multiply_add_fast_f32: "cutlass::arch::OpMultiplyAddFastF32",
|
||||
}
|
||||
|
||||
|
||||
class LayoutType(enum.Enum):
|
||||
ColumnMajor = enum_auto()
|
||||
RowMajor = enum_auto()
|
||||
TensorNHWC = enum_auto()
|
||||
|
||||
|
||||
LayoutTag = {
|
||||
LayoutType.ColumnMajor: "cutlass::layout::ColumnMajor",
|
||||
LayoutType.RowMajor: "cutlass::layout::RowMajor",
|
||||
LayoutType.TensorNHWC: "cutlass::layout::TensorNHWC",
|
||||
}
|
||||
|
||||
|
||||
TransposedLayout = {
|
||||
LayoutType.ColumnMajor: LayoutType.RowMajor,
|
||||
LayoutType.RowMajor: LayoutType.ColumnMajor,
|
||||
LayoutType.TensorNHWC: LayoutType.TensorNHWC,
|
||||
}
|
||||
|
||||
|
||||
ShortLayoutTypeNames = {
|
||||
LayoutType.ColumnMajor: "n",
|
||||
LayoutType.RowMajor: "t",
|
||||
LayoutType.TensorNHWC: "nhwc",
|
||||
}
|
||||
|
||||
|
||||
class OpcodeClass(enum.Enum):
|
||||
Simt = enum_auto()
|
||||
TensorOp = enum_auto()
|
||||
WmmaTensorOp = enum_auto()
|
||||
|
||||
|
||||
OpcodeClassNames = {
|
||||
OpcodeClass.Simt: "simt",
|
||||
OpcodeClass.TensorOp: "tensorop",
|
||||
OpcodeClass.WmmaTensorOp: "wmma_tensorop",
|
||||
}
|
||||
|
||||
OpcodeClassTag = {
|
||||
OpcodeClass.Simt: "cutlass::arch::OpClassSimt",
|
||||
OpcodeClass.TensorOp: "cutlass::arch::OpClassTensorOp",
|
||||
OpcodeClass.WmmaTensorOp: "cutlass::arch::OpClassWmmaTensorOp",
|
||||
}
|
||||
|
||||
|
||||
class OperationKind(enum.Enum):
|
||||
Gemm = enum_auto()
|
||||
Conv2d = enum_auto()
|
||||
|
||||
|
||||
OperationKindNames = {OperationKind.Gemm: "gemm", OperationKind.Conv2d: "conv2d"}
|
||||
|
||||
|
||||
class Target(enum.Enum):
|
||||
library = enum_auto()
|
||||
|
||||
|
||||
def substitute_template(template, values):
|
||||
"""Instantiate a kernel template using `values`."""
|
||||
text = template
|
||||
changed = True
|
||||
while changed:
|
||||
changed = False
|
||||
for key, value in values.items():
|
||||
if isinstance(value, int | IntImm):
|
||||
value = str(int(value))
|
||||
if isinstance(value, float | FloatImm):
|
||||
value = str(float(value))
|
||||
elif isinstance(value, bool):
|
||||
value = str(value).lower()
|
||||
regex = f"\\$\\{{{key}\\}}"
|
||||
newtext = re.sub(regex, value, text)
|
||||
if newtext != text:
|
||||
changed = True
|
||||
text = newtext
|
||||
return text
|
||||
|
||||
|
||||
class GemmKind(enum.Enum):
|
||||
Gemm = enum_auto()
|
||||
|
||||
|
||||
GemmKindNames = {GemmKind.Gemm: "gemm"}
|
||||
|
||||
|
||||
class EpilogueFunctor(enum.Enum):
|
||||
LinearCombination = enum_auto()
|
||||
LinearCombinationRelu = enum_auto()
|
||||
LinearCombinationBias = enum_auto()
|
||||
LinearCombinationGelu = enum_auto()
|
||||
LinearCombinationSigmoid = enum_auto()
|
||||
LinearCombinationSilu = enum_auto()
|
||||
LinearCombinationHardSwish = enum_auto()
|
||||
LinearCombinationResidualBlock = enum_auto()
|
||||
|
||||
|
||||
EpilogueFunctorTag = {
|
||||
EpilogueFunctor.LinearCombination: "cutlass::epilogue::thread::LinearCombination",
|
||||
EpilogueFunctor.LinearCombinationRelu: "cutlass::epilogue::thread::LinearCombinationRelu",
|
||||
EpilogueFunctor.LinearCombinationBias: "cutlass::epilogue::thread::LinearCombination",
|
||||
EpilogueFunctor.LinearCombinationGelu: "cutlass::epilogue::thread::LinearCombinationGELU",
|
||||
EpilogueFunctor.LinearCombinationSigmoid: "cutlass::epilogue::thread::LinearCombinationSigmoid",
|
||||
EpilogueFunctor.LinearCombinationSilu: "cutlass::epilogue::thread::LinearCombinationSilu",
|
||||
EpilogueFunctor.LinearCombinationHardSwish: "cutlass::epilogue::thread::LinearCombinationHardSwish",
|
||||
EpilogueFunctor.LinearCombinationResidualBlock: "cutlass::epilogue::thread::LinearCombinationResidualBlock",
|
||||
}
|
||||
|
||||
|
||||
class SwizzlingFunctor(enum.Enum):
|
||||
Identity1 = enum_auto()
|
||||
Identity2 = enum_auto()
|
||||
Identity4 = enum_auto()
|
||||
Identity8 = enum_auto()
|
||||
Batched = enum_auto()
|
||||
StridedDgradIdentity1 = enum_auto()
|
||||
StridedDgradIdentity4 = enum_auto()
|
||||
|
||||
|
||||
SwizzlingFunctorTag = {
|
||||
SwizzlingFunctor.Identity1: "cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>",
|
||||
SwizzlingFunctor.Identity2: "cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<2>",
|
||||
SwizzlingFunctor.Identity4: "cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<4>",
|
||||
SwizzlingFunctor.Identity8: "cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>",
|
||||
SwizzlingFunctor.Batched: "cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle",
|
||||
SwizzlingFunctor.StridedDgradIdentity1: "cutlass::conv::threadblock::StridedDgradIdentityThreadblockSwizzle<1>",
|
||||
SwizzlingFunctor.StridedDgradIdentity4: "cutlass::conv::threadblock::StridedDgradIdentityThreadblockSwizzle<4>",
|
||||
}
|
||||
|
||||
|
||||
class ConvKind(enum.Enum):
|
||||
Fprop = enum_auto()
|
||||
Dgrad = enum_auto()
|
||||
Wgrad = enum_auto()
|
||||
|
||||
|
||||
ConvKindTag = {
|
||||
ConvKind.Fprop: "cutlass::conv::Operator::kFprop",
|
||||
ConvKind.Dgrad: "cutlass::conv::Operator::kDgrad",
|
||||
ConvKind.Wgrad: "cutlass::conv::Operator::kWgrad",
|
||||
}
|
||||
|
||||
|
||||
ConvKindNames = {ConvKind.Fprop: "fprop", ConvKind.Dgrad: "dgrad", ConvKind.Wgrad: "wgrad"}
|
||||
|
||||
|
||||
class StrideSupport(enum.Enum):
|
||||
Strided = enum_auto()
|
||||
Unity = enum_auto()
|
||||
|
||||
|
||||
StrideSupportTag = {
|
||||
StrideSupport.Strided: "cutlass::conv::StrideSupport::kStrided",
|
||||
StrideSupport.Unity: "cutlass::conv::StrideSupport::kUnity",
|
||||
}
|
||||
|
||||
|
||||
StrideSupportNames = {StrideSupport.Strided: "", StrideSupport.Unity: "unity_stride"}
|
||||
|
||||
|
||||
class IteratorAlgorithm(enum.Enum):
|
||||
Analytic = enum_auto()
|
||||
Optimized = enum_auto()
|
||||
|
||||
|
||||
IteratorAlgorithmTag = {
|
||||
IteratorAlgorithm.Analytic: "cutlass::conv::IteratorAlgorithm::kAnalytic",
|
||||
IteratorAlgorithm.Optimized: "cutlass::conv::IteratorAlgorithm::kOptimized",
|
||||
}
|
||||
|
||||
|
||||
IteratorAlgorithmNames = {
|
||||
IteratorAlgorithm.Analytic: "analytic",
|
||||
IteratorAlgorithm.Optimized: "optimized",
|
||||
}
|
||||
|
||||
|
||||
class MathInstruction:
|
||||
"""Describe characteristics of a math instruction."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
instruction_shape,
|
||||
element_a,
|
||||
element_b,
|
||||
element_c,
|
||||
element_accumulator,
|
||||
opcode_class,
|
||||
math_operation=MathOperation.multiply_add,
|
||||
):
|
||||
self.instruction_shape = instruction_shape
|
||||
self.element_a = element_a
|
||||
self.element_b = element_b
|
||||
self.element_c = element_c
|
||||
self.element_accumulator = element_accumulator
|
||||
self.opcode_class = opcode_class
|
||||
self.math_operation = math_operation
|
||||
|
||||
|
||||
class TileDescription:
|
||||
"""Describe characteristics of a GEMM tile."""
|
||||
|
||||
def __init__(
|
||||
self, threadblock_shape, stages, warp_count, math_instruction, min_compute, max_compute
|
||||
):
|
||||
self.threadblock_shape = threadblock_shape
|
||||
self.stages = stages
|
||||
self.warp_count = warp_count
|
||||
self.math_instruction = math_instruction
|
||||
self.minimum_compute_capability = min_compute
|
||||
self.maximum_compute_capability = max_compute
|
||||
|
||||
def procedural_name(self):
|
||||
return f"{self.threadblock_shape[0]}x{self.threadblock_shape[1]}_{self.threadblock_shape[2]}x{self.stages}"
|
||||
|
||||
|
||||
class TensorDescription:
|
||||
def __init__(self, element, layout, alignment=1):
|
||||
self.element = element
|
||||
self.layout = layout
|
||||
self.alignment = alignment
|
||||
@@ -0,0 +1,47 @@
|
||||
# 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
|
||||
# ruff: noqa: E501
|
||||
"""Generator for CUTLASS rms norm kernels."""
|
||||
|
||||
from .library import substitute_template
|
||||
|
||||
|
||||
def instantiate_rms_norm_template(attrs):
|
||||
"""
|
||||
Return CUTLASS host code for rms norm based on
|
||||
a template and the provided attribute map.
|
||||
"""
|
||||
template = """
|
||||
using data_type = ${data_type};
|
||||
using namespace cutlass::layout;
|
||||
|
||||
int M = ${M};
|
||||
int N = ${N};
|
||||
cutlass::MatrixCoord size(M, N);
|
||||
auto layout_2D = RowMajor::packed(size);
|
||||
auto layout_channels = RowMajor::packed({1, N});
|
||||
|
||||
cutlass::TensorRef<data_type, RowMajor> _input((data_type*)${input}->data, layout_2D);
|
||||
cutlass::TensorRef<data_type, RowMajor> _weight((data_type*)${weight}->data, layout_channels);
|
||||
cutlass::TensorRef<data_type, RowMajor> _output((data_type*)out0->data, layout_2D);
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${input}->device.device_id));
|
||||
|
||||
cutlass::rmsnorm(size, _output, _input, _weight, stream, ${rms_eps});
|
||||
"""
|
||||
return substitute_template(template, attrs)
|
||||
Reference in New Issue
Block a user