Skip to content

Commit

Permalink
[NFC][GPU] Move lowering_config dictionary getter/setters out of attr (
Browse files Browse the repository at this point in the history
…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.
  • Loading branch information
Groverkss authored Nov 25, 2024
1 parent 691b65f commit fdc2d6a
Show file tree
Hide file tree
Showing 11 changed files with 168 additions and 137 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -37,7 +38,7 @@ LogicalResult packToIntrinsic(linalg::LinalgOp linalgOp,
getLoweringConfig<IREE::GPU::LoweringConfigAttr>(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<linalg::ContractionDimensions> contractionDims =
Expand Down Expand Up @@ -78,7 +79,7 @@ struct ConvertToMultiMma final : OpInterfaceRewritePattern<linalg::LinalgOp> {
if (!loweringConfig) {
return failure();
}
IREE::GPU::MmaInterfaceAttr kind = loweringConfig.getMmaKind();
IREE::GPU::MmaInterfaceAttr kind = getMmaKind(loweringConfig);
if (!kind) {
return failure();
}
Expand All @@ -102,7 +103,7 @@ void GPUPackToIntrinsicsPass::runOnOperation() {
if (!loweringConfig) {
return;
}
if (!loweringConfig.getMmaKind()) {
if (!getMmaKind(loweringConfig)) {
return;
}
packingCandidates.push_back(linalgOp);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -93,7 +94,7 @@ struct GPUPromoteMatmulOperandsPass final
}

std::optional<SmallVector<int64_t>> promotedOperands =
loweringConfig.getPromotedOperandList();
getPromotedOperandList(loweringConfig);
if (!promotedOperands) {
return;
}
Expand Down
2 changes: 2 additions & 0 deletions compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ iree_compiler_cc_library(
name = "IREEGPUDialect",
srcs = [
"DerivedConfigUtils.cpp",
"GPULoweringConfigUtils.cpp",
"GPUTileSwizzleUtils.cpp",
"IREEGPUAttrs.cpp",
"IREEGPUDialect.cpp",
Expand All @@ -57,6 +58,7 @@ iree_compiler_cc_library(
],
hdrs = [
"DerivedConfigUtils.h",
"GPULoweringConfigUtils.h",
"GPUTileSwizzleUtils.h",
"IREEGPUAttrs.h",
"IREEGPUDialect.h",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ iree_cc_library(
IREEGPUDialect
HDRS
"DerivedConfigUtils.h"
"GPULoweringConfigUtils.h"
"GPUTileSwizzleUtils.h"
"IREEGPUAttrs.h"
"IREEGPUDialect.h"
Expand All @@ -34,6 +35,7 @@ iree_cc_library(
"IREEGPUOps.h.inc"
SRCS
"DerivedConfigUtils.cpp"
"GPULoweringConfigUtils.cpp"
"GPUTileSwizzleUtils.cpp"
"IREEGPUAttrs.cpp"
"IREEGPUDialect.cpp"
Expand Down
Original file line number Diff line number Diff line change
@@ -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<int64_t> getIntegerVector(ArrayAttr array) {
if (!array || !llvm::all_of(array.getValue(), llvm::IsaPred<IntegerAttr>)) {
return {};
}
return llvm::map_to_vector(array.getValue(), [](Attribute s) -> int64_t {
return cast<IntegerAttr>(s).getInt();
});
}

constexpr StringLiteral kMmaKindName = "mma_kind";

IREE::GPU::MmaInterfaceAttr getMmaKind(LoweringConfigAttr config) {
return config.getAttributes().getAs<IREE::GPU::MmaInterfaceAttr>(
kMmaKindName);
}

void setMmaKind(MLIRContext *context, SmallVectorImpl<NamedAttribute> &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<int64_t> getSubgroupMCount(LoweringConfigAttr config) {
auto subgroup_m_count_attr =
config.getAttributes().getAs<IntegerAttr>(kSubgroupMCountName);
if (!subgroup_m_count_attr) {
return std::nullopt;
}
return subgroup_m_count_attr.getInt();
}

std::optional<int64_t> getSubgroupNCount(LoweringConfigAttr config) {
auto subgroup_n_count_attr =
config.getAttributes().getAs<IntegerAttr>(kSubgroupNCountName);
if (!subgroup_n_count_attr) {
return std::nullopt;
}
return subgroup_n_count_attr.getInt();
}

void setSubgroupMCount(MLIRContext *context,
SmallVectorImpl<NamedAttribute> &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<NamedAttribute> &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<SmallVector<int64_t>>
getPromotedOperandList(LoweringConfigAttr config) {
auto array = config.getAttributes().getAs<ArrayAttr>(kPromoteOperandsName);
if (!array) {
return std::nullopt;
}
return getIntegerVector(array);
}

void setPromotedOperandList(MLIRContext *context,
SmallVectorImpl<NamedAttribute> &attrs,
ArrayRef<int64_t> operands) {
Builder b(context);
attrs.emplace_back(StringAttr::get(context, kPromoteOperandsName),
b.getI64ArrayAttr(operands));
}

} // namespace mlir::iree_compiler::IREE::GPU
Original file line number Diff line number Diff line change
@@ -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<NamedAttribute> &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<int64_t> getSubgroupMCount(LoweringConfigAttr config);
std::optional<int64_t> getSubgroupNCount(LoweringConfigAttr config);
void setSubgroupMCount(MLIRContext *context,
SmallVectorImpl<NamedAttribute> &attrs,
int64_t subgroupMCount);
void setSubgroupNCount(MLIRContext *context,
SmallVectorImpl<NamedAttribute> &attrs,
int64_t subgroupNCount);

/// Helper to retrieve/set a list of operand indices to promote.
std::optional<SmallVector<int64_t>>
getPromotedOperandList(LoweringConfigAttr config);
void setPromotedOperandList(MLIRContext *context,
SmallVectorImpl<NamedAttribute> &attrs,
ArrayRef<int64_t> operands);

} // namespace mlir::iree_compiler::IREE::GPU

#endif // IREE_COMPILER_CODEGEN_DIALECT_GPU_IR_GPULOWERINGCONFIGUTILS_H_
70 changes: 0 additions & 70 deletions compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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<IREE::GPU::MmaInterfaceAttr>(kMmaKindName);
}

void LoweringConfigAttr::setMmaKind(MLIRContext *context,
SmallVectorImpl<NamedAttribute> &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<int64_t> LoweringConfigAttr::getSubgroupMCount() const {
auto subgroup_m_count_attr =
getAttributes().getAs<IntegerAttr>(kSubgroupMCountName);
if (!subgroup_m_count_attr) {
return std::nullopt;
}
return subgroup_m_count_attr.getInt();
}

std::optional<int64_t> LoweringConfigAttr::getSubgroupNCount() const {
auto subgroup_n_count_attr =
getAttributes().getAs<IntegerAttr>(kSubgroupNCountName);
if (!subgroup_n_count_attr) {
return std::nullopt;
}
return subgroup_n_count_attr.getInt();
}

void LoweringConfigAttr::setSubgroupMCount(
MLIRContext *context, SmallVectorImpl<NamedAttribute> &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<NamedAttribute> &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<SmallVector<int64_t>>
LoweringConfigAttr::getPromotedOperandList() const {
auto array = getAttributes().getAs<ArrayAttr>(kPromoteOperandsName);
if (!array) {
return std::nullopt;
}
return getIntegerVector(array);
}

void LoweringConfigAttr::setPromotedOperandList(
MLIRContext *context, SmallVectorImpl<NamedAttribute> &attrs,
ArrayRef<int64_t> operands) {
Builder b(context);
attrs.emplace_back(StringAttr::get(context, kPromoteOperandsName),
b.getI64ArrayAttr(operands));
}

//===----------------------------------------------------------------------===//
// DerivedThreadConfigAttr
//===----------------------------------------------------------------------===//
Expand Down
25 changes: 0 additions & 25 deletions compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td
Original file line number Diff line number Diff line change
Expand Up @@ -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<NamedAttribute> &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<int64_t> getSubgroupMCount() const;
std::optional<int64_t> getSubgroupNCount() const;
static void setSubgroupMCount(MLIRContext *context,
SmallVectorImpl<NamedAttribute> &attrs,
int64_t subgroup_m_count);
static void setSubgroupNCount(MLIRContext *context,
SmallVectorImpl<NamedAttribute> &attrs,
int64_t subgroup_n_count);

/// Helper to retrieve/set a list of operand indices to promote.
std::optional<SmallVector<int64_t>> getPromotedOperandList() const;
static void setPromotedOperandList(MLIRContext *context,
SmallVectorImpl<NamedAttribute> &attrs,
ArrayRef<int64_t> operands);
}];
}

def IREEGPU_DerivedThreadConfig :
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -317,7 +318,7 @@ getMatmulLoweringConfigAndWorkgroupSize(SmallVector<int64_t> 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 =
Expand Down Expand Up @@ -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.
Expand Down
Loading

0 comments on commit fdc2d6a

Please sign in to comment.