[tuner]: add helper function

Signed-off-by: Bangtian Liu <[email protected]>
nod-ai · Dec 11, 2024 · bc88d23 · bc88d23
1 parent 202d04c
commit bc88d23
Show file tree

Hide file tree

Showing 6 changed files with 44 additions and 102 deletions.
diff --git a/tuner/tuner/candidate_gen.py b/tuner/tuner/candidate_gen.py
@@ -322,7 +322,7 @@ def get_transform_function_batch_mmt(
 transform.iree.match.cast_compatible_type %rhs = tensor<{problem_size.rhs_type}> : !transform.any_value
 %config = transform.param.constant #iree_codegen.compilation_info<
     lowering_config = {configuration.lowering_config},
-    translation_info ={configuration.translation_info}
+    translation_info = {configuration.translation_info}
     > -> !transform.any_param
 transform.yield %generic, %config : !transform.any_op, !transform.any_param
 }}

diff --git a/tuner/tuner/candidate_gen_test.py b/tuner/tuner/candidate_gen_test.py
@@ -61,13 +61,7 @@ def test_apply_params_mmt(tuner_ctx: common.TunerContext) -> None:
         iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
     )
     pipeline_options = iree_gpu.PipelineOptionsAttr.get(prefetch_shared_memory=True)
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("8")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            common.LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 8)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [16, 16, 1], 16, config_dict
     )
@@ -136,13 +130,7 @@ def test_apply_params_conv(tuner_ctx: common.TunerContext) -> None:
             iree_gpu.ReorderWorkgroupsStrategy.Transpose
         )
     )
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("2")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            common.LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 2)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [256, 1, 1], 64, config_dict
     )
@@ -216,13 +204,7 @@ def test_apply_params_contract(tuner_ctx: common.TunerContext) -> None:
         iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
     )
     pipeline_options = iree_gpu.PipelineOptionsAttr.get()
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("2")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            common.LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 2)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [256, 1, 1], 64, config_dict
     )
@@ -282,13 +264,7 @@ def test_apply_params_batch_matmul(tuner_ctx: common.TunerContext) -> None:
         iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
     )
     pipeline_options = iree_gpu.PipelineOptionsAttr.get()
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("2")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            common.LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 2)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [128, 2, 1], 64, config_dict
     )
@@ -351,13 +327,7 @@ def test_apply_params_batch_mmt_float(tuner_ctx: common.TunerContext) -> None:
         iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
     )
     pipeline_options = iree_gpu.PipelineOptionsAttr.get()
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("2")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            common.LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 2)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [128, 2, 1], 64, config_dict
     )
@@ -418,13 +388,7 @@ def test_apply_params_batch_mmt_int(tuner_ctx: common.TunerContext) -> None:
         iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
     )
     pipeline_options = iree_gpu.PipelineOptionsAttr.get()
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("4")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            common.LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 4)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [128, 2, 1], 64, config_dict
     )
@@ -509,13 +473,7 @@ def test_apply_params_broadcast_rhs_mmt(tuner_ctx: common.TunerContext) -> None:
         iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
     )
     pipeline_options = iree_gpu.PipelineOptionsAttr.get()
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("4")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            common.LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 4)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [128, 2, 1], 64, config_dict
     )

