diff --git a/tuner/tuner/candidate_gen.py b/tuner/tuner/candidate_gen.py index 41a26a927..a3252130e 100644 --- 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 index 63819e599..7f104bcd9 100644 --- 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 index 2ccb254b4..12b40805e 100644 --- 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 index 26e9e57c9..af1e1bf9a 100644 --- a/tuner/tuner/common_test.py +++ b/tuner/tuner/common_test.py @@ -89,13 +89,7 @@ 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 ) @@ -103,17 +97,13 @@ def test_get_pipeline_config(tuner_ctx: common.TunerContext) -> None: 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 index 8086c75e3..8ba202310 100644 --- 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 index 3ed0faf22..a63576808 100644 --- 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 )