[tuner]: address comments

Signed-off-by: Bangtian Liu <[email protected]>
nod-ai · Dec 10, 2024 · 6db201a · 6db201a
1 parent a0c4e2b
commit 6db201a
Show file tree

Hide file tree

Showing 6 changed files with 59 additions and 32 deletions.
diff --git a/tuner/tuner/candidate_gen.py b/tuner/tuner/candidate_gen.py
@@ -52,7 +52,7 @@ def apply_configuration(
     workgroup_sizes = lowering_config.workgroup_tile_sizes
     reduction_sizes = lowering_config.reduction_tile_sizes
     gpu_pipeline_options = configuration.translation_info.configuration[
-        "gpu_pipeline_options"
+        GPU_PIPELINE_OPTIONS
     ]
     tune_logger.info(f"Applying: {configuration}")
     expr0 = re.compile(

diff --git a/tuner/tuner/candidate_gen_test.py b/tuner/tuner/candidate_gen_test.py
@@ -58,10 +58,12 @@ def test_apply_params_mmt(tuner_ctx: common.TunerContext) -> None:
         subgroup_n_count=16,
     )
     pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
-        iree_codegen.DispatchLoweringPassPipeline.None_
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorize
     )
     pipeline_option = iree_gpu.PipelineOptionsAttr.get(prefetch_shared_memory=True)
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [16, 16, 1], 16, pipeline_option_dict
     )
@@ -124,14 +126,16 @@ def test_apply_params_conv(tuner_ctx: common.TunerContext) -> None:
         subgroup_n_count=4,
     )
     pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
-        iree_codegen.DispatchLoweringPassPipeline.None_
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorize
     )
     pipeline_option = iree_gpu.PipelineOptionsAttr.get(
         reorder_workgroups_strategy=iree_gpu.ReorderWorkgroupsStrategyAttr.get(
             iree_gpu.ReorderWorkgroupsStrategy.Transpose
         )
     )
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [256, 1, 1], 64, pipeline_option_dict
     )
@@ -203,10 +207,12 @@ def test_apply_params_contract(tuner_ctx: common.TunerContext) -> None:
         subgroup_n_count=4,
     )
     pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
-        iree_codegen.DispatchLoweringPassPipeline.None_
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorize
     )
     pipeline_option = iree_gpu.PipelineOptionsAttr.get()
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [256, 1, 1], 64, pipeline_option_dict
     )
@@ -264,10 +270,12 @@ def test_apply_params_batch_matmul(tuner_ctx: common.TunerContext) -> None:
         subgroup_n_count=2,
     )
     pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
-        iree_codegen.DispatchLoweringPassPipeline.None_
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorize
     )
     pipeline_option = iree_gpu.PipelineOptionsAttr.get()
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [128, 2, 1], 64, pipeline_option_dict
     )
@@ -328,10 +336,12 @@ def test_apply_params_batch_mmt_float(tuner_ctx: common.TunerContext) -> None:
         subgroup_n_count=2,
     )
     pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
-        iree_codegen.DispatchLoweringPassPipeline.None_
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorize
     )
     pipeline_option = iree_gpu.PipelineOptionsAttr.get()
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [128, 2, 1], 64, pipeline_option_dict
     )
@@ -390,10 +400,12 @@ def test_apply_params_batch_mmt_int(tuner_ctx: common.TunerContext) -> None:
         subgroup_n_count=2,
     )
     pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
-        iree_codegen.DispatchLoweringPassPipeline.None_
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorize
     )
     pipeline_option = iree_gpu.PipelineOptionsAttr.get()
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [128, 2, 1], 64, pipeline_option_dict
     )
@@ -476,10 +488,12 @@ def test_apply_params_broadcast_rhs_mmt(tuner_ctx: common.TunerContext) -> None:
         subgroup_n_count=2,
     )
     pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
-        iree_codegen.DispatchLoweringPassPipeline.None_
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorize
     )
     pipeline_option = iree_gpu.PipelineOptionsAttr.get()
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [128, 2, 1], 64, pipeline_option_dict
     )

diff --git a/tuner/tuner/common.py b/tuner/tuner/common.py
@@ -118,6 +118,10 @@ class Configuration:
     waves_per_eu: int
 
 
