[MLIR][IRDL] Add a custom constraint operation

[Moxinilian]

Currently, it is not possible to define custom constraints in a portable manner. The c_pred escape hatch exists for TableGen portability reasons, but it is not portable at all and should be treated as a last resort solution.

Instead, we should offer a way to expose custom constraints infrastructure that is as portable as possible. One idea to achieve this would be to offer arbitrary imperative MLIR IR regions that take as input a representation of the IR element to verify and that return whether it is a match. This would be portable in the sense that the dialects used in the region can then offer a lowering to whatever the IRDL target is (maybe LLVM IR for an MLIR-based dialect, SMT dialect for the purposes of formal verification, etc.)

This might depend on introducing an in-IR way to interface with IR constructs. Maybe an mlir MLIR dialect that models the existing IR APIs?

Contact points: @Moxinilian

[PragmaTwice]

Hi @Moxinilian, I want to share some idea about this.

Implementing constraints directly through IR sounds like quite a large project — it seems it would require mapping more MLIR constructs into some dialects, and make dialect loading depend on ExecutionEngine & LLVM JIT. I’ve been thinking about a simpler approach (of course, only as a supplement rather than a replacement — the "in-IR way" is clearly much fancier to me): something similar to pdl.apply_native_constraint, where we could pass in a function object as a constraint (let’s call it irdl.native_pred (or irdl.external_pred?) for now).

Compared to the current irdl.c_pred, this would have a few advantages:

• A major limitation of c_pred is that it only works for irdl_to_cpp, while a native_pred like in PDL could be used with irdl::loadDialects for dynamically loading.
• Similarly, as in PDL, in the Python bindings we could even pass a Python function as a constraint, as long as the relevant APIs (C API) are exposed to Python.

Of course, the downside of this design is that it’s less “declarative”, and such constraints couldn’t be analyzed or modeled within IRDL. But it could still serve as a lightweight option. WDYT?

Here’s a piece of pseudocode to illustrate this idea (in python, since C++ can be more verbose):

module = ModuleOp.parse("""
irdl.dialect @test {
  irdl.operation @op1 {
    %0 = irdl.native_pred "pred1"
    irdl.attributes(%0)
  }
}
""")

def func(attr):
  return attr is an IntegerAttr and attr.value < 10

irdl.load_dialects(module, {"pred1": func})

[Moxinilian]

Hey! Thanks for the interest!

Just as a little piece of context, the reason why we were thinking of doing it like this is for portability purposes. If the semantics of the constraint are defined in terms of "standard dialects", it's very easy to have move them around, and in particular it does not depend on the dialect's embedding context. If a dialect needs to be embedded in a different runtime, you don't need to manually implement the domain-specific constraints of each IRDL-defined dialect in that runtime.

I'm not opposed to they option you are suggesting (the original proposal after all leaves all freedom to users to embed custom constraints that call into linked-in logic), I'm just not sure if there is much of a difference? Wouldn't a custom constraint that just func.calls an external symbol achieve roughly the same objective?

The more I think about this, the more it makes me realize MLIR would benefit a lot from interpretation mechanisms like there are in xDSL...

[PragmaTwice]

Yes, I see where you’re coming from. Expressing constraints in MLIR is a natural idea and has strong portability benefits. But if we dig in, there seem to be quite a few issues to solve before we can implement it:

• We’d need to implement (at least for the Builtin dialect) a whole set of ops to manipulate attributes and types inside predicates: checking an attribute’s type, accessing an attribute’s fields, retrieving child attributes within an attribute, etc. Since we’d have to stay in sync with changes to those dialects, we might also need some automation tooling to keep up—it sounds like a lot of work. (If we consider properties, it could get even more complex; see https://discourse.llvm.org/t/python-bindings-properties/88614.)
• Because MLIR’s ExecutionEngine works via convert-to-llvm + ORC JIT, we may also need to think about ABI (or FFI, arguably). There are two main parts: first, how to lower operations on MLIR attributes/types to the corresponding LLVM IR (e.g., llvm.call, which might involve the MLIR C API?); second, how to invoke JIT’ed functions from a DynamicDialect (for example, how to represent attributes in function parameters).
• In addition, basing this on the ExecutionEngine would make IRDL depend on a relatively heavyweight component. I’m not sure if we could design a pdl_interp-like dialect with its own bytecode to sidestep all the LLVM/JIT-related issues. But that clearly opens another can of worms: if we want it to be sufficiently general-purpose, we have to pay the corresponding implementation cost.

I’d love for us to do some lightweight work to make IRDL more powerful even before all of the above is completed. Your point makes a lot of sense: we could design these “external constraints” while maintaining backward compatibility (so that in the future we can still support defining constraints directly in MLIR). For example, we could do something like this:

irdl.pred_func @pred1(attr: attribute) 
irdl.dialect @test {
  irdl.operation @op1 {
    %0 = irdl.call_pred @pred1 
    irdl.attributes(%0)
  }
}

(The names here are just placeholders—I’d rather not dive into naming bikeshedding yet.)

This would mean we won’t need to invent a completely new system later, and we can first ship a lightweight version (i.e., only support external constraints, just pred_func without regions). What do you think?