Add LowerSpyreOps for lowering scalar math/arith ops to spyreops - #134
Add LowerSpyreOps for lowering scalar math/arith ops to spyreops#134airin711 wants to merge 11 commits into
Conversation
fabianlim
left a comment
fabianlim
left a comment
There was a problem hiding this comment.
rebase this PR to upstream main, and follow suggestions to install pass in _SPYRECODE_STAGE_PASSES, then all the xfails will be handled.
Did you test this on device? Also to note the torch-spyre PRs for handling layer norm as a multi-stage kernel is torch-spyre/torch-spyre#4315
| ConversionPatternRewriter &rewriter) const override { | ||
| if (!isInsideLinalgGeneric(op) || getScalarIntBitWidth(op.getType()) != 32) | ||
| return failure(); | ||
| rewriter.replaceOpWithNewOp<spyreop::MulI32ToI32>( |
There was a problem hiding this comment.
there is no I64 version for this?
There was a problem hiding this comment.
Currently, MulI64ToI64 has not yet been implemented in spyreop.
|
|
||
| /// Whether spyreop's scalar float intrinsics accept this operand type. | ||
| static bool isSpyreOpScalarType(Type type) { | ||
| return isa<Float16Type, Float32Type, spyreop::DF16Type>(type); |
There was a problem hiding this comment.
isSpyreOpScalarType accepts DF16Type, which is unreachable. math::Sqrt/Exp/Rsqrt and arith::DivFOp all constrain operands to builtin FloatType; spyreop::DF16Type isn't one. That arm can never fire. Anyone reading line 51 concludes df16 compute is handled — it isn't, and adding a df16 kernel would produce a confusing "failed to legalize."
| rewriter.replaceOpWithNewOp<spyreop::AddI64ToI64>( | ||
| op, op.getType(), adaptor.getLhs(), adaptor.getRhs()); | ||
| else | ||
| return failure(); |
There was a problem hiding this comment.
this failure path is dead since you marked the width types other than 32 and 64 legal
| // arith.addi/muli are also used for plain index/address arithmetic, so | ||
| // (unlike the ops above) they stay legal everywhere except the one | ||
| // context and bit-width this pass actually converts: scalarized | ||
| // elementwise compute inside a linalg.generic, at a width spyreop has an |
There was a problem hiding this comment.
i see you want to only rewrite those integer ops that lie inside a generic, to detect the compute ops. I guess the assumption that similar scalar ops appearing elsewhere we assume that they are being handled for other purposes (e.g., addressing).
There was a problem hiding this comment.
Yes, that's correct. Integer arithmetic operations are frequently used for purposes such as address calculations, so I impose that restriction.
|
@fabianlim Thank you for all the review comments. I have replied to or just resolved them. |
|
@airin711 thank you. need to ensure that your commits are signed and also have verified signatures. Also if you can do a test on the device that would be great |
| "1D unary spyreop `out = OP(x)` over a fully-static vector, " | ||
| "partitioned across the 32-core grid. Sweeps OP across " | ||
| "sqrt/rsqrt/exp." | ||
| ), |
There was a problem hiding this comment.
did you update ktir-cpu to handle spyreops? I know its good for consistency but im wondering if its worth the effort. Still, if that is already done, then we need to update the hash in setup.py to reflect the new commit. Otherwise if you think its not worth the effort, then we need to discuss how to test these fixtures
There was a problem hiding this comment.
ok i see, based on your comment above, this is not checking spyreop. then it does not make sense to have level A/B/C fixtures
| ``np.sqrt``/``np.exp`` and ``rsqrt`` as ``1.0 / np.sqrt(x)``, confirmed | ||
| empirically (and matched exactly by this file's oracle, below). | ||
|
|
||
| fp32 only, throughout: upstream ``tl.sqrt``/``tl.rsqrt``/``tl.exp`` |
There was a problem hiding this comment.
oh if this is the case, we need to update python/triton/language/math.py to relax this for spyre, see the precedent e.g., #3
| which this fixture does not attempt. | ||
|
|
||
| The Level D variant is a real Spyre binary and does go through | ||
| ``lower_spyre_ops``: that's the whole point of ``check_sqrt_e2e.py`` |
There was a problem hiding this comment.
nit: is `check_sqrt_e2e.py`` a private script?
| generalizes to all three ops. The KTIR a structural test sees is still | ||
| pre-spyrecode, the same as ``elementwise``'s ``1d_device`` and ``reduce``'s | ||
| ``one_tile`` -- so it is not asserted here either; the real, post-lowering | ||
| ``spyreop.<op>`` is what ``check_sqrt_e2e.py`` and the ``ktir-spyreop-sqrt.mlir`` | ||
| lit test cover for ``sqrt``. |
There was a problem hiding this comment.
| generalizes to all three ops. The KTIR a structural test sees is still | |
| pre-spyrecode, the same as ``elementwise``'s ``1d_device`` and ``reduce``'s | |
| ``one_tile`` -- so it is not asserted here either; the real, post-lowering | |
| ``spyreop.<op>`` is what ``check_sqrt_e2e.py`` and the ``ktir-spyreop-sqrt.mlir`` | |
| lit test cover for ``sqrt``. | |
| generalizes to all three ops. |
nit: in general i think lets try to keep the comments to a minimum. The "commentary style" comments can be toned down
| # It fails both halves of the rule above. dbo-opt is the one that needs | ||
| # the spyreop spellings; a kernel that stops at KTIR has no use for them. | ||
| # And the cached .ktir artifact is meant to stay readable by tools built | ||
| # against plain math/arith -- ktir_cpu's numerical oracle among them, | ||
| # which has no MLIRTypeAdapter handler for spyreop.* ops yet (#107) -- | ||
| # so a rewrite only dbo-opt can consume does not belong in that artifact. | ||
| # |
There was a problem hiding this comment.
| # It fails both halves of the rule above. dbo-opt is the one that needs | |
| # the spyreop spellings; a kernel that stops at KTIR has no use for them. | |
| # And the cached .ktir artifact is meant to stay readable by tools built | |
| # against plain math/arith -- ktir_cpu's numerical oracle among them, | |
| # which has no MLIRTypeAdapter handler for spyreop.* ops yet (#107) -- | |
| # so a rewrite only dbo-opt can consume does not belong in that artifact. | |
| # |
| # | ||
| # convert_elementwise_to_linalg and unalias_linalg_outs are what the | ||
| # scheduler inside dbo-opt requires of every kernel it will lower to a | ||
| # binary. lower_spyre_ops is NOT here -- it also depends on | ||
| # convert_elementwise_to_linalg's scalarization, but it belongs to | ||
| # dbo-opt rather than to the IR every compile produces, and it can | ||
| # reject a scalar type spyreop has no intrinsic for (f64, bf16, ...), | ||
| # so it runs later, only for compiles that reach the spyrecode stage. | ||
| # See _SPYRECODE_STAGE_PASSES. |
There was a problem hiding this comment.
| # | |
| # convert_elementwise_to_linalg and unalias_linalg_outs are what the | |
| # scheduler inside dbo-opt requires of every kernel it will lower to a | |
| # binary. lower_spyre_ops is NOT here -- it also depends on | |
| # convert_elementwise_to_linalg's scalarization, but it belongs to | |
| # dbo-opt rather than to the IR every compile produces, and it can | |
| # reject a scalar type spyreop has no intrinsic for (f64, bf16, ...), | |
| # so it runs later, only for compiles that reach the spyrecode stage. | |
| # See _SPYRECODE_STAGE_PASSES. |
| // Not in add_convert_ttir_to_ktdp above: spyreop's scalar intrinsics only | ||
| // accept scalar f16/df16/f32 operands, so this can only fire on a math op | ||
| // that is already scalar -- typically inside a linalg.generic body after | ||
| // convert_elementwise_to_linalg. Reachable individually for now; folding it | ||
| // into the default pipeline is a later ordering decision (anchor on | ||
| // convert_elementwise_to_linalg, same as unalias_linalg_outs). |
There was a problem hiding this comment.
| // Not in add_convert_ttir_to_ktdp above: spyreop's scalar intrinsics only | |
| // accept scalar f16/df16/f32 operands, so this can only fire on a math op | |
| // that is already scalar -- typically inside a linalg.generic body after | |
| // convert_elementwise_to_linalg. Reachable individually for now; folding it | |
| // into the default pipeline is a later ordering decision (anchor on | |
| // convert_elementwise_to_linalg, same as unalias_linalg_outs). |
|
|
||
| Two of that pass's other patterns are deliberately not swept here: | ||
|
|
||
| - ``arith.divf -> spyreop.realdiv`` is binary, not unary -- it needs a second |
There was a problem hiding this comment.
sorry i dont understand, why is spyreop.realdiv not tested? we need to test everything
| already scalarizes a tensor ``/`` into ``arith.divf`` the same way this | ||
| fixture's kernels scalarize into ``math.sqrt``/``math.rsqrt``/``math.exp``, | ||
| so that path already has coverage under a different fixture name). | ||
| - ``arith.addi``/``arith.muli`` (i32/i64) are gated on being inside a |
There was a problem hiding this comment.
we need to figure out how to test this as well, im not sure what the problem really is
| containment predicate, not a type predicate), and, like ``divf``, already | ||
| reachable through ``elementwise``'s int32 add/mul variants. | ||
|
|
||
| Level A/B run *before* ``_make_spyrecode``'s extra stage passes: the tensor |
There was a problem hiding this comment.
i see.. then in this case it does not make sense to have level A/B fixtures, this numerical test with math / arith is irrelevant for spyreop, which is the main focus of this fixture. In that case I suggest lets only do level D here
|
if you do another merge with main, we have removed the docs |
2 participants
Summary
This is PR regarding issue #107, adding a new pass,
LowerSpyreOps, that lowers scalarmath/arithops to theirspyreopdialect intrinsic equivalents, and wires it into the default compile pipeline so every kernel gets it automatically (anchored onrewrite_descriptor_layout, alongside the existingconvert_elementwise_to_linalg/unalias_linalg_outsdefault fixes).Lowering:
Two different legality rules, by op kind:
sqrt/exp/rsqrt/divfare matched unconditionally on scalar type: they only ever appear in real float compute, never address arithmetic, so a scalar operand on an unsupported type (f64, bf16, ...) is reported as illegal rather than silently passed through. Still on a tensor/vector, they're left legal quietly. That's ConvertElementwiseToLinalg's job to scalarize first.addi/muliare scoped: plain scalar integer add/mul is also used for loop indices, offsets, and tile addressing, not just scalarized compute. So these only match inside alinalg.genericbody (the structural signal that it's scalarized elementwise compute) and only at the bit-widths spyreop has an intrinsic for. Anything else is left legal.gelu/sigmoid/silu/softplus/reciprocalare explicitly out of scope: none has a one-to-one source op in TTIR to pattern-match against (they'd need to be expressed as a fused sequence of these primitives instead).Changes
include/Dialect/KTDP/Transforms/Passes.td,include/.../Passes.h,lib/.../CMakeLists.txt,bin/spyre-triton-opt/*lib/Dialect/KTDP/Transforms/LowerSpyreOps.cppConvertMathSqrt/Exp/Rsqrt,ConvertArithDivF/AddI/MulI) plusisSpyreOpScalarType/isInsideLinalgGeneric/getScalarIntBitWidthhelpers.backend/compiler.pylower_spyre_opsas a third defaultrequired_fixesentry, anchored onrewrite_descriptor_layout, afterconvert_elementwise_to_linalgso it sees the scalarizedlinalg.genericbody.test/Conversion/lower-spyre-ops.mlirtest/test_lower_spyre_ops.pylinalg.genericregression tests for the ints) plusDefaultPipelineTestercases proving the pass fires with zero caller-side configuration.test/conftest.py,test/fixtures/{softmax,inter_tile_reduce,elementwise}/meta.pyKnown gap: ktir_cpu has no spyreop support
ktir_cpu (the numerical oracle interpreter used by test_numerical) has no `MLIRTypeAdapter handler for any spyreop.* op. Every fixture that now reaches a spyreop.* op numerically — softmax (all variants), inter_tile_reduce__softmax, elementwise (i32 add/mul, div at any dtype) — would otherwise crash with NotImplementedError.
This PR keeps every conversion wired into the default pipeline rather than gating any off, and marks the affected numerical tests
xfail_numerical(strict-xfail, matching this repo's existing convention for known oracle gaps) with a comment pointing at this issue.Added a new xfail_numerical VariantFactory hook (conftest.py) alongside the existing signature/reference/inputs hooks so the elementwise fixture can condition it per (DTYPE, OP) combo (only 5 of 12 combos are affected); softmax/inter_tile_reduce set it as a static field since they use math.exp unconditionally.
Test plan