[learn from scratch, deep learning compiler] 18. Interfaces in MLIR

0x0. preface

This article is used to learn about Interfaces in MLIR. MLIR is a general and extensible framework, which is composed of different levels of conversations with specific attributes, Operation and Type. Because of the hierarchical design of dialogues, MLIR can express operations at multiple semantic and abstract levels. However, this hierarchical design also has a disadvantage, that is, when performing Operation conversion or Pass at different dialog levels, we need to clarify the specific semantics of each Operation under each dialog, otherwise the conversion or transformation may fail. In fact, readers who have developed based on MLIR should have encountered the possibility of Op conversion failure when combining some MLIR passes to Lower an MLIR file. In order to alleviate this situation, MLIR proposes Interfaces. Actually [learn the compiler from scratch] 13. How to write Pass in MLIR? Here we have used Interfaces to realize inlining and shape derivation Pass. In this section, we will have a deeper understanding of the Interfaces in MLIR. Finally, we combine the useropcompatible interface example in OneFlow IR to further deepen our understanding.

The operations and operations mentioned in this article are the same thing. They are all operations under MLIR Dialect.

0x1. motivation

Interfaces can be translated into interfaces, and MLIR interfaces provide a common way to interact with IR. The design goal of interfaces is to convert and analyze MLIR expressions without invading the specific Operation under a specific Dialect and the specific knowledge of Dialect. In this way, the transformation, analysis and addition of a Dialect and the corresponding Operation can be decoupled, which greatly enhances the scalability of MLIR.

0x2. Dialog interfaces definition (detailed view)

Dialog Interfaces are generally used to Pass or analyze a group of attributes, operations and types. These attributes, operations and types can be defined by different dialog Interfaces. These Interfaces generally cover all levels of conversations and are only used for a few analysis and transformation. Therefore, we should make it clear that the Interface is not the core of Operation, but the core of some general transformations. stay [learn the compiler from scratch] 13. How to write Pass in MLIR? Here is an example of using inline Interface to implement inline Pass. Inlining usually queries high-level information about operations in a dialog, such as cost modeling and legitimacy. These information are usually not specific to an Operation under a dialog and exist separately.

Dialectinterface can be defined by inheriting a CRTP base class dialectinterfacebase:: base < >. For the introduction of CRTP, please refer to: https://zh.wikipedia.org/wiki/ Strange recursive template mode. I understand static polymorphism (CRTP) because there are many dialectinterfaces in MLIR, which should be derived from this dialectinterfacebase:: base < > base class. CRTP is more appropriate for performance considerations. This base class provides some interfaces necessary for the dialog Interface registration, so that they can be referenced in the future. Once the Interface is defined, Dialects can rewrite it with Dialect specific information. Interfaces defined by a dialog are registered through addinterfaces < >, which is similar to the registration mechanism of attribute, Operation and Type. Here is a chestnut:

// Define a basic inline Interface class to allow Dialect to choose to join the inline. 
class DialectInlinerInterface :
    public DialectInterface::Base<DialectInlinerInterface> {
public:
  ///Returns true if the given region 'src' can be inlined into the region.
  ///'dest 'is attached to the Operation registered to the current dialog. 
  ///'valueMapping' contains values from any remapping within the 'src' area. 
  ///For example, this can be used to check which values will replace the entry parameters in the src area. 
  virtual bool isLegalToInline(Region *dest, Region *src,
                               BlockAndValueMapping &valueMapping) const {
    return false;
  }
};

///Override the inline interface to add support for AffineDialect to enable the Operation of inline Affine Dialect. 
struct AffineInlinerInterface : public DialectInlinerInterface {
  ///The Affine structure has specific inline constraints. 
  bool isLegalToInline(Region *dest, Region *src,
                       BlockAndValueMapping &valueMapping) const final {
    ...
  }
};

///Register inline Interfaces under Dialect
AffineDialect::AffineDialect(MLIRContext *context) ... {
  addInterfaces<AffineInlinerInterface>();
}

After these Interfaces are registered, they can be found in the dialog when performing the transformation and analysis of MLIR. There is no need to determine a specific dialog subclass (such as a specific Operation). For example:

Dialect *dialect = ...;
if (DialectInlinerInterface *interface
      = dialect->getRegisteredInterface<DialectInlinerInterface>()) {
  // Dialect provides an implementation of this Interface
  ...
}

