From 6db201af2246012fe7785e81906990c58abb8740 Mon Sep 17 00:00:00 2001 From: Bangtian Liu Date: Tue, 10 Dec 2024 11:23:33 -0600 Subject: [PATCH] [tuner]: address comments Signed-off-by: Bangtian Liu --- tuner/tuner/candidate_gen.py | 2 +- tuner/tuner/candidate_gen_test.py | 42 +++++++++++++++++++---------- tuner/tuner/common.py | 6 ++++- tuner/tuner/common_test.py | 16 +++++++---- tuner/tuner/dispatch_constraints.py | 7 ++--- tuner/tuner/dispatch_parser_test.py | 18 ++++++++----- 6 files changed, 59 insertions(+), 32 deletions(-) diff --git a/tuner/tuner/candidate_gen.py b/tuner/tuner/candidate_gen.py index 775d9d2bb..e81dcdbd9 100644 --- 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 index 5d9807f82..75fc46714 100644 --- 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 index 7434eb5ac..78bfa9eec 100644 --- 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 index 97c7e5485..3bdc93c4e 100644 --- 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 index 9849efb24..33df4a4c9 100644 --- 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 index e86ee0041..ccd49b351 100644 --- 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 )