Page MenuHomePhabricator

[mlir][Affine] Introduce affine memory interfaces
ClosedPublic

Authored by dcaballe on Tue, May 12, 5:15 PM.

Details

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.

Diff Detail

Event Timeline

dcaballe created this revision.Tue, May 12, 5:15 PM
Herald added a project: Restricted Project. · View Herald TranscriptTue, May 12, 5:15 PM

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

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

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 requested changes to this revision.Wed, May 13, 3:20 AM

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.

This revision now requires changes to proceed.Wed, May 13, 3:20 AM

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 retitled this revision from [mlir][Affine][NFCI]: Introduce affine memory interfaces to [mlir][Affine]: Introduce affine memory interfaces.Wed, May 13, 6:23 PM
dcaballe edited the summary of this revision. (Show Details)
dcaballe retitled this revision from [mlir][Affine]: Introduce affine memory interfaces to [mlir][Affine] Introduce affine memory interfaces.Wed, May 13, 6:24 PM
dcaballe updated this revision to Diff 263908.Wed, May 13, 6:56 PM
dcaballe marked 4 inline comments as done.
dcaballe edited the summary of this revision. (Show Details)

Addressing feedback.

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 accepted this revision.Wed, May 13, 8:43 PM
bondhugula added inline comments.
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)

This revision is now accepted and ready to land.Wed, May 13, 8:43 PM
mehdi_amini added inline comments.Wed, May 13, 9:00 PM
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 accepted this revision.Thu, May 14, 1:14 AM
ftynse added inline comments.
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 updated this revision to Diff 264136.Thu, May 14, 6:25 PM

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

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 added a comment.EditedTue, May 19, 5:01 AM

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.

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

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 updated this revision to Diff 264963.Tue, May 19, 10:50 AM

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

This revision was automatically updated to reflect the committed changes.

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.