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

[tuner]: two new utility functions which are more friendly for c binding #19199

Merged
merged 3 commits into from
Nov 19, 2024
Merged

[tuner]: two new utility functions which are more friendly for c binding #19199

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c983227
[tuner] two new functions which are more friendly for c binding
bangtianliu Nov 19, 2024
36d2c39
[tuner] address comments
bangtianliu Nov 19, 2024
4a140db
[tuner] pass by value
bangtianliu Nov 19, 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
11 changes: 6 additions & 5 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,15 +23,16 @@ struct TestLLVMGPUQueryMMAPass final
: impl::TestLLVMGPUQueryMMAPassBase<TestLLVMGPUQueryMMAPass> {
void runOnOperation() override {
ModuleOp moduleOp = getOperation();
llvm::SmallDenseMap<IREE::HAL::ExecutableVariantOp,
SmallVector<IREE::GPU::MMAIntrinsic>>
mmaMap = queryMMAIntrinsics(moduleOp);
for (const auto &[op, mmaAttrs] : mmaMap) {
SmallVector<IREE::HAL::ExecutableVariantOp> executableVariantOps =
getExecutableVariantOps(moduleOp);
for (const auto &op : executableVariantOps) {
bangtianliu marked this conversation as resolved.
Show resolved Hide resolved
llvm::outs() << "Executable Variant Name: "
<< cast<IREE::HAL::ExecutableVariantOp>(*op).getName()
<< "\n";
SmallVector<IREE::GPU::MMAIntrinsic> mmaIntrinsics =
queryMMAIntrinsics(op);
llvm::outs() << "MMA Intrinsics: ";
llvm::interleave(mmaAttrs, llvm::outs(), " ");
llvm::interleave(mmaIntrinsics, llvm::outs(), " ");
llvm::outs() << "\n";
}
}
Expand Down
8 changes: 4 additions & 4 deletions compiler/src/iree/compiler/Codegen/LLVMGPU/test/test_query_mma.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,10 +71,10 @@ module {
}
}

// CHECK-DAG: main_0
// CHECK-DAG: MMA Intrinsics: MFMA_F32_16x16x4_F32 MFMA_F32_16x16x16_F16
// CHECK-DAG: main_1
// CHECK-DAG: MMA Intrinsics: MFMA_F32_32x32x8_F16 MFMA_F32_16x16x16_BF16
// CHECK: main_0
// CHECK-NEXT: MMA Intrinsics: MFMA_F32_16x16x4_F32 MFMA_F32_16x16x16_F16
// CHECK-NEXT: main_1
// CHECK-NEXT: MMA Intrinsics: MFMA_F32_32x32x8_F16 MFMA_F32_16x16x16_BF16

// -----

Expand Down
35 changes: 19 additions & 16 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 @@ -1028,22 +1028,25 @@ std::optional<int> getGPUSubgroupSize(mlir::FunctionOpInterface func) {
return std::nullopt;
}

llvm::SmallDenseMap<IREE::HAL::ExecutableVariantOp,
SmallVector<IREE::GPU::MMAIntrinsic>>
queryMMAIntrinsics(mlir::ModuleOp moduleOp) {
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] = std::move(mmaIntrinsics);
}
});
return mmaAttributesMap;
llvm::SmallVector<IREE::HAL::ExecutableVariantOp>
getExecutableVariantOps(mlir::ModuleOp moduleOp) {
llvm::SmallVector<IREE::HAL::ExecutableVariantOp> executableVariantOps;
bangtianliu marked this conversation as resolved.
Show resolved Hide resolved
moduleOp.walk<WalkOrder::PreOrder>(
bangtianliu marked this conversation as resolved.
Show resolved Hide resolved
[&](IREE::HAL::ExecutableVariantOp executableOp) {
executableVariantOps.push_back(executableOp);
});
return executableVariantOps;
}

llvm::SmallVector<IREE::GPU::MMAIntrinsic>
queryMMAIntrinsics(IREE::HAL::ExecutableVariantOp executableOp) {
llvm::SmallVector<IREE::GPU::MMAIntrinsic> mmaIntrinsics;
if (IREE::GPU::TargetAttr target = getGPUTargetAttr(executableOp)) {
mmaIntrinsics = llvm::map_to_vector(
target.getWgp().getMma(),
[](IREE::GPU::MMAAttr attr) { return attr.getIntrinsic().getValue(); });
}
return mmaIntrinsics;
}

} // namespace mlir::iree_compiler
17 changes: 10 additions & 7 deletions 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 @@ -207,13 +207,16 @@ IREE::GPU::TargetAttr getGPUTargetAttr(Operation *op);
/// Returns std::nullopt if none found.
std::optional<int> getGPUSubgroupSize(mlir::FunctionOpInterface func);

/// Returns a map of supported MMA intrinsic instructions based on the
/// GPU target descriptions in `moduleOp`. Each entry in the map associates
/// an `IREE::HAL::ExecutableVariantOp` with a vector of
/// `IREE::GPU::MMAIntrinsic` attributes.
llvm::SmallDenseMap<IREE::HAL::ExecutableVariantOp,
SmallVector<IREE::GPU::MMAIntrinsic>>
queryMMAIntrinsics(mlir::ModuleOp moduleOp);
/// Returns all `IREE::HAL::ExecutableVariantOp` operations from the
/// given `mlir::ModuleOp`, ensuring they are returned in their original IR
bangtianliu marked this conversation as resolved.
Show resolved Hide resolved
/// order.
llvm::SmallVector<IREE::HAL::ExecutableVariantOp>
bangtianliu marked this conversation as resolved.
Show resolved Hide resolved
getExecutableVariantOps(mlir::ModuleOp moduleOp);

// Returns the MMA intrinsics associated with the given
// `IREE::HAL::ExecutableVariantOp`.
llvm::SmallVector<IREE::GPU::MMAIntrinsic>
bangtianliu marked this conversation as resolved.
Show resolved Hide resolved
queryMMAIntrinsics(IREE::HAL::ExecutableVariantOp executableOp);

} // namespace mlir::iree_compiler

Expand Down
Loading