[Discussion] Recognize arith/math patterns and lower them to spyreop intrinsics

[fabianlim]

Summary

ktir-mlir-frontend#57 adds the spyreop dialect — the IBM Spyre intrinsics the
frontend is allowed to emit. Nothing in this repo emits them yet. Triton's frontend
has no vocabulary for a Spyre intrinsic, so a primitive that Spyre implements as one
instruction reaches KTIR as a chain of arith/math ops.

We need a pass that recognizes those chains and lowers them to the matching
spyreop intrinsic.

The canonical case is precision-widened math. tl.rsqrt on f16 arrives as:

%0 = arith.extf   %a : f16 to f32
%1 = math.rsqrt   %0 : f32
%2 = arith.truncf %1 : f32 to f16

The extf/truncf pair is not arithmetic the user asked for — it is the frontend
expressing "compute this at higher precision". spyreop provides sqrt and
reciprocal operating on !spyreop.df16, so if the intrinsic already carries the
needed internal precision, the widening ops are pure overhead: they cost tile memory
for the intermediate f32 values and hide the single instruction the hardware wants.

What #57 gives us

From the PR's tests, the initial op surface:

spyreop.sqrt        spyreop.reciprocal   spyreop.realdiv     spyreop.exp
spyreop.gelu        spyreop.sigmoid      spyreop.softplus    spyreop.layernormscale
spyreop.addi32toi32 spyreop.addi64toi64  spyreop.muli32toi32 spyreop.idx32toaddr

plus the !spyreop.df16 type.

These are scalar ops, not tensor ops. Every signature in #57 is scalar — e.g.
spyreop.gelu %arg0 : !spyreop.df16. That is the single most important constraint
on this work, and it decides the shape of the pass:

func.func @gelu(%arg0: !spyreop.df16) -> !spyreop.df16 {
  %0 = spyreop.gelu %arg0 : !spyreop.df16
  return %0 : !spyreop.df16
}

So we are not replacing a tensor-level chain with a tensor-level spyreop. The
intrinsic goes inside a linalg.generic region, replacing the scalar body ops
while the surrounding structured op — its iteration space, indexing maps, and
operands — stays as it is. That keeps physicalization, tiling, and the dataflow
scheduler working on a shape they already understand.

Two consequences worth stating early:

• The match happens on generic bodies, not on tensor op chains. This lines up
  well with [Discussion] Normalize to linalg.generic up front so layout physicalization emits no synthesized loops #90 (normalize everything to linalg.generic up front): after that, the
  matcher has one surface instead of several. Ordering between the two should be
  agreed before either is implemented.
• !spyreop.df16 vs f16 needs a decision. A linalg.generic body yielding
  !spyreop.df16 while its output tensor is tensor<...xf16> will not verify
  without a conversion at the region boundary. Whether we convert at the boundary,
  carry tensor<...x!spyreop.df16>, or treat df16 as bitwise-compatible with
  f16 is an open question I can't answer from Remove GIT_PAT requirement from Spyre build #57 alone.

Proposed pass

A pattern set over linalg.generic bodies, after LowerComputeOps (and after #90's
normalization, if that lands first):

• Precision-widening elision — extf → math.op → truncf collapsing to one
  intrinsic where the intrinsic's internal precision matches. sqrt, reciprocal,
  exp are the immediate candidates.
• Composite primitives — chains that map to one intrinsic even without a widen:
  sigmoid, gelu, softplus, layernormscale are each several arith/math ops
  as Triton emits them.
• Guards. Fuse only when every intermediate value is single-use inside the
  region. An extf result read by anything else must not be elided, or that other
  consumer silently loses precision.

Numerics are the real risk

Eliding a widen changes results. extf → rsqrt → truncf computes at f32 and rounds
once; spyreop.sqrt + reciprocal on df16 computes at whatever the hardware does
internally. Those agree only if the intrinsic's internal precision is at least f32,
and df16 is not f32 — so this needs confirming per intrinsic against ktir-cpu or
the hardware spec, not assumed from the op names. A fusion that quietly changes
numerics is worse than no fusion.

