ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 19:18:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix codegen bug

Browse Source
This commit is contained in:
Yanxing-Shi
2025-05-19 14:03:16 +00:00
parent b3caa67694
commit 5b83f76eb0
3 changed files with 35 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tile_engine/ops/gemm/CMakeLists.txt
View File

@@ -3,6 +3,7 @@
execute_process(
COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/gemm_instance_builder.py
--working_path ${CMAKE_CURRENT_BINARY_DIR}
--config_json ${CMAKE_CURRENT_LIST_DIR}/configs/user_provided_config.json
--list_blobs
RESULT_VARIABLE ret
)
@@ -14,10 +15,11 @@ endif()
file(STRINGS ${CMAKE_CURRENT_BINARY_DIR}/gemm_instance_blobs.txt GEMM_CODEGEN_BLOBS)
add_custom_command(
OUTPUT ${GEMM_CODEGEN_BLOBS}
COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/gemm_instance_builder.py
--working_path ${CMAKE_CURRENT_BINARY_DIR}
--gen_blobs
OUTPUT ${GEMM_CODEGEN_BLOBS}
COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/gemm_instance_builder.py
--working_path ${CMAKE_CURRENT_BINARY_DIR}
--config_json ${CMAKE_CURRENT_LIST_DIR}/configs/user_provided_config.json
--gen_blobs
)
set(EXECUTABLE_GEMM_INSTANCE "tile_engine_gemm")

18
tile_engine/ops/gemm/configs/default_config.json
View File

@@ -33,21 +33,21 @@
},
"tile_config": {
"tile_m": {
"max": 512,
"min": 64,
"step": 8,
"max": 256,
"min": 128,
"step": 64,
"exclude": []
},
"tile_n": {
"max": 512,
"min": 64,
"step": 8,
"max": 256,
"min": 128,
"step": 64,
"exclude": []
},
"tile_k": {
"max": 512,
"min": 64,
"step": 8,
"max": 256,
"min": 128,
"step": 64,
"exclude": []
},
"warp_m": {

35
tile_engine/ops/gemm/gemm_instance_builder.py
View File

@@ -438,28 +438,33 @@ struct GemmDispatcher {
get_tile_value(self.config.tile_config.warp_tile_n),
get_tile_value(self.config.tile_config.warp_tile_k),
))
for trait in self.all_trait_names:
tile_valid_params = list(filter(lambda t: self.is_tile_valid(t, trait), tile_params))
content += f""" kernel_map["{trait}"] = {{"""
for tile in tile_params:
if self.is_tile_valid(tile, trait):
content += f"""[&](ck_tile::GemmHostArgs& args, const ck_tile::stream_config& stream) {{ """
content += f""" if(structured_sparsity){{ // SMFMA"""
sparse = self.config.problem.datatype_map['matrix_a'] == 'fp16' and \
self.config.problem.datatype_map['matrix_b'] == 'fp16' and \
self.config.problem.datatype_map['matrix_c'] == 'fp16' and \
((tile[6] == 32 and tile[7] == 32 and tile[8] == 16) or
(tile[6] == 16 and tile[7] == 16 and tile[8] == 32))
content += f"""
return run_kernel<{trait}::GemmKernel<{tile[0]}, {tile[1]}, {tile[2]}, {tile[3]}, {tile[4]}, {tile[5]}, {tile[6]}, {tile[7]}, {tile[8]}, {BOOL_MAP(sparse)}>>(args, stream);"""
content += f"""
for i, tile in enumerate(tile_valid_params):
content += f"""[&](ck_tile::GemmHostArgs& args, const ck_tile::stream_config& stream) {{ """
content += f""" if(structured_sparsity){{ // SMFMA"""
sparse = self.config.problem.datatype_map['matrix_a'] == 'fp16' and \
self.config.problem.datatype_map['matrix_b'] == 'fp16' and \
self.config.problem.datatype_map['matrix_c'] == 'fp16' and \
((tile[6] == 32 and tile[7] == 32 and tile[8] == 16) or
(tile[6] == 16 and tile[7] == 16 and tile[8] == 32))
content += f"""
return run_kernel<{trait}::GemmKernel<{tile[0]}, {tile[1]}, {tile[2]}, {tile[3]}, {tile[4]}, {tile[5]}, {tile[6]}, {tile[7]}, {tile[8]}, {BOOL_MAP(sparse)}>>(args, stream);"""
content += f"""
}} else {{"""
content += f"""
return run_kernel<{trait}::GemmKernel<{tile[0]}, {tile[1]}, {tile[2]}, {tile[3]}, {tile[4]}, {tile[5]}, {tile[6]}, {tile[7]}, {tile[8]}, {BOOL_MAP(False)}>>(args, stream);"""
content += f"""
content += f"""
return run_kernel<{trait}::GemmKernel<{tile[0]}, {tile[1]}, {tile[2]}, {tile[3]}, {tile[4]}, {tile[5]}, {tile[6]}, {tile[7]}, {tile[8]}, {BOOL_MAP(False)}>>(args, stream);"""
content += f"""
}} """
if i == len(tile_valid_params)-1:
content += f"""
}} """
else:
content += f"""
}}, """
content += f"""
}};\n """