mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-12 17:26:00 +00:00
29e58d5b28
Also bundle the CK library and include files with the pip package. The package is pip-installable with `pip install git+https://github.com/tenpercent/composable_kernel@enable-pip` (substitute the repo path and branch if necessary) Testing: `myenv/bin/python3 -m ck4inductor.universal_gemm.gen_instances` (prints a list of instances) `tree myenv/lib/python3.12/site-packages/ck4inductor` (observe the list of sources along the installed package)
571 lines
20 KiB
Python
571 lines
20 KiB
Python
import logging
|
|
import os
|
|
import subprocess
|
|
from dataclasses import fields, replace
|
|
from functools import lru_cache, partial
|
|
from typing import List
|
|
|
|
from ..util import library_path
|
|
|
|
from .op import CKGemmOperation
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def _ck_library_dir():
|
|
gemm_instances_path = os.path.join(
|
|
library_path(), "src", "tensor_operation_instance", "gpu", "gemm_universal"
|
|
)
|
|
if not os.path.exists(gemm_instances_path):
|
|
log.error("CK library path %s does not exist", gemm_instances_path)
|
|
return None
|
|
return gemm_instances_path
|
|
|
|
|
|
def parse_instances(str_instances: List[str]) -> List[CKGemmOperation]:
|
|
"""
|
|
Parse the lines containing Universal Gemm template instances into `CKGemmOperation` instances
|
|
"""
|
|
|
|
def maybe_int(s):
|
|
try:
|
|
return int(s)
|
|
except ValueError:
|
|
return s
|
|
|
|
op_instances = []
|
|
for line in str_instances:
|
|
s_template_args = line.split("DeviceGemm_Xdl_CShuffleV3")[-1].strip("<>, ")
|
|
template_args = []
|
|
i_current = 0
|
|
while i_current < len(s_template_args):
|
|
if s_template_args[i_current] == " ":
|
|
# skip whitespace
|
|
i_current += 1
|
|
continue
|
|
elif s_template_args[i_current : i_current + 2] == "S<":
|
|
# parse template S<Index...>
|
|
i_next = s_template_args.find(">", i_current)
|
|
template_args.append(
|
|
tuple(map(int, s_template_args[i_current + 2 : i_next].split(",")))
|
|
)
|
|
i_current = i_next + 2
|
|
else:
|
|
# all string attributes must be either type aliases or global constants in C++
|
|
i_next = s_template_args.find(",", i_current)
|
|
template_args.append(
|
|
maybe_int(
|
|
s_template_args[i_current : i_next if i_next != -1 else None]
|
|
)
|
|
)
|
|
if i_next != -1:
|
|
i_current = i_next + 1
|
|
if i_next == -1:
|
|
break
|
|
# pad with `None`s for the fields which are not defined in the instance
|
|
new_instance = CKGemmOperation(
|
|
*template_args, # type: ignore[arg-type]
|
|
*((None,) * (len(fields(CKGemmOperation)) - len(template_args))),
|
|
)
|
|
# the last 2 template parameters are optional
|
|
# if they are absent, substitute them with default values from Universal Gemm C++ template declaration
|
|
if new_instance.a_compute_dtype is None:
|
|
new_instance.a_compute_dtype = new_instance.c_element_dtype
|
|
if new_instance.b_compute_dtype is None:
|
|
new_instance.b_compute_dtype = new_instance.c_element_dtype
|
|
|
|
op_instances.append(new_instance)
|
|
return op_instances
|
|
|
|
|
|
def default_instances() -> List[CKGemmOperation]:
|
|
# fallback: known working op instance for problem size M=2240 K=256 N=2048
|
|
# all string attributes must be either type aliases or global constants in C++
|
|
|
|
return [
|
|
CKGemmOperation(
|
|
a_layout="Row",
|
|
b_layout="Row",
|
|
c_layout="Row",
|
|
a_element_dtype="F16",
|
|
b_element_dtype="F16",
|
|
c_element_dtype="F16",
|
|
a_compute_dtype="F16",
|
|
b_compute_dtype="F16",
|
|
acc_dtype="F32",
|
|
c_shuffle_dtype="F16",
|
|
a_elementwise_op="PassThrough",
|
|
b_elementwise_op="PassThrough",
|
|
c_elementwise_op="PassThrough",
|
|
gemm_specialization="GemmSpecialization::Default",
|
|
block_size=256,
|
|
m_per_block=224,
|
|
n_per_block=256,
|
|
k_per_block=64,
|
|
a_k1=8,
|
|
b_k1=2,
|
|
m_per_xdl=16,
|
|
n_per_xdl=16,
|
|
m_xdl_per_wave=7,
|
|
n_xdl_per_wave=8,
|
|
a_block_transfer_thread_cluster_lengths_ak0_m_ak1=(8, 32, 1),
|
|
a_block_transfer_thread_cluster_arrange_order=(1, 0, 2),
|
|
a_block_transfer_src_access_order=(1, 0, 2),
|
|
a_block_transfer_src_vector_dim=2,
|
|
a_block_transfer_src_scalar_per_vector=8,
|
|
a_block_transfer_dst_scalar_per_vector_ak1=8,
|
|
a_block_lds_extra_m=0, # type: ignore[arg-type]
|
|
b_block_transfer_thread_cluster_lengths_bk0_n_bk1=(8, 32, 1),
|
|
b_block_transfer_thread_cluster_arrange_order=(0, 2, 1),
|
|
b_block_transfer_src_access_order=(0, 2, 1),
|
|
b_block_transfer_src_vector_dim=1,
|
|
b_block_transfer_src_scalar_per_vector=8,
|
|
b_block_transfer_dst_scalar_per_vector_bk1=2,
|
|
b_block_lds_extra_n=0, # type: ignore[arg-type]
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=2,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
32,
|
|
1,
|
|
8,
|
|
),
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=8,
|
|
block_gemm_pipeline_scheduler="BlockGemmPipelineScheduler::Intrawave",
|
|
block_gemm_pipeline_version="BlockGemmPipelineVersion::v3",
|
|
)
|
|
]
|
|
|
|
|
|
@lru_cache(None)
|
|
def gen_ops_library() -> List[CKGemmOperation]:
|
|
"""
|
|
Parse the Universal Gemm instances defined in the composable kernel library folder.
|
|
"""
|
|
ck_library_dir = _ck_library_dir()
|
|
if not ck_library_dir:
|
|
return []
|
|
|
|
grep_result = subprocess.run(
|
|
[
|
|
"grep",
|
|
"-inR",
|
|
"DeviceGemm_Xdl_CShuffleV3",
|
|
_ck_library_dir(),
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
|
|
op_instances = parse_instances(grep_result.stdout.strip().split("\n"))
|
|
|
|
log.debug("ck instances from library: %d", len(op_instances))
|
|
|
|
schedulers = [
|
|
"BlockGemmPipelineScheduler::Intrawave",
|
|
"BlockGemmPipelineScheduler::Interwave",
|
|
]
|
|
gemm_specs = [
|
|
"GemmSpecialization::Default",
|
|
"GemmSpecialization::MPadding",
|
|
"GemmSpecialization::NPadding",
|
|
"GemmSpecialization::KPadding",
|
|
"GemmSpecialization::MNPadding",
|
|
"GemmSpecialization::MKPadding",
|
|
"GemmSpecialization::NKPadding",
|
|
"GemmSpecialization::MNKPadding",
|
|
]
|
|
|
|
# substitute templated args by looping through their domains
|
|
substitute_instances = []
|
|
for instance in op_instances:
|
|
sub_scheduler = instance.block_gemm_pipeline_scheduler == "BlkGemmPipeSched"
|
|
sub_spec = instance.gemm_specialization == "GemmSpec"
|
|
schedulers_range = (
|
|
schedulers if sub_scheduler else [instance.block_gemm_pipeline_scheduler]
|
|
)
|
|
spec_range = gemm_specs if sub_spec else [instance.gemm_specialization]
|
|
for scheduler in schedulers_range:
|
|
for spec in spec_range:
|
|
substitute_instances.append(
|
|
replace(
|
|
instance,
|
|
block_gemm_pipeline_scheduler=scheduler,
|
|
gemm_specialization=spec,
|
|
)
|
|
)
|
|
|
|
return substitute_instances
|
|
|
|
|
|
@lru_cache(None)
|
|
def gen_ops_preselected() -> List[CKGemmOperation]:
|
|
"""
|
|
Manually selected (through benchmarking) F16/F16/F16 Row/Col/Row instances
|
|
"""
|
|
ck_gemm_f16_rcr = partial(
|
|
CKGemmOperation,
|
|
a_layout="Row",
|
|
b_layout="Col",
|
|
c_layout="Row",
|
|
a_element_dtype="F16",
|
|
b_element_dtype="F16",
|
|
c_element_dtype="F16",
|
|
acc_dtype="F32",
|
|
c_shuffle_dtype="F16",
|
|
a_elementwise_op="PassThrough",
|
|
b_elementwise_op="PassThrough",
|
|
c_elementwise_op="PassThrough",
|
|
k_per_block=64,
|
|
a_k1=8,
|
|
b_k1=8,
|
|
a_block_transfer_thread_cluster_arrange_order=(1, 0, 2),
|
|
a_block_transfer_src_access_order=(1, 0, 2),
|
|
a_block_transfer_src_vector_dim=2,
|
|
a_block_transfer_src_scalar_per_vector=8,
|
|
a_block_transfer_dst_scalar_per_vector_ak1=8,
|
|
a_block_lds_extra_m=0,
|
|
b_block_transfer_thread_cluster_arrange_order=(1, 0, 2),
|
|
b_block_transfer_src_access_order=(1, 0, 2),
|
|
b_block_transfer_src_vector_dim=2,
|
|
b_block_transfer_src_scalar_per_vector=8,
|
|
b_block_transfer_dst_scalar_per_vector_bk1=8,
|
|
b_block_lds_extra_n=0,
|
|
a_compute_dtype="F16",
|
|
b_compute_dtype="F16",
|
|
)
|
|
ck_gemm_f16_rcr_compute_friendly = partial(
|
|
ck_gemm_f16_rcr,
|
|
block_size=256,
|
|
a_block_transfer_thread_cluster_lengths_ak0_m_ak1=(8, 32, 1),
|
|
b_block_transfer_thread_cluster_lengths_bk0_n_bk1=(8, 32, 1),
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
32,
|
|
1,
|
|
8,
|
|
),
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=8,
|
|
)
|
|
ck_gemm_f16_rcr_memory_friendly = partial(
|
|
ck_gemm_f16_rcr,
|
|
block_size=128,
|
|
a_block_transfer_thread_cluster_lengths_ak0_m_ak1=(8, 16, 1),
|
|
b_block_transfer_thread_cluster_lengths_bk0_n_bk1=(8, 16, 1),
|
|
block_gemm_pipeline_scheduler="BlockGemmPipelineScheduler::Interwave",
|
|
block_gemm_pipeline_version="BlockGemmPipelineVersion::v2",
|
|
)
|
|
ck_gemm_f16_rcr_latency_friendly = partial(
|
|
ck_gemm_f16_rcr,
|
|
gemm_specialization="GemmSpecialization::Default",
|
|
block_size=128,
|
|
m_per_xdl=16,
|
|
n_per_xdl=16,
|
|
m_xdl_per_wave=1,
|
|
n_xdl_per_wave=1,
|
|
a_block_transfer_thread_cluster_lengths_ak0_m_ak1=(8, 16, 1),
|
|
b_block_transfer_thread_cluster_lengths_bk0_n_bk1=(8, 16, 1),
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=4,
|
|
block_gemm_pipeline_scheduler="BlockGemmPipelineScheduler::Intrawave",
|
|
block_gemm_pipeline_version="BlockGemmPipelineVersion::v1",
|
|
)
|
|
return [
|
|
ck_gemm_f16_rcr_compute_friendly(
|
|
gemm_specialization="GemmSpecialization::MNKPadding",
|
|
m_per_block=224,
|
|
n_per_block=256,
|
|
m_per_xdl=16,
|
|
n_per_xdl=16,
|
|
m_xdl_per_wave=7,
|
|
n_xdl_per_wave=8,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=2,
|
|
block_gemm_pipeline_scheduler="BlockGemmPipelineScheduler::Intrawave",
|
|
block_gemm_pipeline_version="BlockGemmPipelineVersion::v3",
|
|
),
|
|
ck_gemm_f16_rcr_compute_friendly(
|
|
gemm_specialization="GemmSpecialization::MNKPadding",
|
|
m_per_block=128,
|
|
n_per_block=128,
|
|
m_per_xdl=32,
|
|
n_per_xdl=32,
|
|
m_xdl_per_wave=2,
|
|
n_xdl_per_wave=2,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
block_gemm_pipeline_scheduler="BlockGemmPipelineScheduler::Intrawave",
|
|
block_gemm_pipeline_version="BlockGemmPipelineVersion::v3",
|
|
),
|
|
ck_gemm_f16_rcr_compute_friendly(
|
|
gemm_specialization="GemmSpecialization::MNKPadding",
|
|
m_per_block=128,
|
|
n_per_block=128,
|
|
m_per_xdl=32,
|
|
n_per_xdl=32,
|
|
m_xdl_per_wave=2,
|
|
n_xdl_per_wave=2,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
block_gemm_pipeline_scheduler="BlockGemmPipelineScheduler::Intrawave",
|
|
block_gemm_pipeline_version="BlockGemmPipelineVersion::v4",
|
|
),
|
|
ck_gemm_f16_rcr_compute_friendly(
|
|
gemm_specialization="GemmSpecialization::MNKPadding",
|
|
m_per_block=128,
|
|
n_per_block=128,
|
|
m_per_xdl=32,
|
|
n_per_xdl=32,
|
|
m_xdl_per_wave=2,
|
|
n_xdl_per_wave=2,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
block_gemm_pipeline_scheduler="BlockGemmPipelineScheduler::Intrawave",
|
|
block_gemm_pipeline_version="BlockGemmPipelineVersion::v5",
|
|
),
|
|
ck_gemm_f16_rcr_compute_friendly(
|
|
gemm_specialization="GemmSpecialization::Default",
|
|
m_per_block=128,
|
|
n_per_block=128,
|
|
m_per_xdl=32,
|
|
n_per_xdl=32,
|
|
m_xdl_per_wave=2,
|
|
n_xdl_per_wave=2,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
block_gemm_pipeline_scheduler="BlockGemmPipelineScheduler::Intrawave",
|
|
block_gemm_pipeline_version="BlockGemmPipelineVersion::v3",
|
|
),
|
|
ck_gemm_f16_rcr_compute_friendly(
|
|
gemm_specialization="GemmSpecialization::Default",
|
|
m_per_block=128,
|
|
n_per_block=128,
|
|
m_per_xdl=32,
|
|
n_per_xdl=32,
|
|
m_xdl_per_wave=2,
|
|
n_xdl_per_wave=2,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
block_gemm_pipeline_scheduler="BlockGemmPipelineScheduler::Intrawave",
|
|
block_gemm_pipeline_version="BlockGemmPipelineVersion::v4",
|
|
),
|
|
ck_gemm_f16_rcr_compute_friendly(
|
|
gemm_specialization="GemmSpecialization::Default",
|
|
m_per_block=128,
|
|
n_per_block=128,
|
|
m_per_xdl=32,
|
|
n_per_xdl=32,
|
|
m_xdl_per_wave=2,
|
|
n_xdl_per_wave=2,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
block_gemm_pipeline_scheduler="BlockGemmPipelineScheduler::Intrawave",
|
|
block_gemm_pipeline_version="BlockGemmPipelineVersion::v5",
|
|
),
|
|
ck_gemm_f16_rcr_memory_friendly(
|
|
gemm_specialization="GemmSpecialization::Default",
|
|
m_per_block=16,
|
|
n_per_block=32,
|
|
m_per_xdl=16,
|
|
n_per_xdl=16,
|
|
m_xdl_per_wave=1,
|
|
n_xdl_per_wave=1,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
16,
|
|
1,
|
|
8,
|
|
),
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=4,
|
|
),
|
|
ck_gemm_f16_rcr_memory_friendly(
|
|
gemm_specialization="GemmSpecialization::MNKPadding",
|
|
m_per_block=16,
|
|
n_per_block=32,
|
|
m_per_xdl=16,
|
|
n_per_xdl=16,
|
|
m_xdl_per_wave=1,
|
|
n_xdl_per_wave=1,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
16,
|
|
1,
|
|
8,
|
|
),
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=4,
|
|
),
|
|
ck_gemm_f16_rcr_memory_friendly(
|
|
gemm_specialization="GemmSpecialization::MNKPadding",
|
|
m_per_block=16,
|
|
n_per_block=64,
|
|
m_per_xdl=16,
|
|
n_per_xdl=16,
|
|
m_xdl_per_wave=1,
|
|
n_xdl_per_wave=2,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=2,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
16,
|
|
1,
|
|
8,
|
|
),
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=8,
|
|
),
|
|
ck_gemm_f16_rcr_memory_friendly(
|
|
gemm_specialization="GemmSpecialization::MNKPadding",
|
|
m_per_block=32,
|
|
n_per_block=64,
|
|
m_per_xdl=32,
|
|
n_per_xdl=32,
|
|
m_xdl_per_wave=1,
|
|
n_xdl_per_wave=1,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
16,
|
|
1,
|
|
8,
|
|
),
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=8,
|
|
),
|
|
ck_gemm_f16_rcr_memory_friendly(
|
|
gemm_specialization="GemmSpecialization::MNKPadding",
|
|
m_per_block=32,
|
|
n_per_block=128,
|
|
m_per_xdl=32,
|
|
n_per_xdl=32,
|
|
m_xdl_per_wave=1,
|
|
n_xdl_per_wave=2,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
16,
|
|
1,
|
|
8,
|
|
),
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=8,
|
|
),
|
|
ck_gemm_f16_rcr_memory_friendly(
|
|
gemm_specialization="GemmSpecialization::Default",
|
|
m_per_block=32,
|
|
n_per_block=16,
|
|
m_per_xdl=16,
|
|
n_per_xdl=16,
|
|
m_xdl_per_wave=1,
|
|
n_xdl_per_wave=1,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
32,
|
|
1,
|
|
4,
|
|
),
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=4,
|
|
),
|
|
ck_gemm_f16_rcr_memory_friendly(
|
|
gemm_specialization="GemmSpecialization::MNKPadding",
|
|
m_per_block=32,
|
|
n_per_block=16,
|
|
m_per_xdl=16,
|
|
n_per_xdl=16,
|
|
m_xdl_per_wave=1,
|
|
n_xdl_per_wave=1,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
32,
|
|
1,
|
|
4,
|
|
),
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=4,
|
|
),
|
|
ck_gemm_f16_rcr_memory_friendly(
|
|
gemm_specialization="GemmSpecialization::MNKPadding",
|
|
m_per_block=64,
|
|
n_per_block=16,
|
|
m_per_xdl=16,
|
|
n_per_xdl=16,
|
|
m_xdl_per_wave=2,
|
|
n_xdl_per_wave=1,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=2,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
64,
|
|
1,
|
|
2,
|
|
),
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=8,
|
|
),
|
|
ck_gemm_f16_rcr_memory_friendly(
|
|
gemm_specialization="GemmSpecialization::MNKPadding",
|
|
m_per_block=64,
|
|
n_per_block=32,
|
|
m_per_xdl=32,
|
|
n_per_xdl=32,
|
|
m_xdl_per_wave=1,
|
|
n_xdl_per_wave=1,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
32,
|
|
1,
|
|
4,
|
|
),
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=8,
|
|
),
|
|
ck_gemm_f16_rcr_memory_friendly(
|
|
gemm_specialization="GemmSpecialization::MNKPadding",
|
|
m_per_block=128,
|
|
n_per_block=32,
|
|
m_per_xdl=32,
|
|
n_per_xdl=32,
|
|
m_xdl_per_wave=2,
|
|
n_xdl_per_wave=1,
|
|
c_shuffle_m_xdl_per_wave_per_shuffle=2,
|
|
c_shuffle_n_xdl_per_wave_per_shuffle=1,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
32,
|
|
1,
|
|
4,
|
|
),
|
|
c_shuffle_block_transfer_scalar_per_vector_n_per_block=8,
|
|
),
|
|
ck_gemm_f16_rcr_latency_friendly(
|
|
m_per_block=16,
|
|
n_per_block=32,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
16,
|
|
1,
|
|
8,
|
|
),
|
|
),
|
|
ck_gemm_f16_rcr_latency_friendly(
|
|
m_per_block=32,
|
|
n_per_block=16,
|
|
c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(
|
|
1,
|
|
32,
|
|
1,
|
|
4,
|
|
),
|
|
),
|
|
]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print(gen_ops_library())
|