Note also that rsqrt has no direct spyreop equivalent: it would be
sqrt + reciprocal, i.e. two intrinsics, with its own rounding question. A
single-intrinsic pattern (exp, or sigmoid) may be the better first target.

Interaction with in-flight work

• [Discussion] Normalize to linalg.generic up front so layout physicalization emits no synthesized loops #90 (normalize to linalg.generic): should probably land first — it gives this
  pass one matching surface instead of several named op forms.
• KTIR lowering admits only one linalg op per compute group; frontend output routinely has several #84 (one linalg op per compute group): fusion reduces op count per group, so it
  relieves pressure rather than adding to it. Worth checking whether collapsing
  chains fixes some of KTIR lowering admits only one linalg op per compute group; frontend output routinely has several #84's violating fixtures outright.
• PR RewriteDescriptorLayout: report consumers no pattern can physicalize #95 / RewriteDescriptorLayout: because the intrinsic sits inside a
  linalg.generic and does not change the structured op's shape, physicalization
  should be unaffected — retypeChain handles a generic by rank-bumping the tensor
  type without touching the body. Worth an explicit test rather than an assumption,
  since Phase 3 will now report a mismatch loudly if that is wrong.
• ktir-mlir-frontend#57 must merge and the submodule pin move before any of this
  can build.

Proposed scope for a first PR

Narrow on purpose, to settle the representation questions before building a table of
patterns:

1. Pick one intrinsic — I'd suggest exp or sigmoid over rsqrt, since both
   are a single intrinsic and avoid the two-op rounding question.
2. Confirm numerically against ktir-cpu that the intrinsic agrees with the chain it
   replaces, at the precision the frontend emits.
3. Resolve the df16/f16 boundary question for that one case.
4. Lit coverage: the fused match; a case where an intermediate has a second use and
   must not fuse; a precision case that must not fuse.
5. A numerical test comparing fused vs unfused.

Explicitly out of scope: a general fusion framework or a full pattern table. One
intrinsic end-to-end tells us whether the df16 boundary and the in-region approach
are right; a framework built first would bake in whatever we guessed.

What I need from reviewers

• The df16 boundary question — convert at the region boundary, carry
  tensor<...x!spyreop.df16>, or treat it as f16-compatible? This is the main
  blocker on writing anything.
• Per-intrinsic internal precision. Is there a spec for what spyreop.sqrt et al.
  do internally, or does this need measuring against ktir-cpu?
• First intrinsic — is exp/sigmoid the right start, or is there one with
  measured impact on a real kernel worth prioritizing?
• Ordering against [Discussion] Normalize to linalg.generic up front so layout physicalization emits no synthesized loops #90.

[KFAFSP]

Two comments:

• There seems to be an rsqrt intinsic in the back-end. However, it's not used at the moment, so it didn't come up in my first survey. It might be an option to add rsqrt to spyreop since there are efficient implementations for it.
• It would be trivial to let the spyreop operations support tensor operands, and then use Elementwise, Scalarizable, Tensorizable traits to allow linalg to convert them into linalg.generic on scalars. @vswagath1989 was against that for now, but maybe you should keep that in mind before building huge work arounds.

[airin711]

I have a few questions and concerns.

Spyreop intrinsics

As noted in the previous comments, DDL (opaque templates) already includes an rsqrt operation. There are also several other operations that have been used in DDL but are not yet implemented as spyreops. Conversely, there may be operations that existed in DDL but should not necessarily be represented in the same way in KTIR.

For example, LayerNorm was split into three operations in DDL: exx2, layernorm_norm, and layernorm_scale. However, my understanding is that within the KTIR compilation pipeline these could be represented as several operations like a DDL, or even without introducing dedicated spyreops at all, using a combination of standard reductions and element-wise operations. I think we should perform a thorough review of all Spyre built-in instructions and carefully evaluate which ones should actually be exposed as spyreops.

