| # Copyright 2023 The IREE Authors | |
| # | |
| # Licensed under the Apache License v2.0 with LLVM Exceptions. | |
| # See https://llvm.org/LICENSE.txt for license information. | |
| # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |
| from iree.compiler import ir | |
| # Test upstream dialects import | |
| from iree.compiler.dialects import ( | |
| affine, | |
| amdgpu, | |
| arith, | |
| builtin, | |
| cf, | |
| complex, | |
| func, | |
| gpu, | |
| # TODO: importing linalg pulls yaml dependency, disable for now | |
| # linalg, | |
| llvm, | |
| math, | |
| memref, | |
| nvvm, | |
| pdl, | |
| rocdl, | |
| scf, | |
| shape, | |
| smt, | |
| tensor, | |
| tosa, | |
| transform, | |
| vector, | |
| ) | |
| # Smoke test for vector transforms | |
| from iree.compiler.dialects.transform import vector as vt | |
| from iree.compiler.dialects.transform import loop | |
| # Make sure that our dialects import. | |
| from iree.compiler.dialects import ( | |
| flow, | |
| hal, | |
| stream, | |
| vm, | |
| util, | |
| iree_codegen, | |
| iree_gpu, | |
| iree_tensor_ext, | |
| preprocessing_transform, | |
| ) | |
| def get_index_attr(val: int) -> ir.IntegerAttr: | |
| return ir.IntegerAttr.get(ir.IndexType.get(), val) | |
| def get_index_array_attr(vals: list[int]) -> ir.ArrayAttr: | |
| return ir.ArrayAttr.get([get_index_attr(val) for val in vals]) | |
| def run(fn): | |
| with ir.Context(), ir.Location.unknown(): | |
| module = ir.Module.create() | |
| with ir.InsertionPoint(module.body): | |
| print("\nTEST:", fn.__name__) | |
| fn() | |
| return fn | |
| # ====================================================================== | |
| # IREE Codegen Dialect | |
| # ====================================================================== | |
| @run | |
| def codegen_vmvx_pipeline_attr(): | |
| pipeline_attr = iree_codegen.VMVXPipelineAttr.get() | |
| assert pipeline_attr is not None | |
| assert isinstance(pipeline_attr, iree_codegen.VMVXPipelineAttr) | |
| assert str(pipeline_attr) == "#iree_codegen.vmvx_pipeline" | |
| @run | |
| def codegen_transform_dialect_codegen_pipeline_attr(): | |
| pipeline_attr = iree_codegen.TransformDialectCodegenPipelineAttr.get() | |
| assert pipeline_attr is not None | |
| assert isinstance(pipeline_attr, iree_codegen.TransformDialectCodegenPipelineAttr) | |
| assert str(pipeline_attr) == "#iree_codegen.transform_dialect_codegen" | |
| @run | |
| def codegen_no_pipeline_attr(): | |
| pipeline_attr = iree_codegen.NoPipelineAttr.get() | |
| assert pipeline_attr is not None | |
| assert isinstance(pipeline_attr, iree_codegen.NoPipelineAttr) | |
| assert str(pipeline_attr) == "#iree_codegen.no_pipeline" | |
| @run | |
| def codegen_pass_pipeline_attr(): | |
| pipeline_attr = iree_codegen.PassPipelineAttr.get("func.func(canonicalize)") | |
| assert pipeline_attr is not None | |
| assert isinstance(pipeline_attr, iree_codegen.PassPipelineAttr) | |
| assert pipeline_attr.pipeline == "func.func(canonicalize)" | |
| assert ( | |
| str(pipeline_attr) == '#iree_codegen.pass_pipeline<"func.func(canonicalize)">' | |
| ) | |
| @run | |
| def gpu_pipeline_attr(): | |
| pipeline_attr = iree_gpu.PipelineAttr.get(iree_gpu.LoweringPipeline.TileAndFuse) | |
| assert pipeline_attr is not None | |
| assert pipeline_attr.value == iree_gpu.LoweringPipeline.TileAndFuse | |
| assert pipeline_attr.raw_value == int(iree_gpu.LoweringPipeline.TileAndFuse) | |
| assert "TileAndFuse" in str(pipeline_attr) | |
| @run | |
| def codegen_translation_info_minimal(): | |
| pipeline_attr = iree_codegen.NoPipelineAttr.get() | |
| translation_info = iree_codegen.TranslationInfoAttr.get(pipeline_attr) | |
| assert translation_info is not None | |
| assert ( | |
| str(translation_info) | |
| == "#iree_codegen.translation_info<pipeline = #iree_codegen.no_pipeline>" | |
| ) | |
| assert translation_info.pass_pipeline == pipeline_attr | |
| assert translation_info.codegen_spec is None | |
| assert translation_info.workgroup_size == [] | |
| assert translation_info.subgroup_size == 0 | |
| assert translation_info.configuration is None | |
| @run | |
| def codegen_translation_info_with_sizes(): | |
| pipeline_attr = iree_codegen.PassPipelineAttr.get("canonicalize") | |
| translation_info = iree_codegen.TranslationInfoAttr.get( | |
| pipeline_attr, None, [64, 4, 1], 32 | |
| ) | |
| assert translation_info is not None | |
| assert translation_info.pass_pipeline == pipeline_attr | |
| assert pipeline_attr.pipeline == "canonicalize" | |
| assert translation_info.codegen_spec is None | |
| assert translation_info.workgroup_size == [64, 4, 1] | |
| assert translation_info.subgroup_size == 32 | |
| assert translation_info.configuration is None | |
| @run | |
| def codegen_translation_info_full(): | |
| pipeline_attr = iree_codegen.TransformDialectCodegenPipelineAttr.get() | |
| foo_symbol = ir.SymbolRefAttr.get(["foo"]) | |
| configuration = ir.DictAttr.get({"A": ir.IntegerAttr.get(ir.IndexType.get(), 42)}) | |
| translation_info = iree_codegen.TranslationInfoAttr.get( | |
| pipeline_attr, foo_symbol, [128], 32, configuration | |
| ) | |
| assert translation_info is not None | |
| assert translation_info.pass_pipeline == pipeline_attr | |
| assert translation_info.codegen_spec == foo_symbol | |
| assert translation_info.workgroup_size == [128] | |
| assert translation_info.subgroup_size == 32 | |
| assert translation_info.configuration == configuration | |
| @run | |
| def codegen_translation_info_with_gpu_pipeline(): | |
| """Test TranslationInfoAttr with GPU PipelineAttr.""" | |
| gpu_pipeline_attr = iree_gpu.PipelineAttr.get( | |
| iree_gpu.LoweringPipeline.VectorDistribute | |
| ) | |
| translation_info = iree_codegen.TranslationInfoAttr.get( | |
| gpu_pipeline_attr, None, [64, 4, 1], 32 | |
| ) | |
| assert translation_info is not None | |
| assert translation_info.pass_pipeline == gpu_pipeline_attr | |
| assert translation_info.codegen_spec is None | |
| assert translation_info.workgroup_size == [64, 4, 1] | |
| assert translation_info.subgroup_size == 32 | |
| assert translation_info.configuration is None | |
| assert "#iree_gpu.pipeline<VectorDistribute>" in str(translation_info) | |
| # ====================================================================== | |
| # IREE GPU Dialect | |
| # ====================================================================== | |
| @run | |
| def gpu_pipeline_options_attr(): | |
| reorder_attr = iree_gpu.ReorderWorkgroupsStrategyAttr.get( | |
| iree_gpu.ReorderWorkgroupsStrategy.Transpose | |
| ) | |
| assert reorder_attr.value == iree_gpu.ReorderWorkgroupsStrategy.Transpose | |
| gpu_attr = iree_gpu.PipelineOptionsAttr.get( | |
| 2, | |
| False, | |
| False, | |
| reorder_attr, | |
| ) | |
| assert type(gpu_attr) is iree_gpu.PipelineOptionsAttr | |
| assert gpu_attr.prefetch_num_stages == 2 | |
| assert not gpu_attr.no_reduce_shared_memory_bank_conflicts | |
| assert not gpu_attr.use_igemm_convolution | |
| gpu_attr = iree_gpu.PipelineOptionsAttr.get( | |
| 0, | |
| True, | |
| True, | |
| iree_gpu.ReorderWorkgroupsStrategyAttr.get( | |
| iree_gpu.ReorderWorkgroupsStrategy.Transpose | |
| ), | |
| ) | |
| assert gpu_attr.prefetch_num_stages == 0 | |
| assert gpu_attr.no_reduce_shared_memory_bank_conflicts | |
| assert gpu_attr.use_igemm_convolution | |
| gpu_attr = iree_gpu.PipelineOptionsAttr.get() | |
| assert ( | |
| gpu_attr.prefetch_num_stages is None | |
| and gpu_attr.no_reduce_shared_memory_bank_conflicts is None | |
| and gpu_attr.use_igemm_convolution is None | |
| and gpu_attr.reorder_workgroups_strategy is None | |
| ) | |
| gpu_attr = iree_gpu.PipelineOptionsAttr.get(2) | |
| assert gpu_attr.prefetch_num_stages == 2 | |
| assert ( | |
| gpu_attr.no_reduce_shared_memory_bank_conflicts is None | |
| and gpu_attr.use_igemm_convolution is None | |
| and gpu_attr.reorder_workgroups_strategy is None | |
| ) | |
| gpu_attr = iree_gpu.PipelineOptionsAttr.get(2, False) | |
| assert ( | |
| gpu_attr.use_igemm_convolution is None | |
| and gpu_attr.reorder_workgroups_strategy is None | |
| ) | |
| gpu_attr = iree_gpu.PipelineOptionsAttr.get(2, False, False) | |
| assert gpu_attr.reorder_workgroups_strategy is None | |
| gpu_attr = iree_gpu.PipelineOptionsAttr.get( | |
| no_reduce_shared_memory_bank_conflicts=False | |
| ) | |
| assert ( | |
| gpu_attr.no_reduce_shared_memory_bank_conflicts is not None | |
| and not gpu_attr.no_reduce_shared_memory_bank_conflicts | |
| ) | |
| assert gpu_attr.prefetch_num_stages is None | |
| assert gpu_attr.use_igemm_convolution is None | |
| assert gpu_attr.reorder_workgroups_strategy is None | |
| gpu_attr = iree_gpu.PipelineOptionsAttr.get( | |
| reorder_workgroups_strategy=reorder_attr | |
| ) | |
| assert gpu_attr.reorder_workgroups_strategy is not None | |
| assert ( | |
| gpu_attr.reorder_workgroups_strategy.value | |
| # unfortunately not `is` | |
| == iree_gpu.ReorderWorkgroupsStrategy.Transpose | |
| ) | |
| @run | |
| def mma_intrinsic_attr(): | |
| mma_intrinsic_attr = iree_gpu.MMAIntrinsicAttr.get( | |
| iree_gpu.MMAIntrinsic.MFMA_F32_32x32x8_F16 | |
| ) | |
| assert mma_intrinsic_attr is not None | |
| assert str(mma_intrinsic_attr) == "#iree_gpu<mma_intrinsic MFMA_F32_32x32x8_F16>" | |
| # Fragment: 0 = lrhs, 1 = rhs, 2 = acc. | |
| fragment = 0 | |
| mma_single_subgroup_layout = iree_gpu.get_single_subgroup_layout( | |
| attr=mma_intrinsic_attr, fragment=fragment | |
| ) | |
| assert isinstance(mma_single_subgroup_layout, iree_gpu.GPUMMASingleSubgroupLayout) | |
| assert mma_single_subgroup_layout.outer == [1, 1] | |
| assert mma_single_subgroup_layout.thread == [32, 2] | |
| assert mma_single_subgroup_layout.tstrides == [1, 32] | |
| assert mma_single_subgroup_layout.element == [1, 4] | |
| raw_value = mma_intrinsic_attr.raw_value | |
| assert raw_value == iree_gpu.MMAIntrinsic.MFMA_F32_32x32x8_F16 | |
| value = mma_intrinsic_attr.value | |
| assert str(value) == "MFMA_F32_32x32x8_F16" | |
| assert int(value) == raw_value | |
| mma_attr = iree_gpu.MMAAttr.get(raw_value) | |
| assert mma_attr is not None | |
| f16 = ir.F16Type.get() | |
| f32 = ir.F32Type.get() | |
| a_type, b_type, c_type = mma_attr.abc_element_types | |
| assert a_type == f16 | |
| assert b_type == f16 | |
| assert c_type == f32 | |
| vec_4xf16 = ir.VectorType.get((4,), f16) | |
| vec_16xf32 = ir.VectorType.get((16,), f32) | |
| a_vec_type, b_vec_type, c_vec_type = mma_attr.abc_vector_types | |
| assert a_vec_type == vec_4xf16 | |
| assert b_vec_type == vec_4xf16 | |
| assert c_vec_type == vec_16xf32 | |
| M, N, K = mma_attr.mnk_shape | |
| assert M == 32 | |
| assert N == 32 | |
| assert K == 8 | |
| assert mma_intrinsic_attr.mma == mma_attr | |
| virtual_mma_intrinsics = mma_attr.get_virtual_intrinsics() | |
| assert isinstance(virtual_mma_intrinsics[0], iree_gpu.VirtualMMAIntrinsic) | |
| assert ( | |
| virtual_mma_intrinsics[0] == iree_gpu.VirtualMMAIntrinsic.VMFMA_F32_32x32x16_F16 | |
| ) | |
| mma_attr = iree_gpu.MMAAttr.get(iree_gpu.MMAIntrinsic.MFMA_F32_16x16x4_F32) | |
| virtual_mma_intrinsics = mma_attr.get_virtual_intrinsics() | |
| assert virtual_mma_intrinsics == [] | |
| mma_attr_col_major = iree_gpu.MMAAttr.get( | |
| iree_gpu.MMAIntrinsic.MFMA_F32_32x32x8_F16, col_major=True | |
| ) | |
| assert mma_attr_col_major is not None | |
| assert "col_major = true" in str(mma_attr_col_major) | |
| assert mma_attr_col_major.col_major | |
| M, N, K = mma_attr_col_major.mnk_shape | |
| assert M == 32 and N == 32 and K == 8 | |
| mma_attr_row_major = iree_gpu.MMAAttr.get( | |
| iree_gpu.MMAIntrinsic.MFMA_F32_32x32x8_F16, col_major=False | |
| ) | |
| assert mma_attr_row_major is not None | |
| assert "col_major" not in str(mma_attr_row_major) | |
| assert not mma_attr_row_major.col_major | |
| mma_attr_default = iree_gpu.MMAAttr.get(iree_gpu.MMAIntrinsic.MFMA_F32_32x32x8_F16) | |
| assert str(mma_attr_default) == str(mma_attr_row_major) | |
| assert not mma_attr_default.col_major | |
| @run | |
| def virtual_mma_intrinsic_attr(): | |
| virtual_mma_intrinsic_attr = iree_gpu.VirtualMMAIntrinsicAttr.get( | |
| iree_gpu.VirtualMMAIntrinsic.VMFMA_F32_16x16x32_F16 | |
| ) | |
| assert virtual_mma_intrinsic_attr is not None | |
| assert ( | |
| str(virtual_mma_intrinsic_attr) | |
| == "#iree_gpu<virtual_mma_intrinsic VMFMA_F32_16x16x32_F16>" | |
| ) | |
| # Fragment: 0 = lhs, 1 = rhs, 2 = acc. | |
| fragment = 0 | |
| virtual_mma_single_subgroup_layout = iree_gpu.get_single_subgroup_layout( | |
| virtual_mma_intrinsic_attr, fragment | |
| ) | |
| assert isinstance( | |
| virtual_mma_single_subgroup_layout, iree_gpu.GPUMMASingleSubgroupLayout | |
| ) | |
| assert virtual_mma_single_subgroup_layout.outer == [1, 1] | |
| assert virtual_mma_single_subgroup_layout.thread == [16, 4] | |
| assert virtual_mma_single_subgroup_layout.tstrides == [ | |
| 1, | |
| 16, | |
| ] | |
| assert virtual_mma_single_subgroup_layout.element == [1, 8] | |
| raw_value = virtual_mma_intrinsic_attr.raw_value | |
| assert raw_value == iree_gpu.VirtualMMAIntrinsic.VMFMA_F32_16x16x32_F16 | |
| value = virtual_mma_intrinsic_attr.value | |
| assert str(value) == "VMFMA_F32_16x16x32_F16" | |
| assert int(value) == raw_value | |
| virtual_mma_attr = iree_gpu.VirtualMMAAttr.get(raw_value) | |
| assert virtual_mma_attr is not None | |
| f16 = ir.F16Type.get() | |
| f32 = ir.F32Type.get() | |
| a_type, b_type, c_type = virtual_mma_attr.abc_element_types | |
| assert a_type == f16 | |
| assert b_type == f16 | |
| assert c_type == f32 | |
| vec_4xf32 = ir.VectorType.get((4,), f32) | |
| vec_8xf16 = ir.VectorType.get((8,), f16) | |
| a_vec_type, b_vec_type, c_vec_type = virtual_mma_attr.abc_vector_types | |
| assert a_vec_type == vec_8xf16 | |
| assert b_vec_type == vec_8xf16 | |
| assert c_vec_type == vec_4xf32 | |
| M, N, K = virtual_mma_attr.mnk_shape | |
| assert M == 16 | |
| assert N == 16 | |
| assert K == 32 | |
| assert virtual_mma_intrinsic_attr.mma == virtual_mma_attr | |
| virtual_mma_attr_col_major = iree_gpu.VirtualMMAAttr.get( | |
| iree_gpu.VirtualMMAIntrinsic.VMFMA_F32_16x16x32_F16, col_major=True | |
| ) | |
| assert virtual_mma_attr_col_major is not None | |
| assert "col_major = true" in str(virtual_mma_attr_col_major) | |
| assert virtual_mma_attr_col_major.col_major | |
| M, N, K = virtual_mma_attr_col_major.mnk_shape | |
| assert M == 16 and N == 16 and K == 32 | |
| virtual_mma_attr_row_major = iree_gpu.VirtualMMAAttr.get( | |
| iree_gpu.VirtualMMAIntrinsic.VMFMA_F32_16x16x32_F16, col_major=False | |
| ) | |
| assert virtual_mma_attr_row_major is not None | |
| assert "col_major" not in str(virtual_mma_attr_row_major) | |
| assert not virtual_mma_attr_row_major.col_major | |
| virtual_mma_attr_default = iree_gpu.VirtualMMAAttr.get( | |
| iree_gpu.VirtualMMAIntrinsic.VMFMA_F32_16x16x32_F16 | |
| ) | |
| assert str(virtual_mma_attr_default) == str(virtual_mma_attr_row_major) | |
| assert not virtual_mma_attr_default.col_major | |
| @run | |
| def lowering_config_attr(): | |
| attributes = ir.DictAttr.get( | |
| { | |
| "reduction": get_index_array_attr([]), | |
| } | |
| ) | |
| lowering_config = iree_gpu.LoweringConfigAttr.get(attributes) | |
| assert lowering_config is not None | |
| assert lowering_config.attributes == attributes | |
| assert lowering_config.workgroup_tile_sizes == [] | |
| assert lowering_config.reduction_tile_sizes == [] | |
| assert lowering_config.mma_kind == None | |
| assert lowering_config.subgroup_basis == ([], []) | |
| mma_intrinsic = iree_gpu.MMAIntrinsic.MFMA_F32_16x16x16_F16 | |
| mma_attr = iree_gpu.MMAAttr.get(mma_intrinsic) | |
| attributes = ir.DictAttr.get( | |
| { | |
| "reduction": get_index_array_attr([1]), | |
| "workgroup": get_index_array_attr([2, 3]), | |
| "mma_kind": mma_attr, | |
| "subgroup_basis": ir.ArrayAttr.get( | |
| [get_index_array_attr([1, 1, 4]), get_index_array_attr([0, 1, 2])] | |
| ), | |
| } | |
| ) | |
| lowering_config = iree_gpu.LoweringConfigAttr.get(attributes) | |
| assert lowering_config.workgroup_tile_sizes == [2, 3] | |
| assert lowering_config.reduction_tile_sizes == [1] | |
| assert lowering_config.mma_kind == mma_attr | |
| assert lowering_config.subgroup_basis == ([1, 1, 4], [0, 1, 2]) | |
| assert ( | |
| str(lowering_config.mma_kind) == "#iree_gpu.mma_layout<MFMA_F32_16x16x16_F16>" | |
| ) | |
| @run | |
| def compilation_info(): | |
| attributes = ir.DictAttr.get({"reduction": get_index_array_attr([])}) | |
| lowering_config = iree_gpu.LoweringConfigAttr.get(attributes) | |
| pipeline_attr = iree_codegen.NoPipelineAttr.get() | |
| translation_info = iree_codegen.TranslationInfoAttr.get(pipeline_attr) | |
| compilation_info = iree_codegen.CompilationInfoAttr.get( | |
| lowering_config, translation_info | |
| ) | |
| assert compilation_info is not None | |
| assert compilation_info.lowering_config == lowering_config | |
| assert compilation_info.translation_info == translation_info | |
| @run | |
| def gpu_target_info_attribute_parsing(): | |
| mlir_string = """ | |
| hal.executable private @main_dispatch_0 { | |
| hal.executable.variant public @rocm_hsaco_fb | |
| target(<"rocm", "rocm-hsaco-fb", | |
| { | |
| abi = "hip", | |
| iree_codegen.target_info = #iree_gpu.target< | |
| arch = "gfx942", | |
| features = "", | |
| wgp = < | |
| compute = fp64, | |
| storage = b64, | |
| subgroup = none, | |
| dot = none, | |
| mma = [<MFMA_F32_16x16x4_F32>, <MFMA_F32_16x16x16_F16>], | |
| subgroup_size_choices = [32, 64], | |
| max_workgroup_sizes = [256, 512, 1024], | |
| max_thread_count_per_workgroup = 1024, | |
| max_workgroup_memory_bytes = 65536, | |
| max_workgroup_counts = [256, 512, 1024], | |
| simds_per_wgp = 4 | |
| >, | |
| chip = <wgp_count = 304, sku = "mi300x"> | |
| > | |
| }> | |
| ) { | |
| } | |
| } | |
| """ | |
| module = ir.Module.parse(mlir_string) | |
| variant_op_list = iree_codegen.get_executable_variant_ops(module) | |
| assert len(variant_op_list) == 1, "Expect one executable variant op" | |
| variant_op = variant_op_list[0] | |
| executable_variant_op = variant_op.opview | |
| target = executable_variant_op.target | |
| gpu_target_info = iree_gpu.TargetInfo.get_gpu_target_info(target) | |
| arch = gpu_target_info.arch | |
| assert arch == "gfx942", f"Expected arch 'gfx942', got '{arch}'" | |
| workgroup_count = gpu_target_info.workgroup_count | |
| simds_per_workgroup = gpu_target_info.simds_per_workgroup | |
| assert ( | |
| workgroup_count == 304 | |
| ), f"Expected workgroup_count 304, got {workgroup_count}" | |
| assert ( | |
| simds_per_workgroup == 4 | |
| ), f"Expected simds_per_workgroup 4, got {simds_per_workgroup}" | |
| subgroup_size_choices = gpu_target_info.subgroup_size_choices | |
| assert subgroup_size_choices == [ | |
| 32, | |
| 64, | |
| ], f"Expected subgroup_size_choice [32, 64], got {subgroup_size_choices}" | |
| max_thread_count = gpu_target_info.max_thread_count_per_workgroup | |
| assert ( | |
| max_thread_count == 1024 | |
| ), f"Expected max_thread_count_per_workgroup 1024, got {max_thread_count}" | |
| max_memory_bytes = gpu_target_info.max_workgroup_memory_bytes | |
| assert ( | |
| max_memory_bytes == 65536 | |
| ), f"Expected max_workgroup_memory_bytes 65536, got {max_memory_bytes}" | |
| max_workgroup_sizes = gpu_target_info.max_workgroup_sizes | |
| assert max_workgroup_sizes == [ | |
| 256, | |
| 512, | |
| 1024, | |
| ], f"Expected max_workgroup_sizes [256, 512, 1024], got {max_workgroup_sizes}" | |
| mma_intrinsics = gpu_target_info.mma_intrinsics | |
| assert mma_intrinsics == [ | |
| iree_gpu.MMAIntrinsic.MFMA_F32_16x16x4_F32, | |
| iree_gpu.MMAIntrinsic.MFMA_F32_16x16x16_F16, | |
| iree_gpu.VirtualMMAIntrinsic.VMFMA_F32_16x16x32_F16, | |
| iree_gpu.VirtualMMAIntrinsic.VDMFMA_F32_8x16x64x2_F16, | |
| ], f"Expected mma_intrinsics [MFMA_F32_16x16x4_F32, MFMA_F32_16x16x16_F16, VMFMA_F32_16x16x32_F16, VDMFMA_F32_8x16x64x2_F16], got {mma_intrinsics}" | |
| @run | |
| def gpu_target_info_constructor(): | |
| context = ir.Context() | |
| target_info = iree_gpu.TargetInfo( | |
| context=context, | |
| arch="gfx942", | |
| subgroup_size_choices=[32, 64], | |
| max_workgroup_sizes=[256, 512, 1024], | |
| max_thread_count_per_workgroup=1024, | |
| max_workgroup_memory_bytes=65536, | |
| workgroup_count=304, | |
| simds_per_workgroup=4, | |
| mma_intrinsics=[ | |
| iree_gpu.MMAIntrinsic.MFMA_F32_16x16x4_F32, | |
| iree_gpu.MMAIntrinsic.MFMA_F32_16x16x16_F16, | |
| iree_gpu.VirtualMMAIntrinsic.VMFMA_F32_16x16x32_F16, | |
| ], | |
| ) | |
| assert ( | |
| target_info.arch == "gfx942" | |
| ), f"Expected arch 'gfx942', got '{target_info.arch}'" | |
| assert target_info.subgroup_size_choices == [ | |
| 32, | |
| 64, | |
| ], f"Expected subgroup_size_choices [32, 64], got {target_info.subgroup_size_choices}" | |
| assert target_info.max_workgroup_sizes == [ | |
| 256, | |
| 512, | |
| 1024, | |
| ], f"Expected max_workgroup_sizes [256, 512, 1024], got {target_info.max_workgroup_sizes}" | |
| assert ( | |
| target_info.max_thread_count_per_workgroup == 1024 | |
| ), f"Expected max_thread_count_per_workgroup 1024, got {target_info.max_thread_count_per_workgroup}" | |
| assert ( | |
| target_info.max_workgroup_memory_bytes == 65536 | |
| ), f"Expected max_workgroup_memory_bytes 65536, got {target_info.max_workgroup_memory_bytes}" | |
| assert ( | |
| target_info.workgroup_count == 304 | |
| ), f"Expected workgroup_count 304, got {target_info.workgroup_count}" | |
| assert ( | |
| target_info.simds_per_workgroup == 4 | |
| ), f"Expected simds_per_workgroup 4, got {target_info.simds_per_workgroup}" | |
| mma_intrinsics = target_info.mma_intrinsics | |
| assert mma_intrinsics == [ | |
| iree_gpu.MMAIntrinsic.MFMA_F32_16x16x4_F32, | |
| iree_gpu.MMAIntrinsic.MFMA_F32_16x16x16_F16, | |
| iree_gpu.VirtualMMAIntrinsic.VMFMA_F32_16x16x32_F16, | |
| ], f"Expected mma_intrinsics [MFMA_F32_16x16x4_F32, MFMA_F32_16x16x16_F16, VMFMA_F32_16x16x32_F16], got {mma_intrinsics}" | |
| assert isinstance(mma_intrinsics[0], iree_gpu.MMAIntrinsic) | |
| assert isinstance(mma_intrinsics[1], iree_gpu.MMAIntrinsic) | |
| assert isinstance(mma_intrinsics[2], iree_gpu.VirtualMMAIntrinsic) | |
| @run | |
| def gpu_target_info_constructor_error_cases(): | |
| context = ir.Context() | |
| try: | |
| iree_gpu.TargetInfo( | |
| context=context, | |
| arch=123, # should be string. | |
| subgroup_size_choices=[32, 64], | |
| max_workgroup_sizes=[256, 512, 1024], | |
| max_thread_count_per_workgroup=1024, | |
| max_workgroup_memory_bytes=65536, | |
| workgroup_count=304, | |
| simds_per_workgroup=4, | |
| mma_intrinsics=[], | |
| ) | |
| assert False, "Expected TypeError for wrong arch type" | |
| except TypeError: | |
| pass | |
| try: | |
| iree_gpu.TargetInfo( | |
| context=context, | |
| arch="gfx942", | |
| subgroup_size_choices=[64.0], # should be list of int. | |
| max_workgroup_sizes=[256, 512, 1024], | |
| max_thread_count_per_workgroup=1024, | |
| max_workgroup_memory_bytes=65536, | |
| workgroup_count=304, | |
| simds_per_workgroup=4, | |
| mma_intrinsics=[], | |
| ) | |
| assert False, "Expected TypeError for wrong subgroup_size_choices type" | |
| except TypeError: | |
| pass | |
| try: | |
| iree_gpu.TargetInfo( | |
| context=context, | |
| arch="gfx942", | |
| subgroup_size_choices=[32, 64], | |
| max_workgroup_sizes=[256.0, 512, 1024], # should be list of int. | |
| max_thread_count_per_workgroup=1024, | |
| max_workgroup_memory_bytes=65536, | |
| workgroup_count=304, | |
| simds_per_workgroup=4, | |
| mma_intrinsics=[], | |
| ) | |
| assert False, "Expected TypeError for wrong max_workgroup_sizes type" | |
| except TypeError: | |
| pass | |
| try: | |
| iree_gpu.TargetInfo( | |
| context=context, | |
| arch="gfx942", | |
| subgroup_size_choices=[32, 64], | |
| max_workgroup_sizes=[256, 512, 1024], | |
| max_thread_count_per_workgroup=1024.0, # should be int. | |
| max_workgroup_memory_bytes=65536, | |
| workgroup_count=304, | |
| simds_per_workgroup=4, | |
| mma_intrinsics=[], | |
| ) | |
| assert False, "Expected TypeError for wrong max_thread_count_per_workgroup type" | |
| except TypeError: | |
| pass | |
| try: | |
| iree_gpu.TargetInfo( | |
| context=context, | |
| arch="gfx942", | |
| subgroup_size_choices=[32, 64], | |
| max_workgroup_sizes=[256, 512, 1024], | |
| max_thread_count_per_workgroup=1024, | |
| max_workgroup_memory_bytes=65536.0, # should be int. | |
| workgroup_count=304, | |
| simds_per_workgroup=4, | |
| mma_intrinsics=[], | |
| ) | |
| assert False, "Expected TypeError for wrong max_workgroup_memory_bytes type" | |
| except TypeError: | |
| pass | |
| try: | |
| iree_gpu.TargetInfo( | |
| context=context, | |
| arch="gfx942", | |
| subgroup_size_choices=[32, 64], | |
| max_workgroup_sizes=[256, 512, 1024], | |
| max_thread_count_per_workgroup=1024, | |
| max_workgroup_memory_bytes=65536, | |
| workgroup_count=304.0, # Should be int. | |
| simds_per_workgroup=4, | |
| mma_intrinsics=[], | |
| ) | |
| assert False, "Expected TypeError for wrong workgroup_count type" | |
| except TypeError: | |
| pass | |
| try: | |
| iree_gpu.TargetInfo( | |
| context=context, | |
| arch="gfx942", | |
| subgroup_size_choices=[32, 64], | |
| max_workgroup_sizes=[256, 512, 1024], | |
| max_thread_count_per_workgroup=1024, | |
| max_workgroup_memory_bytes=65536, | |
| workgroup_count=-304, # Should be non-negative. | |
| simds_per_workgroup=4, | |
| mma_intrinsics=[], | |
| ) | |
| assert False, "Expected ValueError for negative workgroup_count" | |
| except ValueError: | |
| pass | |
| try: | |
| iree_gpu.TargetInfo( | |
| context=context, | |
| arch="gfx942", | |
| subgroup_size_choices=[32, 64], | |
| max_workgroup_sizes=[256, 512, 1024], | |
| max_thread_count_per_workgroup=1024, | |
| max_workgroup_memory_bytes=65536, | |
| workgroup_count=304, | |
| simds_per_workgroup=4.0, # Should be int. | |
| mma_intrinsics=[], | |
| ) | |
| assert False, "Expected TypeError for wrong simds_per_workgroup type" | |
| except TypeError: | |
| pass | |
| try: | |
| iree_gpu.TargetInfo( | |
| context=context, | |
| arch="gfx942", | |
| subgroup_size_choices=[32, 64], | |
| max_workgroup_sizes=[256, 512, 1024], | |
| max_thread_count_per_workgroup=1024, | |
| max_workgroup_memory_bytes=65536, | |
| mma_intrinsics=[123], # should be MMA intrinsic objects. | |
| ) | |
| assert False, "Expected TypeError for wrong MMA intrinsic object type" | |
| except TypeError: | |
| pass | |
| # ====================================================================== | |
| # IREE TensorExt Dialect | |
| # ====================================================================== | |
| @run | |
| def iree_tensor_ext_smoke_test(): | |
| # Make sure that generated op bindings are accessible. | |
| assert issubclass(iree_tensor_ext.ComputeBarrierStartOp, ir.OpView) | |
| # ====================================================================== | |
| # Preprocessing Transform Extensions | |
| # ====================================================================== | |
| @run | |
| def preprocessing_transform_match_contraction_in_named_sequence(): | |
| module_op = ir.Module.create() | |
| module_op.operation.attributes["transform.with_named_sequence"] = ir.UnitAttr.get() | |
| map_lhs = ir.AffineMap.get( | |
| dim_count=3, | |
| symbol_count=0, | |
| exprs=[ir.AffineExpr.get_dim(0), ir.AffineExpr.get_dim(2)], | |
| ) | |
| map_rhs = ir.AffineMap.get( | |
| dim_count=3, | |
| symbol_count=0, | |
| exprs=[ir.AffineExpr.get_dim(2), ir.AffineExpr.get_dim(1)], | |
| ) | |
| map_output = ir.AffineMap.get( | |
| dim_count=3, | |
| symbol_count=0, | |
| exprs=[ir.AffineExpr.get_dim(0), ir.AffineExpr.get_dim(1)], | |
| ) | |
| with ir.InsertionPoint(module_op.body): | |
| named_seq = transform.NamedSequenceOp( | |
| "match_matmul", [transform.AnyOpType.get()], [transform.AnyOpType.get()] | |
| ) | |
| with ir.InsertionPoint(named_seq.body): | |
| batch, m, n, k = preprocessing_transform.MatchContractionOp( | |
| operand_handle=named_seq.bodyTarget, | |
| lhs_type=ir.F32Type.get(), | |
| rhs_type=ir.F32Type.get(), | |
| output_type=ir.F32Type.get(), | |
| indexing_maps=[map_lhs, map_rhs, map_output], | |
| ) | |
| transform.YieldOp([named_seq.bodyTarget]) | |
| module_str = str(module_op) | |
| assert "affine_map<(d0, d1, d2) -> (d0, d2)>" in module_str | |
| assert "affine_map<(d0, d1, d2) -> (d2, d1)>" in module_str | |
| assert "affine_map<(d0, d1, d2) -> (d0, d1)>" in module_str | |
| assert "transform.with_named_sequence" in module_str | |
| assert "transform.named_sequence @match_matmul" in module_str | |
| assert "transform.iree.match.contraction %arg0" in module_str | |
| assert "lhs_type = f32, rhs_type = f32, output_type = f32" in module_str | |
| assert "indexing_maps = [#map, #map1, #map2]" in module_str | |
| @run | |
| def preprocessing_transform_match_convolution_in_named_sequence(): | |
| module_op = ir.Module.create() | |
| module_op.operation.attributes["transform.with_named_sequence"] = ir.UnitAttr.get() | |
| map_input = ir.AffineMap.get( | |
| dim_count=7, | |
| symbol_count=0, | |
| exprs=[ | |
| ir.AffineExpr.get_dim(0), | |
| ir.AffineExpr.get_dim(1) + ir.AffineExpr.get_dim(4), | |
| ir.AffineExpr.get_dim(2) + ir.AffineExpr.get_dim(5), | |
| ir.AffineExpr.get_dim(6), | |
| ], | |
| ) | |
| map_filter = ir.AffineMap.get( | |
| dim_count=7, | |
| symbol_count=0, | |
| exprs=[ | |
| ir.AffineExpr.get_dim(4), | |
| ir.AffineExpr.get_dim(5), | |
| ir.AffineExpr.get_dim(6), | |
| ir.AffineExpr.get_dim(3), | |
| ], | |
| ) | |
| map_output = ir.AffineMap.get( | |
| dim_count=7, | |
| symbol_count=0, | |
| exprs=[ | |
| ir.AffineExpr.get_dim(0), | |
| ir.AffineExpr.get_dim(1), | |
| ir.AffineExpr.get_dim(2), | |
| ir.AffineExpr.get_dim(3), | |
| ], | |
| ) | |
| with ir.InsertionPoint(module_op.body): | |
| named_seq = transform.NamedSequenceOp( | |
| "match_conv", [transform.AnyOpType.get()], [transform.AnyOpType.get()] | |
| ) | |
| with ir.InsertionPoint(named_seq.body): | |
| ( | |
| batch, | |
| out_img, | |
| out_ch, | |
| filt, | |
| in_ch, | |
| depth, | |
| strides, | |
| dilations, | |
| ) = preprocessing_transform.MatchConvolutionOp( | |
| operand_handle=named_seq.bodyTarget, | |
| lhs_type=ir.F32Type.get(), | |
| rhs_type=ir.F32Type.get(), | |
| output_type=ir.F32Type.get(), | |
| indexing_maps=[map_input, map_filter, map_output], | |
| ) | |
| transform.YieldOp([named_seq.bodyTarget]) | |
| module_str = str(module_op) | |
| assert ( | |
| "affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)>" | |
| in module_str | |
| ) | |
| assert "affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)>" in module_str | |
| assert "affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>" in module_str | |
| assert "transform.with_named_sequence" in module_str | |
| assert "transform.named_sequence @match_conv" in module_str | |
| assert "transform.iree.match.convolution %arg0" in module_str | |
| assert "lhs_type = f32, rhs_type = f32, output_type = f32" in module_str | |
| assert "indexing_maps = [#map, #map1, #map2]" in module_str | |
| @run | |
| def preprocessing_transform_match_attention_in_named_sequence(): | |
| module_op = ir.Module.create() | |
| module_op.operation.attributes["transform.with_named_sequence"] = ir.UnitAttr.get() | |
| map_query = ir.AffineMap.get( | |
| dim_count=6, | |
| symbol_count=0, | |
| exprs=[ | |
| ir.AffineExpr.get_dim(0), | |
| ir.AffineExpr.get_dim(1), | |
| ir.AffineExpr.get_dim(2), | |
| ir.AffineExpr.get_dim(4), | |
| ], | |
| ) | |
| map_key = ir.AffineMap.get( | |
| dim_count=6, | |
| symbol_count=0, | |
| exprs=[ | |
| ir.AffineExpr.get_dim(0), | |
| ir.AffineExpr.get_dim(1), | |
| ir.AffineExpr.get_dim(5), | |
| ir.AffineExpr.get_dim(4), | |
| ], | |
| ) | |
| map_value = ir.AffineMap.get( | |
| dim_count=6, | |
| symbol_count=0, | |
| exprs=[ | |
| ir.AffineExpr.get_dim(0), | |
| ir.AffineExpr.get_dim(1), | |
| ir.AffineExpr.get_dim(3), | |
| ir.AffineExpr.get_dim(5), | |
| ], | |
| ) | |
| map_scale = ir.AffineMap.get( | |
| dim_count=6, | |
| symbol_count=0, | |
| exprs=[], | |
| ) | |
| map_output = ir.AffineMap.get( | |
| dim_count=6, | |
| symbol_count=0, | |
| exprs=[ | |
| ir.AffineExpr.get_dim(0), | |
| ir.AffineExpr.get_dim(1), | |
| ir.AffineExpr.get_dim(2), | |
| ir.AffineExpr.get_dim(3), | |
| ], | |
| ) | |
| with ir.InsertionPoint(module_op.body): | |
| named_seq = transform.NamedSequenceOp( | |
| "match_attention", [transform.AnyOpType.get()], [transform.AnyOpType.get()] | |
| ) | |
| with ir.InsertionPoint(named_seq.body): | |
| batch, m, n, k1, k2 = preprocessing_transform.MatchAttentionOp( | |
| operand_handle=named_seq.bodyTarget, | |
| query_type=ir.F16Type.get(), | |
| key_type=ir.F16Type.get(), | |
| value_type=ir.F16Type.get(), | |
| output_type=ir.F16Type.get(), | |
| indexing_maps=[map_query, map_key, map_value, map_scale, map_output], | |
| ) | |
| transform.YieldOp([named_seq.bodyTarget]) | |
| module_str = str(module_op) | |
| assert "affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d4)>" in module_str | |
| assert "affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d4)>" in module_str | |
| assert "affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d5)>" in module_str | |
| assert "affine_map<(d0, d1, d2, d3, d4, d5) -> ()>" in module_str | |
| assert "affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>" in module_str | |
| assert "transform.with_named_sequence" in module_str | |
| assert "transform.named_sequence @match_attention" in module_str | |
| assert "transform.iree.match.attention %arg0" in module_str | |
| assert ( | |
| "query_type = f16, key_type = f16, value_type = f16, output_type = f16" | |
| in module_str | |
| ) | |
| assert "indexing_maps = [#map, #map1, #map2, #map3, #map4]" in module_str |