separate AffineMapAccessInterface from AffineRead/WriteOpInterface

[adstraw]

Summary

Separating the AffineMapAccessInterface from AffineRead/WriteOp interface so that dialects which extend Affine capabilities (e.g. PlaidML PXA = parallel extensions for Affine) can utilize relevant passes (e.g. MemRef normalization).

[ftynse]

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

151

The name should refer to its connection with affine maps. std.load dereferences a memref just fine, but this interfaces wouldn't make sense for it. Something like AffineMapAccessInterface would be more suitable.

167–169

The interface arguments is only used for this assertion. Let's drop both the argument and the assertion...

mlir/lib/Transforms/Utils/Utils.cpp

87

Use cast instead of dyn_cast if you don't check the result of cast for nullity.

Also, avoid using (dyn_)cast immediately after isa. Something like

auto derefInterface = dyn_cast<MemRefDereferencingInterface>(op);
if (!derefInterface) {
  // ...
  return ...;
}
NamedAttribute oldMapAttrPair = derefInterface.getAffineMapAttr();

should work here

[jbruestle]

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

167–169

One of the primary reasons for making this interface is to allow support for memref normalization of other operations not defined in this repo. For our use case, this includes an operation which does affine dereferencing of multiple memrefs, so the argument is required to link the memref to the appropriate AffineMapAttr.

[adstraw]

Addressing code review feedback.

[ftynse]

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

167–169

Then it needs to be appropriately documented. We have no knowledge of downstream projects and we can easily remove the code that is useless in core without warning unless it is clear from the code or documentation that it serves some purpose.

[adstraw]

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

167–169

Note that AffineDmaStartOp uses this interface and provides a non-default implementation of the getAffineMapAttrForMemRef method which returns a different AffineMap depending on whether the src or dst MemRef is passed in. I will update the description to make this more clear, as well.

[dcaballe]

Maybe there's something off on my side but I can't see the diff between the last version of the patch (Diff 3) and Base. When I set up that view in "Revision Contents->History", it shows the Diff against a previous version of the patch, which doesn't include all the changes. Something seems to be off.

[adstraw]

In D96284#2554537, @dcaballe wrote:

Maybe there's something off on my side but I can't see the diff between the last version of the patch (Diff 3) and Base. When I set up that view in "Revision Contents->History", it shows the Diff against a previous version of the patch, which doesn't include all the changes. Something seems to be off.

Apologize for the patching issues. I believe I have corrected the problem now. Diff 322729 shows all my intended changes from base.

[dcaballe]

It works after the last update. Thanks!
LGTM. Some minor comments. Please, wait for @ftynse's final approval.

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

154–155

• double space
• can extended -> can be extended?

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

311

Maybe adding a comment here about this being the implementation of a method of the new AffineMapAccessInterface interface would be useful.

[adstraw]

Thanks for the feedback @dcaballe

[bondhugula]

Mostly minor comments to address.

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

195

Terminate all comments with a full stop please.

309

Likewise.

mlir/lib/Transforms/Utils/Utils.cpp

77–81

Nit: please use braces here (per LLVM coding style).

[bondhugula]

Thanks for making the "memref deferencing ops" thing more systematic here. This is looking good to me. Mostly polishing and doc related comments.

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

153–154

I think we really need a couple more lines explaining the meaning of the interface better (especially for new users). For eg., there is an implication that the memref is being dereferenced / accessed - you actually have it in the name of the interface but this isn't reflected in the description here.

167

Assertion message please.

mlir/lib/Transforms/Utils/Utils.cpp

75–76

Better name here: affMapAcc -> affMapAccInterface (no increase in number of lines due to this).

[bondhugula]

Please also make the commit title more descriptive.

[adstraw]

Thanks for the feedback @ftynse and @bondhugula. Let me know if I missed anything.

[bondhugula]

mlir/lib/Transforms/Utils/Utils.cpp

31

This isn't the only user of such a utility:

/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp:64:static bool isMemRefDereferencingOp(Operation &op) {
./lib/Transforms/LoopFusion.cpp:72:static bool isMemRefDereferencingOp(Operation &op) {

Can we remove those as well?

[adstraw]

Good catch @bondhugula.

[bondhugula]

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

153–154

I still find this documentation not direct/complete enough when read by a new user. It's missing information that clearly specifies which ops should be implementing this interface. Consider rephrasing along:

access a given memref -> dereference and access a given memref operand

154

assume -> assumes

154

default implementation -> current implementation?
Because you'd need to change the API when handling multiple memrefs.

155

to -> on

160–163

You won't need a set method in the future?

168

to match operation memref -> to match memref operand

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

758

Nit: Update comment for the renamed argument.

mlir/lib/Transforms/Utils/Utils.cpp

263–265

Drop trivial braces - unrelated change as well.

267–269

Likewise - inadvertent change unrelated to this revision.

[bondhugula]

LGTM, thanks. Please do resolve the remaining comments.

[adstraw]

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

154

The reference to the "default implementation" in the description was done to address earlier comments regarding the need to pass in the memref to the interface at all + the need for the assert memref == op.getMemRef() in the default implementation. The reason for this is that operations implementing the interface can operate on more than one memref. I have the description to say that --- "Implementers of this interface must operate on at least one memref operand." I hope this clears things up.

160–163

I would prefer to address this in the future, as needed.

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

InterfaceMethod<
       /*args=*/(ins),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{
         ConcreteOp op = cast<ConcreteOp>(this->getOperation());
         return op.getAffineMapAttr().getValue();
       }]
     >,
     InterfaceMethod<
-      /*desc=*/"Returns the AffineMapAttr associated with 'memref'.",
-      /*retTy=*/"NamedAttribute",
-      /*methodName=*/"getAffineMapAttrForMemRef",
-      /*args=*/(ins "Value":$memref),
-      /*methodBody=*/[{}],
-      /*defaultImplementation=*/[{
-        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
-        assert(memref == getMemRef());
-        return {Identifier::get(op.getMapAttrName(), op.getContext()),
-                op.getAffineMapAttr()};
-      }]
-    >,
-    InterfaceMethod<
       /*desc=*/"Returns the value read by this operation.",
       /*retTy=*/"Value",
       /*methodName=*/"getValue",
       /*args=*/(ins),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{
         return cast<ConcreteOp>(this->getOperation());
       }]
InterfaceMethod<
       /*args=*/(ins),
       /*methodName=*/[{}],
       /*defaultImplementation=*/[{
         ConcreteOp op = cast<ConcreteOp>(this->getOperation());
         return op.getAffineMapAttr().getValue();
       }]
     >,
     InterfaceMethod<
-      /*desc=*/"Returns the AffineMapAttr associated with 'memref'.",
-      /*retTy=*/"NamedAttribute",
-      /*methodName=*/"getAffineMapAttrForMemRef",
-      /*args=*/(ins "Value":$memref),
+      /*desc=*/"Returns the value to store.",
+      /*retTy=*/"Value",
+      /*methodName=*/"getValueToStore",
+      /*args=*/(ins),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{
         ConcreteOp op = cast<ConcreteOp>(this->getOperation());
-        assert(memref == getMemRef());
-        return {Identifier::get(op.getMapAttrName(), op.getContext()),
-                op.getAffineMapAttr()};
+        return op.getOperand(op.getStoredValOperandIndex());
       }]
     >,
+  ];
+}
+
+def AffineMapAccessInterface : OpInterface<"AffineMapAccessInterface"> {
+  let description = [{
+      Interface to query the AffineMap associated with a given MemRef.
+      The default implementation assumes an operation with a single MemRef,
+      but the interface can be extended to support multiple MemRefs as well.
+  }];
+
+  let methods = [
     InterfaceMethod<
-      /*desc=*/"Returns the value to store.",
-      /*retTy=*/"Value",
-      /*methodName=*/"getValueToStore",
-      /*args=*/(ins),
+      /*desc=*/"Returns the AffineMapAttr associated with 'memref'.",
+      /*retTy=*/"NamedAttribute",
+      /*methodName=*/"getAffineMapAttrForMemRef",
+      /*args=*/(ins "Value":$memref),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{
         ConcreteOp op = cast<ConcreteOp>(this->getOperation());
-        return op.getOperand(op.getStoredValOperandIndex());
+        assert(memref == op.getMemRef());
+        return {Identifier::get(op.getMapAttrName(), op.getContext()),
+                op.getAffineMapAttr()};
       }]
     >,
   ];
 }
 
 #endif // AFFINEMEMORYOPINTERFACES

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

 //     %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::MemRefsNormalizable,
