Skip to content

Commit fdc2d6a

Browse files
GroverkssGroverkss
authored
[NFC][GPU] Move lowering_config dictionary getter/setters out of attr (iree-org#19268)
These attribute getter/setters are not part of the attr definition and are based on compilation pipelines. Move them to a seperate file so we don't keep accumulating them in the attr definition.
1 parent 691b65f commit fdc2d6a

File tree

11 files changed

+168
-137
lines changed

11 files changed

+168
-137
lines changed

‎compiler/src/iree/compiler/Codegen/Common/GPU/GPUPackToIntrinsics.cpp‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
66

77
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
8+
#include "iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h"
89
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
910
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
1011
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h"
@@ -37,7 +38,7 @@ LogicalResult packToIntrinsic(linalg::LinalgOp linalgOp,
3738
getLoweringConfig<IREE::GPU::LoweringConfigAttr>(linalgOp);
3839
assert(loweringConfig && "Packing unconfigured op");
3940

40-
IREE::GPU::MmaInterfaceAttr kind = loweringConfig.getMmaKind();
41+
IREE::GPU::MmaInterfaceAttr kind = getMmaKind(loweringConfig);
4142
assert(kind && "Packing op without mma kind");
4243

4344
FailureOr<linalg::ContractionDimensions> contractionDims =
@@ -78,7 +79,7 @@ struct ConvertToMultiMma final : OpInterfaceRewritePattern<linalg::LinalgOp> {
7879
if (!loweringConfig) {
7980
return failure();
8081
}
81-
IREE::GPU::MmaInterfaceAttr kind = loweringConfig.getMmaKind();
82+
IREE::GPU::MmaInterfaceAttr kind = getMmaKind(loweringConfig);
8283
if (!kind) {
8384
return failure();
8485
}
@@ -102,7 +103,7 @@ void GPUPackToIntrinsicsPass::runOnOperation() {
102103
if (!loweringConfig) {
103104
return;
104105
}
105-
if (!loweringConfig.getMmaKind()) {
106+
if (!getMmaKind(loweringConfig)) {
106107
return;
107108
}
108109
packingCandidates.push_back(linalgOp);

‎compiler/src/iree/compiler/Codegen/Common/GPU/GPUPromoteMatmulOperands.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
#include "iree/compiler/Codegen/Common/GPU/Passes.h"
88
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
9+
#include "iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h"
910
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
1011
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
1112
#include "iree/compiler/Codegen/Utils/LinalgOpInfo.h"
@@ -93,7 +94,7 @@ struct GPUPromoteMatmulOperandsPass final
9394
}
9495

9596
std::optional<SmallVector<int64_t>> promotedOperands =
96-
loweringConfig.getPromotedOperandList();
97+
getPromotedOperandList(loweringConfig);
9798
if (!promotedOperands) {
9899
return;
99100
}

‎compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ iree_compiler_cc_library(
4848
name = "IREEGPUDialect",
4949
srcs = [
5050
"DerivedConfigUtils.cpp",
51+
"GPULoweringConfigUtils.cpp",
5152
"GPUTileSwizzleUtils.cpp",
5253
"IREEGPUAttrs.cpp",
5354
"IREEGPUDialect.cpp",
@@ -57,6 +58,7 @@ iree_compiler_cc_library(
5758
],
5859
hdrs = [
5960
"DerivedConfigUtils.h",
61+
"GPULoweringConfigUtils.h",
6062
"GPUTileSwizzleUtils.h",
6163
"IREEGPUAttrs.h",
6264
"IREEGPUDialect.h",

‎compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ iree_cc_library(
1515
IREEGPUDialect
1616
HDRS
1717
"DerivedConfigUtils.h"
18+
"GPULoweringConfigUtils.h"
1819
"GPUTileSwizzleUtils.h"
1920
"IREEGPUAttrs.h"
2021
"IREEGPUDialect.h"
@@ -34,6 +35,7 @@ iree_cc_library(
3435
"IREEGPUOps.h.inc"
3536
SRCS
3637
"DerivedConfigUtils.cpp"
38+
"GPULoweringConfigUtils.cpp"
3739
"GPUTileSwizzleUtils.cpp"
3840
"IREEGPUAttrs.cpp"
3941
"IREEGPUDialect.cpp"

‎compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.cpp‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
// Copyright 2024 The IREE Authors
2+
//
3+
// Licensed under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
7+
#include "iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h"
8+
9+
namespace mlir::iree_compiler::IREE::GPU {
10+
11+
static SmallVector<int64_t> getIntegerVector(ArrayAttr array) {
12+
if (!array || !llvm::all_of(array.getValue(), llvm::IsaPred<IntegerAttr>)) {
13+
return {};
14+
}
15+
return llvm::map_to_vector(array.getValue(), [](Attribute s) -> int64_t {
16+
return cast<IntegerAttr>(s).getInt();
17+
});
18+
}
19+
20+
constexpr StringLiteral kMmaKindName = "mma_kind";
21+
22+
IREE::GPU::MmaInterfaceAttr getMmaKind(LoweringConfigAttr config) {
23+
return config.getAttributes().getAs<IREE::GPU::MmaInterfaceAttr>(
24+
kMmaKindName);
25+
}
26+
27+
void setMmaKind(MLIRContext *context, SmallVectorImpl<NamedAttribute> &attrs,
28+
IREE::GPU::MmaInterfaceAttr kind) {
29+
attrs.emplace_back(StringAttr::get(context, kMmaKindName), kind);
30+
}
31+
32+
// TODO: Merge subgroup counts functionality into subgroup tiling level
33+
// lowering, when we have it implemented.
34+
constexpr StringLiteral kSubgroupMCountName = "subgroup_m_count";
35+
constexpr StringLiteral kSubgroupNCountName = "subgroup_n_count";
36+
37+
std::optional<int64_t> getSubgroupMCount(LoweringConfigAttr config) {
38+
auto subgroup_m_count_attr =
39+
config.getAttributes().getAs<IntegerAttr>(kSubgroupMCountName);
40+
if (!subgroup_m_count_attr) {
41+
return std::nullopt;
42+
}
43+
return subgroup_m_count_attr.getInt();
44+
}
45+
46+
std::optional<int64_t> getSubgroupNCount(LoweringConfigAttr config) {
47+
auto subgroup_n_count_attr =
48+
config.getAttributes().getAs<IntegerAttr>(kSubgroupNCountName);
49+
if (!subgroup_n_count_attr) {
50+
return std::nullopt;
51+
}
52+
return subgroup_n_count_attr.getInt();
53+
}
54+
55+
void setSubgroupMCount(MLIRContext *context,
56+
SmallVectorImpl<NamedAttribute> &attrs,
57+
int64_t subgroup_m_count) {
58+
attrs.emplace_back(
59+
StringAttr::get(context, kSubgroupMCountName),
60+
IntegerAttr::get(IntegerType::get(context, 64), subgroup_m_count));
61+
}
62+
63+
void setSubgroupNCount(MLIRContext *context,
64+
SmallVectorImpl<NamedAttribute> &attrs,
65+
int64_t subgroup_n_count) {
66+
attrs.emplace_back(
67+
StringAttr::get(context, kSubgroupNCountName),
68+
IntegerAttr::get(IntegerType::get(context, 64), subgroup_n_count));
69+
}
70+
71+
constexpr StringLiteral kPromoteOperandsName = "promote_operands";
72+
73+
std::optional<SmallVector<int64_t>>
74+
getPromotedOperandList(LoweringConfigAttr config) {
75+
auto array = config.getAttributes().getAs<ArrayAttr>(kPromoteOperandsName);
76+
if (!array) {
77+
return std::nullopt;
78+
}
79+
return getIntegerVector(array);
80+
}
81+
82+
void setPromotedOperandList(MLIRContext *context,
83+
SmallVectorImpl<NamedAttribute> &attrs,
84+
ArrayRef<int64_t> operands) {
85+
Builder b(context);
86+
attrs.emplace_back(StringAttr::get(context, kPromoteOperandsName),
87+
b.getI64ArrayAttr(operands));
88+
}
89+
90+
} // namespace mlir::iree_compiler::IREE::GPU

‎compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
// Copyright 2024 The IREE Authors
2+
//
3+
// Licensed under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
7+
#ifndef IREE_COMPILER_CODEGEN_DIALECT_GPU_IR_GPULOWERINGCONFIGUTILS_H_
8+
#define IREE_COMPILER_CODEGEN_DIALECT_GPU_IR_GPULOWERINGCONFIGUTILS_H_
9+
10+
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
11+
12+
namespace mlir::iree_compiler::IREE::GPU {
13+
14+
/// Helper to retrieve/set a target mma intrinsic.
15+
MmaInterfaceAttr getMmaKind(LoweringConfigAttr config);
16+
void setMmaKind(MLIRContext *context, SmallVectorImpl<NamedAttribute> &attrs,
17+
MmaInterfaceAttr kind);
18+
19+
// TODO: Merge subgroup counts functionality into subgroup tiling level
20+
// lowering, when we have it implemented.
21+
/// Helper to retrieve/set a target subgroup M/N counts.
22+
std::optional<int64_t> getSubgroupMCount(LoweringConfigAttr config);
23+
std::optional<int64_t> getSubgroupNCount(LoweringConfigAttr config);
24+
void setSubgroupMCount(MLIRContext *context,
25+
SmallVectorImpl<NamedAttribute> &attrs,
26+
int64_t subgroupMCount);
27+
void setSubgroupNCount(MLIRContext *context,
28+
SmallVectorImpl<NamedAttribute> &attrs,
29+
int64_t subgroupNCount);
30+
31+
/// Helper to retrieve/set a list of operand indices to promote.
32+
std::optional<SmallVector<int64_t>>
33+
getPromotedOperandList(LoweringConfigAttr config);
34+
void setPromotedOperandList(MLIRContext *context,
35+
SmallVectorImpl<NamedAttribute> &attrs,
36+
ArrayRef<int64_t> operands);
37+
38+
} // namespace mlir::iree_compiler::IREE::GPU
39+
40+
#endif // IREE_COMPILER_CODEGEN_DIALECT_GPU_IR_GPULOWERINGCONFIGUTILS_H_

‎compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp‎

Lines changed: 0 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -1174,76 +1174,6 @@ bool LoweringConfigAttr::hasWorkgroupTilingLevel() const {
11741174
return !getWorkgroupTileSizes().empty();
11751175
}
11761176

1177-
constexpr StringLiteral kMmaKindName = "mma_kind";
1178-
1179-
IREE::GPU::MmaInterfaceAttr LoweringConfigAttr::getMmaKind() const {
1180-
return getAttributes().getAs<IREE::GPU::MmaInterfaceAttr>(kMmaKindName);
1181-
}
1182-
1183-
void LoweringConfigAttr::setMmaKind(MLIRContext *context,
1184-
SmallVectorImpl<NamedAttribute> &attrs,
1185-
IREE::GPU::MmaInterfaceAttr kind) {
1186-
attrs.emplace_back(StringAttr::get(context, kMmaKindName), kind);
1187-
}
1188-
1189-
// TODO: Merge subgroup counts functionality into subgroup tiling level
1190-
// lowering, when we have it implemented.
1191-
constexpr StringLiteral kSubgroupMCountName = "subgroup_m_count";
1192-
constexpr StringLiteral kSubgroupNCountName = "subgroup_n_count";
1193-
1194-
std::optional<int64_t> LoweringConfigAttr::getSubgroupMCount() const {
1195-
auto subgroup_m_count_attr =
1196-
getAttributes().getAs<IntegerAttr>(kSubgroupMCountName);
1197-
if (!subgroup_m_count_attr) {
1198-
return std::nullopt;
1199-
}
1200-
return subgroup_m_count_attr.getInt();
1201-
}
1202-
1203-
std::optional<int64_t> LoweringConfigAttr::getSubgroupNCount() const {
1204-
auto subgroup_n_count_attr =
1205-
getAttributes().getAs<IntegerAttr>(kSubgroupNCountName);
1206-
if (!subgroup_n_count_attr) {
1207-
return std::nullopt;
1208-
}
1209-
return subgroup_n_count_attr.getInt();
1210-
}
1211-
1212-
void LoweringConfigAttr::setSubgroupMCount(
1213-
MLIRContext *context, SmallVectorImpl<NamedAttribute> &attrs,
1214-
int64_t subgroup_m_count) {
1215-
attrs.emplace_back(
1216-
StringAttr::get(context, kSubgroupMCountName),
1217-
IntegerAttr::get(IntegerType::get(context, 64), subgroup_m_count));
1218-
}
1219-
1220-
void LoweringConfigAttr::setSubgroupNCount(
1221-
MLIRContext *context, SmallVectorImpl<NamedAttribute> &attrs,
1222-
int64_t subgroup_n_count) {
1223-
attrs.emplace_back(
1224-
StringAttr::get(context, kSubgroupNCountName),
1225-
IntegerAttr::get(IntegerType::get(context, 64), subgroup_n_count));
1226-
}
1227-
1228-
constexpr StringLiteral kPromoteOperandsName = "promote_operands";
1229-
1230-
std::optional<SmallVector<int64_t>>
1231-
LoweringConfigAttr::getPromotedOperandList() const {
1232-
auto array = getAttributes().getAs<ArrayAttr>(kPromoteOperandsName);
1233-
if (!array) {
1234-
return std::nullopt;
1235-
}
1236-
return getIntegerVector(array);
1237-
}
1238-
1239-
void LoweringConfigAttr::setPromotedOperandList(
1240-
MLIRContext *context, SmallVectorImpl<NamedAttribute> &attrs,
1241-
ArrayRef<int64_t> operands) {
1242-
Builder b(context);
1243-
attrs.emplace_back(StringAttr::get(context, kPromoteOperandsName),
1244-
b.getI64ArrayAttr(operands));
1245-
}
1246-
12471177
//===----------------------------------------------------------------------===//
12481178
// DerivedThreadConfigAttr
12491179
//===----------------------------------------------------------------------===//

‎compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td‎

Lines changed: 0 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -56,31 +56,6 @@ def IREEGPU_LoweringConfigAttr :
5656
AttrParameter<"DictionaryAttr",
5757
"The configured fields, including tiling levels">:$attributes
5858
);
59-
let extraClassDeclaration = [{
60-
/// Helper to retrieve/set a target mma intrinsic.
61-
::mlir::iree_compiler::IREE::GPU::MmaInterfaceAttr getMmaKind() const;
62-
static void setMmaKind(MLIRContext *context,
63-
SmallVectorImpl<NamedAttribute> &attrs,
64-
::mlir::iree_compiler::IREE::GPU::MmaInterfaceAttr kind);
65-
66-
// TODO: Merge subgroup counts functionality into subgroup tiling level
67-
// lowering, when we have it implemented.
68-
/// Helper to retrieve/set a target subgroup M/N counts.
69-
std::optional<int64_t> getSubgroupMCount() const;
70-
std::optional<int64_t> getSubgroupNCount() const;
71-
static void setSubgroupMCount(MLIRContext *context,
72-
SmallVectorImpl<NamedAttribute> &attrs,
73-
int64_t subgroup_m_count);
74-
static void setSubgroupNCount(MLIRContext *context,
75-
SmallVectorImpl<NamedAttribute> &attrs,
76-
int64_t subgroup_n_count);
77-
78-
/// Helper to retrieve/set a list of operand indices to promote.
79-
std::optional<SmallVector<int64_t>> getPromotedOperandList() const;
80-
static void setPromotedOperandList(MLIRContext *context,
81-
SmallVectorImpl<NamedAttribute> &attrs,
82-
ArrayRef<int64_t> operands);
83-
}];
8459
}
8560

8661
def IREEGPU_DerivedThreadConfig :

‎compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
#include "iree/compiler/Codegen/Common/GPU/GPUHeuristics.h"
1010
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
11+
#include "iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h"
1112
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
1213
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.h"
1314
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h"
@@ -81,7 +82,7 @@ setDataTiledMultiMmaLoweringConfig(IREE::GPU::TargetAttr target,
8182
attrs.emplace_back(b.getStringAttr("reduction"),
8283
b.getI64ArrayAttr(reductionTileSizes));
8384
// Promote operands to use shared memory for LHS and RHS.
84-
GPU::LoweringConfigAttr::setPromotedOperandList(context, attrs, {0, 1});
85+
GPU::setPromotedOperandList(context, attrs, {0, 1});
8586
auto configDict = b.getDictionaryAttr(attrs);
8687
auto loweringConfig = IREE::GPU::LoweringConfigAttr::get(context, configDict);
8788

@@ -317,7 +318,7 @@ getMatmulLoweringConfigAndWorkgroupSize(SmallVector<int64_t> bounds,
317318
attrs.emplace_back(StringAttr::get(context, "subgroup"),
318319
b.getI64ArrayAttr(subgroupTileSizes));
319320
attrs.emplace_back(StringAttr::get(context, "mma_kind"), mmaKind);
320-
GPU::LoweringConfigAttr::setPromotedOperandList(context, attrs, {0, 1});
321+
GPU::setPromotedOperandList(context, attrs, {0, 1});
321322
auto configDict = DictionaryAttr::get(context, attrs);
322323
auto loweringConfig = IREE::GPU::LoweringConfigAttr::get(context, configDict);
323324
int64_t flatWorkgroupSize =
@@ -657,7 +658,7 @@ LogicalResult setTileAndFuseLoweringConfig(IREE::GPU::TargetAttr target,
657658
b.getI64ArrayAttr(threadTileSizes));
658659

659660
if (isNonMatvecContraction(linalgOp)) {
660-
GPU::LoweringConfigAttr::setPromotedOperandList(context, attrs, {0, 1});
661+
GPU::setPromotedOperandList(context, attrs, {0, 1});
661662
}
662663

663664
// Heuristic value chosen to limit maximum vector sizes when tiling below.

0 commit comments

Comments
 (0)