diff --git a/compiler/bindings/c/BUILD.bazel b/compiler/bindings/c/BUILD.bazel index 21857d035db0..01c939c528f7 100644 --- a/compiler/bindings/c/BUILD.bazel +++ b/compiler/bindings/c/BUILD.bazel @@ -14,6 +14,7 @@ cc_library( name = "headers", hdrs = [ "iree/compiler/api_support.h", + "iree/compiler/dialects/iree_gpu.h", "iree/compiler/embedding_api.h", "iree/compiler/loader.h", "iree/compiler/mlir_interop.h", diff --git a/compiler/bindings/c/CMakeLists.txt b/compiler/bindings/c/CMakeLists.txt index 2ac9f5e9356c..f13bf426f6c8 100644 --- a/compiler/bindings/c/CMakeLists.txt +++ b/compiler/bindings/c/CMakeLists.txt @@ -22,6 +22,7 @@ iree_cc_library( headers HDRS "iree/compiler/api_support.h" + "iree/compiler/dialects/iree_gpu.h" "iree/compiler/embedding_api.h" "iree/compiler/loader.h" "iree/compiler/mlir_interop.h" diff --git a/compiler/bindings/c/iree/compiler/dialects/iree_gpu.h b/compiler/bindings/c/iree/compiler/dialects/iree_gpu.h new file mode 100644 index 000000000000..6798ed3a45cb --- /dev/null +++ b/compiler/bindings/c/iree/compiler/dialects/iree_gpu.h @@ -0,0 +1,59 @@ +// Copyright 2024 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 + +#ifndef IREE_COMPILER_DIALECTS_IREE_GPU_H +#define IREE_COMPILER_DIALECTS_IREE_GPU_H + +#include "mlir-c/IR.h" +#include "mlir-c/Support.h" + +#ifdef __cplusplus +extern "C" { +#endif + +enum ireeGPUReorderWorkgroupsStrategyEnum { + ireeGPUReorderWorkgroupsStrategyEnumNone = 0, + ireeGPUReorderWorkgroupsStrategyEnumSwizzle = 1, + ireeGPUReorderWorkgroupsStrategyEnumTranspose = 2, +}; + +MLIR_CAPI_EXPORTED bool +ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr(MlirAttribute attr); + +MLIR_CAPI_EXPORTED MlirTypeID +ireeGPUReorderWorkgroupsStrategyAttrGetTypeID(void); + +MLIR_CAPI_EXPORTED MlirAttribute ireeGPUReorderWorkgroupsStrategyAttrGet( + MlirContext mlirCtx, ireeGPUReorderWorkgroupsStrategyEnum value); + +MLIR_CAPI_EXPORTED ireeGPUReorderWorkgroupsStrategyEnum +ireeGPUReorderWorkgroupsStrategyAttrGetValue(MlirAttribute attr); + +MLIR_CAPI_EXPORTED +bool ireeAttributeIsAGPUPipelineOptionsAttr(MlirAttribute attr); + +MLIR_CAPI_EXPORTED MlirAttribute +ireeGPUPipelineOptionsAttrGet(MlirContext mlirCtx, bool *prefetchSharedMemory, + bool *noReduceSharedMemoryBankConflicts, + MlirAttribute *reorderWorkgroupsStrategy); + +MLIR_CAPI_EXPORTED MlirAttribute +ireeGPUPipelineOptionsAttrGetPrefetchSharedMemory(MlirAttribute attr); + +MLIR_CAPI_EXPORTED MlirAttribute +ireeGPUPipelineOptionsAttrGetNoReduceSharedMemoryBankConflicts( + MlirAttribute attr); + +MLIR_CAPI_EXPORTED MlirAttribute +ireeGPUPipelineOptionsAttrGetReorderWorkgroupsStrategy(MlirAttribute attr); + +MLIR_CAPI_EXPORTED MlirTypeID ireeGPUPipelineOptionsAttrGetTypeID(void); + +#ifdef __cplusplus +} +#endif + +#endif // IREE_COMPILER_DIALECTS_IREE_GPU_H diff --git a/compiler/bindings/python/CMakeLists.txt b/compiler/bindings/python/CMakeLists.txt index e27e8e526216..0a0b36474885 100644 --- a/compiler/bindings/python/CMakeLists.txt +++ b/compiler/bindings/python/CMakeLists.txt @@ -100,6 +100,15 @@ declare_mlir_dialect_python_bindings( DIALECT_NAME vm ) +declare_mlir_dialect_python_bindings( + ADD_TO_PARENT IREEPythonSources.Dialects + ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler" + TD_FILE dialects/IREEGPUBinding.td + GEN_ENUM_BINDINGS + SOURCES dialects/iree_gpu.py + DIALECT_NAME iree_gpu +) + declare_mlir_python_sources(IREECompilerAPIPythonCore ADD_TO_PARENT IREEPythonSources ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler" @@ -165,6 +174,17 @@ declare_mlir_python_extension(IREECompilerPythonExtensions.Registration LLVMSupport ) +declare_mlir_python_extension(IREECompilerPythonExtensions.CompilerDialects + MODULE_NAME _ireeCompilerDialects + ADD_TO_PARENT IREECompilerPythonExtensions + SOURCES + IREECompilerDialectsModule.cpp + EMBED_CAPI_LINK_LIBS + iree_compiler_API_SharedImpl + PRIVATE_LINK_LIBS + LLVMSupport +) + ################################################################################ # Generate packages and shared library ################################################################################ diff --git a/compiler/bindings/python/IREECompilerDialectsModule.cpp b/compiler/bindings/python/IREECompilerDialectsModule.cpp new file mode 100644 index 000000000000..f06cd9f41d2e --- /dev/null +++ b/compiler/bindings/python/IREECompilerDialectsModule.cpp @@ -0,0 +1,123 @@ +// Copyright 2024 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 + +#include "iree/compiler/dialects/iree_gpu.h" +#include "mlir-c/BuiltinAttributes.h" +#include "mlir-c/IR.h" +#include "mlir/Bindings/Python/PybindAdaptors.h" + +namespace py = pybind11; +using namespace mlir::python::adaptors; + +PYBIND11_MODULE(_ireeCompilerDialects, m) { + m.doc() = "iree-compiler dialects python extension"; + + auto iree_gpu_module = + m.def_submodule("iree_gpu", "iree_gpu dialect bindings"); + + //===-------------------------------------------------------------------===// + // GPUReorderWorkgroupsStrategyAttr + //===-------------------------------------------------------------------===// + + auto strategyEnum = + py::enum_( + iree_gpu_module, "ReorderWorkgroupsStrategy", py::module_local()) + .value("None_", ireeGPUReorderWorkgroupsStrategyEnumNone) + .value("Swizzle", ireeGPUReorderWorkgroupsStrategyEnumSwizzle) + .value("Transpose", ireeGPUReorderWorkgroupsStrategyEnumTranspose) + .def( + "__str__", + [](ireeGPUReorderWorkgroupsStrategyEnum &self) { + switch (self) { + case ireeGPUReorderWorkgroupsStrategyEnumNone: + return "None"; + case ireeGPUReorderWorkgroupsStrategyEnumSwizzle: + return "Swizzle"; + case ireeGPUReorderWorkgroupsStrategyEnumTranspose: + return "Transpose"; + default: + llvm::report_fatal_error( + "unknown ReorderWorkgroupsStrategy variant"); + } + }, + // pybind overloads are tried in the order they were registered. + // As a result, enums used the default __str__ method instead of + // the custom one. Adding py::prepend() fixes this issue. + py::prepend()); + + mlir_attribute_subclass(iree_gpu_module, "ReorderWorkgroupsStrategyAttr", + ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr, + ireeGPUReorderWorkgroupsStrategyAttrGetTypeID) + .def_classmethod( + "get", + [](const py::object &, ireeGPUReorderWorkgroupsStrategyEnum value, + MlirContext ctx) { + return ireeGPUReorderWorkgroupsStrategyAttrGet(ctx, value); + }, + "cls"_a, "value"_a, "ctx"_a = py::none(), + "Gets a gpu.reorder_workgroups_strategy from parameters.") + .def_property_readonly( + "value", + [](MlirAttribute self) -> ireeGPUReorderWorkgroupsStrategyEnum { + return ireeGPUReorderWorkgroupsStrategyAttrGetValue(self); + }); + + //===-------------------------------------------------------------------===// + // GPUPipelineOptionsAttr + //===-------------------------------------------------------------------===// + + mlir_attribute_subclass(iree_gpu_module, "GPUPipelineOptionsAttr", + ireeAttributeIsAGPUPipelineOptionsAttr, + ireeGPUPipelineOptionsAttrGetTypeID) + .def_classmethod( + "get", + [](const py::object &, std::optional prefetchSharedMemory, + std::optional noReduceSharedMemoryBankConflicts, + std::optional reorderWorkgroupsStrategy, + MlirContext ctx) { + return ireeGPUPipelineOptionsAttrGet( + ctx, + prefetchSharedMemory.has_value() ? &*prefetchSharedMemory + : nullptr, + noReduceSharedMemoryBankConflicts.has_value() + ? &*noReduceSharedMemoryBankConflicts + : nullptr, + reorderWorkgroupsStrategy.has_value() + ? &*reorderWorkgroupsStrategy + : nullptr); + }, + "cls"_a, "prefetch_shared_memory"_a = py::none(), + "no_reduce_shared_memory_bank_conflicts"_a = py::none(), + "reorder_workgroups_strategy"_a = py::none(), py::kw_only(), + "ctx"_a = py::none(), "Gets a gpu.pipeline_options from parameters.") + .def_property_readonly( + "prefetch_shared_memory", + [](MlirAttribute self) -> std::optional { + auto attr = ireeGPUPipelineOptionsAttrGetPrefetchSharedMemory(self); + if (!mlirAttributeIsNull(attr)) + return mlirBoolAttrGetValue(attr); + return std::nullopt; + }) + .def_property_readonly( + "no_reduce_shared_memory_bank_conflicts", + [](MlirAttribute self) -> std::optional { + auto attr = + ireeGPUPipelineOptionsAttrGetNoReduceSharedMemoryBankConflicts( + self); + if (!mlirAttributeIsNull(attr)) + return mlirBoolAttrGetValue(attr); + return std::nullopt; + }) + .def_property_readonly( + "reorder_workgroups_strategy", + [](MlirAttribute self) -> std::optional { + auto attr = + ireeGPUPipelineOptionsAttrGetReorderWorkgroupsStrategy(self); + if (!mlirAttributeIsNull(attr)) + return attr; + return std::nullopt; + }); +} diff --git a/compiler/bindings/python/iree/compiler/dialects/IREEGPUBinding.td b/compiler/bindings/python/iree/compiler/dialects/IREEGPUBinding.td new file mode 100644 index 000000000000..7674967caafd --- /dev/null +++ b/compiler/bindings/python/iree/compiler/dialects/IREEGPUBinding.td @@ -0,0 +1,12 @@ +// Copyright 2024 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 + +#ifndef PYTHON_BINDINGS_IREEGPU_OPS +#define PYTHON_BINDINGS_IREEGPU_OPS + +include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.td" + +#endif // PYTHON_BINDINGS_IREEGPU_OPS diff --git a/compiler/bindings/python/iree/compiler/dialects/iree_gpu.py b/compiler/bindings/python/iree/compiler/dialects/iree_gpu.py new file mode 100644 index 000000000000..9058ec37350b --- /dev/null +++ b/compiler/bindings/python/iree/compiler/dialects/iree_gpu.py @@ -0,0 +1,9 @@ +# Copyright 2024 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_gpu_ops_gen import * +from ._iree_gpu_enum_gen import * +from .._mlir_libs._ireeCompilerDialects.iree_gpu import * diff --git a/compiler/bindings/python/test/ir/dialects_test.py b/compiler/bindings/python/test/ir/dialects_test.py index 63a960360e41..fad5cd9c2e0d 100644 --- a/compiler/bindings/python/test/ir/dialects_test.py +++ b/compiler/bindings/python/test/ir/dialects_test.py @@ -7,10 +7,69 @@ from iree.compiler import ir # Make sure that our dialects import. -from iree.compiler.dialects import ( - flow, - hal, - stream, - vm, - util, -) +from iree.compiler.dialects import flow, hal, stream, vm, util, iree_gpu + + +@lambda _: _() +def gpu_pipeline_options_attr(): + with ir.Context() as ctx, ir.Location.unknown(): + module = ir.Module.create() + with ir.InsertionPoint(module.body): + reorder_attr = iree_gpu.ReorderWorkgroupsStrategyAttr.get( + iree_gpu.ReorderWorkgroupsStrategy.Swizzle, ctx + ) + gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get( + True, + False, + reorder_attr, + ) + assert type(gpu_attr) is iree_gpu.GPUPipelineOptionsAttr + assert gpu_attr.prefetch_shared_memory + assert not gpu_attr.no_reduce_shared_memory_bank_conflicts + + gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get( + False, + True, + iree_gpu.ReorderWorkgroupsStrategyAttr.get( + iree_gpu.ReorderWorkgroupsStrategy.Transpose, ctx + ), + ) + assert not gpu_attr.prefetch_shared_memory + assert gpu_attr.no_reduce_shared_memory_bank_conflicts + + gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get() + assert ( + gpu_attr.prefetch_shared_memory is None + and gpu_attr.no_reduce_shared_memory_bank_conflicts is None + and gpu_attr.reorder_workgroups_strategy is None + ) + + gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get(True) + assert gpu_attr.prefetch_shared_memory + assert ( + gpu_attr.no_reduce_shared_memory_bank_conflicts is None + and gpu_attr.reorder_workgroups_strategy is None + ) + + gpu_attr = iree_gpu.GPUPipelineOptionsAttr.get(True, False) + assert gpu_attr.reorder_workgroups_strategy is None + + gpu_attr = iree_gpu.GPUPipelineOptionsAttr.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_shared_memory is None + assert gpu_attr.reorder_workgroups_strategy is None + + gpu_attr = iree_gpu.GPUPipelineOptionsAttr.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.Swizzle + ) diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel index 2e9d283a7756..a28d668eb043 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel @@ -82,6 +82,7 @@ iree_compiler_cc_library( ":IREEGPUEnums", ":IREEGPUInterfaces", ":IREEGPUOpsGen", + "//compiler/bindings/c:headers", "//compiler/src/iree/compiler/Codegen/Common:TileSwizzle", "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect", "//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect", @@ -92,6 +93,7 @@ iree_compiler_cc_library( "@llvm-project//mlir:AffineDialect", "@llvm-project//mlir:ArithDialect", "@llvm-project//mlir:ArithUtils", + "@llvm-project//mlir:CAPIIR", "@llvm-project//mlir:ControlFlowInterfaces", "@llvm-project//mlir:DialectUtils", "@llvm-project//mlir:IR", diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt index 4c160b495976..907f34bdc9e5 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt @@ -51,6 +51,7 @@ iree_cc_library( MLIRAffineDialect MLIRArithDialect MLIRArithUtils + MLIRCAPIIR MLIRControlFlowInterfaces MLIRIR MLIRLinalgDialect @@ -67,6 +68,7 @@ iree_cc_library( iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect iree::compiler::Codegen::Utils::VectorOpUtils iree::compiler::Dialect::LinalgExt::IR + iree::compiler::bindings::c::headers PUBLIC ) diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp index 471ba8b0c70c..843ae5292fa9 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp @@ -14,12 +14,16 @@ #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h" #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h" #include "iree/compiler/Codegen/Utils/VectorOpUtils.h" +#include "iree/compiler/dialects/iree_gpu.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/STLForwardCompat.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/TypeSwitch.h" #include "llvm/Support/Debug.h" #include "llvm/Support/ErrorHandling.h" +#include "mlir-c/IR.h" +#include "mlir/CAPI/IR.h" +#include "mlir/CAPI/Support.h" #include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -33,7 +37,6 @@ #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/DialectImplementation.h" #include "mlir/IR/OpDefinition.h" -#include "mlir/IR/TypeUtilities.h" #define DEBUG_TYPE "iree-gpu-attrs" #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") @@ -1687,3 +1690,112 @@ void IREEGPUDialect::registerAttributes() { } } // namespace mlir::iree_compiler::IREE::GPU + +bool ireeAttributeIsAGPUPipelineOptionsAttr(MlirAttribute attr) { + return llvm::isa( + unwrap(attr)); +} + +MlirAttribute +ireeGPUPipelineOptionsAttrGet(MlirContext mlirCtx, bool *prefetchSharedMemory, + bool *noReduceSharedMemoryBankConflicts, + MlirAttribute *reorderWorkgroupsStrategy) { + mlir::MLIRContext *ctx = unwrap(mlirCtx); + mlir::Builder b(ctx); + auto prefetchSharedMemoryAttr = mlir::BoolAttr(); + if (prefetchSharedMemory) { + prefetchSharedMemoryAttr = b.getBoolAttr(*prefetchSharedMemory); + } + auto noReduceSharedMemoryBankConflictsAttr = mlir::BoolAttr(); + if (noReduceSharedMemoryBankConflicts) { + noReduceSharedMemoryBankConflictsAttr = + b.getBoolAttr(*noReduceSharedMemoryBankConflicts); + } + auto strategyAttr = + mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategyAttr(); + if (reorderWorkgroupsStrategy) { + strategyAttr = llvm::dyn_cast< + mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategyAttr>( + unwrap(*reorderWorkgroupsStrategy)); + } + return wrap(mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr::get( + ctx, prefetchSharedMemoryAttr, noReduceSharedMemoryBankConflictsAttr, + strategyAttr)); +} + +MlirAttribute +ireeGPUPipelineOptionsAttrGetPrefetchSharedMemory(MlirAttribute attr) { + auto gpuAttr = + llvm::cast( + unwrap(attr)); + return wrap(gpuAttr.getPrefetchSharedMemory()); +} + +MlirAttribute ireeGPUPipelineOptionsAttrGetNoReduceSharedMemoryBankConflicts( + MlirAttribute attr) { + auto gpuAttr = + llvm::cast( + unwrap(attr)); + return wrap(gpuAttr.getNoReduceSharedMemoryBankConflicts()); +} + +MlirAttribute +ireeGPUPipelineOptionsAttrGetReorderWorkgroupsStrategy(MlirAttribute attr) { + auto gpuAttr = + llvm::cast( + unwrap(attr)); + return wrap(gpuAttr.getReorderWorkgroupsStrategy()); +} + +MlirTypeID ireeGPUPipelineOptionsAttrGetTypeID() { + return wrap( + mlir::iree_compiler::IREE::GPU::GPUPipelineOptionsAttr::getTypeID()); +} + +static_assert( + static_cast(ireeGPUReorderWorkgroupsStrategyEnumNone) == + static_cast(mlir::iree_compiler::IREE::GPU:: + ReorderWorkgroupsStrategy::None) && + static_cast(ireeGPUReorderWorkgroupsStrategyEnumSwizzle) == + static_cast(mlir::iree_compiler::IREE::GPU:: + ReorderWorkgroupsStrategy::Swizzle) && + static_cast(ireeGPUReorderWorkgroupsStrategyEnumTranspose) == + static_cast(mlir::iree_compiler::IREE::GPU:: + ReorderWorkgroupsStrategy::Transpose) && + static_cast(ireeGPUReorderWorkgroupsStrategyEnumTranspose) == + mlir::iree_compiler::IREE::GPU:: + getMaxEnumValForReorderWorkgroupsStrategy(), + "ireeGPUReorderWorkgroupsStrategyEnum and " + "mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategy definitions " + "have diverged"); + +bool ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr(MlirAttribute attr) { + return llvm::isa< + mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategyAttr>( + unwrap(attr)); +} + +MlirTypeID ireeGPUReorderWorkgroupsStrategyAttrGetTypeID() { + return wrap(mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategyAttr:: + getTypeID()); +} + +MlirAttribute ireeGPUReorderWorkgroupsStrategyAttrGet( + MlirContext mlirCtx, ireeGPUReorderWorkgroupsStrategyEnum value) { + mlir::MLIRContext *ctx = unwrap(mlirCtx); + return wrap( + mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategyAttr::get( + ctx, static_cast< + mlir::iree_compiler::IREE::GPU::ReorderWorkgroupsStrategy>( + value))); +} + +ireeGPUReorderWorkgroupsStrategyEnum +ireeGPUReorderWorkgroupsStrategyAttrGetValue(MlirAttribute attr) { + assert(ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr(attr) && + "attr is not a GPUReorderWorkgroupsStrategyAttr"); + return static_cast( + llvm::cast( + unwrap(attr)) + .getValue()); +} diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td index 138128988fb1..1f2cad748c08 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td @@ -481,8 +481,8 @@ def IREEGPU_LaneIdAttr : AttrDef { - let assemblyFormat = "``$value"; + EnumAttr { + let assemblyFormat = "`<` $value `>`"; let cppNamespace = "::mlir::iree_compiler::IREE::GPU"; } diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_user_vector_distribute.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_user_vector_distribute.mlir index 44eeb6407be7..c439e042a39a 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_user_vector_distribute.mlir +++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_user_vector_distribute.mlir @@ -83,14 +83,14 @@ hal.executable public @main_0_dispatch_0 { // ----- -// Check that applying the `reorder_workgroups_strategy = Transpose` pipeline option attribute enables workgroup reordering. +// Check that applying the `reorder_workgroups_strategy = ` pipeline option attribute enables workgroup reordering. // OPT-OUT: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info +// OPT-OUT-SAME: gpu_pipeline_options = #iree_gpu.pipeline_options> // OPT-OUT-SAME: mma_schedule = #iree_gpu.mma_schedule, // OPT-IN: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info +// OPT-IN-SAME: gpu_pipeline_options = #iree_gpu.pipeline_options> // OPT-IN-SAME: mma_schedule = #iree_gpu.mma_schedule, #config = #iree_gpu.lowering_config<{workgroup = [128, 128, 0], reduction = [0, 0, 32]}> #pipeline_layout = #hal.pipeline.layout, subgroup_m_count = 2, subgroup_n_count = 2>, - gpu_pipeline_options = #iree_gpu.pipeline_options // enable the 'reorderWorkgroups' pass. + gpu_pipeline_options = #iree_gpu.pipeline_options> // enable the 'reorderWorkgroups' pass. }>} { %cst = arith.constant 0.000000e+00 : f16 %c0 = arith.constant 0 : index @@ -159,10 +159,10 @@ hal.executable public @main_0_dispatch_0 { } // ----- -// Check that applying the `reorder_workgroups_strategy = None` pipeline option disables workgroup reordering. +// Check that applying the `reorder_workgroups_strategy = ` pipeline option disables workgroup reordering. // OPT-OUT: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info +// OPT-OUT-SAME: gpu_pipeline_options = #iree_gpu.pipeline_options> // OPT-OUT-SAME: mma_schedule = #iree_gpu.mma_schedule, #config = #iree_gpu.lowering_config<{workgroup = [128, 128, 0], reduction = [0, 0, 32]}> #pipeline_layout = #hal.pipeline.layout, subgroup_m_count = 2, subgroup_n_count = 2>, - gpu_pipeline_options = #iree_gpu.pipeline_options // Disable the 'reorderWorkgroups' pass. + gpu_pipeline_options = #iree_gpu.pipeline_options> // Disable the 'reorderWorkgroups' pass. }>} { %cst = arith.constant 0.000000e+00 : f16 %c0 = arith.constant 0 : index