Extensions should be first class

[zrho]

Extensibility of the HUGR is not an escape hatch or an attempt at future proofing, but rather a central design goal. Most operations in any HUGR will be in the form an extension. We aren't building a compiler, but rather a library to create compilers. This issue is to discuss some points where I feel the implementation is currently at odds with this goal, and propose some general and implementation ideas on how we can address this.

Egalitarian Data Structure

A benchmark for extensibility is how much of the core operations can (or even are) implemented with the same mechanisms that would be used for operations that are provided by extensions. Perhaps “egalitarian” is a good shorthand name for this, measuring the amount of additional performance and usability that is afforded to the core but not to extensions.

We have an OpType enum which lists the core operations, plus CustomOp for everything else. The custom operations are described by an ExtensionOp (or an OpaqueOp if the extension is not loaded). Unpacking the layers of indirection (both pointer and conceptual), a custom operation therefore consists of its name (a string) and a list of type arguments. Most operations in a practical HUGR will be custom operations, and thus stringly typed in this way.

Moreover the only arguments that can be provided to the operation directly are in the form of type parameters. This would prevent us from implementing custom operations that carry additional information which is not a type (unless we apply hacks such as embedding JSON into the type language). Such additional information is possible to provide to core operations, such as in FuncDecl:

pub struct FuncDecl {
    pub name: String,
    pub signature: PolyFuncType,
}

Neither the name of the function (a string) nor its signature (a type schema) are types themselves, and so could not be passed to a custom operation as type arguments. We therefore could not express FuncDecl as an extension operation at all.

A different point where this happens is for metadata, which is essentially an uninterpreted JSON value. We do need the capacity to deal with uninterpreted metadata so that we can build tools that do not have to know about all extensions. However for metadata of known extensions, we currently have no choice but to serialise and deserialise the metadata every time we want to look at it, or keep it around in external maps. This not only has a performance impact, but is also potentially unergonomic. We do not store the NodeType in the JSON metadata map, and for the same reason we should also allow similar affordances to extension metadata.

Operations, Attributes and Properties

A suggestion on how to move forward: To optimise for the convenience of users of the hugr library that are mainly concerned with extensions outside of the core, we should implement all core operations via extensions. This makes us feel the pain points of the extension mechanism early, and prevents us from taking shortcuts by hardcoding the core stuff. In particular I would suggest that the hugr-core crate is mainly concerned with the HUGR datastructure itself, with any concrete operation living in the hugr crate.

We can annotate each node with an operation name and operation properties. If no extension recognises the operation name, the properties remain uninterpreted as a (JSON) value. When an extension does know the operation, it should be able to deserialise the property into a Rust type and store that type directly in the HUGR. Similarly, node metadata should come in form of attributes, consisting of an attribute name and value. The value starts off as JSON and can be deserialised into the appropriate Rust type by an extension. Once parsed, we should be able to access properties and attributes by their types. When combined with the appropriate checks, this is an instance of Parse, don't validate.

Implementation Suggestion

There are a few ways to implement this. I have experimented with several of these over the last couple of days. I now believe that the following solution makes decent tradeoffs. We already have a design that is similar to an ECS (entity component system) pattern: nodes are already identified by their index, and additional metadata is attached via maps that are keyed with the node index. If you see a node, its operation, its properties and attributes as a row in a table, an ECS inspired data structure would store each column separately.

Taking inspiration from several existing Rust implementations of ECS, we can make a Hugr store the columns in a map that is indexed by its type (see for example the anymap crate). For every attribute type A this type indexed map would contain a store which associates node indices to A s. Depending on the attribute, this could be a hash map, a btreemap, a vector, a bitset or anything else that is appropriate. While this approach does (internally) use trait objects and dynamic dispatch, we only pay this cost whenever we retrieve the store for an attribute and not for every attribute value individually. Moreover any dynamic typing is hidden behind a strictly typed API. This design also works well with the Rust borrow checker, since it allows us to immutably borrow some attributes while writing to another. Potentially this can also help with parallelism in some cases.