-                OpTrait::VariadicOperands, OpTrait::ZeroResult> {
+                OpTrait::VariadicOperands, OpTrait::ZeroResult,
+                AffineMapAccessInterface::Trait> {
 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);
public:
   }
 
   /// Returns the number of elements being transferred by this DMA operation.
   Value getNumElements() {
     return getOperand(getTagMemRefOperandIndex() + 1 +
                       getTagMap().getNumInputs());
   }
 
+  /// Impelements the AffineMapAccessInterface
   /// Returns the AffineMapAttr associated with 'memref'.
   NamedAttribute getAffineMapAttrForMemRef(Value memref) {
     if (memref == getSrcMemRef())
       return {Identifier::get(getSrcMapAttrName(), getContext()),
               getSrcMapAttr()};
     else if (memref == getDstMemRef())
       return {Identifier::get(getDstMapAttrName(), getContext()),
               getDstMapAttr()};

 //   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::MemRefsNormalizable,
-                OpTrait::VariadicOperands, OpTrait::ZeroResult> {
+                OpTrait::VariadicOperands, OpTrait::ZeroResult,
+                AffineMapAccessInterface::Trait> {
 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"; }
 
return {operand_begin() + 1,
             operand_begin() + 1 + getTagMap().getNumInputs()};
   }
 
   // Returns the rank (number of indices) of the tag memref.
   unsigned getTagMemRefRank() {
     return getTagMemRef().getType().cast<MemRefType>().getRank();
   }
 
+  /// Impelements the AffineMapAccessInterface
   /// Returns the AffineMapAttr associated with 'memref'.
   NamedAttribute getAffineMapAttrForMemRef(Value memref) {
     assert(memref == getTagMemRef());
     return {Identifier::get(getTagMapAttrName(), getContext()),
             getTagMapAttr()};
   }
 
   /// Returns the number of elements transferred in the associated DMA op.
   Value getNumElements() { return getOperand(1 + getTagMap().getNumInputs()); }
 

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

def AffineIfOp : Affine_Op<"if",
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 class AffineLoadOpBase<string mnemonic, list<OpTrait> traits = []> :
     Affine_Op<mnemonic, !listconcat(traits,
         [DeclareOpInterfaceMethods<AffineReadOpInterface>,
+        DeclareOpInterfaceMethods<AffineMapAccessInterface>,
         MemRefsNormalizable])> {
   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 getLowerBoundsMapAttrName() { return "lowerBoundsMap"; }
     static StringRef getUpperBoundsMapAttrName() { return "upperBoundsMap"; }
     static StringRef getStepsAttrName() { return "steps"; }
   }];
 
   let hasFolder = 1;
 }
 
