[mlir:Transforms] Move NormalizeMemRefs to MemRef/Transforms/

[rriddle]

Summary

Transforms/ should only contain transformations that are dialect-independent and
this pass interacts with MemRef operations (making it a better fit for living in that
dialect).

[mehdi_amini]

Makes sense, thanks!

[bondhugula]

This will make MLIRMemrefTransforms depend on MLIRAffine -- although the library already depends on it, that's a spurious one and can be removed. (I'll clean it up.) Looks like the test cases haven't been moved:

test/Transforms/normalize-memrefs-ops-dynamic.mlir:// RUN: mlir-opt -normalize-memrefs %s -split-input-file| FileCheck %s
test/Transforms/normalize-memrefs-ops.mlir:// RUN: mlir-opt -normalize-memrefs %s | FileCheck %s
test/Transforms/normalize-memrefs.mlir:// RUN: mlir-opt -normalize-memrefs -allow-unregistered-dialect %s | FileCheck %s

This pass should perhaps move to AffineTransforms or to a new directory meant for cross-dialect transformations where there isn't a natural layering relationship between the involved dialects. MemRefTransforms should not depend on Affine. NormalizeMemRefs creates affine.apply ops and so will depend on the Affine dialect.

[bondhugula]

In D117841#3260230, @bondhugula wrote:

This will make MLIRMemrefTransforms depend on MLIRAffine -- although the library already depends on it, that's a spurious one and can be removed. (I'll clean it up.) Looks like the test cases haven't been moved:

test/Transforms/normalize-memrefs-ops-dynamic.mlir:// RUN: mlir-opt -normalize-memrefs %s -split-input-file| FileCheck %s
test/Transforms/normalize-memrefs-ops.mlir:// RUN: mlir-opt -normalize-memrefs %s | FileCheck %s
test/Transforms/normalize-memrefs.mlir:// RUN: mlir-opt -normalize-memrefs -allow-unregistered-dialect %s | FileCheck %s

This pass should perhaps move to AffineTransforms or to a new directory meant for cross-dialect transformations where there isn't a natural layering relationship between the involved dialects. MemRef should not depend on Affine. NormalizeMemRefs creates affine.apply ops and so will depend on the Affine dialect.

Looks like FoldSubViewOps.cpp and ResolveShapedTypeResultDims.cpp would depend on the Affine dialect - so this move should be fine.

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h

 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
 /// Creates an operation pass to fold memref.subview ops into consumer
 /// load/store ops into `patterns`.
 std::unique_ptr<Pass> createFoldSubViewOpsPass();
 
+/// Creates an interprocedural pass to normalize memrefs to have a trivial
+/// (identity) layout map.
+std::unique_ptr<OperationPass<ModuleOp>> createNormalizeMemRefsPass();
+
 /// Creates an operation pass to resolve `memref.dim` operations with values
 /// that are defined by operations that implement the
 /// `ReifyRankedShapeTypeShapeOpInterface`, in terms of shapes of its input
 /// operands.
 std::unique_ptr<Pass> createResolveRankedShapeTypeResultDimsPass();
 
 /// Creates an operation pass to resolve `memref.dim` operations with values
 /// that are defined by operations that implement the

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td

let description = [{
     from/to the original memref.
   }];
   let constructor = "mlir::memref::createFoldSubViewOpsPass()";
   let dependentDialects = [
       "AffineDialect", "memref::MemRefDialect", "vector::VectorDialect"
   ];
 }
 