+# The key name for GPUPipelineOptionsAttr in the translation info config dictionary.
+GPU_PIPELINE_OPTIONS = "gpu_pipeline_options"
+
+
 def get_lowering_config(
     tuner_ctx: TunerContext,
     **kwargs: Any,
@@ -159,7 +163,7 @@ def get_lowering_config(
 def get_pipeline_config(configuration: Configuration) -> str:
     extra_config = ""
     pipeline_options = configuration.translation_info.configuration[
-        "gpu_pipeline_options"
+        GPU_PIPELINE_OPTIONS
     ]
     if pipeline_options != iree_gpu.PipelineOptionsAttr.get():
         extra_config += f", gpu_pipeline_options = {pipeline_options}"

diff --git a/tuner/tuner/common_test.py b/tuner/tuner/common_test.py
@@ -86,10 +86,12 @@ def test_get_pipeline_config(tuner_ctx: common.TunerContext) -> None:
         subgroup_n_count=1,
     )
     pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
-        iree_codegen.DispatchLoweringPassPipeline.None_
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorize
     )
     pipeline_option = iree_gpu.PipelineOptionsAttr.get()
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [16, 16, 1], 32, pipeline_option_dict
     )
@@ -106,7 +108,9 @@ def test_get_pipeline_config(tuner_ctx: common.TunerContext) -> None:
     assert config2_str == ', llvm_func_attrs = {"amdgpu-waves-per-eu" = "4"}'
 
     pipeline_option = iree_gpu.PipelineOptionsAttr.get(prefetch_shared_memory=True)
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [16, 16, 1], 32, pipeline_option_dict
     )
@@ -222,10 +226,12 @@ def test_get_lowering_config(tuner_ctx: common.TunerContext) -> None:
     )
 
     pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
-        iree_codegen.DispatchLoweringPassPipeline.None_
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorize
     )
     pipeline_option = iree_gpu.PipelineOptionsAttr.get()
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [16, 16, 1], 32, pipeline_option_dict
     )

diff --git a/tuner/tuner/dispatch_constraints.py b/tuner/tuner/dispatch_constraints.py
@@ -248,11 +248,8 @@ def generate_solutions(
             iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorDistribute
         )
         pipeline_option = iree_gpu.PipelineOptionsAttr.get()
-        # TODO: expose the function getDictKeyName() from GPUPipelineOptionsAttr to
-        # get the key name 'gpu_pipeline_options'in the translation info config dictionary
-        pipeline_option_dict = ir.DictAttr.get(
-            {"gpu_pipeline_options": pipeline_option}
-        )
+        pipeline_option_dict = ir.DictAttr.get({GPU_PIPELINE_OPTIONS: pipeline_option})
+
         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
@@ -52,10 +52,12 @@ def test_get_mmt_tile_sizes(tuner_ctx: common.TunerContext) -> None:
         subgroup_n_count=4,
     )
     pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
-        iree_codegen.DispatchLoweringPassPipeline.None_
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorize
     )
     pipeline_option = iree_gpu.PipelineOptionsAttr.get()
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [], 0, pipeline_option_dict
     )
@@ -81,10 +83,12 @@ def test_get_conv_tile_sizes(tuner_ctx: common.TunerContext) -> None:
         subgroup_n_count=4,
     )
     pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
-        iree_codegen.DispatchLoweringPassPipeline.None_
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorize
     )
     pipeline_option = iree_gpu.PipelineOptionsAttr.get()
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [256, 1, 1], 64, pipeline_option_dict
     )
@@ -109,10 +113,12 @@ def test_get_contract_tile_sizes(tuner_ctx: common.TunerContext) -> None:
         subgroup_n_count=1,
     )
     pipeline_attr = iree_codegen.DispatchLoweringPassPipelineAttr.get(
-        iree_codegen.DispatchLoweringPassPipeline.None_
+        iree_codegen.DispatchLoweringPassPipeline.LLVMGPUVectorize
     )
     pipeline_option = iree_gpu.PipelineOptionsAttr.get()
-    pipeline_option_dict = ir.DictAttr.get({"gpu_pipeline_options": pipeline_option})
+    pipeline_option_dict = ir.DictAttr.get(
+        {common.GPU_PIPELINE_OPTIONS: pipeline_option}
+    )
     translation_info = iree_codegen.TranslationInfoAttr.get(
         pipeline_attr, None, [16, 16, 1], 32, pipeline_option_dict
     )