forked from iree-org/iree
-
Notifications
You must be signed in to change notification settings - Fork 11
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[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.
- Loading branch information
Showing
11 changed files
with
168 additions
and
137 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
90 changes: 90 additions & 0 deletions
90
compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.cpp
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
40 changes: 40 additions & 0 deletions
40
compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPULoweringConfigUtils.h
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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_ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.