[MLIR] Updates around MemRef Normalization

[stephenneuendorffer]

Summary

The documentation for the NormalizeMemRefs pass and the associated MemRefsNormalizable
traits was confusing and not on the website. This update clarifies the language
around the difference between a MemRef Type, an operation that accesses the value of
MemRef Type, and better documents the limitations of the current implementation.
This patch also includes some basic debugging information for the pass so people
might have a chance of figuring out why it doesn't work on their code.

[rriddle]

Nice. Please run clang-format though.

mlir/include/mlir/Transforms/Passes.td

375

nit: Add a newline before here.

422

nit: Drop the newline here.

[bondhugula]

mlir/docs/Traits.md

254–256

Hmm... this isn't entirely accurate. We may still need to tweak/improve the original thing. Ops like DeallocOp, CallOp, and ReturnOp also have this trait but they don't access the memref through an arbitrary affine subscript function. The trait is more about when you can normalize the memref operand on the op safely / meaningfully.

[bondhugula]

Thanks for submitting this - esp. the detailed pass description. We should definitely improve the documentation since you didn't find it clear enough.

[AlexEichenberger]

Added suggestion, thanks for attempting to clarify the original text, much appreciated.

mlir/include/mlir/Transforms/Passes.td

327

"must" That is probably too strong, as operations can be handled directly in the pass. I would say what the flag does, namely it will let the pass assume that normalizing that operation is safe.

331

"Currently...." maybe this is too strong a statement as we are adding support, and then we would have to correct this list. I would prefer something weaker. "Initial support includes...."

337

Nit: all passes have to be conservative in general. I would just simply state the current limitation.

Currently the pass is limited to handle functions where all ...

and if you feel like adding this

We may address this limitation if the need arises in the future.

Fine either ways

[stephenneuendorffer]

Thanks for the feedback.

mlir/docs/Traits.md

254–256

I reworded this.

[bondhugula]

A few more minor comments.

mlir/docs/Traits.md

254–260

access -> use? (access could imply dereferencing, which isn't the case with dealloc or call).

258

it's -> its

mlir/include/mlir/IR/OpDefinition.h

1215–1220

Likewise.

1216–1218

Looks like this hasn't been updated.

mlir/include/mlir/Transforms/Passes.td

317–325

Thanks very much for this description.

mlir/docs/Traits.md

 
 This trait is an important structural property of the IR, and enables operations
 to have [passes](PassManagement.md) scheduled under them.
 
 ### MemRefsNormalizable
 
 * `OpTrait::MemRefsNormalizable` -- `MemRefsNormalizable`
 
-This trait is used to flag operations that can accommodate `MemRefs` with
-non-identity memory-layout specifications. This trait indicates that the
-normalization of memory layout can be performed for such operations.
-`MemRefs` normalization consists of replacing an original memory reference
-with layout specifications to an equivalent memory reference where
-the specified memory layout is applied by rewritting accesses and types
-associated with that memory reference.
+This trait is used to flag operations that can access `MemRefs` through an
+arbitrary affine expression. In cases where an accessed MemRef has a
+non-identity memory-layout specificiation, such operations can be
+'normalized' so that the layout of memory is incorporated into the
+index expression of the operation, resulting in a new MemRef type with
+an identity memory-layout specification. See [the -normalize-memrefs pass].
+(https://mlir.llvm.org/docs/Passes/#-normalize-memrefs-normalize-memrefs)
 
 ### Single Block with Implicit Terminator
 
 *   `OpTrait::SingleBlockImplicitTerminator<typename TerminatorOpType>` :
     `SingleBlockImplicitTerminator<string op>`
 
 This trait provides APIs and verifiers for operations with regions that have a
 single block that must terminate with `TerminatorOpType`.

mlir/include/mlir/IR/OpDefinition.h

 /// not have any arguments
 template <typename ConcrentType>
 struct NoRegionArguments : public TraitBase<ConcrentType, NoRegionArguments> {
   static LogicalResult verifyTrait(Operation *op) {
     return ::mlir::OpTrait::impl::verifyNoRegionArguments(op);
   }
 };
 
-/// This trait is used to flag operations that can accommodate MemRefs with
-/// non-identity memory-layout specifications. This trait indicates that the
-/// normalization of memory layout can be performed for such operations.
-/// MemRefs normalization consists of replacing an original memory reference
-/// with layout specifications to an equivalent memory reference where the
-/// specified memory layout is applied by rewritting accesses and types
-/// associated with that memory reference.
+/// This trait is used to flag operations that can access MemRefs through an
+/// arbitrary affine expression. In cases where an accessed MemRef has a
+/// non-identity memory-layout specificiation, such operations can be
+/// 'normalized' so that the layout of memory is incorporated into the
+/// index expression of the operation, resulting in a new MemRef type with
+/// an identity memory-layout specification. 
 // TODO: Right now, the operands of an operation are either all normalizable,
 // or not. In the future, we may want to allow some of the operands to be
 // normalizable.
 template <typename ConcrentType>
 struct MemRefsNormalizable
     : public TraitBase<ConcrentType, MemRefsNormalizable> {};
 
 } // end namespace OpTrait

mlir/include/mlir/Transforms/Passes.td

let description = [{
     }
     ```
   }];
   let constructor = "mlir::createMemRefDataFlowOptPass()";
 }
 
 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. Supported operations must be marked with the
+    [MemRefsNormalizable]
+    (https://mlir.llvm.org/docs/Traits/#memrefsnormalizable) trait. Currently
+    only affine operations, std.alloc, std.dealloc, and std.return are
+    normalizable.
+    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 somewhat conservative: it will 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 = mulf %a, %a : f64
+            affine.store %p, %A[%arg3] : memref<16xf64, #tile>
+      }
+      %c = 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 = mulf %3, %3 : f64
+        affine.store %4, %arg0[%arg3 floordiv 4, %arg3 mod 4]: memref<4x4xf64>
+      }
+      %0 = 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 = constant 8 : index
+      %c0 = constant 0 : index
+      %c1 = 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 = muli %0, %1 : i32
+          %4 = 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 = constant 8 : index
+    %c0 = 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 = muli %0, %1 : i32
+          %4 = addi %2, %3 : i32
+          affine.store %4, %arg2[%arg3 * 8 + %arg4] : memref<64xi32>
+        }
+      }
+    }
+    return
+  }
+  ``` 
+  
+  }];
   let constructor = "mlir::createNormalizeMemRefsPass()";
 }
 
 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",