We can not use anymap directly since we require some additional structure on the attribute stores. In particular we would still like to be able to serialise an attribute store. Luckily this is not too difficult to implement. A proof of concept can be found here: https://github.com/CQCL/hugr/blob/feat/attributes/hugr-core/src/hugr/attributes.rs

For properties we can take a similar approach. I am still experimenting with how this could look like precisely. It should be possible to store properties in such a way that it is very efficient to iterate over all nodes for a particular operation, together with their properties. This is a very common and performance sensitive access pattern (for example every pattern match begins with this).

[doug-q]

Moreover the only arguments that can be provided to the operation directly are in the form of type parameters. This would prevent us from implementing custom operations that carry additional information which is not a type (unless we apply hacks such as embedding JSON into the type language).

Ahem. We can embed JSON into Type Args via TypeArg::Opaque. I think Type Schemas can be too, although I'm not so sure.

However for metadata of known extensions, we currently have no choice but to serialise and deserialise the metadata every time we want to look at it, or keep it around in external maps.

We haven't used metadata too much yet, but do remember that the "origin" of metadata is not necessarily the same extension as the op. For example, guppy may store debug locations in metadata. This means ops cannot expect to be able to interpret all their metadata. My view is that TypeArgs are the right place for op-known data.

A suggestion on how to move forward: To optimise for the convenience of users of the hugr library that are mainly concerned with extensions outside of the core, we should implement all core operations via extensions. This makes us feel the pain points of the extension mechanism early, and prevents us from taking shortcuts by hardcoding the core stuff.

I see your point here, and I agree that there is an argument for leaf ops (MakeTuple,UnpackTuple,Tag,Noop,CallIndirect,possibly Call,FuncDecl,Const,LoadConstant,LoadFunction,Lift) to be extension ops, for the reasons you state. I do not think that argument extends to "Parent" ops (DFG,CFG,FuncDefn,TailLoop,Conditional). Disallowing extension ops to have children is an important design choice that:

• simplifies the extension op interface
• Makes reasoning about control flow simple and "closed" in the sense that we can assume we know everything.

The ops in the second half of the above parenthetical after possibly: There is a tradeoff to be made here. Reasoning about a closed set of operations with only a closed set of ways of loading constants, calling functions, etc. is an important strength of keeping these "hardcoded".

Keeping classes of functionality in "hardcoded" ops means we don't have to expose this functionality to the extension system.

One way of thinking about this, that captures both the "parent" ops and "hardcoded" leaf ops, is: any ops using non-value edges(i.e Static,Const,Function,ControlFlow,Hierarchy) should be "hardcoded".

When an extension does know the operation, it should be able to deserialise the property into a Rust type and store that type directly in the HUGR

I agree that this would be great, but it does present implementation difficulties. This is analogous to how CustomConst currently works. We use the typetag crate which uses linker magic to generate an impl serde::Deserialize for dyn CustomConst. My view is that this is too magic to rely on and that we are going to get bitten by it.

I think the current system, ExtensionOp storing def, signature, and type args is a good balance at avoiding typetag. I do agree that working with ExtensionOp is unpleasant, but I think that problem should be solved by providing "views" onto it via MakeExtensionOp. I think do think our views could be better.

I agree that our current approach does deserialise on every access, which is not great, but I suggest it can be adequately solved by a caching layer outside hugr-core.

[zrho]

Disallowing extension ops to have children is an important design choice

Oh. That quite significantly restricts what extensions can do, and excludes almost everything that I wanted to use HUGR for in the first place. Instead of "MLIR with linear types" we would end up with "LLVM with linear types". The following should not be builtin, but profit immensely from having hierarchical boxes:

