[mlir][Affine] Introduce affine memory interfaces

[dcaballe]

Summary

This patch introduces interfaces for read and write ops with affine
restrictions. I used read/write intead of load/store for the
interfaces so that they can also be implemented by dma ops.
For now, they are only implemented by affine.load, affine.store,
affine.vector_load and affine.vector_store.

For testing purposes, this patch also migrates affine loop fusion and
required analysis to use the new interfaces. No other changes are made
beyond that.

[andydavis1]

Thanks for doing this Diego! Will be very nice to have this. Especially for affine loop fusion and dma generation....

[ftynse]

LGTM in general.

I am wondering whether AffineMemoryOpInterfaces belongs in lib/Interfaces since it is specific to the Affine dialect. I'd rather put it together with the dialect, even though the transformations live in lib/Transforms .

mlir/include/mlir/Dialect/Affine/IR/CMakeLists.txt

2

Please do :)

mlir/lib/Analysis/AffineAnalysis.cpp

665–666

Nit: please make braces symmetric

mlir/lib/Transforms/LoopFusion.cpp

327

Nit: extract cast<AffineWriteLikeOpInterface>(storeOpInst).getMemRef() into a named variable for better formatting here

[bondhugula]

This looks really great! Thanks. CMake dependences and linking can be tricky; please make sure to build/test without LLD as well in case you are using LLD since there are a bunch of dependences / CMakelists.txt that need an update.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2564–2581 ↗

