From ad3ecc24b1db56f87f89b9fd52eb1ad531056894 Mon Sep 17 00:00:00 2001
From: Christian Sigg <csigg@google.com>
Date: Wed, 14 Oct 2020 10:31:08 +0200
Subject: [PATCH] [mlir][gpu] NFC: Make room for more than one GPU rewrite
 pattern.

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

Reviewed By: herhut

Differential Revision: https://reviews.llvm.org/D89370
---
 mlir/include/mlir/Dialect/GPU/Passes.h                       | 12 +++++++++---
 mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp        |  4 ++--
 mlir/test/Dialect/GPU/all-reduce-max.mlir                    |  2 +-
 mlir/test/Dialect/GPU/all-reduce.mlir                        |  2 +-
 mlir/test/lib/Transforms/CMakeLists.txt                      |  2 +-
 .../{TestAllReduceLowering.cpp => TestGpuRewrite.cpp}        | 10 +++++-----
 6 files changed, 19 insertions(+), 13 deletions(-)
 rename mlir/test/lib/Transforms/{TestAllReduceLowering.cpp => TestGpuRewrite.cpp} (81%)

diff --git a/mlir/include/mlir/Dialect/GPU/Passes.h b/mlir/include/mlir/Dialect/GPU/Passes.h
index 64b744b..21c526e 100644
--- a/mlir/include/mlir/Dialect/GPU/Passes.h
+++ b/mlir/include/mlir/Dialect/GPU/Passes.h
@@ -18,9 +18,15 @@
 namespace mlir {
 std::unique_ptr<OperationPass<ModuleOp>> createGpuKernelOutliningPass();
 
-/// Collect a set of patterns to rewrite ops within the GPU dialect.
-void populateGpuRewritePatterns(MLIRContext *context,
-                                OwningRewritePatternList &patterns);
+/// 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.
+inline void populateGpuRewritePatterns(MLIRContext *context,
+                                       OwningRewritePatternList &patterns) {
+  populateGpuAllReducePatterns(context, patterns);
+}
 
 //===----------------------------------------------------------------------===//
 // Registration
diff --git a/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp b/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
index 38df9ef..d3ee3e2 100644
--- a/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
@@ -397,7 +397,7 @@ struct GpuAllReduceConversion : public RewritePattern {
 };
 } // namespace
 
-void mlir::populateGpuRewritePatterns(MLIRContext *context,
-                                      OwningRewritePatternList &patterns) {
+void mlir::populateGpuAllReducePatterns(MLIRContext *context,
+                                        OwningRewritePatternList &patterns) {
   patterns.insert<GpuAllReduceConversion>(context);
 }
diff --git a/mlir/test/Dialect/GPU/all-reduce-max.mlir b/mlir/test/Dialect/GPU/all-reduce-max.mlir
index 142228d..fb5fcaf 100644
--- a/mlir/test/Dialect/GPU/all-reduce-max.mlir
+++ b/mlir/test/Dialect/GPU/all-reduce-max.mlir
@@ -1,4 +1,4 @@
-// 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 {
diff --git a/mlir/test/Dialect/GPU/all-reduce.mlir b/mlir/test/Dialect/GPU/all-reduce.mlir
index 491d9b3..758fb44 100644
--- a/mlir/test/Dialect/GPU/all-reduce.mlir
+++ b/mlir/test/Dialect/GPU/all-reduce.mlir
@@ -1,4 +1,4 @@
-// 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 {
diff --git a/mlir/test/lib/Transforms/CMakeLists.txt b/mlir/test/lib/Transforms/CMakeLists.txt
index effa7e2..e3e82f2 100644
--- a/mlir/test/lib/Transforms/CMakeLists.txt
+++ b/mlir/test/lib/Transforms/CMakeLists.txt
@@ -1,6 +1,5 @@
 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
-  TestAllReduceLowering.cpp
   TestAffineLoopParametricTiling.cpp
   TestBufferPlacement.cpp
   TestExpandTanh.cpp
@@ -15,6 +14,7 @@ add_mlir_library(MLIRTestTransforms
   TestLoopFusion.cpp
   TestGpuMemoryPromotion.cpp
   TestGpuParallelLoopMapping.cpp
+  TestGpuRewrite.cpp
   TestInlining.cpp
   TestLinalgCodegenStrategy.cpp
   TestLinalgFusionTransforms.cpp
diff --git a/mlir/test/lib/Transforms/TestAllReduceLowering.cpp b/mlir/test/lib/Transforms/TestGpuRewrite.cpp
similarity index 81%
rename from mlir/test/lib/Transforms/TestAllReduceLowering.cpp
rename to mlir/test/lib/Transforms/TestGpuRewrite.cpp
index 0c72b6c..13f0d2e 100644
--- a/mlir/test/lib/Transforms/TestAllReduceLowering.cpp
+++ b/mlir/test/lib/Transforms/TestGpuRewrite.cpp
@@ -18,8 +18,8 @@
 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>();
   }
@@ -33,8 +33,8 @@ struct TestAllReduceLoweringPass
 
 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
-- 
2.7.4