diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPackToIntrinsics.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPackToIntrinsics.cpp index ab1805fb98a2..51d4e672af27 100644 --- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPackToIntrinsics.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPackToIntrinsics.cpp @@ -5,6 +5,7 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h" +#include "iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h" @@ -37,7 +38,7 @@ LogicalResult packToIntrinsic(linalg::LinalgOp linalgOp, getLoweringConfig(linalgOp); assert(loweringConfig && "Packing unconfigured op"); - IREE::GPU::MmaInterfaceAttr kind = loweringConfig.getMmaKind(); + IREE::GPU::MmaInterfaceAttr kind = getMmaKind(loweringConfig); assert(kind && "Packing op without mma kind"); FailureOr contractionDims = @@ -78,7 +79,7 @@ struct ConvertToMultiMma final : OpInterfaceRewritePattern { if (!loweringConfig) { return failure(); } - IREE::GPU::MmaInterfaceAttr kind = loweringConfig.getMmaKind(); + IREE::GPU::MmaInterfaceAttr kind = getMmaKind(loweringConfig); if (!kind) { return failure(); } @@ -102,7 +103,7 @@ void GPUPackToIntrinsicsPass::runOnOperation() { if (!loweringConfig) { return; } - if (!loweringConfig.getMmaKind()) { + if (!getMmaKind(loweringConfig)) { return; } packingCandidates.push_back(linalgOp); diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPromoteMatmulOperands.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPromoteMatmulOperands.cpp index 5e50a956bd82..72e30caf709f 100644 --- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPromoteMatmulOperands.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPromoteMatmulOperands.cpp @@ -6,6 +6,7 @@ #include "iree/compiler/Codegen/Common/GPU/Passes.h" #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h" +#include "iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h" #include "iree/compiler/Codegen/Utils/LinalgOpInfo.h" @@ -93,7 +94,7 @@ struct GPUPromoteMatmulOperandsPass final } std::optional> promotedOperands = - loweringConfig.getPromotedOperandList(); + getPromotedOperandList(loweringConfig); if (!promotedOperands) { return; } 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 97ea44550e30..bdd29befadd7 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel @@ -48,6 +48,7 @@ iree_compiler_cc_library( name = "IREEGPUDialect", srcs = [ "DerivedConfigUtils.cpp", + "GPULoweringConfigUtils.cpp", "GPUTileSwizzleUtils.cpp", "IREEGPUAttrs.cpp", "IREEGPUDialect.cpp", @@ -57,6 +58,7 @@ iree_compiler_cc_library( ], hdrs = [ "DerivedConfigUtils.h", + "GPULoweringConfigUtils.h", "GPUTileSwizzleUtils.h", "IREEGPUAttrs.h", "IREEGPUDialect.h", 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 e0df65823007..adaa901e4dfb 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt @@ -15,6 +15,7 @@ iree_cc_library( IREEGPUDialect HDRS "DerivedConfigUtils.h" + "GPULoweringConfigUtils.h" "GPUTileSwizzleUtils.h" "IREEGPUAttrs.h" "IREEGPUDialect.h" @@ -34,6 +35,7 @@ iree_cc_library( "IREEGPUOps.h.inc" SRCS "DerivedConfigUtils.cpp" + "GPULoweringConfigUtils.cpp" "GPUTileSwizzleUtils.cpp" "IREEGPUAttrs.cpp" "IREEGPUDialect.cpp" diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.cpp new file mode 100644 index 000000000000..23e50be4e31b --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.cpp @@ -0,0 +1,90 @@ +// 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/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h" + +namespace mlir::iree_compiler::IREE::GPU { + +static SmallVector getIntegerVector(ArrayAttr array) { + if (!array || !llvm::all_of(array.getValue(), llvm::IsaPred)) { + return {}; + } + return llvm::map_to_vector(array.getValue(), [](Attribute s) -> int64_t { + return cast(s).getInt(); + }); +} + +constexpr StringLiteral kMmaKindName = "mma_kind"; + +IREE::GPU::MmaInterfaceAttr getMmaKind(LoweringConfigAttr config) { + return config.getAttributes().getAs( + kMmaKindName); +} + +void setMmaKind(MLIRContext *context, SmallVectorImpl &attrs, + IREE::GPU::MmaInterfaceAttr kind) { + attrs.emplace_back(StringAttr::get(context, kMmaKindName), kind); +} + +// TODO: Merge subgroup counts functionality into subgroup tiling level +// lowering, when we have it implemented. +constexpr StringLiteral kSubgroupMCountName = "subgroup_m_count"; +constexpr StringLiteral kSubgroupNCountName = "subgroup_n_count"; + +std::optional getSubgroupMCount(LoweringConfigAttr config) { + auto subgroup_m_count_attr = + config.getAttributes().getAs(kSubgroupMCountName); + if (!subgroup_m_count_attr) { + return std::nullopt; + } + return subgroup_m_count_attr.getInt(); +} + +std::optional getSubgroupNCount(LoweringConfigAttr config) { + auto subgroup_n_count_attr = + config.getAttributes().getAs(kSubgroupNCountName); + if (!subgroup_n_count_attr) { + return std::nullopt; + } + return subgroup_n_count_attr.getInt(); +} + +void setSubgroupMCount(MLIRContext *context, + SmallVectorImpl &attrs, + int64_t subgroup_m_count) { + attrs.emplace_back( + StringAttr::get(context, kSubgroupMCountName), + IntegerAttr::get(IntegerType::get(context, 64), subgroup_m_count)); +} + +void setSubgroupNCount(MLIRContext *context, + SmallVectorImpl &attrs, + int64_t subgroup_n_count) { + attrs.emplace_back( + StringAttr::get(context, kSubgroupNCountName), + IntegerAttr::get(IntegerType::get(context, 64), subgroup_n_count)); +} + +constexpr StringLiteral kPromoteOperandsName = "promote_operands"; + +std::optional> +getPromotedOperandList(LoweringConfigAttr config) { + auto array = config.getAttributes().getAs(kPromoteOperandsName); + if (!array) { + return std::nullopt; + } + return getIntegerVector(array); +} + +void setPromotedOperandList(MLIRContext *context, + SmallVectorImpl &attrs, + ArrayRef operands) { + Builder b(context); + attrs.emplace_back(StringAttr::get(context, kPromoteOperandsName), + b.getI64ArrayAttr(operands)); +} + +} // namespace mlir::iree_compiler::IREE::GPU diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h new file mode 100644 index 000000000000..25240907ba4c --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h @@ -0,0 +1,40 @@ +// 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_CODEGEN_DIALECT_GPU_IR_GPULOWERINGCONFIGUTILS_H_ +#define IREE_COMPILER_CODEGEN_DIALECT_GPU_IR_GPULOWERINGCONFIGUTILS_H_ + +#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h" + +namespace mlir::iree_compiler::IREE::GPU { + +/// Helper to retrieve/set a target mma intrinsic. +MmaInterfaceAttr getMmaKind(LoweringConfigAttr config); +void setMmaKind(MLIRContext *context, SmallVectorImpl &attrs, + MmaInterfaceAttr kind); + +// TODO: Merge subgroup counts functionality into subgroup tiling level +// lowering, when we have it implemented. +/// Helper to retrieve/set a target subgroup M/N counts. +std::optional getSubgroupMCount(LoweringConfigAttr config); +std::optional getSubgroupNCount(LoweringConfigAttr config); +void setSubgroupMCount(MLIRContext *context, + SmallVectorImpl &attrs, + int64_t subgroupMCount); +void setSubgroupNCount(MLIRContext *context, + SmallVectorImpl &attrs, + int64_t subgroupNCount); + +/// Helper to retrieve/set a list of operand indices to promote. +std::optional> +getPromotedOperandList(LoweringConfigAttr config); +void setPromotedOperandList(MLIRContext *context, + SmallVectorImpl &attrs, + ArrayRef operands); + +} // namespace mlir::iree_compiler::IREE::GPU + +#endif // IREE_COMPILER_CODEGEN_DIALECT_GPU_IR_GPULOWERINGCONFIGUTILS_H_ 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 2fae167aa987..eaa3f7249c05 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp @@ -1174,76 +1174,6 @@ bool LoweringConfigAttr::hasWorkgroupTilingLevel() const { return !getWorkgroupTileSizes().empty(); } -constexpr StringLiteral kMmaKindName = "mma_kind"; - -IREE::GPU::MmaInterfaceAttr LoweringConfigAttr::getMmaKind() const { - return getAttributes().getAs(kMmaKindName); -} - -void LoweringConfigAttr::setMmaKind(MLIRContext *context, - SmallVectorImpl &attrs, - IREE::GPU::MmaInterfaceAttr kind) { - attrs.emplace_back(StringAttr::get(context, kMmaKindName), kind); -} - -// TODO: Merge subgroup counts functionality into subgroup tiling level -// lowering, when we have it implemented. -constexpr StringLiteral kSubgroupMCountName = "subgroup_m_count"; -constexpr StringLiteral kSubgroupNCountName = "subgroup_n_count"; - -std::optional LoweringConfigAttr::getSubgroupMCount() const { - auto subgroup_m_count_attr = - getAttributes().getAs(kSubgroupMCountName); - if (!subgroup_m_count_attr) { - return std::nullopt; - } - return subgroup_m_count_attr.getInt(); -} - -std::optional LoweringConfigAttr::getSubgroupNCount() const { - auto subgroup_n_count_attr = - getAttributes().getAs(kSubgroupNCountName); - if (!subgroup_n_count_attr) { - return std::nullopt; - } - return subgroup_n_count_attr.getInt(); -} - -void LoweringConfigAttr::setSubgroupMCount( - MLIRContext *context, SmallVectorImpl &attrs, - int64_t subgroup_m_count) { - attrs.emplace_back( - StringAttr::get(context, kSubgroupMCountName), - IntegerAttr::get(IntegerType::get(context, 64), subgroup_m_count)); -} - -void LoweringConfigAttr::setSubgroupNCount( - MLIRContext *context, SmallVectorImpl &attrs, - int64_t subgroup_n_count) { - attrs.emplace_back( - StringAttr::get(context, kSubgroupNCountName), - IntegerAttr::get(IntegerType::get(context, 64), subgroup_n_count)); -} - -constexpr StringLiteral kPromoteOperandsName = "promote_operands"; - -std::optional> -LoweringConfigAttr::getPromotedOperandList() const { - auto array = getAttributes().getAs(kPromoteOperandsName); - if (!array) { - return std::nullopt; - } - return getIntegerVector(array); -} - -void LoweringConfigAttr::setPromotedOperandList( - MLIRContext *context, SmallVectorImpl &attrs, - ArrayRef operands) { - Builder b(context); - attrs.emplace_back(StringAttr::get(context, kPromoteOperandsName), - b.getI64ArrayAttr(operands)); -} - //===----------------------------------------------------------------------===// // DerivedThreadConfigAttr //===----------------------------------------------------------------------===// 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 9a2a4d10e7cd..185caa5ddd4b 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td @@ -56,31 +56,6 @@ def IREEGPU_LoweringConfigAttr : AttrParameter<"DictionaryAttr", "The configured fields, including tiling levels">:$attributes ); - let extraClassDeclaration = [{ - /// Helper to retrieve/set a target mma intrinsic. - ::mlir::iree_compiler::IREE::GPU::MmaInterfaceAttr getMmaKind() const; - static void setMmaKind(MLIRContext *context, - SmallVectorImpl &attrs, - ::mlir::iree_compiler::IREE::GPU::MmaInterfaceAttr kind); - - // TODO: Merge subgroup counts functionality into subgroup tiling level - // lowering, when we have it implemented. - /// Helper to retrieve/set a target subgroup M/N counts. - std::optional getSubgroupMCount() const; - std::optional getSubgroupNCount() const; - static void setSubgroupMCount(MLIRContext *context, - SmallVectorImpl &attrs, - int64_t subgroup_m_count); - static void setSubgroupNCount(MLIRContext *context, - SmallVectorImpl &attrs, - int64_t subgroup_n_count); - - /// Helper to retrieve/set a list of operand indices to promote. - std::optional> getPromotedOperandList() const; - static void setPromotedOperandList(MLIRContext *context, - SmallVectorImpl &attrs, - ArrayRef operands); - }]; } def IREEGPU_DerivedThreadConfig : diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp index 8aaabfa07053..b10a567f1ca2 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp @@ -8,6 +8,7 @@ #include "iree/compiler/Codegen/Common/GPU/GPUHeuristics.h" #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h" +#include "iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h" @@ -81,7 +82,7 @@ setDataTiledMultiMmaLoweringConfig(IREE::GPU::TargetAttr target, attrs.emplace_back(b.getStringAttr("reduction"), b.getI64ArrayAttr(reductionTileSizes)); // Promote operands to use shared memory for LHS and RHS. - GPU::LoweringConfigAttr::setPromotedOperandList(context, attrs, {0, 1}); + GPU::setPromotedOperandList(context, attrs, {0, 1}); auto configDict = b.getDictionaryAttr(attrs); auto loweringConfig = IREE::GPU::LoweringConfigAttr::get(context, configDict); @@ -317,7 +318,7 @@ getMatmulLoweringConfigAndWorkgroupSize(SmallVector bounds, attrs.emplace_back(StringAttr::get(context, "subgroup"), b.getI64ArrayAttr(subgroupTileSizes)); attrs.emplace_back(StringAttr::get(context, "mma_kind"), mmaKind); - GPU::LoweringConfigAttr::setPromotedOperandList(context, attrs, {0, 1}); + GPU::setPromotedOperandList(context, attrs, {0, 1}); auto configDict = DictionaryAttr::get(context, attrs); auto loweringConfig = IREE::GPU::LoweringConfigAttr::get(context, configDict); int64_t flatWorkgroupSize = @@ -657,7 +658,7 @@ LogicalResult setTileAndFuseLoweringConfig(IREE::GPU::TargetAttr target, b.getI64ArrayAttr(threadTileSizes)); if (isNonMatvecContraction(linalgOp)) { - GPU::LoweringConfigAttr::setPromotedOperandList(context, attrs, {0, 1}); + GPU::setPromotedOperandList(context, attrs, {0, 1}); } // Heuristic value chosen to limit maximum vector sizes when tiling below. diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp index 7fc5b80f9fc7..d63fd2d5d258 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp +++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp @@ -12,6 +12,7 @@ #include "iree/compiler/Codegen/Common/GPU/GPUHeuristics.h" #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h" +#include "iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.h" #include "iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.h" @@ -421,13 +422,10 @@ setConvolutionVectorDistributionConfig(IREE::GPU::TargetAttr target, b.getI64ArrayAttr(workgroupTileSizes)); attrs.emplace_back(StringAttr::get(context, "reduction"), b.getI64ArrayAttr(reductionTileSizes)); - IREE::GPU::LoweringConfigAttr::setPromotedOperandList(context, attrs, {0, 1}); - IREE::GPU::LoweringConfigAttr::setMmaKind(context, attrs, - mmaKinds[schedule->index]); - IREE::GPU::LoweringConfigAttr::setSubgroupMCount( - context, attrs, schedule->mSubgroupCounts[0]); - IREE::GPU::LoweringConfigAttr::setSubgroupNCount( - context, attrs, schedule->nSubgroupCounts[0]); + IREE::GPU::setPromotedOperandList(context, attrs, {0, 1}); + IREE::GPU::setMmaKind(context, attrs, mmaKinds[schedule->index]); + IREE::GPU::setSubgroupMCount(context, attrs, schedule->mSubgroupCounts[0]); + IREE::GPU::setSubgroupNCount(context, attrs, schedule->nSubgroupCounts[0]); auto configDict = DictionaryAttr::get(context, attrs); auto loweringConfig = IREE::GPU::LoweringConfigAttr::get(context, configDict); @@ -687,13 +685,10 @@ setMatmulVectorDistributionConfig(IREE::GPU::TargetAttr target, b.getI64ArrayAttr(workgroupTileSizes)); attrs.emplace_back(StringAttr::get(context, "reduction"), b.getI64ArrayAttr(reductionTileSizes)); - IREE::GPU::LoweringConfigAttr::setPromotedOperandList(context, attrs, {0, 1}); - IREE::GPU::LoweringConfigAttr::setMmaKind(context, attrs, - mmaKinds[schedule->index]); - IREE::GPU::LoweringConfigAttr::setSubgroupMCount( - context, attrs, schedule->mSubgroupCounts[0]); - IREE::GPU::LoweringConfigAttr::setSubgroupNCount( - context, attrs, schedule->nSubgroupCounts[0]); + IREE::GPU::setPromotedOperandList(context, attrs, {0, 1}); + IREE::GPU::setMmaKind(context, attrs, mmaKinds[schedule->index]); + IREE::GPU::setSubgroupMCount(context, attrs, schedule->mSubgroupCounts[0]); + IREE::GPU::setSubgroupNCount(context, attrs, schedule->nSubgroupCounts[0]); auto configDict = DictionaryAttr::get(context, attrs); auto loweringConfig = IREE::GPU::LoweringConfigAttr::get(context, configDict); @@ -905,8 +900,7 @@ setAttentionVectorDistributionConfig(IREE::GPU::TargetAttr target, b.getI64ArrayAttr(workgroupTileSizes)); attrs.emplace_back(StringAttr::get(context, "reduction"), b.getI64ArrayAttr(reductionTileSizes)); - IREE::GPU::LoweringConfigAttr::setPromotedOperandList(context, attrs, - {0, 1, 2}); + IREE::GPU::setPromotedOperandList(context, attrs, {0, 1, 2}); SmallVector qkConfig; SmallVector pvConfig; @@ -921,22 +915,16 @@ setAttentionVectorDistributionConfig(IREE::GPU::TargetAttr target, // Configuring for qk matmul. // subgroup_n count for qk matmul is always 1, since we do not tile K1. - IREE::GPU::LoweringConfigAttr::setPromotedOperandList(context, qkConfig, - {0, 1}); - IREE::GPU::LoweringConfigAttr::setMmaKind(context, qkConfig, - mmaKinds[schedule->index]); - IREE::GPU::LoweringConfigAttr::setSubgroupMCount( - context, qkConfig, schedule->mSubgroupCounts[0]); - IREE::GPU::LoweringConfigAttr::setSubgroupNCount(context, qkConfig, 1); + IREE::GPU::setPromotedOperandList(context, qkConfig, {0, 1}); + IREE::GPU::setMmaKind(context, qkConfig, mmaKinds[schedule->index]); + IREE::GPU::setSubgroupMCount(context, qkConfig, schedule->mSubgroupCounts[0]); + IREE::GPU::setSubgroupNCount(context, qkConfig, 1); // Configuring for pv matmul. - IREE::GPU::LoweringConfigAttr::setPromotedOperandList(context, pvConfig, {1}); - IREE::GPU::LoweringConfigAttr::setMmaKind(context, pvConfig, - mmaKinds[schedule->index]); - IREE::GPU::LoweringConfigAttr::setSubgroupMCount( - context, pvConfig, schedule->mSubgroupCounts[0]); - IREE::GPU::LoweringConfigAttr::setSubgroupNCount( - context, pvConfig, schedule->nSubgroupCounts[0]); + IREE::GPU::setPromotedOperandList(context, pvConfig, {1}); + IREE::GPU::setMmaKind(context, pvConfig, mmaKinds[schedule->index]); + IREE::GPU::setSubgroupMCount(context, pvConfig, schedule->mSubgroupCounts[0]); + IREE::GPU::setSubgroupNCount(context, pvConfig, schedule->nSubgroupCounts[0]); SmallVector qkAttrs; SmallVector pvAttrs; diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp index 98b2fb4f817b..41acfa3a4dfe 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp +++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureTensorLayouts.cpp @@ -5,6 +5,7 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h" +#include "iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h" #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h" #include "iree/compiler/Codegen/LLVMGPU/Passes.h" @@ -37,7 +38,7 @@ static SmallVector getPromotedOperands(Operation *op) { } std::optional> promoteConfig = - config.getPromotedOperandList(); + getPromotedOperandList(config); if (!promoteConfig) { return promotedOperands; } @@ -53,7 +54,7 @@ static IREE::GPU::MmaInterfaceAttr getIntrinsic(Operation *op) { auto config = getLoweringConfig(op); assert(config && "Cannot find intrinsic from unconfigured op."); - IREE::GPU::MmaInterfaceAttr mmaIntrinsic = config.getMmaKind(); + IREE::GPU::MmaInterfaceAttr mmaIntrinsic = getMmaKind(config); assert(mmaIntrinsic && "Cannot find intrinsic in lowering config."); return mmaIntrinsic; } @@ -62,14 +63,14 @@ static int64_t getSubgroupMCount(Operation *op) { auto config = getLoweringConfig(op); assert(config && "Cannot find intrinsic from unconfigured op."); - return *config.getSubgroupMCount(); + return getSubgroupMCount(config).value(); } static int64_t getSubgroupNCount(Operation *op) { auto config = getLoweringConfig(op); assert(config && "Cannot find intrinsic from unconfigured op."); - return *config.getSubgroupNCount(); + return getSubgroupNCount(config).value(); } /// Gets a unit vector of the given rank, but fills in the given dimensions @@ -895,7 +896,7 @@ struct LLVMGPUConfigureTensorLayoutsPass final workgroupSize, rewriter); }) .Case([&](IREE::GPU::LoweringConfigAttr config) { - if (config.getMmaKind()) { + if (getMmaKind(config)) { return setIntrinsicLoweringConfigLayout( config, candidate, workgroupSize, rewriter); }