(On Diff #263574)

I missed why we need to introduce these here? Shouldn't they be shared via the op interface method?

[bondhugula]

This commit can't be marked NFC! It is adding functionality - generalizing several utilities to work on both the scalar and vector versions of load/store ops. You can drop the NFC from the title but mention in the commit summary that it isn't changing other functionality beyond migrating passes/utilities to the interface.

[bondhugula]

For the names AffineReadLikeOpInterface, AffineWriteLikeOpInterface, we could even consider dropping "Like" from it. It has become standard to use "Like" in interfaces because the prefixes are often the names of the ops. But here, "read" and "write" are already capturing the "like" part unlike say LoopLikeOp where you have a LoopOp and so it has to be LoopLikeOp. So, if you/others are fine with it, you could go for AffineReadOpInterface, AffineWriteOpInterface. DMAs, load/stores and vector load/stores - all read/write data.

[dcaballe]

Addressing feedback.

[dcaballe]

Thanks for the feedback!

I am wondering whether AffineMemoryOpInterfaces belongs in lib/Interfaces since it is specific to the Affine dialect. I'd rather put it together with the dialect, even though the transformations live in lib/Transforms

Good point. That makes sense. Done.

For the names AffineReadLikeOpInterface, AffineWriteLikeOpInterface, we could even consider dropping "Like" from it. It has become standard to use "Like" in interfaces because the prefixes are often the names of the ops. But here, "read" and "write" are already capturing the "like" part unlike say LoopLikeOp where you have a LoopOp and so it has to be LoopLikeOp. So, if you/others are fine with it, you could go for AffineReadOpInterface, AffineWriteOpInterface. DMAs, load/stores and vector load/stores - all read/write data.

Agreed! Done.

mlir/include/mlir/Dialect/Affine/IR/CMakeLists.txt

2

Sorry, I wanted to ask about it since I didn't see any other tablegen->tablegen dependencies like this one so I'm not sure how to do it. For example, we should also have a dependency with the loop-like interface here, right?

I added a dependency between both targets but this is not enough. I think we need to add a dependency between the dialect target and the interface .td file. Not sure how to do that cleanly.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2564–2581 ↗

(On Diff #263574)

Are you suggesting using the interface also as a base class? I assumed that interfaces shouldn't have any implementation. That could work for some methods but not all of them have the same implementation (e.g., getMapOperands). I could move those with the same implementation.

[bondhugula]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2564–2581 ↗

(On Diff #263574)

You are right that some of the ops that use this interface will not have the same implementation. But for my information, @ftynse, @rriddle, can Op Interfaces have implementations for things that are guaranteed to be shared among all ops using them? I remember seeing a default implementation being provided somewhere but I may be wrong.

Anyway, can these be moved into AffineLoadOpBase, AffineStoreOpBase? (extra class declarations)

[mehdi_amini]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2564–2581 ↗

(On Diff #263574)

I believe OpInterface method are type-erased virtual function, they can have a default implementation in which case an op may or may override the default implementation.

[ftynse]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

2564–2581 ↗

(On Diff #263574)

But for my information, @ftynse, @rriddle, can Op Interfaces have implementations for things that are guaranteed to be shared among all ops using them?

If an interface function declared in ODS has a body, and the ops don't redeclare it with DeclareOpInterfaceMethods, the default implementation is used in all ops. https://mlir.llvm.org/docs/OpDefinitions/#operation-interfaces has some examples.

[dcaballe]

Provide a default implementation for interface methods (not working, see next message)

[dcaballe]

I tried to change all the interface methods to have a default implementation but I'm hitting some problems. After reading again the documentation about interfaces, it's not clear to me what is the difference between providing a methodBody or a defaultImplementation. I had to dig into the generated code to have a bit better understanding but still, I couldn't make it work. I uploaded the patch with both approaches: methods in AffineWriteOpInterface has a methodBody, methods in AffineReadOpInterface has a defaultImplementation.

The problem that I'm hitting with methodBody is that the interface methods are not visible from the concrete op. For example, I cannot call getMemRef() using an AffineStoreOp object. Looking at the generated code, I see that the method getMemRef() is autogenerated for the class AffineWriteOpInterface, which is implemented like this:

Value AffineReadOpInterface::getMemRef() {
      return getImpl()->getMemRef(getOperation());
  }

However, there is no getMemRef() declared/defined in the same way for AffineStoreOp or AffineVectorStoreOp. Shouldn't they have similar autogenerated code? Am I missing something?

For defaultImplementation, I followed the doc example so interface methods are defined like this:

InterfaceMethod<
  /*desc=*/[{ Returns the memref operand to read from. }],
  /*retTy=*/"Value",
  /*methodName=*/"getMemRef",
  /*args=*/(ins),
  /*methodBody*/[{}],
  /*defaultImplementation=*/ [{
    ConcreteOp op = cast<ConcreteOp>(getOperation());
    return op.getOperand(op.getMemRefOperandIndex());
 }]
>,

However, the compiler complains about invoking getOperation() without object:

error: cannot call member function ‘mlir::Operation* mlir::Op<ConcreteT
ype, Traits>::getOperation() [with ConcreteType = mlir::AffineReadOpInterface; Traits = {}]’ without object

I get a bit lost in the internal details but ConcreteType = mlir::AffineReadOpInterface looks suspicious. Shouldn't it be AffineLoadOp/AffineVectorLoadOp?

I would appreciate some help, @ftynse @mehdi_amini @rriddle.

Thanks,
Diego

[ftynse]

TL;DR: "read" part is mostly correct. Use explicit this->getOperation() because class templates+inheritance. Also, cast<MemRefType>(....getType()) is incorrect, use ....getType().cast<MemRefType>() instead.

In D79829#2037697, @dcaballe wrote:

I tried to change all the interface methods to have a default implementation but I'm hitting some problems. After reading again the documentation about interfaces, it's not clear to me what is the difference between providing a methodBody or a defaultImplementation. I had to dig into the generated code to have a bit better understanding but still, I couldn't make it work. I uploaded the patch with both approaches: methods in AffineWriteOpInterface has a methodBody, methods in AffineReadOpInterface has a defaultImplementation.

Let's make a detour to try and untangle this (maybe we can update the doc afterwards). This is probably the most C++-intense part of the code base. I believe OpInterfaces implementation was inspired by the concept-based polymorphism from the inheritance is the base class of evil talk, transposed to LLVM and MLIR infrastructure (in particular LLVM-style casts and MLIR traits). The code example from the talk implements the concept-based polymorphic object that stores a pointer to the underlying object and dispatches polymorphic function calls to free functions differentiated by the type (stored as template parameter) of their first argument.

// This is the "interface" with pure virtual methods.
struct Concept {
  virtual void interfaceMethod() = 0;
};
template <typename Derived>
struct Model {
  void interfaceMethod() override {
    interfaceMethodImpl(data);
  }
  Derived data;
};

// This is a concrete implementation that does not need inheritance.
struct Concrete;
// And this is the implementation of interfaceMethod for Concrete.
void interfaceMethodImpl(Concrete) {}

// This class has polymorphic behavior without having virtual functions itself.
// We can see it as actual user-visible interface for our purposes, users
// will get instances of this class and will be able to work with them opaquely.
struct Polymorphic {
  template <typename Actual>
  Polymorphic(Actual &&a) : container(new Model<Actual>{forward<Actual>(a)}) {}
  void interfaceMethod() { container->interfaceMethod(); }
  std::unique_ptr<Concept> container;
};

Now, in MLIR, we don't actually need to store the underlying object because, as far as interfaces are concerned, there underlying object is an Operation that belongs to a block->region->parent-op->...->top-level-module-op that the caller owns. It would suffice to store a bare non-owning Operation *. Furthermore, we would like interface to integrate with the LLVM-style casting mechanisms. We're in luck, template <ConcereteOp, ...> class Op<ConcreteOp, ...> does exactly that: store an Operation * and provide isa/dyn_cast support. Now we start shifting a bit from the original concept-based idea with more inheritance and into the following:

struct Concept {
  // We will always work on some operation
  virtual void interfaceMethod(Operation *) = 0;
};
template <typename Derived>
struct Model : public Concept {
  void interfaceMethod(Operation *op) override {
    // Call the implementation from the Derived class, which is expected
    // to have a method with the compatible signature. (We could also
    // stick with free functions).
    cast<Derived>(op).interfaceMethod();
  }
  // No need to store the data anymore.
};

struct Polymorphic : public Op<Polymorphic, ...> {
  Poymorphic(Operation *op) : Op<Polymorphic, ...>(op), impl(...) {}
  void interfaceMethod() {
    impl->interfaceMethod(getOperation());
  }
  Concept *impl;
};

The remaining questions are how to we construct an instance of Model<Derived> that we could store in impl and how could we partially share the implementation. That's where the Traits mechanism comes into play. MLIR implements the Traits pattern that allows one to add generalized functions into specific Op classes. The scheme is roughly the following

template <typename Concrete>
struct Trait {
  void traitMethod() { 
    // This has access to actual operation and its type through CRTP.
    // cast<Concrete>(getOperation());
  };

  Operation *getOperation() {
    // In practice, this is bit more complex because it's provided by the base class
    // TraitBase through another layer of CRTP instead of requiring every Trait to
    // reimplement this from scratch.
    return static_cast<Concrete *>(this)->getOperation();
  }
};

// Again, the actual implementation is a more involved because of several layers
// of CRTP and a class variadic template Traits<...> that we use.
struct ConcreteOp : public Trait<ConcreteOp> {
  // This will have "traitMethod" available by inheritance from Trait.
};

Trait mechanism is also used in different places like verifiers and various precondition checkers, but for our purposes here it's basically a collection of templated base class _members_. The key difference with interfaces is that traits are allowed to have (static) data. We can then use this capability to store our instance of Model<Derived> in the trait. As an additional bonus, we get to reuse the Traits<...> machinery already implemented in operations. Now, since we are already giving all Ops that implement the interface a trait to support the interface, we might also use the properties of the trait itself -- in particular provide generic implementations of the methods that will be automatically available in all ops that have the trait (or the interface from which the trait is derived). The final scheme resembles the following

struct Concept {
  virtual void interfaceMethod(Operation *) = 0;
};
template <typename Derived>
struct Model : public Concept {
  void interfaceMethod(Operation *op) override {
    cast<Derived>(op).interfaceMethod();
  };
};
template <typename Derived>
struct Trait {
  static Concept &instance() {
    // This is where the per-class instance of Model lives.
    static Model<Derived> singleton;
    return singleton;
  }

  // We can define (or not!) the interface method implementation that
  // will be available in concrete ops by inheritance.
  void interfaceMethod() {
  }
};
struct Interface : public Op<Interface, ...> {
  Interface(Operation *op) : Op<Interface, ...(op) {
    // We leverage Op registration to store enable looking up Trait::instance in
    // an opaque way.
    impl = op->getAbstractOperation()->getInterfaceFor<Interface>();
  }

  void interfaceMethod() {
    impl->interfaceMethod();
  }
  Concept *impl;
};
struct ConcereteOp : public Op<ConcreteOp, Trait, ...> {
};

As usual, the actual implementation is more involved for practical purposes (e.g., Concept and Model are both nested inside the detail::InterfaceTraits class [not to be confused with Trait], Trait containing the instance is nested inside OpInterface to share the implementation and each specific Interface has a nested Trait that inherits OpInterface::Trait [both are class templates btw] with additional implementations, all Traits inherit from CRTP-base TraitBase, etc). Some of this complexity is hidden by table-genning Concept, Model, Interface and nested Trait classes. Unfortunately, sometimes it bites you back.

Two sources of confusion here are (1) the presence of identically-named functions that may or may not be inheritance-related and (2) the relation between Interfaces and Traits. For (1), concrete Ops are ultimately expected to implement exactly the signature described in Tablegen. These functions will be called by casting Operation * to the specific subclass of Op. Functions with identical signature will be provided in the Interface class, the rest is the dispatching mechanism letting them call similarly-named functions in the specific subclass of Op without directly knowing about them. For (2), Interfaces are extension of Op and are therefore wrappers around Operation * which can be stored or passed around; Traits are a way to share type-parameterized implementations across different Op subclasses. Interfaces rely on Traits to operate. Traits can be used to share the implementation of functions Interfaces require from Ops.

Finally, I can answer the original question about the difference between "methodBody" and "defaultImplementation". "methodBody" is the body of Model::interfaceMethod, if not provided in tablegen, it will just dispatch to ConcreteOp::interfaceMethod. Since there is no inheritance between Model<ConcreteOp> and ConcreteOp, these methods are not automatically available in ConcreteOp; on the contrary, they must be implemented for the entire mechanism to work. "defaultImplementation" is the body of "Trait::interfaceMethod" and will be automatically available in ConcreteOp. Combined with empty "methodBody", it will be the implementation that the interface will call into by default, hence the name. It can, however, be overridden by ConcreteOp and the interface will then call into that because it casts to ConcreteOp.

After this long detour, answers to your questions are relatively straightforward.

The problem that I'm hitting with methodBody is that the interface methods are not visible from the concrete op.

Interface methods are not visible because there is no inheritance from interface. Trait methods (referred to as default implementations) _are_ visible.

For example, I cannot call getMemRef() using an AffineStoreOp object. Looking at the generated code, I see that the method getMemRef() is autogenerated for the class AffineWriteOpInterface, which is implemented like this:

Value AffineReadOpInterface::getMemRef() {
      return getImpl()->getMemRef(getOperation());
  }

However, there is no getMemRef() declared/defined in the same way for AffineStoreOp or AffineVectorStoreOp. Shouldn't they have similar autogenerated code? Am I missing something?

An interface is like OOP interface in e.g. Java, it expects all classes that implement it to define the interface methods themselves. We have the possibility of providing a default implementation for them.

For defaultImplementation, I followed the doc example so interface methods are defined like this:

InterfaceMethod<
  /*desc=*/[{ Returns the memref operand to read from. }],
  /*retTy=*/"Value",
  /*methodName=*/"getMemRef",
  /*args=*/(ins),
  /*methodBody*/[{}],
  /*defaultImplementation=*/ [{
    ConcreteOp op = cast<ConcreteOp>(getOperation());
    return op.getOperand(op.getMemRefOperandIndex());
 }]
>,

However, the compiler complains about invoking getOperation() without object:

error: cannot call member function ‘mlir::Operation* mlir::Op<ConcreteT
ype, Traits>::getOperation() [with ConcreteType = mlir::AffineReadOpInterface; Traits = {}]’ without object

Because default implementation is placed in the Trait class template, which inherits from another class template (template <typename ConcreteOp> struct AffineReadOpInterfaceTrait : public OpInterface<AffineReadOpInterface, detail::AffineReadOpInterfaceInterfaceTraits>::Trait<ConcreteOp>) we need to explicitly disambiguate the call to the parent method (getOperaiton is defined in OpState, inherited by Op, inherited by OpInterface). this->getOperation() will work smoothly.

I get a bit lost in the internal details but ConcreteType = mlir::AffineReadOpInterface looks suspicious. Shouldn't it be AffineLoadOp/AffineVectorLoadOp?

No, it's correct. ConcreteType is a template parameter that points to the derived interface class for downcasting. If it were AffineLoadOp, we could erroneously assume any Op that implements AffineReadOpInterface _isa_ AffineLoadOp.

I would appreciate some help, @ftynse @mehdi_amini @rriddle.

The following diff makes everything compile and pass tests.

diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td b/mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td
index d9e8789a07a..8738000d8d5 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td
@@ -29,7 +29,7 @@ def AffineReadOpInterface : OpInterface<"AffineReadOpInterface"> {
       /*args=*/(ins),
       /*methodBody*/[{}],
       /*defaultImplementation=*/ [{
-        ConcreteOp op = cast<ConcreteOp>(getOperation());
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
         return op.getOperand(op.getMemRefOperandIndex());
       }]
     >,
@@ -40,8 +40,8 @@ def AffineReadOpInterface : OpInterface<"AffineReadOpInterface"> {
       /*args=*/(ins),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{
-        ConcreteOp op = cast<ConcreteOp>(getOperation());
-        return llvm::cast<MemRefType>(op.getMemRef().getType());
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getMemRef().getType().template cast<MemRefType>();
       }]
     >,
     InterfaceMethod<
@@ -51,7 +51,7 @@ def AffineReadOpInterface : OpInterface<"AffineReadOpInterface"> {
       /*args=*/(ins),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{
-        ConcreteOp op = cast<ConcreteOp>(getOperation());
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
         return llvm::drop_begin(op.getOperands(), 1);
       }]
     >,
@@ -63,7 +63,7 @@ def AffineReadOpInterface : OpInterface<"AffineReadOpInterface"> {
       /*args=*/(ins),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{
-        ConcreteOp op = cast<ConcreteOp>(getOperation());
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
         return op.getAffineMapAttr().getValue();
       }]
     >,
@@ -82,21 +82,33 @@ def AffineWriteOpInterface : OpInterface<"AffineWriteOpInterface"> {
       /*retTy=*/"Value",
       /*methodName=*/"getMemRef",
       /*args=*/(ins),
-      /*methodBody=*/ [{ return op.getOperand(op.getMemRefOperandIndex()); }]
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getOperand(op.getMemRefOperandIndex());
+      }]
     >,
     InterfaceMethod<
       /*desc=*/[{ Returns the type of the memref operand. }],
       /*retTy=*/"MemRefType",
       /*methodName=*/"getMemRefType",
       /*args=*/(ins),
-      /*methodBody=*/[{ return llvm::cast<MemRefType>(getMemRef().getType()); }]
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getMemRef().getType().template cast<MemRefType>();
+      }]
     >,
     InterfaceMethod<
       /*desc=*/[{ Returns affine map operands. }],
       /*retTy=*/"Operation::operand_range",
       /*methodName=*/"getMapOperands",
       /*args=*/(ins),
-      /*methodBody=*/[{ return llvm::drop_begin(op.getOperands(), 2); }]
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return llvm::drop_begin(op.getOperands(), 2);
+      }]
     >,
     InterfaceMethod<
       /*desc=*/[{ Returns the affine map used to index the memref for this
@@ -104,7 +116,11 @@ def AffineWriteOpInterface : OpInterface<"AffineWriteOpInterface"> {
       /*retTy=*/"AffineMap",
       /*methodName=*/"getAffineMap",
       /*args=*/(ins),
-      /*methodName=*/[{ return op.getAffineMapAttr().getValue(); }]
+      /*methodName=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getAffineMapAttr().getValue();
+      }]
     >,
   ];
 }

Thanks,
Diego

[dcaballe]

Thank you so much for such a detailed explanation! It's really good! I think we should move this to somewhere in the documentation because it's really elucidating!

Interfaces are extension of Op and are therefore wrappers around Operation * which can be stored or passed around;

Got it! I was missing this key detail.

Interface methods are not visible because there is no inheritance from interface. Trait methods (referred to as default implementations) _are_ visible.
An interface is like OOP interface in e.g. Java, it expects all classes that implement it to define the interface methods themselves. We have the possibility of providing a default implementation for them.

I think the only remaining question for my understanding is why we need both methodBody and defaultImplementation. Wouldn't the latter suffice? I'm trying to think of a use case for the former that can't be implemented with the latter (other than having a default implementation without visibility in the concrete op).

[dcaballe]

Adding Alex's changes. Thanks for the patch!
I'll proceed with the commit if no more comments.

[ftynse]

In D79829#2044547, @dcaballe wrote:

Interface methods are not visible because there is no inheritance from interface. Trait methods (referred to as default implementations) _are_ visible.
An interface is like OOP interface in e.g. Java, it expects all classes that implement it to define the interface methods themselves. We have the possibility of providing a default implementation for them.

I think the only remaining question for my understanding is why we need both methodBody and defaultImplementation. Wouldn't the latter suffice? I'm trying to think of a use case for the former that can't be implemented with the latter (other than having a default implementation without visibility in the concrete op).

I am afraid I don't have an answer for that. Generally, there is some redundancy between interfaces and traits that may be eventually reduced. There may be cases where you want to dispatch the interface function differently than the default approach (wrap/unwrap arguments, call different functions in different classes, etc.); this dispatch may be more expensive and you may not want to pay for it unless you use interfaces. That's the underlying idea of the concept approach: you don't get pay the virtual dispatch cost unless you need it.

[rriddle]

Sorry, I've been OOO but it looks like Alex was able to give an amazing run down. Thanks Alex!

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

14

nit: This is wrong.

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

49

nit: You can use a normal string literal for one line descriptions: "..."

60

nit: These are generally formatted as:

... [{

...

}]

mlir/include/mlir/Dialect/Affine/IR/CMakeLists.txt

4

Seems like this should be using add_mlir_interface?

mlir/lib/Analysis/AffineAnalysis.cpp

666

nit: I would have just inlined the use of the storeOp

mlir/lib/Dialect/Affine/IR/AffineMemoryOpInterfaces.cpp

2

Loop-Like?

[dcaballe]

Thanks @rriddle. I'll take care of this.

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

+//===- AffineMemoryOpInterfaces.h -------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains a set of interfaces for affine memory ops.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_AFFINEMEMORYOPINTERFACES_H_
+#define MLIR_INTERFACES_AFFINEMEMORYOPINTERFACES_H_
+
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/StandardTypes.h"
+
+namespace mlir {
+#include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h.inc"
+} // namespace mlir
+
+#endif // MLIR_INTERFACES_AFFINEMEMORYOPINTERFACES_H_

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

+//===- AffineMemoryOpInterfaces.td -------------------------*- tablegen -*-===//
+//
+// 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 contains a set of interfaces for affine memory ops.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_AFFINEMEMORYOPINTERFACES
+#define MLIR_AFFINEMEMORYOPINTERFACES
+
+include "mlir/IR/OpBase.td"
+
+def AffineReadOpInterface : OpInterface<"AffineReadOpInterface"> {
+  let description = [{
+      Interface to query characteristics of read-like ops with affine
+      restrictions.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{ Returns the memref operand to read from. }],
+      /*retTy=*/"Value",
+      /*methodName=*/"getMemRef",
+      /*args=*/(ins),
+      /*methodBody*/[{}],
+      /*defaultImplementation=*/ [{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getOperand(op.getMemRefOperandIndex());
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{ Returns the type of the memref operand. }],
+      /*retTy=*/"MemRefType",
+      /*methodName=*/"getMemRefType",
+      /*args=*/(ins),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getMemRef().getType().template cast<MemRefType>();
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{ Returns affine map operands. }],
+      /*retTy=*/"Operation::operand_range",
+      /*methodName=*/"getMapOperands",
+      /*args=*/(ins),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return llvm::drop_begin(op.getOperands(), 1);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{ Returns the affine map used to index the memref for this
+                  operation. }],
+      /*retTy=*/"AffineMap",
+      /*methodName=*/"getAffineMap",
+      /*args=*/(ins),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getAffineMapAttr().getValue();
+      }]
+    >,
+  ];
+}
+
+def AffineWriteOpInterface : OpInterface<"AffineWriteOpInterface"> {
+  let description = [{
+      Interface to query characteristics of write-like ops with affine
+      restrictions.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{ Returns the memref operand to write to. }],
+      /*retTy=*/"Value",
+      /*methodName=*/"getMemRef",
+      /*args=*/(ins),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getOperand(op.getMemRefOperandIndex());
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{ Returns the type of the memref operand. }],
+      /*retTy=*/"MemRefType",
+      /*methodName=*/"getMemRefType",
+      /*args=*/(ins),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getMemRef().getType().template cast<MemRefType>();
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{ Returns affine map operands. }],
+      /*retTy=*/"Operation::operand_range",
+      /*methodName=*/"getMapOperands",
+      /*args=*/(ins),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return llvm::drop_begin(op.getOperands(), 2);
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{ Returns the affine map used to index the memref for this
+                  operation. }],
+      /*retTy=*/"AffineMap",
+      /*methodName=*/"getAffineMap",
+      /*args=*/(ins),
+      /*methodName=*/[{}],
+      /*defaultImplementation=*/[{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getAffineMapAttr().getValue();
+      }]
+    >,
+  ];
+}
+
+#endif // MLIR_AFFINEMEMORYOPINTERFACES

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

 //===- AffineOps.h - MLIR Affine Operations -------------------------------===//
 //
 // 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 defines convenience types for working with Affine operations
 // in the MLIR operation set.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_AFFINE_IR_AFFINEOPS_H
 #define MLIR_DIALECT_AFFINE_IR_AFFINEOPS_H
 
