Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[tuner]: Add a utility function to query supported MMA intrinsics #19124

Merged
merged 6 commits into from
Nov 14, 2024
Merged

[tuner]: Add a utility function to query supported MMA intrinsics #19124

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
8d86018
[tuner]:Add a utility function to query supported MMA intrinsics
bangtianliu Nov 12, 2024
7c044d4
[tuner] address comments
bangtianliu Nov 13, 2024
9c43699
[tuner] out of order check
bangtianliu Nov 13, 2024
69cedb8
[tuner] trim down test file and format code
bangtianliu Nov 14, 2024
b07a625
[tuner] format the code again
bangtianliu Nov 14, 2024
77aa485
[tuner] update the code comment
bangtianliu Nov 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 5 additions & 8 deletions compiler/src/iree/compiler/Codegen/LLVMGPU/TestLLVMGPUQueryMMAPass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,16 +23,13 @@ struct TestLLVMGPUQueryMMAPass final
: impl::TestLLVMGPUQueryMMAPassBase<TestLLVMGPUQueryMMAPass> {
void runOnOperation() override {
ModuleOp moduleOp = getOperation();
llvm::SmallDenseMap<Operation *, SmallVector<IREE::GPU::MMAIntrinsic>>
llvm::SmallDenseMap<IREE::HAL::ExecutableVariantOp,
SmallVector<IREE::GPU::MMAIntrinsic>>
mmaMap = queryMMAIntrinsics(moduleOp);
for (const auto &[op, mmaAttrs] : mmaMap) {
if (auto variantOp = llvm::dyn_cast<IREE::HAL::ExecutableVariantOp>(op)) {
llvm::outs() << "Executable Variant Name: " << variantOp.getName()
<< "\n";
} else {
llvm::outs() << "Executable Variant Name: " << "Unnamed Operation"
<< "\n";
}
llvm::outs() << "Executable Variant Name: "
<< cast<IREE::HAL::ExecutableVariantOp>(*op).getName()
bangtianliu marked this conversation as resolved.
Show resolved Hide resolved
<< "\n";
llvm::outs() << "MMA Intrinsics: ";
llvm::interleave(mmaAttrs, llvm::outs(), " ");
llvm::outs() << "\n";
Expand Down
9 changes: 5 additions & 4 deletions compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
#include "llvm/Support/ErrorHandling.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinTypes.h"
Expand Down Expand Up @@ -1029,17 +1028,19 @@ std::optional<int> getGPUSubgroupSize(mlir::FunctionOpInterface func) {
return std::nullopt;
}

llvm::SmallDenseMap<Operation *, SmallVector<IREE::GPU::MMAIntrinsic>>
llvm::SmallDenseMap<IREE::HAL::ExecutableVariantOp,
SmallVector<IREE::GPU::MMAIntrinsic>>
queryMMAIntrinsics(mlir::ModuleOp moduleOp) {
llvm::SmallDenseMap<Operation *, SmallVector<IREE::GPU::MMAIntrinsic>>
llvm::SmallDenseMap<IREE::HAL::ExecutableVariantOp,
SmallVector<IREE::GPU::MMAIntrinsic>>
mmaAttributesMap;
moduleOp.walk([&](IREE::HAL::ExecutableVariantOp executableOp) {
if (IREE::GPU::TargetAttr target = getGPUTargetAttr(executableOp)) {
auto mmaIntrinsics = llvm::map_to_vector(
target.getWgp().getMma(), [](IREE::GPU::MMAAttr attr) {
return attr.getIntrinsic().getValue();
});
mmaAttributesMap[executableOp] = mmaIntrinsics;
mmaAttributesMap[executableOp] = std::move(mmaIntrinsics);
}
});
return mmaAttributesMap;
Expand Down
4 changes: 3 additions & 1 deletion compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#define IREE_COMPILER_CODEGEN_UTILS_GPUUTILS_H_

#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
Expand Down Expand Up @@ -210,7 +211,8 @@ std::optional<int> getGPUSubgroupSize(mlir::FunctionOpInterface func);
/// GPU target descriptions in `moduleOp`. Each entry in the map associates
/// an `Operation*` ( an `IREE::HAL::ExecutableVariantOp`) with a
/// vector of `IREE::GPU::MMAIntrinsic` attributes.
bangtianliu marked this conversation as resolved.
Show resolved Hide resolved
llvm::SmallDenseMap<Operation *, SmallVector<IREE::GPU::MMAIntrinsic>>
llvm::SmallDenseMap<IREE::HAL::ExecutableVariantOp,
SmallVector<IREE::GPU::MMAIntrinsic>>
queryMMAIntrinsics(mlir::ModuleOp moduleOp);

} // namespace mlir::iree_compiler
Expand Down
Loading