diff --git a/tuner/tuner/common.py b/tuner/tuner/common.py
@@ -163,6 +163,31 @@ def get_lowering_config(
     return iree_gpu.LoweringConfigAttr.get(lowering_config_attrs)
 
 
+def get_translation_info_config(
+    pipeline_options: iree_gpu.PipelineOptionsAttr, waves_per_eu: int | str
+) -> ir.DictAttr:
+    if isinstance(waves_per_eu, int):
+        waves_per_eu = str(waves_per_eu)
+    elif not isinstance(waves_per_eu, str):
+        assert (
+            False
+        ), f"waves_per_eu must be an int or str, but got {type(waves_per_eu).__name__}"
+
+    # Create the waves_per_eu dictionary attribute.
+    waves_per_eu_dict = ir.DictAttr.get(
+        {WAVES_PER_EU_KEY: ir.StringAttr.get(waves_per_eu)}
+    )
+
+    config_dict = ir.DictAttr.get(
+        {
+            GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
+            LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
+        }
+    )
+
+    return config_dict
+
+
 def read_input_mlir(filename: str) -> list[str]:
     with open(filename, "r") as f:
         return f.readlines()

diff --git a/tuner/tuner/common_test.py b/tuner/tuner/common_test.py
@@ -89,31 +89,21 @@ def test_get_pipeline_config(tuner_ctx: common.TunerContext) -> None:
         iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
     )
     pipeline_options = iree_gpu.PipelineOptionsAttr.get()
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("2")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            common.LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 2)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [16, 16, 1], 32, config_dict
     )
     config = common.Configuration(
         translation_info=translation_info,
         lowering_config=lowering_config,
     )
-    config1_str: str = str(config.translation_info.configuration["llvm_func_attrs"])
+    config1_str: str = str(
+        config.translation_info.configuration[common.LLVM_FUNC_ATTRS_KEY]
+    )
     assert config1_str == '{"amdgpu-waves-per-eu" = "2"}'
 
     pipeline_options = iree_gpu.PipelineOptionsAttr.get(prefetch_shared_memory=True)
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("4")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            "llvm_func_attrs": waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 4)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [16, 16, 1], 32, config_dict
     )
@@ -231,13 +221,7 @@ def test_get_lowering_config(tuner_ctx: common.TunerContext) -> None:
         iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
     )
     pipeline_options = iree_gpu.PipelineOptionsAttr.get()
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("2")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            "llvm_func_attrs": waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 2)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [16, 16, 1], 32, config_dict
     )

diff --git a/tuner/tuner/dispatch_constraints.py b/tuner/tuner/dispatch_constraints.py
@@ -248,16 +248,9 @@ def generate_solutions(
             iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
         )
         pipeline_options = iree_gpu.PipelineOptionsAttr.get()
-        waves_per_eu_dict = ir.DictAttr.get(
-            {WAVES_PER_EU_KEY: ir.StringAttr.get(str(lookup(waves_per_eu)))}
+        config_dict = get_translation_info_config(
+            pipeline_options, lookup(waves_per_eu)
         )
-        config_dict = ir.DictAttr.get(
-            {
-                GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-                LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
-            }
-        )
-
         translation_info = iree_codegen.TranslationInfoAttr.get(
             pipeline_attr,
             None,

diff --git a/tuner/tuner/dispatch_parser_test.py b/tuner/tuner/dispatch_parser_test.py
@@ -55,13 +55,7 @@ def test_get_mmt_tile_sizes(tuner_ctx: common.TunerContext) -> None:
         iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
     )
     pipeline_options = iree_gpu.PipelineOptionsAttr.get()
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("0")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            common.LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 0)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [], 0, config_dict
     )
@@ -89,13 +83,7 @@ def test_get_conv_tile_sizes(tuner_ctx: common.TunerContext) -> None:
         iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
     )
     pipeline_options = iree_gpu.PipelineOptionsAttr.get()
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("1")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            common.LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 1)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [256, 1, 1], 64, config_dict
     )
@@ -122,13 +110,7 @@ def test_get_contract_tile_sizes(tuner_ctx: common.TunerContext) -> None:
         iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
     )
     pipeline_options = iree_gpu.PipelineOptionsAttr.get()
-    waves_per_eu_dict = ir.DictAttr.get({"amdgpu-waves-per-eu": ir.StringAttr.get("2")})
-    config_dict = ir.DictAttr.get(
-        {
-            common.GPU_PIPELINE_OPTIONS_KEY: pipeline_options,
-            common.LLVM_FUNC_ATTRS_KEY: waves_per_eu_dict,
-        }
-    )
+    config_dict = common.get_translation_info_config(pipeline_options, 2)
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [16, 16, 1], 32, config_dict
     )