diff --git a/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/ContractionSynthesis.cpp b/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/ContractionSynthesis.cpp index 66c638f2ff7e..b2cde43f9d20 100644 --- a/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/ContractionSynthesis.cpp +++ b/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/ContractionSynthesis.cpp @@ -27,6 +27,7 @@ #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/IRMapping.h" #include "mlir/IR/PatternMatch.h" +#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/Debug.h" @@ -1068,6 +1069,88 @@ struct RewriteReducePattern : OpRewritePattern { // retypes, so membership IS the "reachable from a physicalized load" answer. // An op on an unannotated path is simply never in the set, so tt.expand_dims // in softmax (zero tt.spyre_tensor_layout markers) is never retyped. +//===----------------------------------------------------------------------===// +// seedSplatOperands +//===----------------------------------------------------------------------===// + +/// Retype an all-same-value ("splat") constant operand onto `target`'s shape. +/// +/// A Triton kernel that divides by its reduction length emits that length as a +/// whole tensor of one repeated value -- `tl.splat` becomes +/// `arith.constant dense<1.28e2> : tensor<64x128xf32>`. Such a constant has NO +/// operands, and the forward layout analysis walks along operands: it asks +/// "given this physical operand, what type does the result take?". An op with +/// no operands is therefore unreachable by construction, and no analysis rule +/// can claim a physical type for it. Left alone it stays logical while its +/// sibling operand goes physical, and the consuming arithmetic fails its +/// same-type verifier. +/// +/// Rewriting one is sound because every element is identical: the physical form +/// is the same constant at the physical shape, so there is no data to move and +/// no coordinate map to rewrite. `dense<1.28e2> : tensor<64x128xf32>` becomes +/// `dense<1.28e2> : tensor<64x2x64xf32>`. Note this holds even when the split +/// does not divide evenly -- 130 columns become 3 sticks of 64, and the 62 +/// padding lanes get the same value as every real lane, which is exactly what a +/// splat means. +/// +/// Only a splat qualifies. A general `dense<[...]>` constant has per-element +/// data whose stick-tiled placement is a real layout question, so it is left +/// alone for the analysis to reject. +/// +/// In: the op being physicalized, and the shape its physical operands agree on +/// (e.g. `[64, 2, 64]`). +/// Out: true if at least one operand was retyped. Rewrites in place via +/// `rewriter`; a non-splat or already-matching operand is skipped. +static bool seedSplatOperands(Operation *op, ArrayRef target, + PatternRewriter &rewriter) { + bool changed = false; + for (Value o : op->getOperands()) { + auto ty = dyn_cast(o.getType()); + if (!ty || ty.getShape() == target) + continue; + // Only a constant whose value is one repeated element. + auto cst = o.getDefiningOp(); + if (!cst) + continue; + auto dense = dyn_cast(cst.getValue()); + if (!dense || !dense.isSplat()) + continue; + auto newTy = RankedTensorType::get(target, ty.getElementType()); + + // Retype in place when this constant feeds only this op, and mint a fresh + // one otherwise. The distinction is not cosmetic in either direction: a + // constant shared with a consumer that still wants the logical shape must + // NOT be retyped, or that consumer breaks; and retyping the single-use case + // rather than always minting is what keeps the pass from leaving a dead + // logical constant behind for the canonicalizer to sweep. Same + // retype-if-you-can, mint-if-you-must rule rebuildPhysicalInit follows for a + // DPS init. + Value seeded; + if (cst.getResult().hasOneUse()) { + rewriter.modifyOpInPlace(cst, [&]() { + cst.setValueAttr(dense.resizeSplat(newTy)); + cst.getResult().setType(newTy); + }); + seeded = cst.getResult(); + } else { + // Scoped so the caller's insertion point survives: this helper runs + // partway through a pattern that goes on to build more IR. + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(cst); + seeded = arith::ConstantOp::create(rewriter, cst.getLoc(), newTy, + dense.resizeSplat(newTy)) + .getResult(); + } + rewriter.modifyOpInPlace(op, [&]() { + for (OpOperand &use : op->getOpOperands()) + if (use.get() == o) + use.set(seeded); + }); + changed = true; + } + return changed; +} + struct RewriteElementwisePattern : RewritePattern { const PassContext &ctx; RewriteElementwisePattern(MLIRContext *mlirCtx, const PassContext &layoutCtx) @@ -1076,12 +1159,42 @@ struct RewriteElementwisePattern : RewritePattern { LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override { + // Ops that have a rule of their own, excluded by KIND rather than inferred + // from shape. This pattern's shape test -- every tensor operand agrees with + // the others -- is satisfied VACUOUSLY by a single-operand op, so a reshape + // or a broadcast would otherwise be retyped here as though it preserved + // shape, which is precisely what it does not do. Retyping a + // tensor.expand_shape's result to its operand's shape turns the expand into + // a no-op and produces a type Phase 2A never predicted, which + // verifyPhysicalTypeAgreement then reports as the analysis under-claiming. + // + // Structural inference is the wrong instrument here: "elementwise" means + // shape-preserving, and no count of operands or results establishes that. + // Same reasoning that removed isSingleTensorElementwiseOp. + // Excluded by KIND via the shared isShapeChangingOp (Types.h): the analysis + // and the rewrite must agree on which ops are "elementwise". + if (isShapeChangingOp(op)) + return failure(); if (op->getNumResults() != 1) return failure(); auto resTy = dyn_cast(op->getResult(0).getType()); if (!resTy) return failure(); + // A splat constant operand cannot be reached by the operand-driven forward + // analysis (it has no operands), so it arrives here still logical while its + // sibling is physical. Retype it up to the physical shape BEFORE the + // agreement check below, which would otherwise read the disagreement as + // "one side is a guess" and decline. The target comes from an operand the + // analysis actually physicalized, never from another logical operand. + if (auto physIt = llvm::find_if(op->getOperands(), [&](Value o) { + return ctx.physicalValues.contains(o) && + isa(o.getType()); + }); + physIt != op->getOperands().end()) + seedSplatOperands( + op, cast((*physIt).getType()).getShape(), rewriter); + // Every tensor operand must be a RankedTensorType, and they must all // agree on shape -- that agreement is the safety condition: if one // operand is still logical and another already physical, decline rather @@ -1152,6 +1265,87 @@ struct RewriteElementwisePattern : RewritePattern { // Pattern: linalg.transpose (erase and record permutation) //===----------------------------------------------------------------------===// +// Physicalize a linalg.broadcast whose result Phase 2A decided is physical. +// +// A broadcast holds its target shape in two places -- the `outs` operand (a +// tensor.empty) and the `dimensions` attribute naming which OUTPUT positions +// are new -- so retyping its result alone is not enough. Both are rewritten +// here, and both follow mechanically from the marker: +// +// dimensions each added LOGICAL axis contributes one entry per physical dim +// sourced from it: [1] -> [1, 2] when logical axis 1 splits. +// outs applyCoordMap over the logical result shape. +// +// The gate is the analysis, not the operand: unlike RewriteTransposePattern, +// a broadcast's own input is typically logical (its producer is a reduce whose +// result a reshape leaves logical), so gating on physicalValues.contains(input) +// would never fire. Phase 2A already decided this result is physical, and only +// does so when every CARRIED axis is unsplit -- see BroadcastPropagation. +struct RewriteBroadcastPattern : OpRewritePattern { + const PassContext &ctx; + RewriteBroadcastPattern(MLIRContext *mlirCtx, const PassContext &layoutCtx) + : OpRewritePattern(mlirCtx, /*benefit=*/1), ctx(layoutCtx) {} + + LogicalResult matchAndRewrite(linalg::BroadcastOp bc, + PatternRewriter &rewriter) const override { + if (!ctx.physicalTypeAnalysis) + return failure(); + auto it = ctx.physicalTypeAnalysis->find(bc.getResult()[0]); + if (it == ctx.physicalTypeAnalysis->end()) + return failure(); + auto physTy = dyn_cast(it->second.type); + if (!physTy) + return failure(); + + // Idempotence: once the result carries the physical type, stop matching, so + // the greedy driver's re-enqueue cannot re-fire this. + auto resTy = cast(bc.getResult()[0].getType()); + if (resTy.getShape() == physTy.getShape()) + return failure(); + + auto marker = it->second.marker; + if (!marker) + return failure(); + auto physSrc = marker.getPhysSrc(); + + // Renumber `dimensions` from logical output positions to physical ones. An + // added logical axis that splits contributes every physical dim sourced + // from it, which is what keeps linalg.broadcast's rank arithmetic + // (input_rank + |dimensions| == init_rank) true by construction. + llvm::SmallDenseSet addedLogical(bc.getDimensions().begin(), + bc.getDimensions().end()); + llvm::SmallVector newDims; + for (unsigned p = 0; p < physSrc.size(); ++p) + if (addedLogical.contains(physSrc[p])) + newDims.push_back(p); + + Location loc = bc.getLoc(); + Value newInit = rebuildPhysicalInit(rewriter, loc, bc.getInit(), physTy); + if (!newInit) + return failure(); + + auto newBc = linalg::BroadcastOp::create(rewriter, loc, bc.getInput(), + newInit, newDims); + Value newResult = newBc.getResult()[0]; + + // This result IS physical now, under the layout Phase 2A paired it with. + // Record it so a consumer sees a physical value, and carry the analysis's + // decision onto the value that replaces the one it was made about. + ctx.physicalValues[newResult] = PhysicalValueInfo{marker, {}}; + ctx.physicalTypes.carryForward(bc.getResult()[0], newResult); + + rewriter.replaceOp(bc, newResult); + + // REQUIRED, as in RewriteElementwisePattern: the result type its users read + // just changed, and the greedy driver re-enqueues the modified op but not + // its users. A consuming arith.subf whose other operand is already physical + // must be revisited so its own match condition sees agreeing shapes. + for (Operation *user : llvm::make_early_inc_range(newResult.getUsers())) + rewriter.modifyOpInPlace(user, [] {}); + return success(); + } +}; + // The transpose permutation is recorded in ctx.physicalValues (against // `input`, which must already be an entry -- see PhysicalValueInfo) because // this erase happens before dispatchSource can see the transpose: erasing it @@ -1337,6 +1531,7 @@ void populateContractionPatterns(RewritePatternSet &patterns, patterns.add(mlirCtx, ctx); patterns.add(mlirCtx, ctx); + patterns.add(mlirCtx, ctx); patterns.add(mlirCtx, ctx); patterns.add(mlirCtx, ctx); } diff --git a/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/PhysicalTypeAnalysis.cpp b/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/PhysicalTypeAnalysis.cpp index 46a9e2da6a82..30a01e77b3a4 100644 --- a/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/PhysicalTypeAnalysis.cpp +++ b/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/PhysicalTypeAnalysis.cpp @@ -22,6 +22,7 @@ #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/BuiltinTypes.h" +#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/ScopeExit.h" #include "llvm/Support/Debug.h" @@ -51,19 +52,25 @@ struct ElementwisePropagation : PhysicalPropagationPattern { if (op->getNumResults() != 1 || !isa(op->getResult(0).getType())) return false; - ArrayRef commonShape; + // Ops with a rule of their own, excluded by KIND. The shape test below is + // satisfied vacuously by a single-operand op, so a reshape or a broadcast + // would otherwise be claimed here as though it preserved shape. + if (isShapeChangingOp(op)) + return false; + + // Operands are NOT compared to each other. Mid-analysis they routinely + // disagree: Phase 1 physicalizes loads and stops, so an op with one + // load-fed operand and one not is guaranteed to see a mismatch. That state + // is what this pass exists to resolve, not evidence the op is unknown -- + // and `propagate` never reads a sibling's shape, so the comparison gated + // nothing it needed. Requiring agreement here made an arith.subf whose + // operands straddle the split an untaught op, so its result was never + // predicted and verifyPhysicalTypeAgreement reported the analysis as + // under-claiming once Phase 2B retyped it. bool sawTensorOperand = false; - for (Value o : op->getOperands()) { - auto t = dyn_cast(o.getType()); - if (!t) - continue; - if (!sawTensorOperand) { - commonShape = t.getShape(); + for (Value o : op->getOperands()) + if (isa(o.getType())) sawTensorOperand = true; - } else if (t.getShape() != commonShape) { - return false; - } - } return sawTensorOperand; } @@ -267,25 +274,63 @@ struct ReshapePropagation : PhysicalPropagationPattern { propagate(Operation *op, Value result, Value src, const PhysicalTypeInfo &srcInfo, const LayoutRequirement *want) const override { - return failure(); + // The general case declines, for the reason above: a reassociation that + // fuses two real axes leaves no physical dim for one of them. + // + // One narrow case does carry, and it is the one that matters here. When the + // value being reshaped is ALREADY at its physical shape -- srcInfo.type + // equals the operand's own type -- then physical and logical coincide on + // this value, nothing about it is stick-split, and the reshape's own result + // type is the physical type. The marker rides along unchanged. + // + // The test is on the VALUE, not on the marker. The marker a reduce result + // carries is the STORE's, which is typically split (phys_op = id/floordiv/ + // mod); that says nothing about whether this rank-1 value is split. Testing + // the marker's ops would reject exactly the case this exists to allow. + // + // This is what carries a reduce's result past the expand_shape/collapse_shape + // pair LowerComputeOps emits between a reduce and a broadcast (rules A3 and + // A4 lower tt.expand_dims and tt.broadcast independently, so A3 expands + // 1 -> 1x1 and A4 immediately collapses it back). Without it the forward + // walk stops there and the broadcast is never asked. + auto marker = srcInfo.marker; + if (!marker) + return failure(); + if (!srcInfo.transposePerm.empty()) + return failure(); + + auto srcTy = dyn_cast(srcInfo.type); + auto operandTy = dyn_cast(src.getType()); + auto resTy = dyn_cast(result.getType()); + if (!srcTy || !operandTy || !resTy) + return failure(); + // Physical == logical for this value, or there is a split to rewrite and + // no coordinate map can express it across a reassociation. + if (srcTy != operandTy) + return failure(); + return PhysicalTypeInfo{resTy, marker, {}}; } }; -/// linalg.broadcast: no physical result, for a different reason than the reshape -/// family above. +/// linalg.broadcast: recompute the target shape against the physical rank. /// -/// A broadcast's result shape is not derived from its operand at all: it is a -/// target shape fixed when the op was built, which LowerComputeOps builds from +/// A broadcast's result shape is not derived from its operand: it is a target +/// shape fixed when the op was built, which LowerComputeOps builds from /// tt.broadcast against the operand's LOGICAL rank. Given a physical operand -/// that target shape is simply stale -- it describes a tensor of the wrong rank, -/// and a consuming elementwise op then fails its same-type constraint. +/// that shape is stale -- it describes a tensor of the wrong rank, so a +/// consuming elementwise op fails its same-type constraint. That is the failure +/// a reduce -> broadcast -> elementwise chain hits, the shape softmax and +/// layernorm are written in. /// -/// Unlike a reshape's coordinate map, this is recoverable in principle: -/// recomputing the target shape and `dimensions` against the physical rank would -/// let the op work on physical operands directly, which is what a reduce -> -/// broadcast -> elementwise chain needs -- the shape softmax and layernorm are -/// written in. Until that exists, declining is the safe answer and such an -/// operand is rejected downstream rather than guessed at. +/// This rule recomputes the shape by pushing the result's logical extents +/// through the marker's coordinate map, and the paired rewrite pattern +/// renumbers `dimensions` to name the new physical positions. +/// +/// It handles the case where every axis the broadcast CARRIES is unsplit; an +/// added axis may split freely, since it is new in the output and so changes +/// only `outs` and `dimensions`, never the input. A split carried axis needs the +/// INPUT retyped too, which a rule producing only a result type cannot express, +/// so that case declines and is repaired at the consuming op instead. struct BroadcastPropagation : PhysicalPropagationPattern { bool match(Operation *op) const override { return isa(op); @@ -295,7 +340,106 @@ struct BroadcastPropagation : PhysicalPropagationPattern { propagate(Operation *op, Value result, Value src, const PhysicalTypeInfo &srcInfo, const LayoutRequirement *want) const override { - return failure(); + auto bc = cast(op); + + // The rule reads `srcInfo` as the INPUT's layout, so it must actually be + // the input. findPhysicalTensorOperand returns the first physical tensor + // operand in operand order, and a broadcast's operands are (input, init) -- + // so a physicalized init would otherwise be read as though it were the + // input, taking its marker and its transposePerm from the wrong chain. + // BroadcastRequirement::induce deliberately hands the init the requirement + // whole, which is exactly what makes the init a candidate here. + // ReducePropagation guards the same way against its own inits. + if (src != bc.getInput()) + return failure(); + + // The emission rebuilds the broadcast's init, and rebuildPhysicalInit only + // handles a tensor.empty or a fill over one. Gate on the same predicate the + // emission uses so this decision and that one cannot disagree: claiming a + // physical result the rewrite then declines to produce would leave the value + // logical while its consumers were predicted physical. + if (!canRebuildPhysicalInit(bc.getInit())) + return failure(); + + auto marker = srcInfo.marker; + if (!marker) + return failure(); + // An erased transpose reorders the marker's dims relative to the tile, so + // "the physical dims in the marker's order" is not the order the op sees. + if (!srcInfo.transposePerm.empty()) + return failure(); + + auto resTy = dyn_cast(result.getType()); + if (!resTy) + return failure(); + + auto physSrc = marker.getPhysSrc(); + auto physOp = marker.getPhysOp(); + auto physArg = marker.getPhysArg(); + + // The marker describes the RESULT's layout, so its logical rank must be the + // result's rank -- `dimensions` indexes logical output positions. + unsigned logicalRank = 0; + for (int64_t d : physSrc) + logicalRank = std::max(logicalRank, (unsigned)(d + 1)); + if (logicalRank != (unsigned)resTy.getRank()) + return failure(); + + // A maximum is not a set cover: `phys_src = [0, 2]` on a rank-3 result has + // logicalRank 3 and passes the check above while logical dim 1 is + // unrepresented. applyCoordMap would then silently return a rank-2 shape, + // dropping dim 1's extent. Require every logical dim to be named. + { + llvm::SmallVector covered(logicalRank, false); + for (int64_t d : physSrc) + covered[d] = true; + if (llvm::is_contained(covered, false)) + return failure(); + } + + // Which logical output axes does the broadcast ADD, and which does it CARRY + // from its input? `dimensions` names the added ones. + llvm::SmallDenseSet added; + for (int64_t d : bc.getDimensions()) { + if (d < 0 || d >= (int64_t)logicalRank) + return failure(); + added.insert(d); + } + + // A CARRIED axis that is stick-split would need the INPUT retyped too -- + // linalg.broadcast matches its input against the non-broadcast init dims + // positionally, so a split carried axis demands a higher-rank input. A + // propagation rule produces a type for the RESULT only, so decline. The + // consuming elementwise op is where that case has to be repaired. + for (unsigned p = 0; p < physSrc.size(); ++p) + if (!added.contains(physSrc[p]) && + static_cast(physOp[p]) != CoordOp::Identity) + return failure(); + + // The emission renumbers `dimensions` by scanning physical dims in + // ascending order, which only names the input's axes correctly when the + // CARRIED axes appear in the same relative order physically as logically. + // A permuting marker (e.g. phys_src = [2, 0, 1], every op Identity) passes + // every check above -- the marker verifier does not require monotonicity -- + // and would transpose the carried data while still satisfying the broadcast + // verifier. Nothing downstream can catch that, so reject it here. + int64_t prevCarried = -1; + for (unsigned p = 0; p < physSrc.size(); ++p) { + if (added.contains(physSrc[p])) + continue; + if (physSrc[p] <= prevCarried) + return failure(); + prevCarried = physSrc[p]; + } + + // Every added axis may split freely: it is new in the output, so splitting + // it changes only `outs` and `dimensions`, never the input. + llvm::SmallVector physShape; + if (!applyCoordMap(resTy.getShape(), physSrc, physOp, physArg, physShape)) + return failure(); + + return PhysicalTypeInfo{ + RankedTensorType::get(physShape, resTy.getElementType()), marker, {}}; } }; diff --git a/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/RequirementAnalysis.cpp b/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/RequirementAnalysis.cpp index ef904cd89688..363ce663e0e2 100644 --- a/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/RequirementAnalysis.cpp +++ b/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/RequirementAnalysis.cpp @@ -81,19 +81,21 @@ struct ElementwiseRequirement : RequirementBackwardPattern { if (op->getNumResults() != 1 || !isa(op->getResult(0).getType())) return false; - ArrayRef commonShape; + // Deliberately does NOT compare operand shapes to each other. Mid-analysis + // the IR is half-retyped -- Phase 1 physicalizes loads and stops -- so a + // sibling operand being a different shape is the state this pass exists to + // resolve, not evidence that the op is unknown. An elementwise op's + // requirement is the same on every operand regardless (induce is `return + // req`), so the comparison gates nothing it needs. + // + // Shape-changing ops are excluded by KIND; isShapeChangingOp (Types.h) is + // the one list all three elementwise predicates share, and states why. + if (isShapeChangingOp(op)) + return false; bool sawTensorOperand = false; - for (Value o : op->getOperands()) { - auto t = dyn_cast(o.getType()); - if (!t) - continue; - if (!sawTensorOperand) { - commonShape = t.getShape(); + for (Value o : op->getOperands()) + if (isa(o.getType())) sawTensorOperand = true; - } else if (t.getShape() != commonShape) { - return false; - } - } return sawTensorOperand; } @@ -167,13 +169,29 @@ struct MatmulRequirement : RequirementBackwardPattern { } }; -/// tensor.expand_shape / collapse_shape / reshape: the requirement terminates. -/// The physical dim count changes across a reassociation map, so a -/// per-physical-dim requirement on the result says nothing about the operand -- -/// the mirror of why ReshapePropagation declines forward. +/// tensor.expand_shape / collapse_shape / reshape: the requirement crosses only +/// a reshape that adds or removes SIZE-1 dims, and terminates otherwise. +/// +/// The general case cannot cross, for the reason the forward rule states: a +/// requirement is indexed per PHYSICAL dim, and a reassociation that fuses two +/// real axes leaves no dim for one of them to map to. Collapsing physical +/// [1, 64, 64] under phys_op = [floor, id, mod] via [[0, 1], [2]] fuses the +/// stick index with the row, and no coordinate map describes the result. /// -/// Registered ahead of the structural elementwise rule; see -/// populateRequirementBackwardPatterns. +/// A reshape that only inserts or drops size-1 dims is different: it touches no +/// real axis, and a size-1 dim carries no coordinate information. So the +/// requirement crosses with phys_src renumbered. That narrow class is what +/// LowerComputeOps emits between a reduce and a broadcast -- rules A3 and A4 +/// lower tt.expand_dims and tt.broadcast independently, so A3 expands 1 -> 1x1 +/// and A4 immediately collapses it back. +/// +/// The reassociation groups always index the HIGHER-rank side: the result for +/// expand_shape, the operand for collapse_shape. So the two directions are not +/// symmetric and are handled separately below. +/// +/// The safety condition is on the marker, not on extents: a stick index can +/// itself have extent 1 while still carrying coordinate meaning, so a group of +/// size > 1 must contain no floordiv or mod dim. struct ReshapeRequirement : RequirementBackwardPattern { bool match(Operation *op) const override { return isa induce(Operation *op, Value result, Value operand, const LayoutRequirement &req) const override { - return failure(); + // tensor.reshape takes a runtime shape operand, so there is no static + // reassociation to reason about. + auto expand = dyn_cast(op); + auto collapse = dyn_cast(op); + if (!expand && !collapse) + return failure(); + + auto inTy = dyn_cast(operand.getType()); + auto resTy = dyn_cast(result.getType()); + if (!inTy || !resTy) + return failure(); + + // The requirement is stated over the RESULT's logical dims. + unsigned logicalRank = 0; + for (int64_t d : req.physSrc) + logicalRank = std::max(logicalRank, (unsigned)(d + 1)); + if (logicalRank != (unsigned)resTy.getRank()) + return failure(); + + auto reassoc = expand ? expand.getReassociationIndices() + : collapse.getReassociationIndices(); + // Groups index the higher-rank side. Verify that side's rank matches, so a + // malformed pairing is refused rather than mis-indexed. + llvm::ArrayRef bigShape = + expand ? resTy.getShape() : inTy.getShape(); + if (reassoc.size() != (unsigned)(expand ? inTy.getRank() + : resTy.getRank())) + return failure(); + + // remap[result dim] -> operand dim, or -1 when the dim disappears. + llvm::SmallVector remap(logicalRank, -1); + for (unsigned g = 0; g < reassoc.size(); ++g) { + // Within a group, at most one dim of the higher-rank side may be + // non-unit; it is the one the lower-rank side's dim corresponds to. + int64_t nonUnit = -1; + for (int64_t d : reassoc[g]) { + if (d < 0 || d >= (int64_t)bigShape.size()) + return failure(); + if (bigShape[d] != 1) { + if (nonUnit >= 0) + return failure(); // two real axes in one group: cannot cross + nonUnit = d; + } + } + // A floordiv/mod dim inside a multi-dim group would be fused or split. + if (reassoc[g].size() > 1) + for (int64_t d : reassoc[g]) { + int64_t reqDim = expand ? d : (int64_t)g; + for (unsigned p = 0; p < req.physSrc.size(); ++p) + if (req.physSrc[p] == reqDim && + static_cast(req.physOp[p]) != CoordOp::Identity) + return failure(); + } + + if (expand) { + // Groups index the RESULT. Group g corresponds to operand dim g, and + // the surviving result dim within it is the non-unit one (or the first). + int64_t keep = nonUnit >= 0 ? nonUnit : reassoc[g].front(); + remap[keep] = (int64_t)g; + } else { + // Groups index the OPERAND. Result dim g corresponds to the group's + // non-unit operand dim (or its first). + int64_t keep = nonUnit >= 0 ? nonUnit : reassoc[g].front(); + if ((unsigned)g >= logicalRank) + return failure(); + remap[g] = keep; + } + } + + // Build the operand-side requirement. The two directions differ in whether + // entries are dropped or added: + // + // expand_shape the operand has FEWER dims, so a requirement entry whose + // logical dim is a newly inserted size-1 dim is dropped. + // collapse_shape the operand has MORE dims, so an entry must be ADDED for + // each size-1 operand dim the collapse removed. Such a dim + // is Identity with extent 1 -- it carries no coordinate + // information, which is exactly why crossing is sound. + // + // Either way the result is indexed per operand dim, so it is assembled by + // walking the operand's dims rather than the requirement's. + LayoutRequirement out; + out.marker = req.marker; + llvm::SmallVector reqDimForOperandDim(inTy.getRank(), -1); + for (unsigned d = 0; d < logicalRank; ++d) + if (remap[d] >= 0 && remap[d] < inTy.getRank()) + reqDimForOperandDim[remap[d]] = (int64_t)d; + + for (int64_t od = 0; od < inTy.getRank(); ++od) { + int64_t reqDim = reqDimForOperandDim[od]; + if (reqDim < 0) { + // A size-1 operand dim the reshape removed. Only sound because it is + // size 1; refuse anything else rather than invent a coordinate for it. + if (inTy.getDimSize(od) != 1) + return failure(); + out.physSrc.push_back(od); + out.physOp.push_back((int64_t)CoordOp::Identity); + out.physArg.push_back(0); + out.physExtents.push_back(1); + continue; + } + // Carry every requirement entry naming this logical dim. + bool found = false; + for (unsigned p = 0; p < req.physSrc.size(); ++p) { + if (req.physSrc[p] != reqDim) + continue; + out.physSrc.push_back(od); + out.physOp.push_back(req.physOp[p]); + out.physArg.push_back(req.physArg[p]); + out.physExtents.push_back(req.physExtents[p]); + found = true; + } + if (!found) + return failure(); + } + if ((int64_t)out.physSrc.size() != inTy.getRank()) + return failure(); + return out; } }; -/// linalg.broadcast: the requirement terminates. The result has dims the operand -/// does not, so no per-dim requirement on it constrains the operand. This is -/// also what keeps softmax's and layernorm's reduce results out of the map: they -/// reach their store only through a broadcast. +/// linalg.broadcast: the requirement PROJECTS onto the carried axes. +/// +/// The result has dims the operand does not -- that is what a broadcast is -- +/// so the requirement cannot cross unchanged. But it does project: keep the +/// physical dims whose phys_src names a logical axis the operand CARRIES, drop +/// those naming an axis the broadcast ADDS. What survives is a requirement of +/// exactly the operand's logical rank. +/// +/// For softmax that projection is trivial -- the surviving dim is Identity, so +/// it asks nothing a rank-1 logical value does not already satisfy. It is still +/// load-bearing: it is what carries the requirement past this op to the reduce, +/// which needs `want` non-null to take the Physical space. Once the reduce's +/// result is in the forward map, the reshapes and then this broadcast acquire a +/// physical operand, which is what lets BroadcastPropagation be asked at all. +/// The real shape work happens there, from the store's marker. +/// +/// Declines when a CARRIED axis is split across two physical dims. Projecting +/// would then demand an operand of higher rank than it has -- linalg.broadcast +/// matches its input against the non-broadcast init dims positionally -- so +/// there is no requirement the operand could satisfy. That is the case the +/// consuming elementwise op has to repair instead. struct BroadcastRequirement : RequirementBackwardPattern { bool match(Operation *op) const override { return isa(op); @@ -199,7 +351,61 @@ struct BroadcastRequirement : RequirementBackwardPattern { llvm::FailureOr induce(Operation *op, Value result, Value operand, const LayoutRequirement &req) const override { - return failure(); + auto bc = cast(op); + // The init carries the RESULT's shape, so it gets the requirement whole -- + // the same split TransposeRequirement makes for its own init. + if (operand == bc.getInit()) + return req; + if (operand != bc.getInput()) + return failure(); + + // `dimensions` names the added LOGICAL output axes, so the requirement's + // logical rank must be the result's rank for the two to be comparable. + unsigned logicalRank = 0; + for (int64_t d : req.physSrc) + logicalRank = std::max(logicalRank, (unsigned)(d + 1)); + auto resTy = dyn_cast(result.getType()); + if (!resTy || logicalRank != (unsigned)resTy.getRank()) + return failure(); + + llvm::SmallDenseSet added; + for (int64_t d : bc.getDimensions()) { + if (d < 0 || d >= (int64_t)logicalRank) + return failure(); + added.insert(d); + } + + // A carried axis split across two physical dims cannot be projected: the + // operand would need a rank it does not have. + for (unsigned p = 0; p < req.physSrc.size(); ++p) + if (!added.contains(req.physSrc[p]) && + static_cast(req.physOp[p]) != CoordOp::Identity) + return failure(); + + // Renumber the carried logical axes down, since the added ones are gone. + llvm::SmallVector logicalRemap(logicalRank, -1); + int64_t next = 0; + for (unsigned d = 0; d < logicalRank; ++d) + if (!added.contains((int64_t)d)) + logicalRemap[d] = next++; + + LayoutRequirement out; + out.marker = req.marker; + for (unsigned p = 0; p < req.physSrc.size(); ++p) { + int64_t mapped = logicalRemap[req.physSrc[p]]; + if (mapped < 0) + continue; // a dim of an added axis: dropped + out.physSrc.push_back(mapped); + out.physOp.push_back(req.physOp[p]); + out.physArg.push_back(req.physArg[p]); + out.physExtents.push_back(req.physExtents[p]); + } + // The projection must have the operand's own rank, or it describes a + // different value than the one it is about. + auto inTy = dyn_cast(operand.getType()); + if (!inTy || (int64_t)out.physSrc.size() != inTy.getRank()) + return failure(); + return out; } }; diff --git a/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/Types.h b/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/Types.h index b76a11cd6f9c..61aa313700f2 100644 --- a/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/Types.h +++ b/third_party/spyre/lib/Dialect/KTDP/Transforms/RewriteDescriptorLayout/Types.h @@ -3,6 +3,8 @@ #include "RewriteDescriptorLayout/PermutationUtils.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Value.h" @@ -13,6 +15,35 @@ namespace mlir::triton::ktdp { +/// True when `op` changes tensor shape and therefore has (or needs) a layout +/// rule of its own, rather than being treated as shape-preserving. +/// +/// The three "elementwise" predicates -- ElementwiseRequirement::match, +/// ElementwisePropagation::match and RewriteElementwisePattern -- each ask this +/// question, and they must agree: if the analysis treats an op as elementwise +/// and the rewrite does not (or vice versa), one claims a value the other never +/// predicted and verifyPhysicalTypeAgreement reports the disagreement. +/// +/// Excluding by KIND rather than by comparing operand shapes is deliberate. +/// Phase 1 physicalizes loads and stops, so mid-pass the IR is half-retyped by +/// design; an operand disagreeing with its sibling is the state this pass exists +/// to resolve, not evidence the op is unknown. And a shape comparison is +/// satisfied VACUOUSLY by a single-operand op, so a reshape or a broadcast would +/// pass it while being exactly what it is meant to catch. +/// +/// Add an op here whenever a shape-changing op is added to the pipeline. Ops +/// with a rule of their own are belt-and-braces (a named rule is asked first); +/// the load-bearing entries are the ones with NO rule, such as tensor::PadOp and +/// tensor::ConcatOp, which would otherwise be crossed as shape-preserving and +/// silently retyped. +inline bool isShapeChangingOp(Operation *op) { + return isa(op); +} + + /// What is known about one physical value (a physicalized ktdp.load result, /// or a value retyped to physical by Phase 2's elementwise pattern). struct PhysicalValueInfo { diff --git a/third_party/spyre/test/Conversion/rewrite-descriptor-layout-broadcast-splat.mlir b/third_party/spyre/test/Conversion/rewrite-descriptor-layout-broadcast-splat.mlir new file mode 100644 index 000000000000..0df8962d9954 --- /dev/null +++ b/third_party/spyre/test/Conversion/rewrite-descriptor-layout-broadcast-splat.mlir @@ -0,0 +1,175 @@ +// RUN: spyre-triton-opt %s --lower-descriptor-memory --lower-scalar-load --lower-compute-ops --rewrite-descriptor-layout -split-input-file | FileCheck %s + +// Seeding a splat constant so layernorm can carry a layout. +// Design: docs/designs/rewrite-broadcast-pattern.md. +// +// WHAT THESE TESTS ASK OF THE PASS: in the layernorm shape -- reduce, broadcast, +// then divide by the element count N -- the divisor must reach the physical shape +// along with the broadcast, so the divide's operands agree. +// +// The remaining cases bound that: which constants qualify, and what must happen +// when one is shared with a consumer that still wants the logical shape. +// +// WHY THE DIVISOR NEEDED SEPARATE WORK: `tl.splat`/a Python scalar becomes an +// `arith.constant dense<...>` -- a value with ZERO operands. The forward layout +// analysis walks along operands, asking "given this physical operand, what type +// does the result take?", so an op with no operands is unreachable by +// construction and no analysis rule can claim a type for it. Left logical while +// its sibling went physical, it failed `arith.divf`'s same-type verifier. +// +// The fix is a SEED, not a propagation rule: when an elementwise op has one +// physical operand and one splat-constant operand, the constant is rebuilt at +// the physical shape. Sound because every element is identical -- there is no +// data to move and no coordinate map to rewrite. + +// CHECK-LABEL: func @layernorm_splat_divisor +// The divisor carries the same repeated value at the physical shape. It feeds +// only this divide, so it is retyped IN PLACE -- no canonicalizer needed to +// clear a leftover logical copy, and CHECK-NOT below proves none is left. +// CHECK: %[[N:.*]] = arith.constant dense<1.280000e+02> : tensor<64x2x64xf32> +// CHECK-NOT: arith.constant dense<1.280000e+02> : tensor<64x128xf32> +// CHECK: ktdp.load {{.*}} -> tensor<64x2x64xf32> +// CHECK: linalg.reduce +// CHECK-SAME: dimensions = [1, 2] +// CHECK: linalg.broadcast +// CHECK-SAME: outs(%{{.*}} : tensor<64x2x64xf32>) +// CHECK-SAME: dimensions = [1, 2] +// Both arithmetic ops now have agreeing operands. +// CHECK: arith.divf %{{.*}}, %[[N]] : tensor<64x2x64xf32> +// CHECK: arith.subf %{{.*}}, %{{.*}} : tensor<64x2x64xf32> +// CHECK: ktdp.store %{{.*}}, %{{.*}} : tensor<64x2x64xf32> +module { +tt.func @layernorm_splat_divisor(%a_ptr: !tt.ptr, %o_ptr: !tt.ptr) { + %c0 = arith.constant 0 : i32 + %c64 = arith.constant 64 : i32 + %c128 = arith.constant 128 : i32 + %s128 = arith.constant 128 : i64 + %s1 = arith.constant 1 : i64 + // The element count along the reduced axis, as a splat: 128 columns. + %n = arith.constant dense<1.280000e+02> : tensor<64x128xf32> + %ad = tt.make_tensor_descriptor %a_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %ad {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + %a = tt.descriptor_load %ad[%c0, %c0] : !tt.tensordesc<64x128xf32> -> tensor<64x128xf32> + // Sum along the columns, then fan the row-sums back out: the mean's numerator. + %s = "tt.reduce"(%a) ({ ^bb0(%x: f32, %y: f32): %sm = arith.addf %x, %y : f32 + tt.reduce.return %sm : f32 }) {axis = 1 : i32} : (tensor<64x128xf32>) -> tensor<64xf32> + %se = tt.expand_dims %s {axis = 1 : i32} : tensor<64xf32> -> tensor<64x1xf32> + %sb = tt.broadcast %se : tensor<64x1xf32> -> tensor<64x128xf32> + %mean = arith.divf %sb, %n : tensor<64x128xf32> + %d = arith.subf %a, %mean : tensor<64x128xf32> + %od = tt.make_tensor_descriptor %o_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %od {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + tt.descriptor_store %od[%c0, %c0], %d : !tt.tensordesc<64x128xf32>, tensor<64x128xf32> + tt.return +} +} + +// ----- + +// Non-divisible: 130 columns become 3 sticks of 64, i.e. 192 lanes for 130 real +// ones. Every lane including the 62 padding ones gets the same value, which is +// what "splat" means -- so unlike the store bridge, seeding a splat is correct +// on a non-divisible split. This also exercises resizeSplat rather than +// reshape: the element count CHANGES (130 -> 192), which reshape forbids. + +// CHECK-LABEL: func @splat_non_divisible +// CHECK: arith.constant dense<1.300000e+02> : tensor<64x3x64xf32> +// CHECK: arith.divf %{{.*}}, %{{.*}} : tensor<64x3x64xf32> +module { +tt.func @splat_non_divisible(%a_ptr: !tt.ptr, %o_ptr: !tt.ptr) { + %c0 = arith.constant 0 : i32 + %c64 = arith.constant 64 : i32 + %c130 = arith.constant 130 : i32 + %s130 = arith.constant 130 : i64 + %s1 = arith.constant 1 : i64 + %n = arith.constant dense<1.300000e+02> : tensor<64x130xf32> + %ad = tt.make_tensor_descriptor %a_ptr, [%c64, %c130], [%s130, %s1] : !tt.ptr, !tt.tensordesc<64x130xf32> + tt.spyre_tensor_layout %ad {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x130xf32> + %a = tt.descriptor_load %ad[%c0, %c0] : !tt.tensordesc<64x130xf32> -> tensor<64x130xf32> + %d = arith.divf %a, %n : tensor<64x130xf32> + %od = tt.make_tensor_descriptor %o_ptr, [%c64, %c130], [%s130, %s1] : !tt.ptr, !tt.tensordesc<64x130xf32> + tt.spyre_tensor_layout %od {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x130xf32> + tt.descriptor_store %od[%c0, %c0], %d : !tt.tensordesc<64x130xf32>, tensor<64x130xf32> + tt.return +} +} + +// ----- + +// One splat feeding TWO consumers: an annotated path that goes physical and an +// unannotated one that must stay logical. The seed mints a FRESH constant +// rather than retyping in place, so the logical use is untouched. Retyping in +// place would corrupt it -- which is why sharing is checked. + +// CHECK-LABEL: func @shared_splat_two_consumers +// The physical consumer gets a new constant at the physical shape... +// CHECK: arith.constant dense<1.280000e+02> : tensor<64x2x64xf32> +// ...and the original survives for the logical consumer. +// CHECK: arith.constant dense<1.280000e+02> : tensor<64x128xf32> +// CHECK: arith.divf %{{.*}}, %{{.*}} : tensor<64x2x64xf32> +// CHECK: arith.mulf %{{.*}}, %{{.*}} : tensor<64x128xf32> +module { +tt.func @shared_splat_two_consumers(%a_ptr: !tt.ptr, %o_ptr: !tt.ptr, %p_ptr: !tt.ptr) { + %c0 = arith.constant 0 : i32 + %c64 = arith.constant 64 : i32 + %c128 = arith.constant 128 : i32 + %s128 = arith.constant 128 : i64 + %s1 = arith.constant 1 : i64 + %n = arith.constant dense<1.280000e+02> : tensor<64x128xf32> + %ad = tt.make_tensor_descriptor %a_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %ad {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + %a = tt.descriptor_load %ad[%c0, %c0] : !tt.tensordesc<64x128xf32> -> tensor<64x128xf32> + %d = arith.divf %a, %n : tensor<64x128xf32> + %od = tt.make_tensor_descriptor %o_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %od {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + tt.descriptor_store %od[%c0, %c0], %d : !tt.tensordesc<64x128xf32>, tensor<64x128xf32> + %pd = tt.make_tensor_descriptor %p_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + %q = tt.descriptor_load %pd[%c0, %c0] : !tt.tensordesc<64x128xf32> -> tensor<64x128xf32> + %r = arith.mulf %q, %n : tensor<64x128xf32> + tt.descriptor_store %pd[%c0, %c0], %r : !tt.tensordesc<64x128xf32>, tensor<64x128xf32> + tt.return +} +} + +// ----- + +// The layernorm shape under the STICK-OUTERMOST marker, phys_src=[1,0,1], which +// is the ordering most fixtures in this suite use. The split axis owns physical +// dims 0 and 2 here rather than 1 and 2, so the seeded constant has to land at +// <2x64x64> and not <64x2x64>. +// +// A splat is indifferent to that order -- every element is the same value, so +// resizeSplat only needs the total shape -- which is exactly what makes this +// worth pinning: the seed must take its shape from the physical operand beside +// it, never reconstruct one from the marker. + +// CHECK-LABEL: func @layernorm_splat_stick_outermost +// CHECK: %[[N:.*]] = arith.constant dense<1.280000e+02> : tensor<2x64x64xf32> +// CHECK-NOT: arith.constant dense<1.280000e+02> : tensor<64x128xf32> +// CHECK: linalg.reduce +// CHECK-SAME: dimensions = [0, 2] +// CHECK: arith.divf %{{.*}}, %[[N]] : tensor<2x64x64xf32> +// CHECK: ktdp.store %{{.*}}, %{{.*}} : tensor<2x64x64xf32> +module { +tt.func @layernorm_splat_stick_outermost(%a_ptr: !tt.ptr, %o_ptr: !tt.ptr) { + %c0 = arith.constant 0 : i32 + %c64 = arith.constant 64 : i32 + %c128 = arith.constant 128 : i32 + %s128 = arith.constant 128 : i64 + %s1 = arith.constant 1 : i64 + %n = arith.constant dense<1.280000e+02> : tensor<64x128xf32> + %ad = tt.make_tensor_descriptor %a_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %ad {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + %a = tt.descriptor_load %ad[%c0, %c0] : !tt.tensordesc<64x128xf32> -> tensor<64x128xf32> + %s = "tt.reduce"(%a) ({ ^bb0(%x: f32, %y: f32): %sm = arith.addf %x, %y : f32 + tt.reduce.return %sm : f32 }) {axis = 1 : i32} : (tensor<64x128xf32>) -> tensor<64xf32> + %se = tt.expand_dims %s {axis = 1 : i32} : tensor<64xf32> -> tensor<64x1xf32> + %sb = tt.broadcast %se : tensor<64x1xf32> -> tensor<64x128xf32> + %mean = arith.divf %sb, %n : tensor<64x128xf32> + %d = arith.subf %a, %mean : tensor<64x128xf32> + %od = tt.make_tensor_descriptor %o_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %od {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + tt.descriptor_store %od[%c0, %c0], %d : !tt.tensordesc<64x128xf32>, tensor<64x128xf32> + tt.return +} +} diff --git a/third_party/spyre/test/Conversion/rewrite-descriptor-layout-broadcast.mlir b/third_party/spyre/test/Conversion/rewrite-descriptor-layout-broadcast.mlir new file mode 100644 index 000000000000..d50a92765e73 --- /dev/null +++ b/third_party/spyre/test/Conversion/rewrite-descriptor-layout-broadcast.mlir @@ -0,0 +1,170 @@ +// RUN: spyre-triton-opt %s --lower-descriptor-memory --lower-scalar-load --lower-compute-ops --rewrite-descriptor-layout -split-input-file | FileCheck %s + +// Physicalizing `linalg.broadcast` (issue #91). +// +// A `tt.spyre_tensor_layout` annotation asks for the Spyre "stick" layout: one +// logical axis splits into a stick index (floordiv) and a lane (mod), so the +// rank goes up by one. phys_src says which logical axis each physical dim comes +// from, phys_op is 0=identity / 1=floordiv / 2=mod, and phys_arg is the stick +// WIDTH (64), not a dim index. +// +// The marker also fixes the ORDER of the physical dims, and that order decides +// which dims the reduce and the broadcast name. Both orders appear below, +// because the renumbering must not assume either: +// +// phys_src=[1,0,1] phys_op=[floordiv,identity,mod] [stick, row, lane] +// logical <64x128> -> physical <2x64x64>; split axis owns dims 0 and 2 +// phys_src=[0,1,1] phys_op=[identity,floordiv,mod] [row, stick, lane] +// logical <64x128> -> physical <64x2x64>; split axis owns dims 1 and 2 +// +// WHAT THESE TESTS ASK OF THE PASS: when the layout splits an axis the reduce +// CONSUMES, every axis the broadcast CARRIES stays whole, and the broadcast must +// follow its operands to the physical shape -- its `dimensions` list renumbered +// to name the new physical dims, its `outs` rebuilt at the new rank. +// +// The second case asks the opposite: with no reduce in front of it, nothing +// requires the broadcast to be physical, and it must be LEFT ALONE. + +// ----- + +// A reduce feeding a broadcast feeding elementwise arithmetic -- the softmax +// shape, and the case this whole file exists for. +// +// Axis 1 (the columns, length 128) is split at 64 and is also the axis the +// reduce consumes. Required outcome: the reduce absorbs BOTH physical dims of +// that axis, the broadcast re-adds the same two, and the subtraction downstream +// ends up with operands of one shape. + +// CHECK-LABEL: func @softmax_broadcast_split_reduced_axis +// The load is physical: rank 3, the split axis became 2 sticks of 64. +// CHECK: ktdp.load {{.*}} -> tensor<64x2x64xf32> +// The reduce absorbs BOTH physical dims of the split axis, so dimensions = [1, 2] +// and the result is rank 1. +// CHECK: linalg.reduce +// CHECK-SAME: outs(%{{.*}} : tensor<64xf32>) +// CHECK-SAME: dimensions = [1, 2] +// The broadcast re-adds those same two dims. +// CHECK: linalg.broadcast +// CHECK-SAME: outs(%{{.*}} : tensor<64x2x64xf32>) +// CHECK-SAME: dimensions = [1, 2] +// With the broadcast physical, the subtraction's operands finally agree. +// CHECK: arith.subf %{{.*}}, %{{.*}} : tensor<64x2x64xf32> +// CHECK: math.exp %{{.*}} : tensor<64x2x64xf32> +// CHECK: ktdp.store %{{.*}}, %{{.*}} : tensor<64x2x64xf32> +module { +tt.func @softmax_broadcast_split_reduced_axis(%a_ptr: !tt.ptr, %o_ptr: !tt.ptr) { + %c0 = arith.constant 0 : i32 + %c64 = arith.constant 64 : i32 + %c128 = arith.constant 128 : i32 + %s128 = arith.constant 128 : i64 + %s1 = arith.constant 1 : i64 + %ad = tt.make_tensor_descriptor %a_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %ad {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + %a = tt.descriptor_load %ad[%c0, %c0] : !tt.tensordesc<64x128xf32> -> tensor<64x128xf32> + %m = "tt.reduce"(%a) ({ ^bb0(%x: f32, %y: f32): %mx = arith.maximumf %x, %y : f32 + tt.reduce.return %mx : f32 }) {axis = 1 : i32} : (tensor<64x128xf32>) -> tensor<64xf32> + // tt.expand_dims then tt.broadcast lower to a size-1 reshape plus a + // linalg.broadcast; the reshape is a cancelling pair the analysis must cross. + %me = tt.expand_dims %m {axis = 1 : i32} : tensor<64xf32> -> tensor<64x1xf32> + %bc = tt.broadcast %me : tensor<64x1xf32> -> tensor<64x128xf32> + %d = arith.subf %a, %bc : tensor<64x128xf32> + %e = math.exp %d : tensor<64x128xf32> + %od = tt.make_tensor_descriptor %o_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %od {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + tt.descriptor_store %od[%c0, %c0], %e : !tt.tensordesc<64x128xf32>, tensor<64x128xf32> + tt.return +} +} + +// ----- + +// The same broadcast with NO reduce in front of it. Nothing upstream of the +// broadcast reads the split axis, so the backward analysis -- the pass walking +// from each store toward its producers, asking what layout each value must +// have -- never asks the broadcast for one, and it stays LOGICAL. The pass instead bridges at the +// store: an scf.for slices the logical <64x128> into 2 sticks of 64 and +// assembles the physical <64x2x64>. Both outcomes are correct; which one you +// get depends on whether a physical value reaches the broadcast, so this test +// pins the no-reduce half of that choice. + +// CHECK-LABEL: func @broadcast_only_split_added_axis +// The broadcast is left logical -- rank 2, dimensions = [1]. +// CHECK: linalg.broadcast +// CHECK-SAME: outs(%{{.*}} : tensor<64x128xf32>) +// CHECK-SAME: dimensions = [1] +// CHECK: math.exp %{{.*}} : tensor<64x128xf32> +// The store bridge: one trip per stick (128/64 = 2), each copying a 64-wide slice. +// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<64x2x64xf32> +// CHECK: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[ACC:.*]] = %[[EMPTY]]) +// CHECK: tensor.extract_slice %{{.*}}[0, %{{.*}}] [64, 64] [1, 1] : tensor<64x128xf32> to tensor<64x64xf32> +// CHECK: tensor.insert_slice %{{.*}} into %[[ACC]][0, %[[IV]], 0] [64, 1, 64] [1, 1, 1] +// CHECK: ktdp.store %{{.*}}, %{{.*}} : tensor<64x2x64xf32> +module { +tt.func @broadcast_only_split_added_axis(%a_ptr: !tt.ptr, %o_ptr: !tt.ptr) { + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + %c64 = arith.constant 64 : i32 + %c128 = arith.constant 128 : i32 + %s128 = arith.constant 128 : i64 + %s1 = arith.constant 1 : i64 + // Load a single column so the value fanned out is genuinely rank 1 (length 64) + // and the broadcast ADDS the split axis rather than carrying it. + %ad = tt.make_tensor_descriptor %a_ptr, [%c64, %c1], [%s1, %s1] : !tt.ptr, !tt.tensordesc<64x1xf32> + %a = tt.descriptor_load %ad[%c0, %c0] : !tt.tensordesc<64x1xf32> -> tensor<64x1xf32> + %bc = tt.broadcast %a : tensor<64x1xf32> -> tensor<64x128xf32> + %e = math.exp %bc : tensor<64x128xf32> + %od = tt.make_tensor_descriptor %o_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %od {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + tt.descriptor_store %od[%c0, %c0], %e : !tt.tensordesc<64x128xf32>, tensor<64x128xf32> + tt.return +} +} + +// ----- + +// The same softmax shape under the STICK-OUTERMOST marker, which is the ordering +// most fixtures in this suite use: phys_src=[1,0,1] puts the stick index at +// physical dim 0, the untouched row at dim 1, and the lane at dim 2. The split +// axis therefore owns dims 0 and 2, NOT 1 and 2. +// +// This is the case the renumbering could plausibly get wrong: it maps a logical +// axis to every physical dim sourced from it, and those dims are not adjacent +// here. A rule that assumed the split axis contributed a contiguous pair, or +// that read the physical order off the logical order, would produce [1, 2] and +// silently reduce the wrong axes. + +// CHECK-LABEL: func @softmax_broadcast_stick_outermost +// Stick index leads, so the load is <2x64x64> rather than <64x2x64>. +// CHECK: ktdp.load {{.*}} -> tensor<2x64x64xf32> +// Both dims of the split axis are absorbed, and they are 0 and 2. +// CHECK: linalg.reduce +// CHECK-SAME: outs(%{{.*}} : tensor<64xf32>) +// CHECK-SAME: dimensions = [0, 2] +// The broadcast re-adds the same two, in the same positions. +// CHECK: linalg.broadcast +// CHECK-SAME: outs(%{{.*}} : tensor<2x64x64xf32>) +// CHECK-SAME: dimensions = [0, 2] +// CHECK: arith.subf %{{.*}}, %{{.*}} : tensor<2x64x64xf32> +// CHECK: ktdp.store %{{.*}}, %{{.*}} : tensor<2x64x64xf32> +module { +tt.func @softmax_broadcast_stick_outermost(%a_ptr: !tt.ptr, %o_ptr: !tt.ptr) { + %c0 = arith.constant 0 : i32 + %c64 = arith.constant 64 : i32 + %c128 = arith.constant 128 : i32 + %s128 = arith.constant 128 : i64 + %s1 = arith.constant 1 : i64 + %ad = tt.make_tensor_descriptor %a_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %ad {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + %a = tt.descriptor_load %ad[%c0, %c0] : !tt.tensordesc<64x128xf32> -> tensor<64x128xf32> + %m = "tt.reduce"(%a) ({ ^bb0(%x: f32, %y: f32): %mx = arith.maximumf %x, %y : f32 + tt.reduce.return %mx : f32 }) {axis = 1 : i32} : (tensor<64x128xf32>) -> tensor<64xf32> + %me = tt.expand_dims %m {axis = 1 : i32} : tensor<64xf32> -> tensor<64x1xf32> + %bc = tt.broadcast %me : tensor<64x1xf32> -> tensor<64x128xf32> + %d = arith.subf %a, %bc : tensor<64x128xf32> + %e = math.exp %d : tensor<64x128xf32> + %od = tt.make_tensor_descriptor %o_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %od {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + tt.descriptor_store %od[%c0, %c0], %e : !tt.tensordesc<64x128xf32>, tensor<64x128xf32> + tt.return +} +} diff --git a/third_party/spyre/test/Conversion/rewrite-descriptor-layout-reshape.mlir b/third_party/spyre/test/Conversion/rewrite-descriptor-layout-reshape.mlir new file mode 100644 index 000000000000..247ab67c7a57 --- /dev/null +++ b/third_party/spyre/test/Conversion/rewrite-descriptor-layout-reshape.mlir @@ -0,0 +1,61 @@ +// RUN: spyre-triton-opt %s --lower-descriptor-memory --lower-scalar-load --lower-compute-ops --rewrite-descriptor-layout -split-input-file | FileCheck %s + +// Crossing size-1 reshapes when physicalizing a layout. +// +// A `tt.spyre_tensor_layout` annotation asks for the Spyre "stick" layout: one +// logical axis splits into a stick index (floordiv) and a lane (mod), so a +// logical <64x128> becomes a physical <64x2x64> and the rank goes 2 -> 3. +// phys_src says which logical axis each physical dim comes from, phys_op is +// 0=identity / 1=floordiv / 2=mod, and phys_arg is the stick WIDTH (64), not a +// dim index. +// +// `tt.expand_dims` inserts a length-1 axis and `tt.broadcast` fans out along +// it. LowerComputeOps lowers them independently, so expand-then-broadcast +// leaves a CANCELLING PAIR: expand 64 -> 64x1, then collapse 64x1 -> 64 right +// back. No real extent changes -- only size-1 axes are added and dropped. +// +// WHAT THIS TEST ASKS OF THE PASS: the layout must survive that round trip. A +// reshape that only adds or drops size-1 axes moves no data, so it must not stop +// the layout propagating -- if it does, the chain breaks before the broadcast +// that needs it and the kernel fails to compile. + +// ----- + +// The cancelling pair, reached from a store that wants a physical layout. +// Axis 1 (length 128) is split into 2 sticks of 64 and is also the reduced +// axis, so the reduce absorbs both physical dims and the chain +// reduce -> collapse -> broadcast must stay connected across the reshape. + +// CHECK-LABEL: func @reshape_cancelling_pair_crossed +// CHECK: ktdp.load {{.*}} -> tensor<64x2x64xf32> +// CHECK: linalg.reduce +// CHECK-SAME: outs(%{{.*}} : tensor<64xf32>) +// CHECK-SAME: dimensions = [1, 2] +// The reshape survives as a collapse of the size-1 axis and does NOT block the +// layout: the broadcast downstream of it is physical. +// CHECK: linalg.broadcast +// CHECK-SAME: outs(%{{.*}} : tensor<64x2x64xf32>) +// CHECK-SAME: dimensions = [1, 2] +// CHECK: ktdp.store %{{.*}}, %{{.*}} : tensor<64x2x64xf32> +module { +tt.func @reshape_cancelling_pair_crossed(%a_ptr: !tt.ptr, %o_ptr: !tt.ptr) { + %c0 = arith.constant 0 : i32 + %c64 = arith.constant 64 : i32 + %c128 = arith.constant 128 : i32 + %s128 = arith.constant 128 : i64 + %s1 = arith.constant 1 : i64 + %ad = tt.make_tensor_descriptor %a_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %ad {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + %a = tt.descriptor_load %ad[%c0, %c0] : !tt.tensordesc<64x128xf32> -> tensor<64x128xf32> + %m = "tt.reduce"(%a) ({ ^bb0(%x: f32, %y: f32): %sm = arith.addf %x, %y : f32 + tt.reduce.return %sm : f32 }) {axis = 1 : i32} : (tensor<64x128xf32>) -> tensor<64xf32> + // expand then broadcast: the expand's 64x1 is collapsed straight back to 64. + %me = tt.expand_dims %m {axis = 1 : i32} : tensor<64xf32> -> tensor<64x1xf32> + %bc = tt.broadcast %me : tensor<64x1xf32> -> tensor<64x128xf32> + %d = arith.subf %a, %bc : tensor<64x128xf32> + %od = tt.make_tensor_descriptor %o_ptr, [%c64, %c128], [%s128, %s1] : !tt.ptr, !tt.tensordesc<64x128xf32> + tt.spyre_tensor_layout %od {phys_src = array, phys_op = array, phys_arg = array} : !tt.tensordesc<64x128xf32> + tt.descriptor_store %od[%c0, %c0], %d : !tt.tensordesc<64x128xf32>, tensor<64x128xf32> + tt.return +} +}