We propose a minimal integration of the SPV_EXT_replicated_composites extension into the MLIR SPIR-V dialect. This will enable the SPIR-V serialization layer to emit OpConstantCompositeReplicateEXT and OpSpecConstantCompositeReplicateEXT when appropriate, based on existing spirv.Constant and spirv.SpecConstantComposite operations that represent splat values.
No new MLIR operations will be introduced as part of this proposal.
OpConstantCompositeReplicateEXT and OpSpecConstantCompositeReplicateEXT allow for defining replicated (splat) constants and spec constants in a more compact and semantically explicit manner.
Proposal
We propose extending the MLIR to SPIR-V serialization logic to detect splat composite constants and emit the new extension instructions when:
A spirv.Constant represents a splat value → emit OpConstantCompositeReplicateEXT.
A spirv.SpecConstantComposite represents a splat value → emit OpSpecConstantCompositeReplicateEXT.
Note:OpCompositeConstructReplicateEXT instruction is meant for runtime construction of splat composites from SSA values, whereas this proposal only addresses constant and spec constant lowering. As such, there is no need to implement or consider OpCompositeConstructReplicateEXT for this RFC.
Constraints
This extension only supports full replication where all constituents are identical.
The pass will conditionally emit extension ops only when the module enables SPV_EXT_replicated_composites extension and ReplicatedCompositesEXT capability, that is, SPIR-V modules using this transformation must explicitly declare the following otherwise fallbacks to emiting legacy OpConstantComposite / OpSpecConstantComposite.
The OpConstantCompositeReplicateEXT and OpSpecConstantCompositeReplicateEXT instructions require the result type to be a composite type of Vector, Array, Tensors or Struct. Any other usage will be considered invalid.
The replicated operand used in these instructions must be Exactly the same type as each constituent element in the composite type. Especially for structs, where all members must be of the same type for replication to be valid.
Only SPIR-V binary generation (serialization) will be affected.
No new MLIR ops will be added.
Motivation
Improves SPIR-V size and readability when using repeated values.
Preserves MLIR’s existing simplicity and avoids dialect bloat.
Support includes:
Round-tripping between MLIR and SPIR-V binary.
Versioning
The SPV_EXT_replicated_composites extension requires SPIR-V version 1.0
Testing
Tests are included to verify roundtrip conversion between MLIR and SPIR-V binary where OpConstantCompositeReplicateEXT and OpSpecConstantCompositeReplicateEXT are serialized and deserialized.
Wouldn’t it make more sense to represent those two in MLIR (maybe it could be as simple as adding a splat attribute to the existing constant), have a proper analysis and transformation pass to convert spirv.Constant where legal, and then just have a simple serialization? That way serializer remains a “dumb” layer that just emits instructions but doesn’t do any analysis (beyond what’s necessary). But, also, I can see this being an overkill. Another, minor argument, for not doing it in serialization would be that any users of the dialect that don’t go through serialization can benefit from this change (although not sure how useful it would be). That’s my two cents.
+1 to what @IgWod said, I’d prefer to keep the (de)serializer as simple as possible. It’s much easier to maintain and debug IR-level transformations.
We should only emit those when users explicitly opt into using the SPV_EXT_replicated_composites extension. A new pass can turn splat constants into new ops that mirror Op*CompositeReplicateEXT.
Just to make sure I understand your suggestion correctly, you’re proposing, as one of the suggestions, that instead of detecting replicated patterns in the serializer, we should have a dedicated MLIR transformation pass that runs at the SPIRV dialect level, after lowering, and looks for splat spirv.Constant or spirv.SpecConstantComposite and then rewrites them to new ops like spirv.ConstantCompositeReplicated / spirv.SpecConstantCompositeReplicated, assuming the extension is enabled. Then the serializer would just emit those directly to OpConstantCompositeReplicatedEXT/ OpSpecConstantCompositeReplicatedEXT, respectively, is that right?
Assuming that’s the case, I have a few questions/notes around this approach:
@kuhar We still can have the logic of conditional emitting of OpConstantCompositeReplicatedEXT/ OpSpecConstantCompositeReplicatedEXT only when the SPV_EXT_replicated_composites extension is explicitly enabled during serialization.
Since, for example OpConstantCompositeReplicatedEXT and OpConstantComposite are semantically identical, I wonder what would be the gain by having two distinct ops in the IR, if we take this approach, considering that new ops would require extra plumbing, e.g., parsing, verification, printing, etc. Also I wonder if introducing new ops may confuse the users in what is the practical difference of new ops with the existing ops.
As @IgWod mentioned I’m also wondering if there are consumers that would actually benefit from seeing new replicated ops? Otherwise, we might be exposing a detail that is mostly relevant to binary size, not analysis or transformation.
maybe it could be as simple as adding a splat attribute to the existing constant
I see a splat attribute to only express encoding hint, not a semantic distinction. It wouldn’t change the meaning or behavior of the constant, so I don’t see how IR consumers would benefit from it, plus it calls for extra pluming as well as possible user confusion (e.g., thinking the attribute has semantic effects).
Appreciate your thoughts on the above, so that I can gauge better the trade-off with the additionally complexity that your suggestion may introduce.
We still can have the logic of conditional emitting of OpConstantCompositeReplicatedEXT / OpSpecConstantCompositeReplicatedEXT only when the SPV_EXT_replicated_composites extension is explicitly enabled during serialization .
I would strongly discourage this, IMO the simpler the serialization logic the better. The reason is that it’s much easier to debug IR → IR transformations than IR → binary serialization that may result in malformed binaries that tools can’t read / print.
You’re right, it was just something that crossed my mind when writing, but probably I should’ve given it a bit more thought. My reasoning was that it would simplify spirv.Constant [0, 0, 0] : !spirv.array<3 x i64> so it only specifies one element, but I’ve just realised you can already do that in MLIR: spirv.Constant dense<0> : vector<3xi32>. Let’s forget I said it
Most of it is handled by tablegen, so plumbing should be minimal. But I can see opportunities. We can now add verification (again, with some coming at minimal cost from tablegen) to check if those new constants are correct, so we can catch problems at IR level. It is also in the spirit of the dialect:
To summarize, the SPIR-V dialect follows the following design principles:
Stay as the same semantic level as the SPIR-V specification by having one-to-one mapping for most concepts and entities.
Adopt SPIR-V specification’s syntax if possible, but deviate intentionally to utilize MLIR mechanisms if it results in better representation and benefits transformation.
Be straightforward to serialize into and deserialize from the SPIR-V binary format.
Finally, I admit you make some good points and it’ll always be a trade-off, but I think keeping things in MLIR and keeping the serializer simple is important.
Many thanks again @kuhar / @IgWod for your clarifying points.
@kuhar I have a question on your suggestion of having new ops. For the case of prospective op spirv.ConstantCompositeReplicated I can think of two possible ways for the op’s format which would affect parsing/printing and serialization:
SSA-style operand which mirrors how their corresponding SPIRV ops are presented in binary
The problem with this approach would be that it wont work for a less trivial constant type like the one below, as one would need to specify the type of the splat value (e.g. here the splat type would be vector<2xi32> not i32 - Also in this example there is no easy way to deduce type of dense<[1, 2]>)
It’s up to you. You can support both options at the same time – you can find examples of ops that take either attributes or constant operands in the vector dialect.
Personally, I think that this extension seems very niche and I’d suggest you to implement whatever you find actually useful.
