MemRef Normalization for Dialects

[AlexEichenberger]

Summary

When dealing with dialects that will results in function calls to external libraries, it is important to be able to handle maps as some dialects may require mapped data. Before this patch, the detection of whether normalization can apply or not, operations are compared to an explicit list of operations (alloc, dealloc, return) or to the presence of specific operation interfaces (AffineReadOpInterface, AffineWriteOpInterface, AffineDMAStartOp, or AffineDMAWaitOp).

This patch add a trait, MemRefsNormalizable to determine if an operation can have its memrefs normalized.

This trait can be used in turn by dialects to assert that such operations are compatible with normalization of memrefs with nontrivial memory layout specification. An example is given in the literal tests.

[AlexEichenberger]

remove spurious commented code

[bondhugula]

Thanks for contributing this. You are missing the documentation and the definition of it in doc/Traits.md. A reviewer would typically start form there :-) - to make sure the semantics are understood. Also, the way you've defined it in the commit summary, it doesn't immediately look like DeallocOp and ReturnOp or anything return like would fall into that class and those would still have to be hardcoded / treated specially the way it is. We probably need to iterate/refine on the formal definition of this trait.

Commit summary:
may requires -> may require

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

661

Please keep the trait names in sorted order.

1392

Likewise.

mlir/include/mlir/IR/OpBase.td

1701

All comments are to be terminated with a period.

mlir/include/mlir/IR/OpDefinition.h

1215

Typo here. a the

1215

Please spell out the entire definition of this trait - here and in doc/Traits.md.

[AlexEichenberger]

Added requested doc, edited the summary to be more specific with respect to the contribution of this patch, and removed a prior TODO with respect to recognizing operations that can be normalized by systematically using the new trait.

[AlexEichenberger]

Added revisions suggested for comments.

[AlexEichenberger]

added proper order for traits

[AlexEichenberger]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

661

reordered systematically.

1392

done everywhere. Some had interface and traits, still ordered alphabetically. Let me know if this is not how it is done.

mlir/include/mlir/IR/OpDefinition.h

1215

Added the refs in doc/Traits and fixed the comment.

[bondhugula]

Looks good to me. Some superficial comments - first round.

mlir/docs/Traits.md

279

Naturalizable -> Normalizable

283

Nit: Spell fix: accommodate

285

Nit: proceed -> be performed

287

-> with non-identity layout ..?

mlir/include/mlir/IR/OpBase.td

1701

Nit: in -> in the

mlir/test/lib/Dialect/Test/TestOps.td

630

Nit: drop extra line.

[bondhugula]

1. Stretch question: do you see a scenario where some operand memrefs are normalizable but not others?

Since this 'MemRefsNormalizable' trait is more of "can/may be normalized" and is a first pass check because there are other things that could prevent it. For eg., CallOp is marked with this trait, but if a memref argument for the callee isn't normalizable due to a user op in that callee where it's not normalizable, that call op won't be normalized and transitively other ops won't be as well.

In the future or as a TODO, I think the right way to model this is using SideEffects, in particular, MemoryEffects (I briefly mentioned this on thread). This allows you to specify an effect at operand granularity. The design will require more thought though.

[bondhugula]

mlir/test/mlir-tblgen/op-memrefs-norm.mlir

2

The test case is in the wrong directory!

5

Please append these test cases to the existing ones for normalize-memrefs, or at least in the same directory.

[bondhugula]

mlir/test/mlir-tblgen/op-memrefs-norm.mlir

3

Please have a line or two here on what these test cases are testing.

11–15

Please don't use numbered SSA values in test cases. They can't be maintained - either %{{.*}} if you don't care about capturing and comparing, or capture it and use it. On the LHS, if you aren't capturing, you can even drop the part to the left of = ....

24–26

Likewise.

29

It would be good to have a line saying what this is testing for readability in the future.

36–38

No numbered SSA ids in check lines please.

[avarmapml]

Looks good to me.

As suggested by @bondhugula the transitive dependency of normalizable memrefs needs to be taken into account because currently that's how the normalize-memrefs pass works.
Can be added as TODO.

[AlexEichenberger]

Implemented all of the requested changes. Introduced at TODO for partially normalizable operations.

[AlexEichenberger]

mlir/docs/Traits.md

287

Implemented all of the requested text changes.

mlir/include/mlir/IR/OpBase.td

1701

changed the comment.

mlir/include/mlir/IR/OpDefinition.h

1215

Added the text of doc/Traits.md here as well, if that is what you asked.

mlir/test/lib/Dialect/Test/TestOps.td

630

Removed the extra line.

