diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/TestLLVMGPUQueryMMAPass.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/TestLLVMGPUQueryMMAPass.cpp index 355187fd6f2b..2e65f81320ae 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/TestLLVMGPUQueryMMAPass.cpp +++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/TestLLVMGPUQueryMMAPass.cpp @@ -23,15 +23,16 @@ struct TestLLVMGPUQueryMMAPass final : impl::TestLLVMGPUQueryMMAPassBase { void runOnOperation() override { ModuleOp moduleOp = getOperation(); - llvm::SmallDenseMap> - mmaMap = queryMMAIntrinsics(moduleOp); - for (const auto &[op, mmaAttrs] : mmaMap) { + SmallVector executableVariantOps = + getExecutableVariantOps(moduleOp); + for (IREE::HAL::ExecutableVariantOp op : executableVariantOps) { llvm::outs() << "Executable Variant Name: " << cast(*op).getName() << "\n"; + SmallVector mmaIntrinsics = + queryMMAIntrinsics(op); llvm::outs() << "MMA Intrinsics: "; - llvm::interleave(mmaAttrs, llvm::outs(), " "); + llvm::interleave(mmaIntrinsics, llvm::outs(), " "); llvm::outs() << "\n"; } } diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/test_query_mma.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/test_query_mma.mlir index 070355f7085b..4863d38c8cc8 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/test_query_mma.mlir +++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/test_query_mma.mlir @@ -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 // ----- diff --git a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp index 1ad0bd397f9a..97c12a0e19f2 100644 --- a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp +++ b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp @@ -1028,22 +1028,24 @@ std::optional getGPUSubgroupSize(mlir::FunctionOpInterface func) { return std::nullopt; } -llvm::SmallDenseMap> -queryMMAIntrinsics(mlir::ModuleOp moduleOp) { - llvm::SmallDenseMap> - mmaAttributesMap; +SmallVector +getExecutableVariantOps(mlir::ModuleOp moduleOp) { + llvm::SmallVector executableVariantOps; 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); - } + executableVariantOps.push_back(executableOp); }); - return mmaAttributesMap; + return executableVariantOps; +} + +SmallVector +queryMMAIntrinsics(IREE::HAL::ExecutableVariantOp executableOp) { + llvm::SmallVector 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 diff --git a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h index ead0fc42e01c..133d7246a5fa 100644 --- a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h +++ b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h @@ -207,13 +207,16 @@ IREE::GPU::TargetAttr getGPUTargetAttr(Operation *op); /// Returns std::nullopt if none found. std::optional 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> -queryMMAIntrinsics(mlir::ModuleOp moduleOp); +/// Returns all `IREE::HAL::ExecutableVariantOp` operations from the +/// given `mlir::ModuleOp`, ensuring they are returned in their original IR +/// order. +SmallVector +getExecutableVariantOps(mlir::ModuleOp moduleOp); + +// Returns the MMA intrinsics associated with the given +// `IREE::HAL::ExecutableVariantOp`. +SmallVector +queryMMAIntrinsics(IREE::HAL::ExecutableVariantOp executableOp); } // namespace mlir::iree_compiler