diff --git a/tile_engine/ops/gemm/gemm_instance_builder.py b/tile_engine/ops/gemm/gemm_instance_builder.py index 23675e73e8..c43797f3e0 100755 --- a/tile_engine/ops/gemm/gemm_instance_builder.py +++ b/tile_engine/ops/gemm/gemm_instance_builder.py @@ -72,8 +72,16 @@ class GemmCodeGenerator: for trait, tile_valid_params in self.valid_trait_tile_combinations.items(): for tile in tile_valid_params: for tile_m, tile_n, tile_k, warp_m, warp_n, warp_k, warp_tile_m, warp_tile_n, warp_tile_k in tile: + 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 \ + ((warp_tile_m == 32 and warp_tile_n == 32 and warp_tile_k == 16) or + (warp_tile_m == 16 and warp_tile_n == 16 and warp_tile_k == 32)) + if sparse: + f.write(str( + w_p / f"gemm_{trait}_{tile_m}x{tile_n}x{tile_k}_{warp_m}x{warp_n}x{warp_k}_{warp_tile_m}x{warp_tile_n}x{warp_tile_k}_true.cpp") + "\n") f.write(str( - w_p / f"gemm_{trait}_{tile_m}x{tile_n}x{tile_k}_{warp_m}x{warp_n}x{warp_k}_{warp_tile_m}x{warp_tile_n}x{warp_tile_k}.cpp") + "\n") + w_p / f"gemm_{trait}_{tile_m}x{tile_n}x{tile_k}_{warp_m}x{warp_n}x{warp_k}_{warp_tile_m}x{warp_tile_n}x{warp_tile_k}_false.cpp") + "\n") def _generate_all_traits(self): """Generate all possible kernel traits names.""" @@ -326,7 +334,7 @@ struct GemmKernel {{ # Parameter validity check invalid_params = [] - if (warp_m, warp_n, warp_k) not in [(1,4,1), (2,2,1), (4,1,1)]: + if (warp_m, warp_n, warp_k) not in [(1, 4, 1), (2, 2, 1), (4, 1, 1)]: invalid_params.append( f"warp_m({warp_m}) * warp_n({warp_n}) * warp_k({warp_k})") if (warp_m * warp_tile_m) == 0: @@ -466,15 +474,17 @@ struct GemmKernel {{ ((warp_tile_m == 32 and warp_tile_n == 32 and warp_tile_k == 16) or (warp_tile_m == 16 and warp_tile_n == 16 and warp_tile_k == 32)) if sparse: - content += f""" + sparse_content = content + f""" template struct {trait}::GemmKernel<{tile_m}, {tile_n}, {tile_k}, {warp_m}, {warp_n}, {warp_k}, {warp_tile_m}, {warp_tile_n}, {warp_tile_k}, true>; """ - content += f""" + (self.output_dir / + f"gemm_{trait}_{tile_m}x{tile_n}x{tile_k}_{warp_m}x{warp_n}x{warp_k}_{warp_tile_m}x{warp_tile_n}x{warp_tile_k}_true.cpp").write_text(sparse_content) + + no_sparse_content = content + f""" template struct {trait}::GemmKernel<{tile_m}, {tile_n}, {tile_k}, {warp_m}, {warp_n}, {warp_k}, {warp_tile_m}, {warp_tile_n}, {warp_tile_k}, false>; """ - (self.output_dir / - f"gemm_{trait}_{tile_m}x{tile_n}x{tile_k}_{warp_m}x{warp_n}x{warp_k}_{warp_tile_m}x{warp_tile_n}x{warp_tile_k}.cpp").write_text(content) + f"gemm_{trait}_{tile_m}x{tile_n}x{tile_k}_{warp_m}x{warp_n}x{warp_k}_{warp_tile_m}x{warp_tile_n}x{warp_tile_k}_false.cpp").write_text(no_sparse_content) def _generate_dispatcher_file(self): """Generate the code block of dispatch mechanism.""" @@ -528,7 +538,8 @@ struct GemmDispatcher { content += f""" kernel_map["{trait}"] = {{""" for _, tile in enumerate(tile_valid_params): for j in range(len(tile)): - tile_m, tile_n, tile_k, warp_m, warp_n, warp_k, warp_tile_m, warp_tile_n, warp_tile_k = tile[j] + tile_m, tile_n, tile_k, warp_m, warp_n, warp_k, warp_tile_m, warp_tile_n, warp_tile_k = tile[ + j] content += f"""[=](ck_tile::GemmHostArgs& args, const ck_tile::stream_config& stream) {{ """ content += f""" if(structured_sparsity){{ // SMFMA""" @@ -545,7 +556,7 @@ struct GemmDispatcher { return run_kernel<{trait}::GemmKernel<{tile_m}, {tile_n}, {tile_k}, {warp_m}, {warp_n}, {warp_k}, {warp_tile_m}, {warp_tile_n}, {warp_tile_k}, {BOOL_MAP(False)}>>(args, stream);""" content += f""" }} """ - + if j == len(tile)-1: content += f""" }} """