Interfaces and their implementation

Interfaces

pliron is an extensible IR framework. This means that the framework does not restrict the kind of Operations (Ops / Opcodes), Attributes or Types that can be represented in the IR. A user is free to define new Ops, Attributes or Types and assign semantics to them.

For sanity to prevail, to not have to specify analysis and transformation details for each of these individually, common behaviour can be grouped into an Interface. MLIR uses Interfaces and Traits to achieve this. There is some difference between the two.

In pliron, we wrap around Rust traits to provide Interfaces. In the following discussion, we use Op and Op interfaces to explain the workings of interfaces, but it equally applies to Attributes and Types.

As an example, here's an Interface for Ops that always have exactly one region:

The OneRegionInterface:

/// [Op]s that have exactly one region.
#[op_interface]
pub trait OneRegionInterface {
    /// Get the single region that this [Op] has.
    fn get_region(&self, ctx: &Context) -> Ptr<Region> {
        self.get_operation().deref(ctx).get_region(0)
    }

    /// Checks that the operation has exactly one region.
    fn verify(op: &dyn Op, ctx: &Context) -> Result<()>
    where
        Self: Sized,
    {
        let self_op = op.get_operation().deref(ctx);
        if self_op.regions.len() != 1 {
            return verify_err!(self_op.loc(), OneRegionVerifyErr(op.get_opid().to_string()));
        }
        Ok(())
    }
}

Here, the #[op_interface] annotates the Rust trait as an Op Interface. Complementary to it, we also have an #[op_interface_impl] macro that helps an Op implement an interface.

#[op_interface]
pub trait MyOpInterface {
  fn gubbi(&self);
}

#[op_interface_impl]
impl MyOpInterface for MyOp {
    fn gubbi(&self) { println!("gubbi"); }
}

These two macros enable the following additional functionality:

1. Casting from a dyn Op to a dyn Interface object, thus enabling dynamic checking and use of interfaces.
2. Automatic verification of Interface properties. For example, a verifier for the OneRegionInterface above must ensure that every Op that implements the interface must have exactly one region. These macros ensure that such an Interface verifier is automatically called when the verifier for the Op is called. Also, if there are super-interfaces specified, the verifiers for those are called prior to calling the verifier for the Interface itself.

We now go on to discuss how these two functionality work, in more detail.

Interface casts

Interfaces convey that an Op implementing it provides certain functionality or holds certain properties. When working on the IR, it becomes essential to, given a dyn Op, determine whether the Op implements an interface, and if it does, to be able to call the methods defined by that interface. i.e., we want a function

pub fn op_cast<T: ?Sized + Op>(op: &dyn Op) -> Option<&T>

that can be called for any interface (Rust trait) object dyn T.

To enable this, pliron provides infrastructure for casting from dyn Any to dyn Trait, for traits that the type contained by the Any object implements. This utility works by maintaining a static table TRAIT_CASTERS_MAP that maps TypeIds of concrete types to each interface (Rust trait) that it implements, and querying the table when performing the actual cast. The #[op_interface_impl] macro is responsible for making an entry in this static table so that future queries can perform the actual cast. We use the downcast_rs crate to help upcast dyn Op objects to an Any object so that we can use the above trait_cast utility.

Verifying Interfaces

An Op implementing an interface must satisfy certain constraints defined by the interface itself. While an Op's verifier can manually call the verifier for each interface that it implements, it is certainly more convenient and safer to have that be automatic.

MLIR achieves automatic verification of an Op's interfaces and traits via variadic template magic.

template <typename ConcreteType, template <typename T> class... Traits>
class Op : public OpState, public Traits<ConcreteType>... {
  ...
  /// Implementation of `VerifyInvariantsFn` OperationName hook.
  static LogicalResult verifyInvariants(Operation *op) {
    static_assert(hasNoDataMembers(),
                  "Op class shouldn't define new data members");
    return failure(
        failed(op_definition_impl::verifyTraits<Traits<ConcreteType>...>(op)) ||
        failed(cast<ConcreteType>(op).verify()));
  }
  ...
}

Every trait or interface that an Op implements is listed as a type parameter during the class definition of that Op. A behaviour different from the interface's default definition will be overridden in the class definition. A minor downside to this approach is that, implementing a new interface for an Op requires changes to the Op's definition. This is especially a problem when the interface is defined in client code (i.e., not part of the library in which the Op is defined).

To get automatic verification of interfaces for pliron, we need to generate, at compile time, a function that has calls to every interface implemented by an Op. A less desirable, but simpler solution was to have the definition of an Op explicitly list the interfaces that it implements. While this has the same downside as MLIR's solution, there is also the possibility that a user implements an interface but forgets to add it to this list.

An additional requirement is that we want the verifiers of super-interfaces to run before the sub-interfaces. This allows every verifier implementation to conveniently assume that it's super-interface is already verified.

To get this working, we declare a static OP_INTERFACE_VERIFIERS_MAP that maps Ops to an ordered (as per the super-interface dependence) list of verifiers. So when an Op needs to be verified, each verifier for that Op in this list is called, in order. This is again made convenient by the #[op_interface_impl] macro that helps register verifiers of Op interfaces into this map. The map is Lazy and hence is built on the first access, which involves sorting a list of dependence pairs, to arrive at the final order.

Distributed updates of static tables

In both the topics above, we maintain tables TRAIT_CASTERS_MAP and OP_INTERFACE_VERIFIERS_MAP that need to be updated at each application of #[op_interface_impl], but all at compile time.

The linkme crate allows us to declare a static array distributed_slice of const expressions, whose contents can be any static value of that type, marked by a provided macro. Every marked element is finally linked into one distributed_slice.

Conclusion

Our solution has the added advantage of not limiting the user from implementing his own interface for an externally defined Op, or for implementing an externally defined interface for his Op. The only restriction is that both the Op and the interface cannot be externally defined, and this comes directly from Rust's restriction on either a trait or a type that implements the trait to be in the current crate.