For example, in llvm / MLIR / lib / IR / Dialect The registerDelayedInterfaces function in the cpp file shows the above usage. This function is used to register the loaded dialog Interfaces:

void DialectRegistry::registerDelayedInterfaces(Dialect *dialect) const {
  auto it = interfaces.find(dialect->getTypeID());
  if (it == interfaces.end())
    return;

  // Add an interface if it is not already present.
  for (const auto &kvp : it->getSecond().dialectInterfaces) {
    if (dialect->getRegisteredInterface(kvp.first))
      continue;
    dialect->addInterface(kvp.second(dialect));
  }

  // Add attribute, operation and type interfaces.
  for (const auto &info : it->getSecond().objectInterfaces)
    std::get<2>(info)(dialect->getContext());
}

0x3. DialectInterfaceCollection (I haven't used it yet)

Dialect interfacecollection provides an additional utility. This class allows the collection of all dialects that have registered a given Interface in the MLIRContext instance. This is useful for hiding and optimizing the lookup of registered dialog interfaces.

class InlinerInterface : public
    DialectInterfaceCollection<DialectInlinerInterface> {
  // The hook of this class is the hook image of DialectInlinerInterface. By default, it is implemented to call the hook on a given Dialect Interface. 
  virtual bool isLegalToInline(Region *dest, Region *src,
                               BlockAndValueMapping &valueMapping) const {
    auto *handler = getInterfaceFor(dest->getContainingOp());
    return handler ? handler->isLegalToInline(dest, src, valueMapping) : false;
  }
};

MLIRContext *ctx = ...;
InlinerInterface interface(ctx);
if(!interface.isLegalToInline(...))
   ...

0x4. Property, operation, type Interfaces (optional)

As the name implies, property / action / type interfaces are those registered at a specific property / action / type level. These interfaces provide access to derived objects by providing virtual interfaces that must be implemented. For example, many analysts and transformations want to know the side effects of Operation to improve performance and correctness. The side effects of an Operation are usually related to the semantics of a specific Operation, such as affinity Load Operation has a read effect (as the name suggests).

These interfaces are defined by CRTP classes covering specific IR entities; AttrInterface, OpInterface, or TypeInterface. These classes use the Traits class that defines the Concept and Model classes as template parameters. These classes provide an implementation of Concept based polymorphism, in which Concept defines a set of virtual methods that are covered by a Model templated on a specific entity type. It should be noted that these classes should be pure and should not contain non static data members or other variable data. To attach an Interface to an object, the base class provides a trail class that can be attached to the list of features of the object (skip the example code below to see the explanation).

struct ExampleOpInterfaceTraits {
  ///Define a basic Concept class and specify the virtual interface to be implemented. 
  struct Concept {
    virtual ~Concept();

    ///This is an example of a non static hook for an Operation 
    virtual unsigned exampleInterfaceHook(Operation *op) const = 0;

    ///This is an example of a static hook for Operation. Static hooks do not require a specific instance of Operation. The implementation is a virtual hook, as in the non static case, because the implementation of the hook itself still needs to be implemented indirectly. 
    virtual unsigned exampleStaticInterfaceHook() const = 0;
  };

  ///Define a model class for the concept of the given Operation type 
  template <typename ConcreteOp>
  struct Model : public Concept {
    ///Override methods to be distributed on specific operations 
    unsigned exampleInterfaceHook(Operation *op) const final {
      return llvm::cast<ConcreteOp>(op).exampleInterfaceHook();
    }

    ///Override methods to be distributed on specific operations
    unsigned exampleStaticInterfaceHook() const final {
      return ConcreteOp::exampleStaticInterfaceHook();
    }
  };
};

///Defines the main Interface class with which the analysis and transformation will interact. 
class ExampleOpInterface : public OpInterface<ExampleOpInterface,
                                              ExampleOpInterfaceTraits> {
public:
  ///Inherit the base class constructor to support LLVM style conversion. 
  using OpInterface<ExampleOpInterface, ExampleOpInterfaceTraits>::OpInterface;

  ///Interface is distributed to 'getImpl()'. This is a method provided by the basic 'OpInterface' class, which returns an instance of concept. 
  unsigned exampleInterfaceHook() const {
    return getImpl()->exampleInterfaceHook(getOperation());
  }
  unsigned exampleStaticInterfaceHook() const {
    return getImpl()->exampleStaticInterfaceHook(getOperation()->getName());
  }
};

Once the Interface is defined, it can be registered in the operation by adding the provided feature ExampleOpInterface::Trait, as described earlier. Using this Interface is like using any other derived operation type, that is, cast:

///When defining an Operation, the Interface is registered through the nested 'trail' class provided by the 'opinterface < >' base class.
class MyOp : public Op<MyOp, ExampleOpInterface::Trait> {
public:
  /// The definition of the interface method on the derived operation.
  unsigned exampleInterfaceHook() { return ...; }
  static unsigned exampleStaticInterfaceHook() { return ...; }
};

///Later, we can query whether a specific Operation (such as "MyOp") overrides a given Interface. 
Operation *op = ...;
if (ExampleOpInterface example = dyn_cast<ExampleOpInterface>(op))
  llvm::errs() << "hook returned = " << example.exampleInterfaceHook() << "\n";

If you read here and read it before [learn the compiler from scratch] 13. How to write Pass in MLIR? Using the example of inline Interface, I believe you can better understand the steps of registering inline Pass under Toy Dialect.

0x5. External Model of property, operation and type Interfaces (optional)

This may require providing an interface implementation for an IR object without modifying the definition of the object. It is worth noting that this allows interfaces for properties, operations, and types to be implemented outside the dialog that defines them, for example, by providing interfaces for built-in types. This is achieved by using two Concept based polymorphic Model extension classes derived from Concept, as shown below (note the note):

struct ExampleTypeInterfaceTraits {
  struct Concept {
    virtual unsigned exampleInterfaceHook(Type type) const = 0;
    virtual unsigned exampleStaticInterfaceHook() const = 0;
  };

  template <typename ConcreteType>
  struct Model : public Concept { /*...*/ };

  ///Unlike Model, FallbackModel passes type objects to
  ///Hook to make it accessible in the method body, even if the method is not defined in the class itself,
  ///Therefore, there is no "this" access. ODS automatically generates this class for all Interfaces. 
  template <typename ConcreteType>
  struct FallbackModel : public Concept {
    unsigned exampleInterfaceHook(Type type) const override {
      getImpl()->exampleInterfaceHook(type);
    }
    unsigned exampleStaticInterfaceHook() const override {
      ConcreteType::exampleStaticInterfaceHook();
    }
  };

  ///'ExternalModel' provides a location for the default implementation of the Interface method by explicitly separating the model class that implements the Interface from the type class that implements the Interface. Then you can use ` cast < ConcreteType > ` to define the default implementation. If 'ConcreteType' does not provide the API required by the default implementation, the user-defined implementation can directly use 'FallbackModel' to override the default implementation. Located in the class template, it will never be instantiated and will not cause compilation errors. ODS automatically generates this class and puts the default method implementation into it. 
  template <typename ConcreteModel, typename ConcreteType>
  struct ExternalModel : public FallbackModel<ConcreteModel> {
    unsigned exampleInterfaceHook(Type type) const override {
      // Default implementation can be provided here.
      return type.cast<ConcreteType>().callSomeTypeSpecificMethod();
    }
  };
};

You can provide external Models for property, operation, and type interfaces by deriving FallbackMode or ExternalModel and registering the Model class with the relevant class in a given context. Unless registered, other contexts will not see the Interface. For example:

///The external Interface implementation of the concrete class. This does not require modifying the definition of the type class itself. 
struct ExternalModelExample
    : public ExampleTypeInterface::ExternalModel<ExternalModelExample,
                                                 IntegerType> {
  static unsigned exampleStaticInterfaceHook() {
    // The implementation is provided here
    return IntegerType::someStaticMethod();
  }
  // There is no need to define an "exampleInterfaceHook" with a default implementation in the "external model". But it can be overwritten if needed. 
}

int main() {
  MLIRContext context;
  /* ... */;

  // Attach the interface model to the type in the given context before use. It is expected that the dialog containing this type will be loaded at this time. 
  IntegerType::attachInterface<ExternalModelExample>(context);
}

Finally, the document also puts forward a suggestion that we should use this mechanism only when we "own" the interface of external applications. This prevents both the owner of the dialog containing the object and the owner of the interface from being aware of the interface implementation, which may lead to repeated or divergent implementations. I haven't encountered the need to use this mechanism, so I won't go further here.

0x6. Dial fallback of opinterface (optional)

Some dialogues have an open ecosystem and do not register all possible operations. In this case, the OpInterface that implements these operations can still be supported. When the Operation is not registered or no Interface implementation is provided, the query will fallback to the Dialect itself.

The second Model is used in such cases and is automatically generated when ODS uses a Model named FallbackModel (see below). You can implement this Model for a specific Dialect:

// This is the implementation of a dialect fallback for `ExampleOpInterface`.
struct FallbackExampleOpInterface
    : public ExampleOpInterface::FallbackModel<
          FallbackExampleOpInterface> {
  static bool classof(Operation *op) { return true; }

  unsigned exampleInterfaceHook(Operation *op) const;
  unsigned exampleStaticInterfaceHook() const;
};

Dialect can then instantiate this implementation and return it on a specific Operation by overriding the getRegisteredInterfaceForOp method:

void *TestDialect::getRegisteredInterfaceForOp(TypeID typeID,
                                               StringAttr opName) {
  if (typeID == TypeID::get<ExampleOpInterface>()) {
    if (isSupported(opName))
      return fallbackExampleOpInterface;
    return nullptr;
  }
  return nullptr;
}

I don't understand this section very well. Let's record it first. Through the above introduction to Interfaces, we can leave some basic impressions. I think it should be enough. Next, we will talk about how to define Interfaces based on ODS, which is the focus of this article.

0x7. Using ODS framework to define Interface (important)

As mentioned above, the Interface allows properties, operations, and types to expose the calling method without the caller knowing the specific derived Type. The disadvantage of this infrastructure is that it requires some templates to connect all parts together. MLIR provides a mechanism for defining interfaces declaratively in ODS and automatically generating C + + definitions. For example, chestnuts:

def ExampleOpInterface : OpInterface<"ExampleOpInterface"> {
  let description = [{
    This is an example interface definition.
  }];

  let methods = [
    InterfaceMethod<
      "This is an example of a non-static hook to an operation.",
      "unsigned", "exampleInterfaceHook"
    >,
    StaticInterfaceMethod<
      "This is an example of a static hook to an operation.",
      "unsigned", "exampleStaticInterfaceHook"
    >,
  ];
}

Providing the definition of AttrInterface, OpInterface, or TypeInterface class will automatically generate the C + + class of the interface. The interface consists of the following components:

  • C + + class name (provided through template parameters).
  • Description. (description).
  • C++ Namespace. (cppNamespace). That is, under which C + + namespace the Interface class should be generated.
  • Methods(methods). List of Interfaces hook methods defined by the IR object.
  • Extra Class Declarations. Optional: extraClassDeclaration. Additional C + + code generated in the declaration of the Interface class. This allows you to define methods, etc. on user oriented Interface classes without hooking to IR entities. These declarations are not implicitly visible in the default implementation of Interface methods, but static declarations can be accessed using full name qualification.

The OpInterface class may also include an additional Verifier (verify). It is a C + + code block with additional validation applied to the Operation attached to the Interface. The structure of this code block corresponds to structure 1-1 of the trail:: verifytrail method.

There are two types of methods that can be used with Interface, InterfaceMethod and StaticInterfaceMethod. They all consist of the same core components, except that StaticInterfaceMethod models static methods on derived IR objects. The Interface method has the following components:

  • Description: description information of the method, a string.
  • ReturnType: the string corresponding to the C + + return type of the method.
  • MethodName: the string corresponding to the C + + name of the method.
  • Arguments (Optional): strings corresponding to C + + type and variable name respectively.
  • MethodBody (Optional) and DefaultImplementation. I haven't used it yet. I need to check it later.

If the Operation specifies the Interface using declaraeopinterfacemethods, ODS also allows the generation of declarations for the InterfaceMethods of the Operation (see the following example).

ef MyInterface : OpInterface<"MyInterface"> {
  let description = [{
    This is the description of the interface. It provides concrete information
    on the semantics of the interface, and how it may be used by the compiler.
  }];

  let methods = [
    InterfaceMethod<[{
      This method represents a simple non-static interface method with no
      inputs, and a void return type. This method is required to be implemented
      by all operations implementing this interface. This method roughly
      correlates to the following on an operation implementing this interface:

      ```c++
      class ConcreteOp ... {
      public:
        void nonStaticMethod();
      };
      ```
    }], "void", "nonStaticMethod"
    >,

    InterfaceMethod<[{
      This method represents a non-static interface method with a non-void
      return value, as well as an `unsigned` input named `i`. This method is
      required to be implemented by all operations implementing this interface.
      This method roughly correlates to the following on an operation
      implementing this interface:

      ```c++
      class ConcreteOp ... {
      public:
        Value nonStaticMethod(unsigned i);
      };
      ```
    }], "Value", "nonStaticMethodWithParams", (ins "unsigned":$i)
    >,

    StaticInterfaceMethod<[{
      This method represents a static interface method with no inputs, and a
      void return type. This method is required to be implemented by all
      operations implementing this interface. This method roughly correlates
      to the following on an operation implementing this interface:

      ```c++
      class ConcreteOp ... {
      public:
        static void staticMethod();
      };
      ```
    }], "void", "staticMethod"
    >,

    StaticInterfaceMethod<[{
      This method corresponds to a static interface method that has an explicit
      implementation of the method body. Given that the method body has been
      explicitly implemented, this method should not be defined by the operation
      implementing this method. This method merely takes advantage of properties
      already available on the operation, in this case its `build` methods. This
      method roughly correlates to the following on the interface `Model` class:

      ```c++
      struct InterfaceTraits {
        /// ... The `Concept` class is elided here ...

        template <typename ConcreteOp>
        struct Model : public Concept {
          Operation *create(OpBuilder &builder, Location loc) const override {
            return builder.create<ConcreteOp>(loc);
          }
        }
      };
      ```

      Note above how no modification is required for operations implementing an
      interface with this method.
    }],
      "Operation *", "create", (ins "OpBuilder &":$builder, "Location":$loc),
      /*methodBody=*/[{
        return builder.create<ConcreteOp>(loc);
    }]>,

    InterfaceMethod<[{
      This method represents a non-static method that has an explicit
      implementation of the method body. Given that the method body has been
      explicitly implemented, this method should not be defined by the operation
      implementing this method. This method merely takes advantage of properties
      already available on the operation, in this case its `build` methods. This
      method roughly correlates to the following on the interface `Model` class:

      ```c++
      struct InterfaceTraits {
        /// ... The `Concept` class is elided here ...

        template <typename ConcreteOp>
        struct Model : public Concept {
          Operation *create(Operation *opaqueOp, OpBuilder &builder,
                            Location loc) const override {
            ConcreteOp op = cast<ConcreteOp>(opaqueOp);
            return op.getNumInputs() + op.getNumOutputs();
          }
        }
      };
      ```

      Note above how no modification is required for operations implementing an
      interface with this method.
    }],
      "unsigned", "getNumInputsAndOutputs", (ins), /*methodBody=*/[{
        return $_op.getNumInputs() + $_op.getNumOutputs();
    }]>,

    InterfaceMethod<[{
      This method represents a non-static method that has a default
      implementation of the method body. This means that the implementation
      defined here will be placed in the trait class that is attached to every
      operation that implements this interface. This has no effect on the
      generated `Concept` and `Model` class. This method roughly correlates to
      the following on the interface `Trait` class:

      ```c++
      template <typename ConcreteOp>
      class MyTrait : public OpTrait::TraitBase<ConcreteType, MyTrait> {
      public:
        bool isSafeToTransform() {
          ConcreteOp op = cast<ConcreteOp>(this->getOperation());
          return op.getNumInputs() + op.getNumOutputs();
        }
      };
      ```

      As detailed in [Traits](Traits.md), given that each operation implementing
      this interface will also add the interface trait, the methods on this
      interface are inherited by the derived operation. This allows for
      injecting a default implementation of this method into each operation that
      implements this interface, without changing the interface class itself. If
      an operation wants to override this default implementation, it merely
      needs to implement the method and the derived implementation will be
      picked up transparently by the interface class.

      ```c++
      class ConcreteOp ... {
      public:
        bool isSafeToTransform() {
          // Here we can override the default implementation of the hook
          // provided by the trait.
        }
      };
      ```
    }],
      "bool", "isSafeToTransform", (ins), /*methodBody=*/[{}],
      /*defaultImplementation=*/[{
    }]>,
  ];
}

// Operation interfaces can optionally be wrapped inside
// DeclareOpInterfaceMethods. This would result in autogenerating declarations
// for members `foo`, `bar` and `fooStatic`. Methods with bodies are not
// declared inside the op declaration but instead handled by the op interface
// trait directly.
def OpWithInferTypeInterfaceOp : Op<...
    [DeclareOpInterfaceMethods<MyInterface>]> { ... }

// Methods that have a default implementation do not have declarations
// generated. If an operation wishes to override the default behavior, it can
// explicitly specify the method that it wishes to override. This will force
// the generation of a declaration for those methods.
def OpWithOverrideInferTypeInterfaceOp : Op<...
    [DeclareOpInterfaceMethods<MyInterface, ["getNumWithDefault"]>]> { ... }

Note: in the ODS framework, the existing Operation interface defined in C + + can be accessed through the opinterfacetrain class.

0x8. Operation Interface list

MLIR includes standard interfaces that provide functionality that may be common in many different operations. Below is a list of key interfaces that can be used directly by any dialect. The title format of each Interface section is as follows:

  • Interface class name
    • (C++ class – ODS class(if applicable))

The screenshot of the standard Interface is as follows:

0x9. OneFlow's Interface

Next, take OneFlow IR as an example to see which interfaces are defined in OneFlow Dialect. The code is in: OneFlow / IR / include / OneFlow / oneflowinterfaces Here. OneFlow defines UserOpCompatibleInterface, ControlEdgeCompatibleInterface, NoGrad and other Interface types based on ODS. Let's take UserOpCompatibleInterface as an example to see its implementation:

def UserOpCompatibleInterface : OpInterface<"UserOpCompatible"> {
  let description = [{
    Interface to getting the hard-coded bn
  }];

  let methods = [
    StaticInterfaceMethod<"",
        "const std::vector<std::string>*", "inputKeys", (ins), [{
        static std::vector<std::string> val(mlir::oneflow::support::GetInputKeys(ConcreteOp::getOperationName().split('.').second.str()));
        return &val;
    }]>,
    StaticInterfaceMethod<"",
        "const std::vector<std::string>*", "outputKeys", (ins), [{
        static std::vector<std::string> val(mlir::oneflow::support::GetOutputKeys(ConcreteOp::getOperationName().split('.').second.str()));
        return &val;
    }]>,
    InterfaceMethod<"",
        "std::pair<unsigned, unsigned>", "getODSOperandIndexAndLength", (ins "unsigned":$index), [{
        return $_op.getODSOperandIndexAndLength(index);
    }]>,
    InterfaceMethod<"",
        "std::pair<unsigned, unsigned>", "getODSResultIndexAndLength", (ins "unsigned":$index), [{
        return $_op.getODSResultIndexAndLength(index);
    }]>
  ];
}

You can see that UserOpCompatibleInterface uses the StaticInterfaceMethod and InterfaceMethod in the Interface ODS specification in Section 7 to specify the methods for obtaining the Operation input operand name, output operand name, operand and length, result and length for this Interface. Then in OneFlow's OneFlow / IR / include / OneFlow / oneflowuserops TD uses declaraeopinterfacemethods < UserOpCompatibleInterface > to specify interfaces for various operations, and this Interface declaration will be carried in the generated Operation code.

So what are the benefits of doing so? The first point is that the UserOp of OneFlow is equipped with UserOpCompatibleInterface. As long as we implement a general GetInputKeys function for the UserOp of OneFlow, all operations derived from UserOp have the function of this function because they are equipped with UserOpCompatibleInterface.

A more general example is to develop some general Pass based on InterFace, such as inline and shape derivation Pass. see [learn the compiler from scratch] 13. How to write Pass in MLIR?

0x10. summary

This article mainly introduces the Interface of MLIR, adds some understanding and description on the basis of MLIR documents, and shows an example of OneFlow to illustrate the benefits of Interface and how to use ODS to write Interface.

Keywords: AI Deep Learning

Added by datoshway on Thu, 06 Jan 2022 03:56:37 +0200