[LLVMGPU] Switch to new pass generation tablegen definitions. (#18213)

This is mostly an NFC change. The revision applies a little cleanups: - Switch a couple of passes to follow `create.*Pass` naming convention. - Move the detailed comments from `Passes.h` to `Passes.td`. --------- Signed-off-by: hanhanW <[email protected]>
iree-org · Aug 14, 2024 · a72e78b · a72e78b
1 parent 0c2f51b
commit a72e78b
Show file tree

Hide file tree

Showing 31 changed files with 224 additions and 389 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/Common/Passes.h b/compiler/src/iree/compiler/Codegen/Common/Passes.h
@@ -88,11 +88,6 @@ createTileAndDistributeToWorkgroupsPass(
     int32_t maxWorkgroupParallelDims,
     linalg::DistributionMethod distributionMethod);
 
-// TODO(hanchung): Move it where it is defined (i.e., Codegen/LLVMGPU).
-// Extract address computations (including the ones with GPU instructions) into
-// their own separate instructions.
-std::unique_ptr<Pass> createExtractAddressComputationGPUPass();
-
 //----------------------------------------------------------------------------//
 // CodeGen Common Patterns
 //----------------------------------------------------------------------------//

diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/AMDGPUChainedMatmulPass.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/AMDGPUChainedMatmulPass.cpp
@@ -5,14 +5,17 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 #include <numeric>
-#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
+
 #include "iree/compiler/Codegen/LLVMGPU/Passes.h"
 #include "iree/compiler/Codegen/Utils/VectorOpUtils.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 
 namespace mlir::iree_compiler {
 
+#define GEN_PASS_DEF_AMDGPUPREPAREFORCHAINEDMATMULPASS
+#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"
+
 using VectorValue = TypedValue<VectorType>;
 
 namespace {
@@ -59,8 +62,8 @@ namespace {
 ///   C = A @ B --> C.T = B.T @ A.T
 /// is only defined on standard "@" function, it may be a different
 /// transformation for other indexing maps.
-struct AMDGPUPrepareForChainedMatmulPass
-    : public AMDGPUPrepareForChainedMatmulBase<
+struct AMDGPUPrepareForChainedMatmulPass final
+    : impl::AMDGPUPrepareForChainedMatmulPassBase<
           AMDGPUPrepareForChainedMatmulPass> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<vector::VectorDialect>();
@@ -255,10 +258,4 @@ struct AMDGPUPrepareForChainedMatmulPass
 };
 
 } // namespace
-
-std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
-createAMDGPUPrepareForChainedMatmulPass() {
-  return std::make_unique<AMDGPUPrepareForChainedMatmulPass>();
-}
-
 } // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMGPU/BUILD.bazel
@@ -32,7 +32,6 @@ iree_gentbl_cc_library(
 iree_compiler_cc_library(
     name = "PassHeaders",
     hdrs = [
-        "PassDetail.h",
         "Passes.h",
         "Passes.h.inc",
     ],
@@ -69,7 +68,6 @@ iree_gentbl_cc_library(
 iree_compiler_cc_library(
     name = "ROCDLPassHeaders",
     hdrs = [
-        "ROCDLPassDetail.h",
         "ROCDLPasses.h",
         "ROCDLPasses.h.inc",
     ],

diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt
@@ -23,7 +23,6 @@ iree_cc_library(
   NAME
     PassHeaders
   HDRS
-    "PassDetail.h"
     "Passes.h"
     "Passes.h.inc"
   DEPS
@@ -52,7 +51,6 @@ iree_cc_library(
   NAME
     ROCDLPassHeaders
   HDRS
-    "ROCDLPassDetail.h"
     "ROCDLPasses.h"
     "ROCDLPasses.h.inc"
   DEPS

diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp
@@ -6,7 +6,6 @@
 
 #include "iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.h"
 
-#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
 #include "iree/compiler/Codegen/LLVMGPU/Passes.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
 #include "iree/compiler/Codegen/Utils/Utils.h"
@@ -25,6 +24,9 @@
 
 namespace mlir::iree_compiler {
 
+#define GEN_PASS_DEF_TESTLLVMGPUSCALARIZEMATHOPPASS
+#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"
+
 void ConvertToDynamicSharedMemory(ModuleOp moduleOp) {
   SymbolTableCollection symbolTableCollection;
   // Collect all the adressOfOps to static shared memory globals.
@@ -183,8 +185,9 @@ struct ConvertSharedMemAllocOp : public OpRewritePattern<memref::AllocOp> {
 
 /// Pass to test in dialect transformation used to legalize the IR before
 /// convertToNVVM/ConvertToROCDL.
-class TestLLVMGPULegalizeOpPass
-    : public TestLLVMGPUScalarizeMathOpBase<TestLLVMGPULegalizeOpPass> {
+class TestLLVMGPULegalizeOpPass final
+    : public impl::TestLLVMGPUScalarizeMathOpPassBase<
+          TestLLVMGPULegalizeOpPass> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<vector::VectorDialect>();
   }
@@ -542,10 +545,6 @@ void populateLowerHALInterfaceOp(RewritePatternSet &patterns) {
       patterns.getContext());
 }
 
-std::unique_ptr<OperationPass<ModuleOp>> createTestLLVMGPULegalizePass() {
-  return std::make_unique<TestLLVMGPULegalizeOpPass>();
-}
-
 static IntegerAttr wrapNumericMemorySpace(MLIRContext *ctx, unsigned space) {
   return IntegerAttr::get(IntegerType::get(ctx, 64), space);
 }

diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToNVVM.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToNVVM.cpp
@@ -8,7 +8,6 @@
 #include "iree/compiler/Codegen/Common/Transforms.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
 #include "iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.h"
-#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
 #include "iree/compiler/Codegen/LLVMGPU/Passes.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
@@ -24,6 +23,7 @@
 #include "mlir/Conversion/NVGPUToNVVM/NVGPUToNVVM.h"
 #include "mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/Transforms/Passes.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
@@ -35,14 +35,20 @@
 
 namespace mlir::iree_compiler {
 
+#define GEN_PASS_DEF_CONVERTTONVVMPASS
+#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"
+
 namespace {
 
 /// A pass that replaces all occurrences of GPU device operations with their
 /// corresponding NVVM equivalent.
 ///
 /// This pass only handles device code and is not meant to be run on GPU host
 /// code.
-struct ConvertToNVVMPass : public ConvertToNVVMBase<ConvertToNVVMPass> {
+struct ConvertToNVVMPass final
+    : impl::ConvertToNVVMPassBase<ConvertToNVVMPass> {
+  using impl::ConvertToNVVMPassBase<ConvertToNVVMPass>::ConvertToNVVMPassBase;
+
   void getDependentDialects(DialectRegistry &registry) const override {
     registry
         .insert<gpu::GPUDialect, IREE::GPU::IREEGPUDialect, LLVM::LLVMDialect,
@@ -180,9 +186,4 @@ struct ConvertToNVVMPass : public ConvertToNVVMBase<ConvertToNVVMPass> {
 };
 
 } // namespace
-
-std::unique_ptr<OperationPass<ModuleOp>> createConvertToNVVMPass() {
-  return std::make_unique<ConvertToNVVMPass>();
-}
-
 } // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToROCDL.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToROCDL.cpp
@@ -8,7 +8,6 @@
 #include "iree/compiler/Codegen/Common/Transforms.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
 #include "iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.h"
-#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
 #include "iree/compiler/Codegen/LLVMGPU/Passes.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
 #include "mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h"
@@ -40,6 +39,9 @@
 
 namespace mlir::iree_compiler {
 
+#define GEN_PASS_DEF_CONVERTTOROCDLPASS
+#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"
+
 static llvm::cl::opt<int>
     clROCMIndexingBits("iree-rocm-index-bits",
                        llvm::cl::desc("Set the bit width of indices in ROCm."),
@@ -75,7 +77,11 @@ static void populateConvertGPUToAMDGPUPatterns(RewritePatternSet &patterns) {
 ///
 /// This pass only handles device code and is not meant to be run on GPU host
 /// code.
-struct ConvertToROCDLPass : public ConvertToROCDLBase<ConvertToROCDLPass> {
+struct ConvertToROCDLPass final
+    : impl::ConvertToROCDLPassBase<ConvertToROCDLPass> {
+  using impl::ConvertToROCDLPassBase<
+      ConvertToROCDLPass>::ConvertToROCDLPassBase;
+
   void getDependentDialects(DialectRegistry &registry) const override {
     registry
         .insert<IREE::GPU::IREEGPUDialect, LLVM::LLVMDialect,
@@ -203,9 +209,4 @@ struct ConvertToROCDLPass : public ConvertToROCDLBase<ConvertToROCDLPass> {
     LDBG("After converting to dynamic shared memory\n" << m);
   }
 };
-
-std::unique_ptr<OperationPass<ModuleOp>> createConvertToROCDLPass() {
-  return std::make_unique<ConvertToROCDLPass>();
-}
-
 } // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/ExtractAddressComputationGPUPass.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/ExtractAddressComputationGPUPass.cpp
@@ -5,20 +5,22 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 #include "iree/compiler/Codegen/Common/ExtractAddressComputation.h"
-#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
 #include "iree/compiler/Codegen/LLVMGPU/Passes.h"
 #include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #define DEBUG_TYPE "extract-address-computation-gpu"
 
-using namespace mlir;
-
 namespace mlir::iree_compiler {
 
+#define GEN_PASS_DEF_EXTRACTADDRESSCOMPUTATIONGPUPASS
+#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"
+
 //===----------------------------------------------------------------------===//
 // Helper functions for the `load base[off0...]`
 //  => `load (subview base[off0...])[0...]` pattern.
@@ -80,8 +82,8 @@ populateExtractAddressComputationGPUPatterns(RewritePatternSet &patterns) {
 // Pass registration
 //===----------------------------------------------------------------------===//
 namespace {
-struct ExtractAddressComputationGPUPass
-    : public ExtractAddressComputationGPUBase<
+struct ExtractAddressComputationGPUPass final
+    : impl::ExtractAddressComputationGPUPassBase<
           ExtractAddressComputationGPUPass> {
   void runOnOperation() override;
 };
@@ -96,7 +98,4 @@ void ExtractAddressComputationGPUPass::runOnOperation() {
   }
 }
 
-std::unique_ptr<Pass> createExtractAddressComputationGPUPass() {
-  return std::make_unique<ExtractAddressComputationGPUPass>();
-}
 } // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUCastAddressSpaceFunction.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUCastAddressSpaceFunction.cpp
@@ -4,23 +4,26 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
 #include "iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
+#include "mlir/Pass/Pass.h"
 
 #define DEBUG_TYPE "iree-llvmgpu-cast-address-space-function"
 
 namespace mlir::iree_compiler {
 
+#define GEN_PASS_DEF_LLVMGPUCASTADDRESSSPACEFUNCTIONPASS
+#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"
+
 namespace {
 
-struct LLVMGPUCastAddressSpaceFunctionPass
-    : public LLVMGPUCastAddressSpaceFunctionBase<
+struct LLVMGPUCastAddressSpaceFunctionPass final
+    : impl::LLVMGPUCastAddressSpaceFunctionPassBase<
           LLVMGPUCastAddressSpaceFunctionPass> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<affine::AffineDialect, gpu::GPUDialect>();
@@ -75,10 +78,4 @@ struct LLVMGPUCastAddressSpaceFunctionPass
 };
 
 } // namespace
-
-std::unique_ptr<OperationPass<ModuleOp>>
-createLLVMGPUCastAddressSpaceFunction() {
-  return std::make_unique<LLVMGPUCastAddressSpaceFunctionPass>();
-}
-
 } // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUCastTypeToFitMMA.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUCastTypeToFitMMA.cpp
@@ -6,7 +6,6 @@
 
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h"
-#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
 #include "iree/compiler/Codegen/LLVMGPU/Passes.h"
 #include "iree/compiler/Codegen/Utils/VectorOpUtils.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
@@ -21,6 +20,9 @@
 
 namespace mlir::iree_compiler {
 
+#define GEN_PASS_DEF_LLVMGPUCASTTYPETOFITMMAPASS
+#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"
+
 namespace {
 
 struct UpcastContractOutput final : OpRewritePattern<vector::ContractionOp> {
@@ -72,9 +74,8 @@ struct UpcastContractOutput final : OpRewritePattern<vector::ContractionOp> {
   }
 };
 
-struct LLVMGPUCastTypeToFitMMAPass
-    : public LLVMGPUCastTypeToFitMMABase<LLVMGPUCastTypeToFitMMAPass> {
-public:
+struct LLVMGPUCastTypeToFitMMAPass final
+    : impl::LLVMGPUCastTypeToFitMMAPassBase<LLVMGPUCastTypeToFitMMAPass> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<vector::VectorDialect>();
     registry.insert<arith::ArithDialect>();
@@ -114,9 +115,4 @@ struct LLVMGPUCastTypeToFitMMAPass
   }
 };
 } // namespace
-std::unique_ptr<InterfacePass<FunctionOpInterface>>
-createLLVMGPUCastTypeToFitMMAPass() {
-  return std::make_unique<LLVMGPUCastTypeToFitMMAPass>();
-}
-
 } // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureVectorLayouts.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPUConfigureVectorLayouts.cpp
@@ -9,7 +9,6 @@
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
 #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
-#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
 #include "iree/compiler/Codegen/LLVMGPU/Passes.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
 #include "llvm/ADT/SetVector.h"
@@ -27,6 +26,9 @@
 
 namespace mlir::iree_compiler {
 
+#define GEN_PASS_DEF_LLVMGPUCONFIGUREVECTORLAYOUTSPASS
+#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"
+
 namespace {
 
 // Sets an anchoring layout for the given contraction op. Looks for a
@@ -279,10 +281,9 @@ LogicalResult setTransferReadAnchor(ArrayRef<int64_t> workgroupSize,
   return success();
 }
 
-struct LLVMGPUConfigureVectorLayoutsPass
-    : public LLVMGPUConfigureVectorLayoutsBase<
+struct LLVMGPUConfigureVectorLayoutsPass final
+    : impl::LLVMGPUConfigureVectorLayoutsPassBase<
           LLVMGPUConfigureVectorLayoutsPass> {
-public:
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<IREE::VectorExt::IREEVectorExtDialect>();
     registry.insert<vector::VectorDialect>();
@@ -360,10 +361,4 @@ struct LLVMGPUConfigureVectorLayoutsPass
   }
 };
 } // namespace
-
-std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
-createLLVMGPUConfigureVectorLayouts() {
-  return std::make_unique<LLVMGPUConfigureVectorLayoutsPass>();
-}
-
 } // namespace mlir::iree_compiler