Precision

Regarding numerical precision, many operations already support fp32 at the DDL (opaque template) level. In other words, if the frontend explicitly selects fp32 because additional precision is genuinely required, then it should also be possible to process those operations as fp32 at the spyreop level.

Is it really correct to eliminate precision-widening in all cases? From previous experiments, we know that df16 precision is sufficient for many element-wise operations, but there are also important cases, particularly operations involving accumulation such as LayerNorm, where fp32 computation is necessary to achieve the desired numerical accuracy.

Similar to GPUs, I believe the model should explicitly specify the precision appropriate for the hardware. The frontend should generate precision-widening code accordingly, and that intent should be represented in KTIR. If that is the case, then those precision-widening operations should not be removed indiscriminately during optimization. What are your thoughts on this?

[KFAFSP]

For example, LayerNorm was split into three operations in DDL: exx2, layernorm_norm, and layernorm_scale. However, my understanding is that within the KTIR compilation pipeline these could be represented as several operations like a DDL, or even without introducing dedicated spyreops at all, using a combination of standard reductions and element-wise operations. I think we should perform a thorough review of all Spyre built-in instructions and carefully evaluate which ones should actually be exposed as spyreops.

Eventually, we might not need to expose any operations at all, but with the current method of doing things (e.g., some decompositions happen in the front-end), they are needed to cross the boundary.

Additionally, until things like linalg get proper support for custom data types, we will always be stuck with having to roll our own ops for the dl16 data type.

Is it really correct to eliminate precision-widening in all cases?

Where did this argument come from? As far as I understand, we're currently pretending that f16 == dl16, which is the main issue mentioned here. So far, spyreop allows IEEE binary32 wherever an intrinsic implementation can be provided. In the cases where f32 isn't an allowed type, that's because it can't be implemented in that way.

We can agree on widening the types on spyreop and have the scheduler choose alternatives, however, that might end up just hiding the inaccuracy in a later stage.

Similar to GPUs, I believe the model should explicitly specify the precision appropriate for the hardware.

With model, do you mean the KTIR programming model, or the ML model?

[airin711]

I don't think spyreop will ever become completely unnecessary. For example, operations such as exp, log, sqrt, and realdiv are expressed as special combinations of Spyre ISA instructions, so I agree that they should be provided as intrinsic operations in spyreop.
On the other hand, computations such as sigmoid, gelu, softplus, and layernorm can be expressed as combinations of the primary operations mentioned above. Shouldn't we carefully evaluate whether these should really be bundled into single spyreop operations?

Regarding precision, I was concerned because I saw a statement suggesting that precision widening would be eliminated even in a case like extf → rsqrt → truncf where FP32 computation is explicitly indicated. If I misunderstood the intended meaning, I apologize.

By "model," I was referring to the ML model. For example, on GPUs, I believe models from frameworks such as Hugging Face explicitly specify that most operations should be computed in bf16, while certain operations should be computed in fp32. In the same way, if the frontend model explicitly specifies that a computation should be performed in FP32, then I think the dl16 → fp32 precision widening should be performed explicitly.
On the other hand, if no such specification exists in the frontend, and the widening to FP32 is being inserted automatically during the conversion process to Triton or KTIR, then I agree that those inserted precision expansions should be removed.

As for the f16 → dl16 issue, I agree that explicit precision conversion annotations should be eliminated.

[fabianlim]

I think its time to revist this now that we have concrete demonstrations of how this is done in torch-spyre torch-spyre/torch-spyre#4125

• key point is that these ops accept both f16/f32/i32 dtypes upfront, so for these types it is straightforward we do not need type-casts
• but for the dl16 yes perhaps we may need conversions (but currently the distinction between f16 and fl16 is cosmetic)