+#include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 

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

 //===- AffineOps.td - Affine operation definitions ---------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Defines MLIR affine operations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef AFFINE_OPS
 #define AFFINE_OPS
 
 include "mlir/Dialect/Affine/IR/AffineOpsBase.td"
+include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def Affine_Dialect : Dialect {
   let name = "affine";
   let cppNamespace = "";
   let hasConstantMaterializer = 1;
 }
let extraClassDeclaration = [{
     }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 class AffineLoadOpBase<string mnemonic, list<OpTrait> traits = []> :
-    Affine_Op<mnemonic, traits> {
+    Affine_Op<mnemonic, !listconcat(traits,
+        [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; }
 
-    /// Get memref operand.
-    Value getMemRef() { return getOperand(getMemRefOperandIndex()); }
     void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }
-    MemRefType getMemRefType() {
-      return getMemRef().getType().cast<MemRefType>();
-    }
-
-    /// Get affine map operands.
-    operand_range getMapOperands() { return llvm::drop_begin(getOperands(), 1); }
 
     /// Returns the affine map used to index the memref for this operation.
-    AffineMap getAffineMap() { return getAffineMapAttr().getValue(); }
     AffineMapAttr getAffineMapAttr() {
       return 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 AffineLoadOp : AffineLoadOpBase<"load", []> {
+def AffineLoadOp : AffineLoadOpBase<"load"> {
   let summary = "affine load operation";
   let description = [{
     The "affine.load" op reads an element from a memref, where the index
     for each memref dimension is an affine expression of loop induction
     variables and symbols. The output of 'affine.load' is a new value with the
     same type as the elements of the memref. An affine expression of loop IVs
     and symbols must be specified for each dimension of the memref. The keyword
     'symbol' can be used to indicate SSA identifiers which are symbolic.
let extraClassDeclaration = [{
     static StringRef getIsDataCacheAttrName() { return "isDataCache"; }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 class AffineStoreOpBase<string mnemonic, list<OpTrait> traits = []> :
-    Affine_Op<mnemonic, traits> {
-
+    Affine_Op<mnemonic, !listconcat(traits,
+        [DeclareOpInterfaceMethods<AffineWriteOpInterface>])> {
   code extraClassDeclarationBase = [{
     /// Get value to be stored by store operation.
     Value getValueToStore() { return getOperand(0); }
 
     /// Returns the operand index of the memref.
     unsigned getMemRefOperandIndex() { return 1; }
 
-    /// Get memref operand.
-    Value getMemRef() { return getOperand(getMemRefOperandIndex()); }
     void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }
 
-    MemRefType getMemRefType() {
-      return getMemRef().getType().cast<MemRefType>();
-    }
-
-    /// Get affine map operands.
-    operand_range getMapOperands() { return llvm::drop_begin(getOperands(), 2); }
-
     /// Returns the affine map used to index the memref for this operation.
-    AffineMap getAffineMap() { return getAffineMapAttr().getValue(); }
     AffineMapAttr getAffineMapAttr() {
       return 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", []> {
+def AffineStoreOp : AffineStoreOpBase<"store"> {
   let summary = "affine store operation";
   let description = [{
     The "affine.store" op writes an element to a memref, where the index
     for each memref dimension is an affine expression of loop induction
     variables and symbols. The 'affine.store' op stores a new value which is the
     same type as the elements of the memref. An affine expression of loop IVs
     and symbols must be specified for each dimension of the memref. The keyword
     'symbol' can be used to indicate SSA identifiers which are symbolic.
def AffineTerminatorOp :
   // No custom parsing/printing form.
   let parser = ?;
   let printer = ?;
 
   // Fully specified by traits.
   let verifier = ?;
 }
 
-def AffineVectorLoadOp : AffineLoadOpBase<"vector_load", []> {
+def AffineVectorLoadOp : AffineLoadOpBase<"vector_load"> {
   let summary = "affine vector load operation";
   let description = [{
     The "affine.vector_load" is the vector counterpart of
     [affine.load](#affineload-operation). It reads a slice from a
     [MemRef](../LangRef.md#memref-type), supplied as its first operand,
     into a [vector](../LangRef.md#vector-type) of the same base elemental type.
     The index for each memref dimension is an affine expression of loop induction
     variables and symbols. These indices determine the start position of the read
def AffineVectorLoadOp : AffineLoadOpBase<"vector_load"> {
 
   let extraClassDeclaration = extraClassDeclarationBase # [{
     VectorType getVectorType() {
       return result().getType().cast<VectorType>();
     }
   }];
 }
 
-def AffineVectorStoreOp : AffineStoreOpBase<"vector_store", []> {
+def AffineVectorStoreOp : AffineStoreOpBase<"vector_store"> {
   let summary = "affine vector store operation";
   let description = [{
     The "affine.vector_store" is the vector counterpart of
     [affine.store](#affinestore-affinestoreop). It writes a
     [vector](../LangRef.md#vector-type), supplied as its first operand,
     into a slice within a [MemRef](../LangRef.md#memref-type) of the same base
     elemental type, supplied as its second operand.
     The index for each memref dimension is an affine expression of loop

mlir/include/mlir/Dialect/Affine/IR/CMakeLists.txt

 add_mlir_dialect(AffineOps affine)
 add_mlir_doc(AffineOps -gen-op-doc AffineOps Dialects/)
+
+set(LLVM_TARGET_DEFINITIONS AffineMemoryOpInterfaces.td)
+mlir_tablegen(AffineMemoryOpInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(AffineMemoryOpInterfaces.cpp.inc -gen-op-interface-defs)
+add_public_tablegen_target(MLIRAffineMemoryOpInterfacesIncGen)
+add_dependencies(mlir-generic-headers MLIRAffineMemoryOpInterfacesIncGen)
+
+add_dependencies(MLIRAffineOpsIncGen MLIRAffineMemoryOpInterfacesIncGen)

mlir/lib/Analysis/AffineAnalysis.cpp

static void computeDirectionVector(
   }
 }
 
 // Populates 'accessMap' with composition of AffineApplyOps reachable from
 // indices of MemRefAccess.
 void MemRefAccess::getAccessMap(AffineValueMap *accessMap) const {
   // Get affine map from AffineLoad/Store.
   AffineMap map;
-  if (auto loadOp = dyn_cast<AffineLoadOp>(opInst))
+  if (auto loadOp = dyn_cast<AffineReadOpInterface>(opInst)) {
     map = loadOp.getAffineMap();
-  else if (auto storeOp = dyn_cast<AffineStoreOp>(opInst))
+  } else {
+    auto storeOp = cast<AffineWriteOpInterface>(opInst);
     map = storeOp.getAffineMap();
+  }
   SmallVector<Value, 8> operands(indices.begin(), indices.end());
   fullyComposeAffineMapAndOperands(&map, &operands);
   map = simplifyAffineMap(map);
   canonicalizeMapAndOperands(&map, &operands);
   accessMap->reset(map, operands);
 }
 
 // Builds a flat affine constraint system to check if there exists a dependence
LLVM_DEBUG(llvm::dbgs() << "Checking for dependence at depth: "
                           << Twine(loopDepth) << " between:\n";);
   LLVM_DEBUG(srcAccess.opInst->dump(););
   LLVM_DEBUG(dstAccess.opInst->dump(););
 
   // Return 'NoDependence' if these accesses do not access the same memref.
   if (srcAccess.memref != dstAccess.memref)
     return DependenceResult::NoDependence;
 
-  // Return 'NoDependence' if one of these accesses is not an AffineStoreOp.
-  if (!allowRAR && !isa<AffineStoreOp>(srcAccess.opInst) &&
-      !isa<AffineStoreOp>(dstAccess.opInst))
+  // Return 'NoDependence' if one of these accesses is not an
+  // AffineWriteOpInterface.
+  if (!allowRAR && !isa<AffineWriteOpInterface>(srcAccess.opInst) &&
+      !isa<AffineWriteOpInterface>(dstAccess.opInst))
     return DependenceResult::NoDependence;
 
   // Get composed access function for 'srcAccess'.
   AffineValueMap srcAccessMap;
   srcAccess.getAccessMap(&srcAccessMap);
 
   // Get composed access function for 'dstAccess'.
   AffineValueMap dstAccessMap;

 /// Gathers dependence components for dependences between all ops in loop nest
 /// rooted at 'forOp' at loop depths in range [1, maxLoopDepth].
 void mlir::getDependenceComponents(
     AffineForOp forOp, unsigned maxLoopDepth,
     std::vector<SmallVector<DependenceComponent, 2>> *depCompsVec) {
   // Collect all load and store ops in loop nest rooted at 'forOp'.
   SmallVector<Operation *, 8> loadAndStoreOpInsts;
   forOp.getOperation()->walk([&](Operation *opInst) {
-    if (isa<AffineLoadOp>(opInst) || isa<AffineStoreOp>(opInst))
+    if (isa<AffineReadOpInterface>(opInst) ||
+        isa<AffineWriteOpInterface>(opInst))
       loadAndStoreOpInsts.push_back(opInst);
   });
 
   unsigned numOps = loadAndStoreOpInsts.size();
   for (unsigned d = 1; d <= maxLoopDepth; ++d) {
     for (unsigned i = 0; i < numOps; ++i) {
       auto *srcOpInst = loadAndStoreOpInsts[i];
       MemRefAccess srcAccess(srcOpInst);

mlir/lib/Analysis/Utils.cpp

 // region:  {memref = %A, write = false, {%i <= m0 <= %i + 7} }
 // The last field is a 2-d FlatAffineConstraints symbolic in %i.
 //
 // TODO(bondhugula): extend this to any other memref dereferencing ops
 // (dma_start, dma_wait).
 LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
                                     ComputationSliceState *sliceState,
                                     bool addMemRefDimBounds) {
-  assert((isa<AffineLoadOp>(op) || isa<AffineStoreOp>(op)) &&
-         "affine load/store op expected");
+  assert((isa<AffineReadOpInterface>(op) || isa<AffineWriteOpInterface>(op)) &&
+         "affine read/write op expected");
 
   MemRefAccess access(op);
   memref = access.memref;
   write = access.isStore();
 
   unsigned rank = access.getRank();
 
   LLVM_DEBUG(llvm::dbgs() << "MemRefRegion::compute: " << *op
for (unsigned i = 0, e = memRefType.getRank(); i < e; i++) {
     sizeInBytes = sizeInBytes * memRefType.getDimSize(i);
   }
   return sizeInBytes;
 }
 
 template <typename LoadOrStoreOp>
 LogicalResult mlir::boundCheckLoadOrStoreOp(LoadOrStoreOp loadOrStoreOp,
                                             bool emitError) {
-  static_assert(
-      llvm::is_one_of<LoadOrStoreOp, AffineLoadOp, AffineStoreOp>::value,
-      "argument should be either a AffineLoadOp or a AffineStoreOp");
+  static_assert(llvm::is_one_of<LoadOrStoreOp, AffineReadOpInterface,
+                                AffineWriteOpInterface>::value,
+                "argument should be either a AffineReadOpInterface or a "
+                "AffineWriteOpInterface");
 
   Operation *op = loadOrStoreOp.getOperation();
   MemRefRegion region(op->getLoc());
   if (failed(region.compute(op, /*loopDepth=*/0, /*sliceState=*/nullptr,
                             /*addMemRefDimBounds=*/false)))
     return success();
 
   LLVM_DEBUG(llvm::dbgs() << "Memory region");
if (outOfBounds && emitError) {
       loadOrStoreOp.emitOpError()
           << "memref out of lower bound access along dimension #" << (r + 1);
     }
   }
   return failure(outOfBounds);
 }
 
 // Explicitly instantiate the template so that the compiler knows we need them!
-template LogicalResult mlir::boundCheckLoadOrStoreOp(AffineLoadOp loadOp,
-                                                     bool emitError);
-template LogicalResult mlir::boundCheckLoadOrStoreOp(AffineStoreOp storeOp,
-                                                     bool emitError);
+template LogicalResult
+mlir::boundCheckLoadOrStoreOp(AffineReadOpInterface loadOp, bool emitError);
+template LogicalResult
+mlir::boundCheckLoadOrStoreOp(AffineWriteOpInterface storeOp, bool emitError);
 
 // Returns in 'positions' the Block positions of 'op' in each ancestor
 // Block from the Block containing operation, stopping at 'limitBlock'.
 static void findInstPosition(Operation *op, Block *limitBlock,
                              SmallVectorImpl<unsigned> *positions) {
   Block *block = op->getBlock();
   while (block != limitBlock) {
     // FIXME: This algorithm is unnecessarily O(n) and should be improved to not
for (unsigned j = 0, numOpsB = opsB.size(); j < numOpsB; ++j) {
         continue;
       // Check if 'loopDepth' exceeds nesting depth of src/dst ops.
       if ((!isBackwardSlice && loopDepth > getNestingDepth(opsA[i])) ||
           (isBackwardSlice && loopDepth > getNestingDepth(opsB[j]))) {
         LLVM_DEBUG(llvm::dbgs() << "Invalid loop depth\n");
         return failure();
       }
 
-      bool readReadAccesses = isa<AffineLoadOp>(srcAccess.opInst) &&
-                              isa<AffineLoadOp>(dstAccess.opInst);
+      bool readReadAccesses = isa<AffineReadOpInterface>(srcAccess.opInst) &&
+                              isa<AffineReadOpInterface>(dstAccess.opInst);
       FlatAffineConstraints dependenceConstraints;
       // Check dependence between 'srcAccess' and 'dstAccess'.
       DependenceResult result = checkMemrefAccessDependence(
           srcAccess, dstAccess, /*loopDepth=*/numCommonLoops + 1,
           &dependenceConstraints, /*dependenceComponents=*/nullptr,
           /*allowRAR=*/readReadAccesses);
       if (result.value == DependenceResult::Failure) {
         LLVM_DEBUG(llvm::dbgs() << "Dependence check failed\n");
void mlir::getComputationSliceState(
   sliceState->ubOperands.resize(numSliceLoopIVs, sliceBoundOperands);
 
   // Set destination loop nest insertion point to block start at 'dstLoopDepth'.
   sliceState->insertPoint =
       isBackwardSlice ? dstLoopIVs[loopDepth - 1].getBody()->begin()
                       : std::prev(srcLoopIVs[loopDepth - 1].getBody()->end());
 
   llvm::SmallDenseSet<Value, 8> sequentialLoops;
-  if (isa<AffineLoadOp>(depSourceOp) && isa<AffineLoadOp>(depSinkOp)) {
+  if (isa<AffineReadOpInterface>(depSourceOp) &&
+      isa<AffineReadOpInterface>(depSinkOp)) {
     // For read-read access pairs, clear any slice bounds on sequential loops.
     // Get sequential loops in loop nest rooted at 'srcLoopIVs[0]'.
     getSequentialLoops(isBackwardSlice ? srcLoopIVs[0] : dstLoopIVs[0],
                        &sequentialLoops);
   }
   // Clear all sliced loop bounds beginning at the first sequential loop, or
   // first loop with a slice fusion barrier attribute..
   // TODO(andydavis, bondhugula) Use MemRef read/write regions instead of
if (AffineMap ubMap = sliceState->ubs[i])
       forOp.setUpperBound(sliceState->ubOperands[i], ubMap);
   }
   return sliceLoopNest;
 }
 
 // Constructs  MemRefAccess populating it with the memref, its indices and
 // opinst from 'loadOrStoreOpInst'.
 MemRefAccess::MemRefAccess(Operation *loadOrStoreOpInst) {
-  if (auto loadOp = dyn_cast<AffineLoadOp>(loadOrStoreOpInst)) {
+  if (auto loadOp = dyn_cast<AffineReadOpInterface>(loadOrStoreOpInst)) {
     memref = loadOp.getMemRef();
     opInst = loadOrStoreOpInst;
     auto loadMemrefType = loadOp.getMemRefType();
     indices.reserve(loadMemrefType.getRank());
     for (auto index : loadOp.getMapOperands()) {
       indices.push_back(index);
     }
   } else {
-    assert(isa<AffineStoreOp>(loadOrStoreOpInst) && "load/store op expected");
-    auto storeOp = dyn_cast<AffineStoreOp>(loadOrStoreOpInst);
+    assert(isa<AffineWriteOpInterface>(loadOrStoreOpInst) &&
+           "Affine read/write op expected");
+    auto storeOp = cast<AffineWriteOpInterface>(loadOrStoreOpInst);
     opInst = loadOrStoreOpInst;
     memref = storeOp.getMemRef();
     auto storeMemrefType = storeOp.getMemRefType();
     indices.reserve(storeMemrefType.getRank());
     for (auto index : storeOp.getMapOperands()) {
       indices.push_back(index);
     }
   }
 }
 
 unsigned MemRefAccess::getRank() const {
   return memref.getType().cast<MemRefType>().getRank();
 }
 
-bool MemRefAccess::isStore() const { return isa<AffineStoreOp>(opInst); }
+bool MemRefAccess::isStore() const {
+  return isa<AffineWriteOpInterface>(opInst);
+}
 
 /// Returns the nesting depth of this statement, i.e., the number of loops
 /// surrounding this statement.
 unsigned mlir::getNestingDepth(Operation *op) {
   Operation *currOp = op;
   unsigned depth = 0;
   while ((currOp = currOp->getParentOp())) {
     if (isa<AffineForOp>(currOp))

 static Optional<int64_t> getMemoryFootprintBytes(Block &block,
                                                  Block::iterator start,
                                                  Block::iterator end,
                                                  int memorySpace) {
   SmallDenseMap<Value, std::unique_ptr<MemRefRegion>, 4> regions;
 
   // Walk this 'affine.for' operation to gather all memory regions.
   auto result = block.walk(start, end, [&](Operation *opInst) -> WalkResult {
-    if (!isa<AffineLoadOp>(opInst) && !isa<AffineStoreOp>(opInst)) {
+    if (!isa<AffineReadOpInterface>(opInst) &&
+        !isa<AffineWriteOpInterface>(opInst)) {
       // Neither load nor a store op.
       return WalkResult::advance();
     }
 
     // Compute the memref region symbolic in any IVs enclosing this block.
     auto region = std::make_unique<MemRefRegion>(opInst->getLoc());
     if (failed(
             region->compute(opInst,
void mlir::getSequentialLoops(AffineForOp forOp,
   });
 }
 
 /// Returns true if 'forOp' is parallel.
 bool mlir::isLoopParallel(AffineForOp forOp) {
   // Collect all load and store ops in loop nest rooted at 'forOp'.
   SmallVector<Operation *, 8> loadAndStoreOpInsts;
   auto walkResult = forOp.walk([&](Operation *opInst) -> WalkResult {
-    if (isa<AffineLoadOp>(opInst) || isa<AffineStoreOp>(opInst))
+    if (isa<AffineReadOpInterface>(opInst) ||
+        isa<AffineWriteOpInterface>(opInst))
       loadAndStoreOpInsts.push_back(opInst);
     else if (!isa<AffineForOp>(opInst) && !isa<AffineTerminatorOp>(opInst) &&
              !isa<AffineIfOp>(opInst) &&
              !MemoryEffectOpInterface::hasNoEffect(opInst))
       return WalkResult::interrupt();
 
     return WalkResult::advance();
   });

mlir/lib/Dialect/Affine/IR/AffineMemoryOpInterfaces.cpp

+//===- AffineMemoryOpInterfaces.cpp - Loop-like operations in MLIR --------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Affine Memory Op Interfaces
+//===----------------------------------------------------------------------===//
+
+/// Include the definitions of the affine memory op interfaces.
+#include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.cpp.inc"

mlir/lib/Dialect/Affine/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRAffineOps
+  AffineMemoryOpInterfaces.cpp
   AffineOps.cpp
   AffineValueMap.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
 
   DEPENDS
+  MLIRAffineMemoryOpInterfacesIncGen
   MLIRAffineOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIREDSC
   MLIRIR
   MLIRLoopLikeInterface
   MLIRSideEffectInterfaces
   MLIRStandardOps
   )

mlir/lib/Transforms/LoopFusion.cpp

 mlir::createLoopFusionPass(unsigned fastMemorySpace,
                            uint64_t localBufSizeThreshold, bool maximalFusion) {
   return std::make_unique<LoopFusion>(fastMemorySpace, localBufSizeThreshold,
                                       maximalFusion);
 }
 
 // TODO(b/117228571) Replace when this is modeled through side-effects/op traits
 static bool isMemRefDereferencingOp(Operation &op) {
-  if (isa<AffineLoadOp>(op) || isa<AffineStoreOp>(op) ||
+  if (isa<AffineReadOpInterface>(op) || isa<AffineWriteOpInterface>(op) ||
       isa<AffineDmaStartOp>(op) || isa<AffineDmaWaitOp>(op))
     return true;
   return false;
 }
 
 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 hasNonForRegion = false;
 
   void collect(Operation *opToWalk) {
     opToWalk->walk([&](Operation *op) {
       if (isa<AffineForOp>(op))
         forOps.push_back(cast<AffineForOp>(op));
       else if (op->getNumRegions() != 0)
         hasNonForRegion = true;
-      else if (isa<AffineLoadOp>(op))
+      else if (isa<AffineReadOpInterface>(op))
         loadOpInsts.push_back(op);
-      else if (isa<AffineStoreOp>(op))
+      else if (isa<AffineWriteOpInterface>(op))
         storeOpInsts.push_back(op);
     });
   }
 };
 
 // MemRefDependenceGraph is a graph data structure where graph nodes are
 // top-level operations in a FuncOp which contain load/store ops, and edges
 // are memref dependences between the nodes.
struct Node {
     // List of store op insts.
     SmallVector<Operation *, 4> stores;
     Node(unsigned id, Operation *op) : id(id), op(op) {}
 
     // Returns the load op count for 'memref'.
     unsigned getLoadOpCount(Value memref) {
       unsigned loadOpCount = 0;
       for (auto *loadOpInst : loads) {
-        if (memref == cast<AffineLoadOp>(loadOpInst).getMemRef())
+        if (memref == cast<AffineReadOpInterface>(loadOpInst).getMemRef())
           ++loadOpCount;
       }
       return loadOpCount;
     }
 
     // Returns the store op count for 'memref'.
     unsigned getStoreOpCount(Value memref) {
       unsigned storeOpCount = 0;
       for (auto *storeOpInst : stores) {
-        if (memref == cast<AffineStoreOp>(storeOpInst).getMemRef())
+        if (memref == cast<AffineWriteOpInterface>(storeOpInst).getMemRef())
           ++storeOpCount;
       }
       return storeOpCount;
     }
 
     // Returns all store ops in 'storeOps' which access 'memref'.
     void getStoreOpsForMemref(Value memref,
                               SmallVectorImpl<Operation *> *storeOps) {
       for (auto *storeOpInst : stores) {
-        if (memref == cast<AffineStoreOp>(storeOpInst).getMemRef())
+        if (memref == cast<AffineWriteOpInterface>(storeOpInst).getMemRef())
           storeOps->push_back(storeOpInst);
       }
     }
 
     // Returns all load ops in 'loadOps' which access 'memref'.
     void getLoadOpsForMemref(Value memref,
                              SmallVectorImpl<Operation *> *loadOps) {
       for (auto *loadOpInst : loads) {
-        if (memref == cast<AffineLoadOp>(loadOpInst).getMemRef())
+        if (memref == cast<AffineReadOpInterface>(loadOpInst).getMemRef())
           loadOps->push_back(loadOpInst);
       }
     }
 
     // Returns all memrefs in 'loadAndStoreMemrefSet' for which this node
     // has at least one load and store operation.
     void getLoadAndStoreMemrefSet(DenseSet<Value> *loadAndStoreMemrefSet) {
       llvm::SmallDenseSet<Value, 2> loadMemrefs;
       for (auto *loadOpInst : loads) {
-        loadMemrefs.insert(cast<AffineLoadOp>(loadOpInst).getMemRef());
+        loadMemrefs.insert(cast<AffineReadOpInterface>(loadOpInst).getMemRef());
       }
       for (auto *storeOpInst : stores) {
-        auto memref = cast<AffineStoreOp>(storeOpInst).getMemRef();
+        auto memref = cast<AffineWriteOpInterface>(storeOpInst).getMemRef();
         if (loadMemrefs.count(memref) > 0)
           loadAndStoreMemrefSet->insert(memref);
       }
     }
   };
 
   // Edge represents a data dependence between nodes in the graph.
   struct Edge {
void removeNode(unsigned id) {
     nodes.erase(id);
   }
 
   // Returns true if node 'id' writes to any memref which escapes (or is an
   // argument to) the function/block. Returns false otherwise.
   bool writesToLiveInOrEscapingMemrefs(unsigned id) {
     Node *node = getNode(id);
     for (auto *storeOpInst : node->stores) {
-      auto memref = cast<AffineStoreOp>(storeOpInst).getMemRef();
+      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))
           return true;
     }
     return false;
   }
 
-  // Returns the unique AffineStoreOp in `node` that meets all the following:
+  // Returns the unique AffineWriteOpInterface in `node` that meets all the
+  // following:
   //   *) store is the only one that writes to a function-local memref live out
   //      of `node`,
   //   *) store is not the source of a self-dependence on `node`.
-  // Otherwise, returns a null AffineStoreOp.
-  AffineStoreOp getUniqueOutgoingStore(Node *node) {
-    AffineStoreOp uniqueStore;
+  // Otherwise, returns a null AffineWriteOpInterface.
+  AffineWriteOpInterface getUniqueOutgoingStore(Node *node) {
+    AffineWriteOpInterface uniqueStore;
 
     // Return null if `node` doesn't have any outgoing edges.
     auto outEdgeIt = outEdges.find(node->id);
     if (outEdgeIt == outEdges.end())
       return nullptr;
 
     const auto &nodeOutEdges = outEdgeIt->second;
     for (auto *op : node->stores) {
-      auto storeOp = cast<AffineStoreOp>(op);
+      auto storeOp = cast<AffineWriteOpInterface>(op);
       auto memref = storeOp.getMemRef();
       // Skip this store if there are no dependences on its memref. This means
       // that store either:
       // *) writes to a memref that is only read within the same loop nest
       //    (self-dependence edges are not represented in graph at the moment),
       // *) writes to a function live out memref (function parameter), or
       // *) is dead.
       if (llvm::all_of(nodeOutEdges, [=](const Edge &edge) {
public:
   //    function/block argument).
   // *) The node has no successors in the dependence graph.
   bool canRemoveNode(unsigned id) {
     if (writesToLiveInOrEscapingMemrefs(id))
       return false;
     Node *node = getNode(id);
     for (auto *storeOpInst : node->stores) {
       // Return false if there exist out edges from 'id' on 'memref'.
-      if (getOutEdgeCount(id, cast<AffineStoreOp>(storeOpInst).getMemRef()) > 0)
+      auto storeMemref = cast<AffineWriteOpInterface>(storeOpInst).getMemRef();
+      if (getOutEdgeCount(id, storeMemref) > 0)
         return false;
     }
     return true;
   }
 
   // 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.
if (auto forOp = dyn_cast<AffineForOp>(op)) {
       collector.collect(&op);
       // Return false if a non 'affine.for' region was found (not currently
       // supported).
       if (collector.hasNonForRegion)
         return false;
       Node node(nextNodeId++, &op);
       for (auto *opInst : collector.loadOpInsts) {
         node.loads.push_back(opInst);
-        auto memref = cast<AffineLoadOp>(opInst).getMemRef();
+        auto memref = cast<AffineReadOpInterface>(opInst).getMemRef();
         memrefAccesses[memref].insert(node.id);
       }
       for (auto *opInst : collector.storeOpInsts) {
         node.stores.push_back(opInst);
-        auto memref = cast<AffineStoreOp>(opInst).getMemRef();
+        auto memref = cast<AffineWriteOpInterface>(opInst).getMemRef();
         memrefAccesses[memref].insert(node.id);
       }
       forToNodeMap[&op] = node.id;
       nodes.insert({node.id, node});
-    } else if (auto loadOp = dyn_cast<AffineLoadOp>(op)) {
+    } else if (auto loadOp = dyn_cast<AffineReadOpInterface>(op)) {
       // Create graph node for top-level load op.
       Node node(nextNodeId++, &op);
       node.loads.push_back(&op);
-      auto memref = cast<AffineLoadOp>(op).getMemRef();
+      auto memref = cast<AffineReadOpInterface>(op).getMemRef();
       memrefAccesses[memref].insert(node.id);
       nodes.insert({node.id, node});
-    } else if (auto storeOp = dyn_cast<AffineStoreOp>(op)) {
+    } else if (auto storeOp = dyn_cast<AffineWriteOpInterface>(op)) {
       // Create graph node for top-level store op.
       Node node(nextNodeId++, &op);
       node.stores.push_back(&op);
-      auto memref = cast<AffineStoreOp>(op).getMemRef();
+      auto memref = cast<AffineWriteOpInterface>(op).getMemRef();
       memrefAccesses[memref].insert(node.id);
       nodes.insert({node.id, node});
     } else if (op.getNumRegions() != 0) {
       // Return false if another region is found (not currently supported).
       return false;
     } else if (op.getNumResults() > 0 && !op.use_empty()) {
       // Create graph node for top-level producer of SSA values, which
       // could be used by loop nest nodes.

 // Removes load operations from 'srcLoads' which operate on 'memref', and
 // adds them to 'dstLoads'.
 static void moveLoadsAccessingMemrefTo(Value memref,
                                        SmallVectorImpl<Operation *> *srcLoads,
                                        SmallVectorImpl<Operation *> *dstLoads) {
   dstLoads->clear();
   SmallVector<Operation *, 4> srcLoadsToKeep;
   for (auto *load : *srcLoads) {
-    if (cast<AffineLoadOp>(load).getMemRef() == memref)
+    if (cast<AffineReadOpInterface>(load).getMemRef() == memref)
       dstLoads->push_back(load);
     else
       srcLoadsToKeep.push_back(load);
   }
   srcLoads->swap(srcLoadsToKeep);
 }
 
 // Returns the innermost common loop depth for the set of operations in 'ops'.
static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
                                  uint64_t localBufSizeThreshold) {
   auto *forInst = forOp.getOperation();
 
   // Create builder to insert alloc op just before 'forOp'.
   OpBuilder b(forInst);
   // Builder to create constants at the top level.
   OpBuilder top(forInst->getParentOfType<FuncOp>().getBody());
   // Create new memref type based on slice bounds.
-  auto oldMemRef = cast<AffineStoreOp>(srcStoreOpInst).getMemRef();
+  auto oldMemRef = cast<AffineWriteOpInterface>(srcStoreOpInst).getMemRef();
   auto oldMemRefType = oldMemRef.getType().cast<MemRefType>();
   unsigned rank = oldMemRefType.getRank();
 
   // Compute MemRefRegion for 'srcStoreOpInst' at depth 'dstLoopDepth'.
   MemRefRegion region(srcStoreOpInst->getLoc());
   bool validRegion = succeeded(region.compute(srcStoreOpInst, dstLoopDepth));
   (void)validRegion;
   assert(validRegion && "unexpected memref region failure");

 }
 
 // Checks if node 'srcId' can be safely fused into node 'dstId'. Node 'srcId'
 // may write to multiple memrefs but it is required that only one of them,
 // 'srcLiveOutStoreOp', has output edges.
 // Returns true if 'dstNode's read/write region to 'memref' is a super set of
 // 'srcNode's write region to 'memref' and 'srcId' has only one output edge.
 // TODO(andydavis) Generalize this to handle more live in/out cases.
-static bool canFuseSrcWhichWritesToLiveOut(unsigned srcId, unsigned dstId,
-                                           AffineStoreOp srcLiveOutStoreOp,
-                                           MemRefDependenceGraph *mdg) {
+static bool
+canFuseSrcWhichWritesToLiveOut(unsigned srcId, unsigned dstId,
+                               AffineWriteOpInterface srcLiveOutStoreOp,
   assert(srcLiveOutStoreOp && "Expected a valid store op");
   auto *dstNode = mdg->getNode(dstId);
   Value memref = srcLiveOutStoreOp.getMemRef();
   // Return false if 'srcNode' has more than one output edge on 'memref'.
   if (mdg->getOutEdgeCount(srcId, memref) > 1)
     return false;
 
   // Compute MemRefRegion 'srcWriteRegion' for 'srcStoreOp' on 'memref'.
while (!worklist.empty()) {
       // consumer loop nest.
       sinkSequentialLoops(dstNode);
 
       SmallVector<Operation *, 4> loads = dstNode->loads;
       SmallVector<Operation *, 4> dstLoadOpInsts;
       DenseSet<Value> visitedMemrefs;
       while (!loads.empty()) {
         // Get memref of load on top of the stack.
-        auto memref = cast<AffineLoadOp>(loads.back()).getMemRef();
+        auto memref = cast<AffineReadOpInterface>(loads.back()).getMemRef();
         if (visitedMemrefs.count(memref) > 0)
           continue;
         visitedMemrefs.insert(memref);
         // Move all loads in 'loads' accessing 'memref' to 'dstLoadOpInsts'.
         moveLoadsAccessingMemrefTo(memref, &loads, &dstLoadOpInsts);
         // Skip if no input edges along which to fuse.
         if (mdg->inEdges.count(dstId) == 0)
           continue;
while (!worklist.empty()) {
           // fusion.
           auto srcStoreOp = mdg->getUniqueOutgoingStore(srcNode);
           if (!srcStoreOp) {
             // Get the src store op at the deepest loop depth.
             // We will use 'LoopFusionUtils::canFuseLoops' to check fusion
             // feasibility for loops with multiple stores.
             unsigned maxLoopDepth = 0;
             for (auto *op : srcNode->stores) {
-              auto storeOp = cast<AffineStoreOp>(op);
+              auto storeOp = cast<AffineWriteOpInterface>(op);
               if (storeOp.getMemRef() != memref) {
                 srcStoreOp = nullptr;
                 break;
               }
               unsigned loopDepth = getNestingDepth(storeOp);
               if (loopDepth > maxLoopDepth) {
                 maxLoopDepth = loopDepth;
                 srcStoreOp = storeOp;
while (!worklist.empty()) {
 
           // Skip if fusion is not feasible at all loop depths.
           if (!canFuse)
             continue;
 
           // Gather 'dstNode' store ops to 'memref'.
           SmallVector<Operation *, 2> dstStoreOpInsts;
           for (auto *storeOpInst : dstNode->stores)
-            if (cast<AffineStoreOp>(storeOpInst).getMemRef() == memref)
+            if (cast<AffineWriteOpInterface>(storeOpInst).getMemRef() == memref)
               dstStoreOpInsts.push_back(storeOpInst);
 
           unsigned bestDstLoopDepth;
           mlir::ComputationSliceState sliceState;
           // Check if fusion would be profitable.
           if (!isFusionProfitable(srcStoreOp, srcStoreOp, dstLoadOpInsts,
                                   dstStoreOpInsts, &sliceState,
                                   &bestDstLoopDepth, maximalFusion,
while (!worklist.empty()) {
             // Promote single iteration slice loops to single IV value.
             for (auto forOp : sliceCollector.forOps) {
               promoteIfSingleIteration(forOp);
             }
             if (createPrivateMemref) {
               // Create private memref for 'memref' in 'dstAffineForOp'.
               SmallVector<Operation *, 4> storesForMemref;
               for (auto *storeOpInst : sliceCollector.storeOpInsts) {
-                if (cast<AffineStoreOp>(storeOpInst).getMemRef() == memref)
+                if (cast<AffineWriteOpInterface>(storeOpInst).getMemRef() ==
+                    memref)
                   storesForMemref.push_back(storeOpInst);
               }
               // TODO(andydavis) Use union of memref write regions to compute
               // private memref footprint.
               auto newMemRef = createPrivateMemRef(
                   dstAffineForOp, storesForMemref[0], bestDstLoopDepth,
                   fastMemorySpace, localBufSizeThreshold);
               visitedMemrefs.insert(newMemRef);
               // Create new node in dependence graph for 'newMemRef' alloc op.
               unsigned newMemRefNodeId =
                   mdg->addNode(newMemRef.getDefiningOp());
               // Add edge from 'newMemRef' node to dstNode.
               mdg->addEdge(newMemRefNodeId, dstId, newMemRef);
             }
 
             // Collect dst loop stats after memref privatization transformation.
             LoopNestStateCollector dstLoopCollector;
             dstLoopCollector.collect(dstAffineForOp.getOperation());
 
             // Add new load ops to current Node load op list 'loads' to
             // continue fusing based on new operands.
             for (auto *loadOpInst : dstLoopCollector.loadOpInsts) {
-              auto loadMemRef = cast<AffineLoadOp>(loadOpInst).getMemRef();
+              auto loadMemRef =
+                  cast<AffineReadOpInterface>(loadOpInst).getMemRef();
               // NOTE: Change 'loads' to a hash set in case efficiency is an
               // issue. We still use a vector since it's expected to be small.
               if (visitedMemrefs.count(loadMemRef) == 0 &&
                   !llvm::is_contained(loads, loadOpInst))
                 loads.push_back(loadOpInst);
             }
 
             // Clear and add back loads and stores.
auto canFuseWithSibNode = [&](Node *sibNode, Value memref) {
       if (llvm::any_of(loadAndStoreMemrefSet, [=](Value memref) {
             return mdg->getIncomingMemRefAccesses(sibNode->id, memref) > 0;
           }))
         return false;
 
       // Check that all stores are to the same memref.
       DenseSet<Value> storeMemrefs;
       for (auto *storeOpInst : sibNode->stores) {
-        storeMemrefs.insert(cast<AffineStoreOp>(storeOpInst).getMemRef());
+        storeMemrefs.insert(
+            cast<AffineWriteOpInterface>(storeOpInst).getMemRef());
       }
       if (storeMemrefs.size() != 1)
         return false;
       return true;
     };
 
     // Search for siblings which load the same memref function argument.
     auto fn = dstNode->op->getParentOfType<FuncOp>();
     for (unsigned i = 0, e = fn.getNumArguments(); i != e; ++i) {
       for (auto *user : fn.getArgument(i).getUsers()) {
-        if (auto loadOp = dyn_cast<AffineLoadOp>(user)) {
+        if (auto loadOp = dyn_cast<AffineReadOpInterface>(user)) {
           // Gather loops surrounding 'use'.
           SmallVector<AffineForOp, 4> loops;
           getLoopIVs(*user, &loops);
           // Skip 'use' if it is not within a loop nest.
           if (loops.empty())
             continue;
           Node *sibNode = mdg->getForOpNode(loops[0]);
           assert(sibNode != nullptr);

mlir/test/lib/Transforms/TestMemRefBoundCheck.cpp

struct TestMemRefBoundCheck
     : public PassWrapper<TestMemRefBoundCheck, FunctionPass> {
   void runOnFunction() override;
 };
 
 } // end anonymous namespace
 
 void TestMemRefBoundCheck::runOnFunction() {
   getFunction().walk([](Operation *opInst) {
-    TypeSwitch<Operation *>(opInst).Case<AffineLoadOp, AffineStoreOp>(
+    TypeSwitch<Operation *>(opInst)
+        .Case<AffineReadOpInterface, AffineWriteOpInterface>(
             [](auto op) { boundCheckLoadOrStoreOp(op); });
 
     // TODO(bondhugula): do this for DMA ops as well.
   });
 }
 
 namespace mlir {
 void registerMemRefBoundCheck() {
   PassRegistration<TestMemRefBoundCheck>(
       "test-memref-bound-check", "Check memref access bounds in a Function");
 }
 } // namespace mlir