-def AffinePrefetchOp : Affine_Op<"prefetch"> {
+def AffinePrefetchOp : Affine_Op<"prefetch",
+  [DeclareOpInterfaceMethods<AffineMapAccessInterface>]> {
   let summary = "affine prefetch operation";
   let description = [{
     The "affine.prefetch" op prefetches data from a memref location described
     with an affine subscript similar to affine.load, and has three attributes:
     a read/write specifier, a locality hint, and a cache type specifier as shown
     below:
 
     ```mlir
let extraClassDeclaration = [{
     }
 
     /// Returns the affine map used to index the memref for this operation.
     AffineMap getAffineMap() { return getAffineMapAttr().getValue(); }
     AffineMapAttr getAffineMapAttr() {
       return (*this)->getAttr(getMapAttrName()).cast<AffineMapAttr>();
     }
 
+    /// Impelements the AffineMapAccessInterface
     /// Returns the AffineMapAttr associated with 'memref'.
     NamedAttribute getAffineMapAttrForMemRef(Value mref) {
       assert(mref == memref());
       return {Identifier::get(getMapAttrName(), getContext()),
         getAffineMapAttr()};
     }
 
     /// Get affine map operands.
     operand_range getMapOperands() {
       return {operand_begin() + 1, operand_end()};
     }
 
     static StringRef getMapAttrName() { return "map"; }
     static StringRef getLocalityHintAttrName() { return "localityHint"; }
     static StringRef getIsWriteAttrName() { return "isWrite"; }
     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>,
+    DeclareOpInterfaceMethods<AffineMapAccessInterface>,
     MemRefsNormalizable])> {
   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); }
 
     /// Returns the affine map used to index the memref for this operation.
     AffineMapAttr getAffineMapAttr() {
       return (*this)->getAttr(getMapAttrName()).cast<AffineMapAttr>();
     }
 
+    /// Impelements the AffineMapAccessInterface
     /// Returns the AffineMapAttr associated with 'memref'.
     NamedAttribute getAffineMapAttrForMemRef(Value memref) {
       assert(memref == getMemRef());
       return {Identifier::get(getMapAttrName(), getContext()),
               getAffineMapAttr()};
     }
 
     static StringRef getMapAttrName() { return "map"; }

mlir/lib/Transforms/Utils/Utils.cpp

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/TypeSwitch.h"
 using namespace mlir;
 
-/// Return true if this operation dereferences one or more memref's.
-// Temporary utility: will be replaced when this is modeled through
-// side-effects/op traits. TODO
-static bool isMemRefDereferencingOp(Operation &op) {
-  return isa<AffineReadOpInterface, AffineWriteOpInterface, AffineDmaStartOp,
-             AffineDmaWaitOp>(op);
-}
-
-/// Return the AffineMapAttr associated with memory 'op' on 'memref'.
-static NamedAttribute getAffineMapAttrForMemRef(Operation *op, Value memref) {
-  return TypeSwitch<Operation *, NamedAttribute>(op)
-      .Case<AffineDmaStartOp, AffineReadOpInterface, AffinePrefetchOp,
-            AffineWriteOpInterface, AffineDmaWaitOp>(
-          [=](auto op) { return op.getAffineMapAttrForMemRef(memref); });
-}
-
 // Perform the replacement in `op`.
 LogicalResult mlir::replaceAllMemRefUsesWith(Value oldMemRef, Value newMemRef,
                                              Operation *op,
                                              ArrayRef<Value> extraIndices,
                                              AffineMap indexRemap,
                                              ArrayRef<Value> extraOperands,
                                              ArrayRef<Value> symbolOperands,
                                              bool allowNonDereferencingOps) {
if (usePositions.size() > 1) {
     return failure();
   }
 
   unsigned memRefOperandPos = usePositions.front();
 
   OpBuilder builder(op);
   // The following checks if op is dereferencing memref and performs the access
   // index rewrites.
-  if (!isMemRefDereferencingOp(*op)) {
+  auto affMapAcc = dyn_cast<AffineMapAccessInterface>(op);
+  if (!affMapAcc) {
     if (!allowNonDereferencingOps)
       // Failure: memref used in a non-dereferencing context (potentially
       // escapes); no replacement in these cases unless allowNonDereferencingOps
       // is set.
       return failure();
     op->setOperand(memRefOperandPos, newMemRef);
     return success();
   }
   // Perform index rewrites for the dereferencing op and then replace the op
-  NamedAttribute oldMapAttrPair = getAffineMapAttrForMemRef(op, oldMemRef);
+  NamedAttribute oldMapAttrPair =
+      affMapAcc.getAffineMapAttrForMemRef(oldMemRef);
   AffineMap oldMap = oldMapAttrPair.second.cast<AffineMapAttr>().getValue();
   unsigned oldMapNumInputs = oldMap.getNumInputs();
   SmallVector<Value, 4> oldMapOperands(
       op->operand_begin() + memRefOperandPos + 1,
       op->operand_begin() + memRefOperandPos + 1 + oldMapNumInputs);
 
   // Apply 'oldMemRefOperands = oldMap(oldMapOperands)'.
   SmallVector<Value, 4> oldMemRefOperands;
for (auto *op : oldMemRef.getUsers()) {
     // Skip dealloc's - no replacement is necessary, and a memref replacement
     // at other uses doesn't hurt these dealloc's.
     if (isa<DeallocOp>(op) && !replaceInDeallocOp)
       continue;
 
     // Check if the memref was used in a non-dereferencing context. It is fine
     // for the memref to be used in a non-dereferencing way outside of the
     // region where this replacement is happening.
-    if (!isMemRefDereferencingOp(*op)) {
+    if (!isa<AffineMapAccessInterface>(*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 (!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);
   }
 
   for (auto *op : opsToReplace) {
     if (failed(replaceAllMemRefUsesWith(