1. A "run somewhere else node", for instance for tierkreis
2. An accumulating loop construct, like used by Weld. I have planned to use something like this for the supercomputing project by compiling high level parallel code to MPI.
3. Automatic differentiation as a box (done here classically and for quantum here)
4. The "simd boxes" for machine learning that the Oxford office is using in their notation
5. Fixed point computations for streams, which are useful to express incremental computation (see DBSP)
6. Polyhedral forms for loop optimisation

If we exclude operations with children from HUGR, those would not be possible. In particular, the Oxford office would be unable to use HUGR for their plans and would have to build their own thing.

Let's for now define "General HUGR" as basically just the data structure, with an infrastructure for rewrites, analysis, serialisation, visualisation etc. Then there would be a "Restricted HUGR" that is a subset of that, including the "core" operations and simple extensions. If we build "Restricted HUGR" we would have to build "General HUGR" as well. If we build "General HUGR", we can still write analyses that only work on the restricted subset. These passes would then also be applicable to general hugr modules after lowering. This is not premature generalisation; I have very concrete use cases in mind already that I simply wouldn't be able to do.

---

I agree that the dynamic typing approach is somewhat clunky, but it is not essential. There is another option that is growing on me: We want to have a declarative specification of operations anyway, including operation names, types, docs, etc. There is also already code for this. Yet there also are quite a few MakeOpDef implementations that express this information in code. At that point we might as well generate that code from the declarative specification, making it a single source of truth. That would come with several advantages:

1. We'd automatically keep the Rust implementation and the declarative specification in sync.
2. We'd be able to generate nice Python bindings as well.
3. We can experiment more with concrete implementations since they become more interchangeable (in particular a datalog based one)
4. We can generate documentation from the specification and be sure that it is complete.
5. By reducing the amount of special casing we are guided towards coherent design.

I can experiment with how that would look like.

[doug-q]

Disallowing extension ops to have children is an important design choice

Oh. That quite significantly restricts what extensions can do, and excludes almost everything that I wanted to use HUGR for in the first place. Instead of "MLIR with linear types" we would end up with "LLVM with linear types". The following should not be builtin, but profit immensely from having hierarchical boxes:

I agree those (hmm, the subset I understand) are all good use cases.

I do expect that they can all be achieved by Const nodes containing Value::Function nodes(which contain Hugrs). (Other CustomConsts are free to contain anything, including Hugrs. Value::Function should likely be refactored into an extension CustomConst). Users would rewrite LoadConstant targeting such nodes to do whatever would be needed.

If we were to add parent extension ops I would suggest that we should very carefully consider whether we should instead be building on MLIR. We do have some progress in this direction, although it is a bit stale https://github.com/CQCL/hugr-mlir.

In any case, I do think we should all urgently agree on this(parent extension ops) design goal.

We want to have a declarative specification of operations anyway, including operation names, types, docs

My personal view is that a yaml (or similar) declarative specification is not likely to work out well, although I am surely in the minority on this. I think that enough ops will need enough rust code that you might as well do your declarative specification in rust and generate docs/bindings from that. I don't understand what you mean by "alternative implementations"(I heart datalog though). I admit that we do not generate docs from our rust specs now, but we certainly shoud.

[zrho]

What I had in mind with HUGR is indeed very similar to MLIR. If the MLIR project would have started with a specification, of which their implementation is just one of potentially many, I would be pushing hard to adopt it. The idea is amazing. But unfortunately the project is very intimately tied to the one massive C++ implementation; even the declarative parts leak a lot of implementation details. You can't even parse an MLIR file without linking in custom parsing code for each dialect. If you do MLIR, you have no choice but to link in over half a million lines of code. MLIR can be very interesting as a compile target. But my prejudices tell me that maintaining a binding would be more effort than writing a simple toolchain ourselves that fits our needs.