[AlexEichenberger]

In D86236#2231069, @bondhugula wrote:

1. Stretch question: do you see a scenario where some operand memrefs are normalizable but not others?

Since this 'MemRefsNormalizable' trait is more of "can/may be normalized" and is a first pass check because there are other things that could prevent it. For eg., CallOp is marked with this trait, but if a memref argument for the callee isn't normalizable due to a user op in that callee where it's not normalizable, that call op won't be normalized and transitively other ops won't be as well.

In the future or as a TODO, I think the right way to model this is using SideEffects, in particular, MemoryEffects (I briefly mentioned this on thread). This allows you to specify an effect at operand granularity. The design will require more thought though.

Good suggestion; I am leaving it as a TODO as this patch is already extending prior functionality to other dialects. Your approach should guide us when we will want to handle such operations.

[rriddle]

mlir/docs/Traits.md

279

Keep these sorted.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2129

Please format this such that the traits are aligned.

mlir/lib/Transforms/Utils/Utils.cpp

282

nit: Looks like an extra space here.

mlir/test/Transforms/normalize-memrefs-ops.mlir

1 ↗

(On Diff #287475)

Why are there multiple options here? You should generally only have one to keep the test focused, and not rely on the compounded behavior.

[AlexEichenberger]

Update to reflect latest comments.

[AlexEichenberger]

mlir/docs/Traits.md

279

Make sense, done.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

2129

aligned them.

mlir/lib/Transforms/Utils/Utils.cpp

282

Good catch, removed.

mlir/test/Transforms/normalize-memrefs-ops.mlir

1 ↗

(On Diff #287475)

I mistakenly believed that I needed the other option because the test uses a test op. Indeed it is not needed. Removed.

[bondhugula]

Some more minor comments - mostly on documentation and comments. Can you please fix?

mlir/docs/Traits.md

257–258

Not sure you should be using MemRefs in plural here. Is this for cross-referencing? I think you'll have to use a different syntax for that with a #.

259

access pattern -> access

-> .. is applied by rewriting accesses ...

260

this memory references -> those memref references or that memory reference.

mlir/include/mlir/IR/OpBase.td

1701

-> with non-identity maps ? Memrefs always have maps.

mlir/include/mlir/IR/OpDefinition.h

1220–1221

To be fixed in line with the trait doc.

mlir/test/lib/Dialect/Test/TestOps.td

626

Comment here please.

[AlexEichenberger]

performed all requested changes but one. Left a clarification request for the suggestion that I am not sure to understand properly.

mlir/docs/Traits.md

257–258

@bondhugula Are you asking to rename the trait to MemRefNormalizable where the "s" is removed ? I just want to make sure I understood your suggestion.

I am fine either way. In general, I don't like "s" in the middle. If it was up to my choice, I would have preferred no "s" to MemRef either. The trait property is on the operation so some of them have multiple MemRefs.

Just let me know how you want to go and I will make necessary changes.

260

performed the changes

mlir/include/mlir/IR/OpBase.td

1701

good catch, changed.

mlir/include/mlir/IR/OpDefinition.h

1220–1221

copied it over

mlir/test/lib/Dialect/Test/TestOps.td

626

Added as requested.

[bondhugula]

mlir/docs/Traits.md

257–258

I wasn't actually thinking about the trait renaming - in fact, I didn't think about that part. I was referring to MemRefs you were using further below which don't parse well in their sentences

an original MemRefs

an equivalent MemRefs

Are you referring to plural or singular? These are broken and have to be fixed.

[AlexEichenberger]

Completed all current requests for changes.

mlir/docs/Traits.md

257–258

make sense; done.

[bondhugula]

Thanks!

[bondhugula]

@AlexEichenberger Let me know if you are able to commit or if you'd like me to commit this for you.

[AlexEichenberger]

@bondhugula I don't believe I have commit privilege yet. Please commit for me, if you don't mind.

Thanks to all for your helpful comments, especially for guiding me toward the right approach. Happy to add the TODOs when the need arises.

[bondhugula]

In D86236#2238898, @AlexEichenberger wrote:

@bondhugula I don't believe I have commit privilege yet. Please commit for me, if you don't mind.

Sure. To confirm, should I be using these credentials on the commit: "Alexandre E. Eichenberger <[e-mail]>"? In previous ones, I see another address as well.

[AlexEichenberger]

yes please.

mlir/docs/Traits.md

 
 ### 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`.
 
 ### Symbol
 
 *   `OpTrait::Symbol` -- `Symbol`
 
 This trait is used for operations that define a
 [`Symbol`](SymbolsAndSymbolTables.md#symbol).
 
 ### SymbolTable
 
 *   `OpTrait::SymbolTable` -- `SymbolTable`
 
 This trait is used for operations that define a
 [`SymbolTable`](SymbolsAndSymbolTables.md#symbol-table).
 
 ### Terminator
 
 *   `OpTrait::IsTerminator` -- `Terminator`
 
 This trait provides verification and functionality for operations that are known
 to be [terminators](LangRef.md#terminator-operations).
+
+### MemRefsNaturalizable
+
+* `OpTrait::MemRefsNormalizable` -- `MemRefsNormalizable`
+
+This trait is used to flag operations that can accomodate `MemRefs` with
+nontrivial memory-layout specifications. This trait indicates that the
+normalization of memory layout can proceed for such operations.
+`MemRefs` normalization consists of replacing the original `MemRefs`
+with layout specifications to an equivalent `MemRefs` where the specified
+memory layout is applied to the access pattern and the type associated with
+this memory references.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

 //   affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%idx], %num_elements,
 //     %stride, %num_elt_per_stride : ...
 //
 // TODO: add additional operands to allow source and destination striding, and
 // multiple stride levels (possibly using AffineMaps to specify multiple levels
 // of striding).
 // TODO: Consider replacing src/dst memref indices with view memrefs.
 class AffineDmaStartOp : public Op<AffineDmaStartOp, OpTrait::VariadicOperands,
-                                   OpTrait::ZeroResult> {
+    OpTrait::ZeroResult, OpTrait::MemRefsNormalizable> {
 public:
   using Op::Op;
 
   static void build(OpBuilder &builder, OperationState &result, Value srcMemRef,
                     AffineMap srcMap, ValueRange srcIndices, Value destMemRef,
                     AffineMap dstMap, ValueRange destIndices, Value tagMemRef,
                     AffineMap tagMap, ValueRange tagIndices, Value numElements,
                     Value stride = nullptr, Value elementsPerStride = nullptr);

 //
 //   affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%index], %num_elements :
 //     memref<2048xf32, 0>, memref<256xf32, 1>, memref<1xi32, 2>
 //   ...
 //   ...
 //   affine.dma_wait %tag[%index], %num_elements : memref<1xi32, 2>
 //
 class AffineDmaWaitOp : public Op<AffineDmaWaitOp, OpTrait::VariadicOperands,
-                                  OpTrait::ZeroResult> {
+    OpTrait::ZeroResult, OpTrait::MemRefsNormalizable> {
 public:
   using Op::Op;
 
   static void build(OpBuilder &builder, OperationState &result, Value tagMemRef,
                     AffineMap tagMap, ValueRange tagIndices, Value numElements);
 
   static StringRef getOperationName() { return "affine.dma_wait"; }
 

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

let extraClassDeclaration = [{
     }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 class AffineLoadOpBase<string mnemonic, list<OpTrait> traits = []> :
-    Affine_Op<mnemonic, !listconcat(traits,
-        [DeclareOpInterfaceMethods<AffineReadOpInterface>])> {
+    Affine_Op<mnemonic, !listconcat(traits, [MemRefsNormalizable,
+        DeclareOpInterfaceMethods<AffineReadOpInterface>])> {
   let arguments = (ins Arg<AnyMemRef, "the reference to load from",
       [MemRead]>:$memref,
       Variadic<Index>:$indices);
 
   code extraClassDeclarationBase = [{
     /// Returns the operand index of the memref.
     unsigned getMemRefOperandIndex() { return 0; }
 
let extraClassDeclaration = [{
     static StringRef getIsDataCacheAttrName() { return "isDataCache"; }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 class AffineStoreOpBase<string mnemonic, list<OpTrait> traits = []> :
-    Affine_Op<mnemonic, !listconcat(traits,
-        [DeclareOpInterfaceMethods<AffineWriteOpInterface>])> {
+    Affine_Op<mnemonic, !listconcat(traits, [MemRefsNormalizable,
+        DeclareOpInterfaceMethods<AffineWriteOpInterface>])> {
   code extraClassDeclarationBase = [{
     /// Returns the operand index of the value to be stored.
     unsigned getStoredValOperandIndex() { return 0; }
 
     /// Returns the operand index of the memref.
     unsigned getMemRefOperandIndex() { return 1; }
 
     void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let assemblyFormat = [{
     $dest (`(` $destOperands^ `:` type($destOperands) `)`)? attr-dict
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // CallOp
 //===----------------------------------------------------------------------===//
 
-def CallOp : Std_Op<"call", [CallOpInterface]> {
+def CallOp : Std_Op<"call", [MemRefsNormalizable, CallOpInterface]> {
   let summary = "call operation";
   let description = [{
     The `call` operation represents a direct call to a function that is within
     the same symbol scope as the call. The operands and result types of the
     call must match the specified function type. The callee is encoded as a
     symbol reference attribute named "callee".
 
     Example:
def SinOp : FloatUnaryOp<"sin"> {
   }];
 }
 
 
 //===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
 
-def DeallocOp : Std_Op<"dealloc", [MemoryEffects<[MemFree]>]> {
+def DeallocOp : Std_Op<"dealloc",
+    [MemRefsNormalizable, MemoryEffects<[MemFree]>]> {
   let summary = "memory deallocation operation";
   let description = [{
     The `dealloc` operation frees the region of memory referenced by a memref
     which was originally created by the `alloc` operation.
     The `dealloc` operation should not be called on memrefs which alias an
     alloc'd memref (e.g. memrefs returned by `view` operations).
 
     Example:
def RemFOp : FloatArithmeticOp<"remf"> {
   let summary = "floating point division remainder operation";
 }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
 def ReturnOp : Std_Op<"return", [NoSideEffect, HasParent<"FuncOp">, ReturnLike,
-                                 Terminator]> {
+    Terminator, MemRefsNormalizable]> {
   let summary = "return operation";
   let description = [{
     The `return` operation represents a return operation within a function.
     The operation takes variable number of operands and produces no results.
     The operand number and types must match the signature of the function
     that contains the operation.
 
     Example:

mlir/include/mlir/IR/OpBase.td

 def SameOperandsAndResultType : NativeOpTrait<"SameOperandsAndResultType">;
 // Op has the same element type (or type itself, if scalar) for all operands.
 def SameOperandsElementType : NativeOpTrait<"SameOperandsElementType">;
 // Op has the same operand and result element type (or type itself, if scalar).
 def SameOperandsAndResultElementType :
   NativeOpTrait<"SameOperandsAndResultElementType">;
 // Op is a terminator.
 def Terminator : NativeOpTrait<"IsTerminator">;
+// Op can be safely normalized in presence of MemRefs with maps
+def MemRefsNormalizable : NativeOpTrait<"MemRefsNormalizable">;
 
 // Op's regions have a single block with the specified terminator.
 class SingleBlockImplicitTerminator<string op>
     : ParamNativeOpTrait<"SingleBlockImplicitTerminator", op>;
 
 // Op's parent operation is the provided one.
 class HasParent<string op>
     : ParamNativeOpTrait<"HasParent", op>;

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 provides a the definition for MemRefsNormalizable
+template <typename ConcrentType>
+struct MemRefsNormalizable
+    : public TraitBase<ConcrentType, MemRefsNormalizable> {
+};
+
 } // end namespace OpTrait
 
 //===----------------------------------------------------------------------===//
 // Operation Definition classes
 //===----------------------------------------------------------------------===//
 
 /// This provides public APIs that all operations should have.  The template
 /// argument 'ConcreteType' should be the concrete type by CRTP and the others
 /// are base classes by the policy pattern.

mlir/lib/Transforms/NormalizeMemRefs.cpp

moduleOp.walk([&](FuncOp funcOp) {
     }
   });
 
   // Those functions which can be normalized are subjected to normalization.
   for (FuncOp &funcOp : normalizableFuncs)
     normalizeFuncOpMemRefs(funcOp, moduleOp);
 }
 
-/// Return true if this operation dereferences one or more memref's.
-/// TODO: Temporary utility, will be replaced when this is modeled through
-/// side-effects/op traits.
-static bool isMemRefDereferencingOp(Operation &op) {
-  return isa<AffineReadOpInterface, AffineWriteOpInterface, AffineDmaStartOp,
-             AffineDmaWaitOp>(op);
-}
-
 /// 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.
 /// TODO: Extend this for DimOps.
 static bool isMemRefNormalizable(Value::user_range opUsers) {
   if (llvm::any_of(opUsers, [](Operation *op) {
-        if (isMemRefDereferencingOp(*op))
+        if (op->hasTrait<OpTrait::MemRefsNormalizable>())
           return false;
-        return !isa<DeallocOp, CallOp, ReturnOp>(*op);
+        return true;
       }))
     return false;
   return true;
 }
 
 /// Set all the calling functions and the callees of the function as not
 /// normalizable.
 void NormalizeMemRefs::setCalleesAndCallersNonNormalizable(

mlir/lib/Transforms/Utils/Utils.cpp

for (auto *op : oldMemRef.getUsers()) {
     // for the memref to be used in a non-dereferencing way outside of the
     // region where this replacement is happening.
     if (!isMemRefDereferencingOp(*op)) {
       if (!allowNonDereferencingOps)
         return failure();
       // Currently we support the following non-dereferencing ops to be a
       // candidate for replacement: Dealloc, CallOp and ReturnOp.
       // TODO: Add support for other kinds of ops.
-      if (!isa<DeallocOp, CallOp, ReturnOp>(*op))
+      if (! op->hasTrait<OpTrait::MemRefsNormalizable>())
         return failure();
     }
 
     // We'll first collect and then replace --- since replacement erases the op
     // that has the use, and that op could be postDomFilter or domFilter itself!
     opsToReplace.insert(op);
   }
 

mlir/test/lib/Dialect/Test/TestOps.td

 def OpK : TEST_Op<"op_k">, Arguments<(ins)>, Results<(outs I32)>;
 def : Pat<(OpJ), (OpK)>;
 
 // Test that natives calls are only called once during rewrites.
 def OpM : TEST_Op<"op_m"> {
   let arguments = (ins I32, OptionalAttr<I32Attr>:$optional_attr);
   let results = (outs I32);
 }
+
+// Test for memrefs normalization of an op with normalizable memrefs.
+def OpNorm : TEST_Op<"op_norm", [MemRefsNormalizable]> {
+  let arguments = (ins AnyMemRef:$X, AnyMemRef:$Y);
+}
+def OpNonNorm : TEST_Op<"op_nonnorm"> {
+  let arguments = (ins AnyMemRef:$X, AnyMemRef:$Y);
+}
+
+
 // Pattern add the argument plus a increasing static number hidden in
 // OpMTest function. That value is set into the optional argument.
 // That way, we will know if operations is called once or twice.
 def OpMGetNullAttr : NativeCodeCall<"Attribute()">;
 def OpMAttributeIsNull : Constraint<CPred<"! ($_self)">, "Attribute is null">;
 def OpMVal : NativeCodeCall<"OpMTest($_builder, $0)">;
 def : Pat<(OpM $attr, $optAttr), (OpM $attr, (OpMVal $attr) ),
     [(OpMAttributeIsNull:$optAttr)]>;

mlir/test/mlir-tblgen/op-memrefs-norm.mlir

+// RUN: mlir-opt -test-patterns -normalize-memrefs %s | FileCheck %s
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 64, d2 mod 32, d3 mod 64)>
+
+// CHECK-LABEL: test_norm
+func @test_norm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
+    %0 = alloc() : memref<1x16x14x14xf32, #map0>
+    "test.op_norm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32, #map0>) -> ()
+    dealloc %0 :  memref<1x16x14x14xf32, #map0>
+
+    // CHECK: %0 = alloc() : memref<1x16x1x1x32x64xf32>
+    // CHECK: "test.op_norm"(%arg0, %0) : (memref<1x16x1x1x32x64xf32>, memref<1x16x1x1x32x64xf32>) -> ()
+    // CHECK: dealloc %0 : memref<1x16x1x1x32x64xf32>
+    return
+}
+
+// CHECK-LABEL: test_nonnorm
+func @test_nonnorm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
+    %0 = alloc() : memref<1x16x14x14xf32, #map0>
+    "test.op_nonnorm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32, #map0>) -> ()
+    dealloc %0 :  memref<1x16x14x14xf32, #map0>
+
+    // CHECK: %0 = alloc() : memref<1x16x14x14xf32, #map0>
+    // CHECK: "test.op_nonnorm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32, #map0>) -> ()
+    // CHECK: dealloc %0 : memref<1x16x14x14xf32, #map0>
+    return
+}
+
+// CHECK-LABEL: test_norm_mix
+func @test_norm_mix(%arg0 : memref<1x16x1x1x32x64xf32>) -> () {
+    %0 = alloc() : memref<1x16x14x14xf32, #map0>
+    "test.op_norm"(%arg0, %0) : (memref<1x16x1x1x32x64xf32>, memref<1x16x14x14xf32, #map0>) -> ()
+    dealloc %0 :  memref<1x16x14x14xf32, #map0>
+
+    // CHECK: %0 = alloc() : memref<1x16x1x1x32x64xf32>
+    // CHECK: "test.op_norm"(%arg0, %0) : (memref<1x16x1x1x32x64xf32>, memref<1x16x1x1x32x64xf32>) -> ()
+    // CHECK: dealloc %0 : memref<1x16x1x1x32x64xf32>
+    return
+}