[tuner] two new utility functions which are more friendly for c bindi…

…ng (iree-org#19199) Follow-up PR after iree-org#19124. TODO: add binding in this PR. --------- Signed-off-by: Bangtian Liu <[email protected]>
nod-ai · Nov 19, 2024 · 495adc6 · 495adc6
1 parent 5b0740c
commit 495adc6
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 30 deletions.
diff --git 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<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 (IREE::HAL::ExecutableVariantOp op : executableVariantOps) {
       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";
     }
   }

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
@@ -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
@@ -1028,22 +1028,24 @@ 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;
+SmallVector<IREE::HAL::ExecutableVariantOp>
+getExecutableVariantOps(mlir::ModuleOp moduleOp) {
+  llvm::SmallVector<IREE::HAL::ExecutableVariantOp> 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<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
diff --git 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<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
+/// order.
+SmallVector<IREE::HAL::ExecutableVariantOp>
+getExecutableVariantOps(mlir::ModuleOp moduleOp);
+
+// Returns the MMA intrinsics associated with the given
+// `IREE::HAL::ExecutableVariantOp`.
+SmallVector<IREE::GPU::MMAIntrinsic>
+queryMMAIntrinsics(IREE::HAL::ExecutableVariantOp executableOp);
 
 } // namespace mlir::iree_compiler