If you take the ideas of MLIR and make them into a generic file format with some tooling around it, you would be very very close to how I imagine HUGR. In my mind I have been using "extension" and "dialect" synonymously. The most significant differences are an orientation towards purely functional/data focused dialects over MLIR's imperative ones, as well as linear types. MLIR instructions have an ordering that is significant, and have some effect annotations that tell the optimiser what instructions can be reordered. In HUGR we express mostly data dependencies, which makes sense for a compiler of high-level programs. Quantum computing forces linear types on us already, but we can make use of them to handle mutability and resources in many cases. We can therefore do optimisations that would be difficult with general effectful programs. We can also structure the IR around linearity, instead of having to carefully preserve it in MLIR.

One reason why I am bullish on the declarative side is precisely to avoid the single implementation problem of MLIR. By having HUGR be implementation agnostic, we can experiment with stuff like a datalog backend or data structures with different performance tradeoffs while retaining compatibility.

I do not know how many ops will need to break out of the declarative garden but I suspect that we can do a lot. Consider for example the tierkreis eval node, which was special cased at first but with row types becomes just a normal node. Now row polymorphism is used for other things in tierkreis as well, and these things work well together since the goal of having complete type inference forced a coherent design. With rows, type level lists, a sufficiently powerful system of type constraints I expect we can go far.

Even for ops that do need special treatment after all, we might still have a declarative overapproximation that comes with a validator. We can then manipulate, produce, print, pattern match and rewrite that overapproximated op without any specialised Rust code, with the understanding that the validity of this is not entirely captured by the type system. What the validator should do can be documented alongside the declarative spec.

[doug-q]

I do expect that they can all be achieved by Const nodes containing Value::Function nodes(which contain Hugrs). (Other CustomConsts are free to contain anything, including Hugrs. Value::Function should likely be refactored into an extension CustomConst). Users would rewrite LoadConstant targeting such nodes to do whatever would be needed.

Some further thoughts on this.

Being a bit more concrete, how to "run code on another tierkreis worker" :

• define a Custom Type "TierkriesCode" which contains a Hugr, has a TypeArg with the FunctionType of that Hugr, and information needed to compile (i.e. codegen target etc.)
• a Const Node containing an instance of TierkriessCode<FT>
• a LoadConstant pointing at that node gives you a wire of type TierkriessCode<FT>
• an extension op "run" with signature [TierkreisCode<FT>] -> [FT]
• Define the on-the-wire representation of TierkriesCode<FT> to be an elf shared library and an entry symbol.
• Rewrite the LoadConstant: compile the Hugr inside the Const into a shared library, embed that as a binary blob and load the binary blob onto a wire.
• Lower the "run" to perform networking calls to some remote agent.

For Automatic Differentiation I think you can use the same method, triggered by LoadConstanting a specific CustomType to run some Hugr -> Hugr pass.

I have realised that there is nothing preventing one from putting a Hugr in a TypeArg, and one could use this to implement "parent" ops. For example, an extension Conditional could take Case Hugrs as TypeArgs.

[doug-q]

You can't even parse an MLIR file without linking in custom parsing code for each dialect. If you do MLIR, you have no choice but to link in over half a million lines of code. MLIR can be very interesting as a compile target. But my prejudices tell me that maintaining a binding would be more effort than writing a simple toolchain ourselves that fits our needs.

I agree that MLIR is large and C++ and imperfect.

I'm not sure exactly what "maintaining a binding" means, but my prejudices tell me the opposite: reusing an existing imperfect tool would be less effort than rewriting much of it in our own imperfect tool.

One can of course parse a "generic" MLIR file(text or bitcode) without custom parsing code for each dialect. MLIR providing pretty printing/parsing is a great strength. It's very hard to work with HUGR while we have no good textual format. Our best way to pretty print HUGRs is to generate and print mermaid diagram source code. One expects this will not scale well with larger HUGRs.

[ss2165]

Some general responses to above discussions:

• While extensibility is indeed a core goal, that does not necessarily mean "MLIR with linear types". The extensibility is implicitly limited in scope to satisfy our expected use cases, in order to limit implementation effort and improve ease of use (in brief, it should be as simple as possible while satisfying the use cases). Now, if we've got our set of expected use cases not quite right, we can reassess what needs to happen to satisfy them.
• "Closed" control flow over the core operations by restricting hierarchy to the core ops is indeed a design intent. Anything beyond this required by extensions is intended to be covered by higher order inputs. Perhaps we should think about an implementation optimisation for Constant Hugrs that adds them to the hierarchy of the parent hugr?
• I like the argument for a better declarative extension specification that would allow us to move more extensions to be defined outside Rust. Indeed, even for our own use cases HUGR touches enough varied projects that we are at risk of overly tying the spec to one implementation, and the more we can do in an "implementation-agnostic" fashion, the more we benefit in the long run. At the outset of the implementation, I was not convinced that doing things like Rust code generation from a declarative specification was sensible - it was likely to result in a difficult to maintain code base. I am interested to see what recent expansions in the capabilities of the extension system (broader type schemes and row variables) make possible within declarative extensions, so encourage exploratory experiments. Note that such experiments should think about how to ergonomically use such extensions from at least Rust and Python.
• A scalable pretty print of HUGRs is a necessary feature. How that is achieved can be debated but MLIR is a good option from the outset.

[zrho]

I've found a way to recontextualise the operation properties as type arguments approach that makes me a bit less skeptical.

We have been encoding values as types in order to pass them as properties to operations. This reminds me of datatype promotion in Haskell (see Giving Haskell a Promotion or DataKinds). Any type is given the kind Type, while every data type is promoted to a kind itself. Datatype promotion allows to express types that depend on values, such as a fixed length vector type whose type constructor has the kind Vec : Int -> Type -> Type. Note that this is not full dependent typing.

The signature of an operation is a type schema, consisting of

• a sequence of type variables together with their kinds
• a collection of type constraints
• a function type

A prefix of the type variables of the operation's signature are to be supplied as properties to the operation. The operation specifies which type variables belong to this prefix. These can be types, but due to datatype promotion, they can also be values. When a type variable belongs to the operation's properties, it must be explicitly specified. The remaining type variables are inferred.

• This way the type checker can be used to verify that an operation's properties have the correct type.
• The operation's properties can influence the type of the operation, in a way that can be specified declaratively.
• This is compatible with how properties are handled as type arguments in HUGR currently.

Type Constraints as Escape Hatches for Custom Logic

There are operations whose types depend on the operation properties in complex ways. In these situations we can use the type constraint language as an "escape hatch". When an extension requires custom typing, it can declaratively specify a collection
of type constraints and express the signatures of its operations using these constraints. The extension can then supply Rust code as a type checker plugin which can use arbitrary logic to solve the constraints. This is similar to the Programmable Matching used by Cranelift's ISLE to extend their declarative matching system with arbitrary logic in Rust.

• Every operation will be expressible declaratively, and so can take part in codegen for documentation, bindings etc.
• While the type checker plugin for an extension's constraints is not declarative itself, the constraint system provides a clean interface. Alternative implementations of the type checker (say in other languages or using different ideas) can implement the same constraints.
• Moving custom typing logic from the operation to the constraint language also provides opportunities for composition and reuse. With a rich and well-designed set of builtin constraints, I expect that many operations will already be typable without any custom logic.

This story does not yet address how to capture operations with nesting. There are some approaches which I can think of, but I'd need a bit to mull it over and see how it can be made coherent.

Code Generation

I typically share concerns about maintainability whenever the standard build process is messed with. However, it appears to me that in our case codegen would lead to more maintainable code. We could add an extension and simultaneously generate documentation, specialised Rust code, and convenient Python bindings. We could change implementation details without having to update every extension implementation (be it builtin or custom). I expect the number of operations that we will deal with to realistically be in the low hundreds, so that would be a non-trivial amount of tedious work when done manually.

As precedent, I want to point out that Cranelift uses code generation for their opcode and instruction enums (see here). Because Cranelift does not have an extensible instruction set, this is done merely for the convenience of it. With extensibility coming into the mix, the case for codegen would be even stronger.