+def NormalizeMemRefs : Pass<"normalize-memrefs", "ModuleOp"> {
+  let summary = "Normalize memrefs";
+   let description = [{
+    This pass transforms memref types with a non-trivial
+    [layout map](https://mlir.llvm.org/docs/LangRef/#layout-map) into
+    memref types with an identity layout map, e.g. (i, j) -> (i, j). This
+    pass is inter-procedural, in the sense that it can modify function
+    interfaces and call sites that pass memref types. In order to modify
+    memref types while preserving the original behavior, users of those
+    memref types are also modified to incorporate the resulting layout map.
+    For instance, an [AffineLoadOp]
+    (https://mlir.llvm.org/docs/Dialects/Affine/#affineload-affineloadop)
+    will be updated to compose the layout map with with the affine expression
+    contained in the op. Operations marked with the [MemRefsNormalizable]
+    (https://mlir.llvm.org/docs/Traits/#memrefsnormalizable) trait are
+    expected to be normalizable. Supported operations include affine
+    operations, memref.alloc, memref.dealloc, and std.return.
+
+    Given an appropriate layout map specified in the code, this transformation
+    can express tiled or linearized access to multi-dimensional data
+    structures, but will not modify memref types without an explicit layout
+    map.
+
+    Currently this pass is limited to only modify
+    functions where all memref types can be normalized. If a function
+    contains any operations that are not MemRefNormalizable, then the function
+    and any functions that call or call it will not be modified.
+
+    Input
+
+    ```mlir
+    #tile = affine_map<(i) -> (i floordiv 4, i mod 4)>
+    func @matmul(%A: memref<16xf64, #tile>,
+                 %B: index, %C: memref<16xf64>) -> (memref<16xf64, #tile>) {
+      affine.for %arg3 = 0 to 16 {
+            %a = affine.load %A[%arg3] : memref<16xf64, #tile>
+            %p = arith.mulf %a, %a : f64
+            affine.store %p, %A[%arg3] : memref<16xf64, #tile>
+      }
+      %c = memref.alloc() : memref<16xf64, #tile>
+      %d = affine.load %c[0] : memref<16xf64, #tile>
+      return %A: memref<16xf64, #tile>
+    }
+    ```
+
+    Output
+
+    ```mlir
+    func @matmul(%arg0: memref<4x4xf64>, %arg1: index, %arg2: memref<16xf64>)
+      -> memref<4x4xf64> {
+      affine.for %arg3 = 0 to 16 {
+        %3 = affine.load %arg0[%arg3 floordiv 4, %arg3 mod 4]: memref<4x4xf64>
+        %4 = arith.mulf %3, %3 : f64
+        affine.store %4, %arg0[%arg3 floordiv 4, %arg3 mod 4]: memref<4x4xf64>
+      }
+      %0 = memref.alloc() : memref<4x4xf64>
+      %1 = affine.apply #map1()
+      %2 = affine.load %0[0, 0] : memref<4x4xf64>
+      return %arg0 : memref<4x4xf64>
+    }
+    ```
+
+    Input
+
+    ```
+    #linear8 = affine_map<(i, j) -> (i * 8 + j)>
+    func @linearize(%arg0: memref<8x8xi32, #linear8>,
+                    %arg1: memref<8x8xi32, #linear8>,
+                    %arg2: memref<8x8xi32, #linear8>) {
+      %c8 = arith.constant 8 : index
+      %c0 = arith.constant 0 : index
+      %c1 = arith.constant 1 : index
+      affine.for %arg3 = %c0 to %c8  {
+      affine.for %arg4 = %c0 to %c8  {
+        affine.for %arg5 = %c0 to %c8 {
+          %0 = affine.load %arg0[%arg3, %arg5] : memref<8x8xi32, #linear8>
+          %1 = affine.load %arg1[%arg5, %arg4] : memref<8x8xi32, #linear8>
+          %2 = affine.load %arg2[%arg3, %arg4] : memref<8x8xi32, #linear8>
+          %3 = arith.muli %0, %1 : i32
+          %4 = arith.addi %2, %3 : i32
+          affine.store %4, %arg2[%arg3, %arg4] : memref<8x8xi32, #linear8>
+        }
+      }
+      }
+      return
+    }
+    ```
+
+    Output
+
+    ```mlir
+    func @linearize(%arg0: memref<64xi32>,
+                    %arg1: memref<64xi32>,
+                    %arg2: memref<64xi32>) {
+    %c8 = arith.constant 8 : index
+    %c0 = arith.constant 0 : index
+    affine.for %arg3 = %c0 to %c8 {
+      affine.for %arg4 = %c0 to %c8 {
+        affine.for %arg5 = %c0 to %c8 {
+          %0 = affine.load %arg0[%arg3 * 8 + %arg5] : memref<64xi32>
+          %1 = affine.load %arg1[%arg5 * 8 + %arg4] : memref<64xi32>
+          %2 = affine.load %arg2[%arg3 * 8 + %arg4] : memref<64xi32>
+          %3 = arith.muli %0, %1 : i32
+          %4 = arith.addi %2, %3 : i32
+          affine.store %4, %arg2[%arg3 * 8 + %arg4] : memref<64xi32>
+        }
+      }
+    }
+    return
+  }
+  ```
+  }];
+  let constructor = "mlir::memref::createNormalizeMemRefsPass()";
+  let dependentDialects = ["AffineDialect"];
+}
+
 def ResolveRankedShapeTypeResultDims :
     Pass<"resolve-ranked-shaped-type-result-dims"> {
   let summary = "Resolve memref.dim of result values of ranked shape type";
   let description = [{
     The pass resolves memref.dim of result of operations that
     implement the `ReifyRankedShapedTypeOpInterface` in terms of
     shapes of its operands.
   }];

mlir/include/mlir/Transforms/Passes.h

 /// Creates a pass which performs sparse conditional constant propagation over
 /// nested operations.
 std::unique_ptr<Pass> createSCCPPass();
 
 /// Creates a pass which delete symbol operations that are unreachable. This
 /// pass may *only* be scheduled on an operation that defines a SymbolTable.
 std::unique_ptr<Pass> createSymbolDCEPass();
 
-/// Creates an interprocedural pass to normalize memrefs to have a trivial
-/// (identity) layout map.
-std::unique_ptr<OperationPass<ModuleOp>> createNormalizeMemRefsPass();
-
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Transforms/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_PASSES_H

mlir/include/mlir/Transforms/Passes.td

def LoopCoalescing : Pass<"loop-coalescing", "FuncOp"> {
   let dependentDialects = ["arith::ArithmeticDialect"];
 }
 
 def LoopInvariantCodeMotion : Pass<"loop-invariant-code-motion"> {
   let summary = "Hoist loop invariant instructions outside of the loop";
   let constructor = "mlir::createLoopInvariantCodeMotionPass()";
 }
 
-def NormalizeMemRefs : Pass<"normalize-memrefs", "ModuleOp"> {
-  let summary = "Normalize memrefs";
-   let description = [{
-    This pass transforms memref types with a non-trivial
-    [layout map](https://mlir.llvm.org/docs/LangRef/#layout-map) into
-    memref types with an identity layout map, e.g. (i, j) -> (i, j). This
-    pass is inter-procedural, in the sense that it can modify function
-    interfaces and call sites that pass memref types. In order to modify
-    memref types while preserving the original behavior, users of those
-    memref types are also modified to incorporate the resulting layout map.
-    For instance, an [AffineLoadOp]
-    (https://mlir.llvm.org/docs/Dialects/Affine/#affineload-affineloadop)
-    will be updated to compose the layout map with with the affine expression
-    contained in the op. Operations marked with the [MemRefsNormalizable]
-    (https://mlir.llvm.org/docs/Traits/#memrefsnormalizable) trait are
-    expected to be normalizable. Supported operations include affine
-    operations, memref.alloc, memref.dealloc, and std.return.
-
-    Given an appropriate layout map specified in the code, this transformation
-    can express tiled or linearized access to multi-dimensional data
-    structures, but will not modify memref types without an explicit layout
-    map.
-
-    Currently this pass is limited to only modify
-    functions where all memref types can be normalized. If a function
-    contains any operations that are not MemRefNormalizable, then the function
-    and any functions that call or call it will not be modified.
-
-    Input
-
-    ```mlir
-    #tile = affine_map<(i) -> (i floordiv 4, i mod 4)>
-    func @matmul(%A: memref<16xf64, #tile>,
-                 %B: index, %C: memref<16xf64>) -> (memref<16xf64, #tile>) {
-      affine.for %arg3 = 0 to 16 {
-            %a = affine.load %A[%arg3] : memref<16xf64, #tile>
-            %p = arith.mulf %a, %a : f64
-            affine.store %p, %A[%arg3] : memref<16xf64, #tile>
-      }
-      %c = memref.alloc() : memref<16xf64, #tile>
-      %d = affine.load %c[0] : memref<16xf64, #tile>
-      return %A: memref<16xf64, #tile>
-    }
-    ```
-
-    Output
-
-    ```mlir
-    func @matmul(%arg0: memref<4x4xf64>, %arg1: index, %arg2: memref<16xf64>)
-      -> memref<4x4xf64> {
-      affine.for %arg3 = 0 to 16 {
-        %3 = affine.load %arg0[%arg3 floordiv 4, %arg3 mod 4]: memref<4x4xf64>
-        %4 = arith.mulf %3, %3 : f64
-        affine.store %4, %arg0[%arg3 floordiv 4, %arg3 mod 4]: memref<4x4xf64>
-      }
-      %0 = memref.alloc() : memref<4x4xf64>
-      %1 = affine.apply #map1()
-      %2 = affine.load %0[0, 0] : memref<4x4xf64>
-      return %arg0 : memref<4x4xf64>
-    }
-    ```
-
-    Input
-
-    ```
-    #linear8 = affine_map<(i, j) -> (i * 8 + j)>
-    func @linearize(%arg0: memref<8x8xi32, #linear8>,
-                    %arg1: memref<8x8xi32, #linear8>,
-                    %arg2: memref<8x8xi32, #linear8>) {
-      %c8 = arith.constant 8 : index
-      %c0 = arith.constant 0 : index
-      %c1 = arith.constant 1 : index
-      affine.for %arg3 = %c0 to %c8  {
-      affine.for %arg4 = %c0 to %c8  {
-        affine.for %arg5 = %c0 to %c8 {
-          %0 = affine.load %arg0[%arg3, %arg5] : memref<8x8xi32, #linear8>
-          %1 = affine.load %arg1[%arg5, %arg4] : memref<8x8xi32, #linear8>
-          %2 = affine.load %arg2[%arg3, %arg4] : memref<8x8xi32, #linear8>
-          %3 = arith.muli %0, %1 : i32
-          %4 = arith.addi %2, %3 : i32
-          affine.store %4, %arg2[%arg3, %arg4] : memref<8x8xi32, #linear8>
-        }
-      }
-      }
-      return
-    }
-    ```
-
-    Output
-
-    ```mlir
-    func @linearize(%arg0: memref<64xi32>,
-                    %arg1: memref<64xi32>,
-                    %arg2: memref<64xi32>) {
-    %c8 = arith.constant 8 : index
-    %c0 = arith.constant 0 : index
-    affine.for %arg3 = %c0 to %c8 {
-      affine.for %arg4 = %c0 to %c8 {
-        affine.for %arg5 = %c0 to %c8 {
-          %0 = affine.load %arg0[%arg3 * 8 + %arg5] : memref<64xi32>
-          %1 = affine.load %arg1[%arg5 * 8 + %arg4] : memref<64xi32>
-          %2 = affine.load %arg2[%arg3 * 8 + %arg4] : memref<64xi32>
-          %3 = arith.muli %0, %1 : i32
-          %4 = arith.addi %2, %3 : i32
-          affine.store %4, %arg2[%arg3 * 8 + %arg4] : memref<64xi32>
-        }
-      }
-    }
-    return
-  }
-  ```
-  }];
-  let constructor = "mlir::createNormalizeMemRefsPass()";
-  let dependentDialects = ["AffineDialect"];
-}
-
 def ParallelLoopCollapsing : Pass<"parallel-loop-collapsing"> {
   let summary = "Collapse parallel loops to use less induction variables";
   let constructor = "mlir::createParallelLoopCollapsingPass()";
   let options = [
     ListOption<"clCollapsedIndices0", "collapsed-indices-0", "unsigned",
                "Which loop indices to combine 0th loop index",
                "llvm::cl::MiscFlags::CommaSeparated">,
     ListOption<"clCollapsedIndices1", "collapsed-indices-1", "unsigned",

mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefTransforms
   FoldSubViewOps.cpp
+  NormalizeMemRefs.cpp
   ResolveShapedTypeResultDims.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef
 
   DEPENDS
   MLIRMemRefPassIncGen
 

mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp

 //===- NormalizeMemRefs.cpp -----------------------------------------------===//
 //
 // 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 an interprocedural pass to normalize memrefs to have
 // identity layout maps.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Transforms/Passes.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Transforms/Utils.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "normalize-memrefs"
 
 using namespace mlir;
 
struct NormalizeMemRefs : public NormalizeMemRefsBase<NormalizeMemRefs> {
   void updateFunctionSignature(FuncOp funcOp, ModuleOp moduleOp);
   void setCalleesAndCallersNonNormalizable(FuncOp funcOp, ModuleOp moduleOp,
                                            DenseSet<FuncOp> &normalizableFuncs);
   Operation *createOpResultsNormalized(FuncOp funcOp, Operation *oldOp);
 };
 
 } // namespace
 
-std::unique_ptr<OperationPass<ModuleOp>> mlir::createNormalizeMemRefsPass() {
+std::unique_ptr<OperationPass<ModuleOp>>
+mlir::memref::createNormalizeMemRefsPass() {
   return std::make_unique<NormalizeMemRefs>();
 }
 
 void NormalizeMemRefs::runOnOperation() {
   LLVM_DEBUG(llvm::dbgs() << "Normalizing Memrefs...\n");
   ModuleOp moduleOp = getOperation();
   // We maintain all normalizable FuncOps in a DenseSet. It is initialized
   // with all the functions within a module and then functions which are not

mlir/lib/Dialect/MemRef/Transforms/PassDetail.h

+//===- PassDetail.h - MemRef Pass class details -----------------*- 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_MEMREF_TRANSFORMS_PASSDETAIL_H_
+#define DIALECT_MEMREF_TRANSFORMS_PASSDETAIL_H_
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+class AffineDialect;
+
+// Forward declaration from Dialect.h
+template <typename ConcreteDialect>
+void registerDialect(DialectRegistry &registry);
+
+namespace arith {
+class ArithmeticDialect;
+} // namespace arith
+
+namespace memref {
+class MemRefDialect;
+} // namespace memref
+
+namespace tensor {
+class TensorDialect;
+} // namespace tensor
+
+namespace vector {
+class VectorDialect;
+} // namespace vector
+
+#define GEN_PASS_CLASSES
+#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
+
+} // namespace mlir
+
+#endif // DIALECT_MEMREF_TRANSFORMS_PASSDETAIL_H_

mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp

 //===- ResolveShapedTypeResultDims.cpp - Resolve dim ops of result values -===//
 //
 // 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 pass resolves `memref.dim` operations of result values in terms of
 // shapes of their operands using the `InferShapedTypeOpInterface`.
 //
 //===----------------------------------------------------------------------===//
 
+#include "PassDetail.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"

 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//
 
 namespace {
-#define GEN_PASS_CLASSES
-#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
-
 struct ResolveRankedShapeTypeResultDimsPass final
     : public ResolveRankedShapeTypeResultDimsBase<
           ResolveRankedShapeTypeResultDimsPass> {
   void runOnOperation() override;
 };
 
 struct ResolveShapedTypeResultDimsPass final
     : public ResolveShapedTypeResultDimsBase<ResolveShapedTypeResultDimsPass> {

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   Canonicalizer.cpp
   ControlFlowSink.cpp
   CSE.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopCoalescing.cpp
   LoopFusion.cpp
   LoopInvariantCodeMotion.cpp
-  NormalizeMemRefs.cpp
   OpStats.cpp
   ParallelLoopCollapsing.cpp
   PipelineDataTransfer.cpp
   SCCP.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
   ViewOpGraph.cpp