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.

168–171

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

mlir/lib/Transforms/Utils/Utils.cpp

88

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

168–171

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

168–171

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

168–171

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

75–82

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.

168–171

Assertion message please.

mlir/lib/Transforms/Utils/Utils.cpp

74–77

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

161–164

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

169

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

264–266

Drop trivial braces - unrelated change as well.

268–270

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.

161–164

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 used to dereference and access a given
+      memref. Implementers of this interface must operate on at least one
+      memref operand.  The memref argument given to this interface much match
+      one of those memref operands.
+  }];
+
+  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() &&
+               "Expected memref argument to match memref operand");
+        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());
+      assert(mref == memref() &&
+             "Expected mref argument to match memref operand");
       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>();
     }
 
-    /// Returns the AffineMapAttr associated with 'memref'.
-    NamedAttribute getAffineMapAttrForMemRef(Value memref) {
-      assert(memref == getMemRef());
-      return {Identifier::get(getMapAttrName(), getContext()),
-              getAffineMapAttr()};
-    }
-
     static StringRef getMapAttrName() { return "map"; }
   }];
 }
 
 def AffineStoreOp : AffineStoreOpBase<"store"> {
   let summary = "affine store operation";
   let description = [{
     The "affine.store" op writes an element to a memref, where the index

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

static bool isOpLoopInvariant(Operation &op, Value indVar,
                               SmallPtrSetImpl<Operation *> &definedOps,
                               SmallPtrSetImpl<Operation *> &opsToHoist);
 
 static bool
 areAllOpsInTheBlockListInvariant(Region &blockList, Value indVar,
                                  SmallPtrSetImpl<Operation *> &definedOps,
                                  SmallPtrSetImpl<Operation *> &opsToHoist);
 
-static bool isMemRefDereferencingOp(Operation &op) {
-  // TODO: Support DMA Ops.
-  return isa<AffineReadOpInterface, AffineWriteOpInterface>(op);
-}
-
 // Returns true if the individual op is loop invariant.
 bool isOpLoopInvariant(Operation &op, Value indVar,
                        SmallPtrSetImpl<Operation *> &definedOps,
                        SmallPtrSetImpl<Operation *> &opsToHoist) {
   LLVM_DEBUG(llvm::dbgs() << "iterating on op: " << op;);
 
   if (isa<AffineIfOp>(op)) {
     if (!checkInvarianceOfNestedIfOps(&op, indVar, definedOps, opsToHoist)) {
       return false;
     }
   } else if (isa<AffineForOp>(op)) {
     // If the body of a predicated region has a for loop, we don't hoist the
     // 'affine.if'.
     return false;
   } else if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op)) {
     // TODO: Support DMA ops.
     return false;
   } else if (!isa<ConstantOp>(op)) {
     // Register op in the set of ops defined inside the loop. This set is used
     // to prevent hoisting ops that depend on other ops defined inside the loop
     // which are themselves not being hoisted.
     definedOps.insert(&op);
 
-    if (isMemRefDereferencingOp(op)) {
+    if (isa<AffineMapAccessInterface>(op)) {
       Value memref = isa<AffineReadOpInterface>(op)
                          ? cast<AffineReadOpInterface>(op).getMemRef()
                          : cast<AffineWriteOpInterface>(op).getMemRef();
       for (auto *user : memref.getUsers()) {
         // If this memref has a user that is a DMA, give up because these
         // operations write to this memref.
         if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op)) {
           return false;

mlir/lib/Transforms/LoopFusion.cpp

 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLoopFusionPass(unsigned fastMemorySpace,
                            uint64_t localBufSizeThreshold, bool maximalFusion) {
   return std::make_unique<LoopFusion>(fastMemorySpace, localBufSizeThreshold,
                                       maximalFusion);
 }
 
-// TODO: Replace when this is modeled through side-effects/op traits
-static bool isMemRefDereferencingOp(Operation &op) {
-  return isa<AffineReadOpInterface, AffineWriteOpInterface, AffineDmaStartOp,
-             AffineDmaWaitOp>(op);
-}
-
 namespace {
 
 // LoopNestStateCollector walks loop nests and collects load and store
 // operations, and whether or not an IfInst was encountered in the loop nest.
 struct LoopNestStateCollector {
   SmallVector<AffineForOp, 4> forOps;
   SmallVector<Operation *, 4> loadOpInsts;
   SmallVector<Operation *, 4> storeOpInsts;
bool writesToLiveInOrEscapingMemrefs(unsigned id) {
     for (auto *storeOpInst : node->stores) {
       auto memref = cast<AffineWriteOpInterface>(storeOpInst).getMemRef();
       auto *op = memref.getDefiningOp();
       // Return true if 'memref' is a block argument.
       if (!op)
         return true;
       // Return true if any use of 'memref' escapes the function.
       for (auto *user : memref.getUsers())
-        if (!isMemRefDereferencingOp(*user))
+        if (!isa<AffineMapAccessInterface>(*user))
           return true;
     }
     return false;
   }
 
   // Returns true iff there is an edge from node 'srcId' to node 'dstId' which
   // is for 'value' if non-null, or for any value otherwise. Returns false
   // otherwise.
for (auto *storeOpInst : node->stores) {
     // Check if 'memref' escapes because it's a block argument.
     if (memref.isa<BlockArgument>()) {
       escapingMemRefs.insert(memref);
       continue;
     }
     // Check if 'memref' escapes through a non-affine op (e.g., std load/store,
     // call op, etc.).
     for (Operation *user : memref.getUsers())
-      if (!isMemRefDereferencingOp(*user))
+      if (!isa<AffineMapAccessInterface>(*user))
         escapingMemRefs.insert(memref);
   }
 }
 
 } // end anonymous namespace
 
 // Initializes the data dependence graph by walking operations in 'f'.
 // Assigns each node in the graph a node id based on program order in 'f'.
static bool hasNonAffineUsersOnThePath(unsigned srcId, unsigned dstId,
   for (auto &idAndNode : mdg->nodes) {
     Operation *op = idAndNode.second.op;
     // Take care of operations between 'srcNode' and 'dstNode'.
     if (srcNode->op->isBeforeInBlock(op) && op->isBeforeInBlock(dstNode->op)) {
       // Walk inside the operation to find any use of the memref.
       // Interrupt the walk if found.
       auto walkResult = op->walk([&](Operation *user) {
         // Skip affine ops.
-        if (isMemRefDereferencingOp(*user))
+        if (isa<AffineMapAccessInterface>(*user))
           return WalkResult::advance();
         // Find a non-affine op that uses the memref.
         if (llvm::is_contained(users, user))
           return WalkResult::interrupt();
         return WalkResult::advance();
       });
       if (walkResult.wasInterrupted())
         return true;

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) {
     assert(false && "multiple dereferencing uses in a single op not supported");
     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)) {
-    if (!allowNonDereferencingOps)
+  auto affMapAccInterface = dyn_cast<AffineMapAccessInterface>(op);
+  if (!affMapAccInterface) {
+    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 =
+      affMapAccInterface.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(