mlir/lib/Transforms/NormalizeMemRefs.cpp

 namespace {
 
 /// All memrefs passed across functions with non-trivial layout maps are
 /// converted to ones with trivial identity layout ones.
 /// If all the memref types/uses in a function are normalizable, we treat
 /// such functions as normalizable. Also, if a normalizable function is known
 /// to call a non-normalizable function, we treat that function as
 /// non-normalizable as well. We assume external functions to be normalizable.
-///
-/// Input :-
-/// #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 = mulf %a, %a : f64
-///         affine.store %p, %A[%arg3] : memref<16xf64, #tile>
-///   }
-///   %c = alloc() : memref<16xf64, #tile>
-///   %d = affine.load %c[0] : memref<16xf64, #tile>
-///   return %A: memref<16xf64, #tile>
-/// }
-///
-/// Output :-
-///   func @matmul(%arg0: memref<4x4xf64>, %arg1: index, %arg2: memref<16xf64>)
-///   -> memref<4x4xf64> {
-///     affine.for %arg3 = 0 to 16 {
-///       %2 = affine.load %arg0[%arg3 floordiv 4, %arg3 mod 4] :
-///       memref<4x4xf64> %3 = mulf %2, %2 : f64 affine.store %3, %arg0[%arg3
-///       floordiv 4, %arg3 mod 4] : memref<4x4xf64>
-///     }
-///     %0 = alloc() : memref<16xf64, #map0>
-///     %1 = affine.load %0[0] : memref<16xf64, #map0>
-///     return %arg0 : memref<4x4xf64>
-///   }
-///
 struct NormalizeMemRefs : public NormalizeMemRefsBase<NormalizeMemRefs> {
   void runOnOperation() override;
   void normalizeFuncOpMemRefs(FuncOp funcOp, ModuleOp moduleOp);
   bool areMemRefsNormalizable(FuncOp funcOp);
   void updateFunctionSignature(FuncOp funcOp, ModuleOp moduleOp);
   void setCalleesAndCallersNonNormalizable(FuncOp funcOp, ModuleOp moduleOp,
                                            DenseSet<FuncOp> &normalizableFuncs);
 };
 
 } // end anonymous namespace
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::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
   // normalizable are removed from this set.
   // TODO: Change this to work on FuncLikeOp once there is an operation
   // interface for it.
   DenseSet<FuncOp> normalizableFuncs;
   // Initialize `normalizableFuncs` with all the functions within a module.
   moduleOp.walk([&](FuncOp funcOp) { normalizableFuncs.insert(funcOp); });
 
   // Traverse through all the functions applying a filter which determines
   // whether that function is normalizable or not. All callers/callees of
   // a non-normalizable function will also become non-normalizable even if
   // they aren't passing any or specific non-normalizable memrefs. So,
   // functions which calls or get called by a non-normalizable becomes non-
   // normalizable functions themselves.
   moduleOp.walk([&](FuncOp funcOp) {
     if (normalizableFuncs.contains(funcOp)) {
       if (!areMemRefsNormalizable(funcOp)) {
+        LLVM_DEBUG(llvm::dbgs() << "@" << funcOp.getName() << " contains ops that cannot normalize MemRefs\n");
         // Since this function is not normalizable, we set all the caller
         // functions and the callees of this function as not normalizable.
         // TODO: Drop this conservative assumption in the future.
         setCalleesAndCallersNonNormalizable(funcOp, moduleOp,
                                             normalizableFuncs);
       }
     }
   });
 
+  LLVM_DEBUG(llvm::dbgs() << "Normalizing " << normalizableFuncs.size() << " functions\n");
   // Those functions which can be normalized are subjected to normalization.
   for (FuncOp &funcOp : normalizableFuncs)
     normalizeFuncOpMemRefs(funcOp, moduleOp);
 }
 
 /// Check whether all the uses of oldMemRef are either dereferencing uses or the
 /// op is of type : DeallocOp, CallOp or ReturnOp. Only if these constraints
 /// are satisfied will the value become a candidate for replacement.

 
 /// Set all the calling functions and the callees of the function as not
 /// normalizable.
 void NormalizeMemRefs::setCalleesAndCallersNonNormalizable(
     FuncOp funcOp, ModuleOp moduleOp, DenseSet<FuncOp> &normalizableFuncs) {
   if (!normalizableFuncs.contains(funcOp))
     return;
 
+  LLVM_DEBUG(llvm::dbgs() << "@" << funcOp.getName() << " calls or is called by non-normalizable function\n");
   normalizableFuncs.erase(funcOp);
   // Caller of the function.
   Optional<SymbolTable::UseRange> symbolUses = funcOp.getSymbolUses(moduleOp);
   for (SymbolTable::SymbolUse symbolUse : *symbolUses) {
     // TODO: Extend this for ops that are FunctionLike. This would require
     // creating an OpInterface for FunctionLike ops.
     FuncOp parentFuncOp = symbolUse.getUser()->getParentOfType<FuncOp>();
     for (FuncOp &funcOp : normalizableFuncs) {