[mlir][gpu] NFC: Make room for more than one GPU rewrite pattern.

[csigg]

Summary

AllReduceLowering is currently the only GPU rewrite pattern, but more are coming. This is a preparation change.

[csigg]

mlir/test/lib/Transforms/TestGpuRewrite.cpp

37

[csigg]

Fix pass description.

[csigg]

Fix.

mlir/include/mlir/Dialect/GPU/Passes.h

 #ifndef MLIR_DIALECT_GPU_PASSES_H_
 #define MLIR_DIALECT_GPU_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 std::unique_ptr<OperationPass<ModuleOp>> createGpuKernelOutliningPass();
 
-/// Collect a set of patterns to rewrite ops within the GPU dialect.
+/// Collect a set of patterns to rewrite all-reduce ops within the GPU dialect.
+void populateGpuAllReducePatterns(MLIRContext *context,
+                                  OwningRewritePatternList &patterns);
+
+/// Collect all patterns to rewrite ops within the GPU dialect.
 void populateGpuRewritePatterns(MLIRContext *context,
                                 OwningRewritePatternList &patterns);
 
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/GPU/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_GPU_PASSES_H_

mlir/lib/Dialect/GPU/CMakeLists.txt

 add_mlir_dialect_library(MLIRGPU
   IR/GPUDialect.cpp
   Transforms/AllReduceLowering.cpp
   Transforms/KernelOutlining.cpp
   Transforms/MemoryPromotion.cpp
   Transforms/ParallelLoopMapper.cpp
+  Transforms/Passes.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/GPU
 
   DEPENDS
   MLIRGPUOpsIncGen
   MLIRGPUOpInterfacesIncGen
   MLIRGPUPassIncGen

mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp

LogicalResult matchAndRewrite(Operation *op,
     };
     while (funcOp.walk(callback).wasInterrupted()) {
     }
     return success();
   }
 };
 } // namespace
 
-void mlir::populateGpuRewritePatterns(MLIRContext *context,
+void mlir::populateGpuAllReducePatterns(MLIRContext *context,
                                         OwningRewritePatternList &patterns) {
   patterns.insert<GpuAllReduceConversion>(context);
 }

mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp

 //===- KernelOutlining.cpp - Implementation of GPU kernel outlining -------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the GPU dialect kernel outlining pass.
 //
 //===----------------------------------------------------------------------===//
 
-#include "PassDetail.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/GPU/Utils.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Support/LLVM.h"

mlir/lib/Dialect/GPU/Transforms/Passes.cpp

-//===- PassDetail.h - GPU Pass class details --------------------*- C++ -*-===//
+//===- Passes.cpp - GPU passes registration ---------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef DIALECT_GPU_TRANSFORMS_PASSDETAIL_H_
 #define DIALECT_GPU_TRANSFORMS_PASSDETAIL_H_
 
+#include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 #define GEN_PASS_CLASSES
 #include "mlir/Dialect/GPU/Passes.h.inc"
 
+void populateGpuRewritePatterns(MLIRContext *context,
+                                OwningRewritePatternList &patterns) {
+  populateGpuAllReducePatterns(context, patterns);
+}
+
 } // end namespace mlir
 
 #endif // DIALECT_GPU_TRANSFORMS_PASSDETAIL_H_

mlir/test/Dialect/GPU/all-reduce-max.mlir

-// RUN: mlir-opt -test-all-reduce-lowering %s | FileCheck %s
+// RUN: mlir-opt -test-gpu-rewrite %s | FileCheck %s
 
 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
 // CHECK: gpu.module @kernels {
 gpu.module @kernels {
 
   // CHECK-LABEL: gpu.func @kernel(
   // CHECK-SAME: [[VAL_0:%.*]]: f32) workgroup([[VAL_1:%.*]] : memref<32xf32, 3>) kernel {
   gpu.func @kernel(%arg0 : f32) kernel {

mlir/test/Dialect/GPU/all-reduce.mlir

-// RUN: mlir-opt -test-all-reduce-lowering %s | FileCheck %s
+// RUN: mlir-opt -test-gpu-rewrite %s | FileCheck %s
 
 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
 // CHECK: gpu.module @kernels {
 gpu.module @kernels {
 
   // CHECK-LABEL: gpu.func @kernel(
   // CHECK-SAME: [[VAL_0:%.*]]: f32) workgroup([[VAL_1:%.*]] : memref<32xf32, 3>) kernel {
   gpu.func @kernel(%arg0 : f32) kernel {

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
-  TestAllReduceLowering.cpp
   TestAffineLoopParametricTiling.cpp
   TestBufferPlacement.cpp
   TestExpandTanh.cpp
   TestCallGraph.cpp
   TestConstantFold.cpp
   TestConvVectorization.cpp
   TestConvertCallOp.cpp
   TestConvertGPUKernelToCubin.cpp
   TestConvertGPUKernelToHsaco.cpp
   TestDominance.cpp
   TestDynamicPipeline.cpp
   TestLoopFusion.cpp
   TestGpuMemoryPromotion.cpp
   TestGpuParallelLoopMapping.cpp
+  TestGpuRewrite.cpp
   TestInlining.cpp
   TestLinalgFusionTransforms.cpp
   TestLinalgHoisting.cpp
   TestLinalgTransforms.cpp
   TestLiveness.cpp
   TestLoopMapping.cpp
   TestLoopParametricTiling.cpp
   TestLoopUnrolling.cpp

mlir/test/lib/Transforms/TestGpuRewrite.cpp

 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
-struct TestAllReduceLoweringPass
-    : public PassWrapper<TestAllReduceLoweringPass, OperationPass<ModuleOp>> {
+struct TestGpuRewritePass
+    : public PassWrapper<TestGpuRewritePass, OperationPass<ModuleOp>> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<StandardOpsDialect>();
   }
   void runOnOperation() override {
     OwningRewritePatternList patterns;
     populateGpuRewritePatterns(&getContext(), patterns);
     applyPatternsAndFoldGreedily(getOperation(), patterns);
   }
 };
 } // namespace
 
 namespace mlir {
 void registerTestAllReduceLoweringPass() {
-  PassRegistration<TestAllReduceLoweringPass> pass(
-      "test-all-reduce-lowering",
-      "Lowers gpu.all-reduce ops within the GPU dialect.");
+  PassRegistration<TestGpuRewritePass> pass(
+      "test-gpu-rewrite",
+      "Applies all rewrite patterns within the GPU dialect.");
 }
 } // namespace mlir