Handle zero-extent rms_norm decomposition - #4678
Conversation
ccb1065 to
21f9de4
Compare
b1fe74a to
0cce3b3
Compare
| @@ -8184,6 +8184,14 @@ class DecomposeAtenRMSLayerNormOp : public OpRewritePattern<AtenRmsNormOp> { | |||
| if (!outputTy.hasDtype()) | |||
| return rewriter.notifyMatchFailure(op, "output should have a dtype."); | |||
|
|
|||
There was a problem hiding this comment.
If RMS norm input has a zero dimension aten.rms_norm operator is directly replaced with the input tensor.
| ValueTensorType::get(context, reducedShape, inputTy.getDtype()); | ||
| // x^2 | ||
| Value inputSquared = AtenSquareOp::create(rewriter, loc, inputTy, input); | ||
| // x^2. Emit multiplication directly instead of aten.square so zero-extent |
There was a problem hiding this comment.
This avoids emitting aten.square operator which later becomes torch.aten.pow.Tensor_Scalar. The TOSA legalization was failing on zero-extent pow.Tensor_Scalar.
| }; | ||
| } // namespace | ||
|
|
||
| namespace { |
There was a problem hiding this comment.
Separate fold patterns are needed because the pipeline may see RMSNorm either before or after it has already been decomposed. The direct aten.rms_norm rewrite handles the intact op, but already-decomposed zero-extent cases can end at different ops depending on dtype.
FoldZeroExtentRmsNormDecompositionFromMul handles the float32 decomposition ending in aten.mul.Tensor. FoldZeroExtentRmsNormDecompositionFromConvert handles paths ending in prims.convert_element_type. FoldZeroExtentRmsNormDecompositionFromAtenToDtype handles fp16/bf16 eager-opmath paths ending in aten.to.dtype.
All three apply the same rule: if the matched chain is just RMSNorm computation over a statically zero-sized input, replace the chain with the original empty input so illegal zero-extent pow or cast ops do not reach TOSA legalization.
2faffcc to
48a9d8d
Compare
48a9d8d to
5b279cb
Compare
Lallapallooza
left a comment
Lallapallooza
left a comment
There was a problem hiding this comment.
Thanks for the patch, few comments.
| if (llvm::is_contained(inputTy.getSizes(), 0)) { | ||
| if (input.getType() != op.getType()) | ||
| return rewriter.notifyMatchFailure( | ||
| op, "zero-extent input and output types should match."); | ||
| rewriter.replaceOp(op, input); | ||
| return success(); | ||
| } |
There was a problem hiding this comment.
Could we move this fast path behind the existing normalized-shape gate and the other statically visible RMSNorm checks? Please validate normalized shape, known weight shape, and supported dtype before folding, with negative lit coverage.
| return success(); | ||
| } | ||
|
|
||
| static LogicalResult matchMeanOfPow2(Value value, Value input) { |
There was a problem hiding this comment.
The current path for this eps=0.5 case emits pow -> mean -> add(eps) -> rsqrt -> mul -> mul, but matchMeanOfPow2 accepts only a mean or sum/div directly below rsqrt.
| return input; | ||
| } | ||
|
|
||
| static void eraseDeadRmsNormDecompositionOp(Operation *op, |
There was a problem hiding this comment.
eraseDeadRmsNormDecompositionFromOperands can walk a saved Value after another operand branch has erased its defining op.
Fixes zero-extent RMSNorm folding by validating normalized shape, weight shape/dtype, and floating dtypes before replacing the op with the input. Also handles already-decomposed RMSNorm graphs with eps by matching the add.Scalar before rsqrt, and makes dead decomposition cleanup safe when shared branches are erased. Fixes #4678
Fixes zero-extent RMSNorm folding by validating normalized shape, weight shape/dtype, and floating dtypes before replacing the op with the input. Also handles already-decomposed RMSNorm graphs with eps by matching the add.Scalar before rsqrt, and makes dead decomposition cleanup safe when shared branches are erased. Fixes llvm#4678
Fixes zero-extent RMSNorm folding by validating normalized shape, weight shape/dtype, and floating dtypes before replacing the op with the input. Also handles already-decomposed RMSNorm graphs with eps by matching the add.Scalar before rsqrt, and makes dead decomposition cleanup safe when shared branches are erased. Fixes llvm#4678
96c1e38 to
183474f
Compare
Fixes zero-extent RMSNorm folding by validating normalized shape, weight shape/dtype, and floating dtypes before replacing the op with the input. Also handles already-decomposed RMSNorm graphs with eps by matching the add.Scalar before rsqrt, and makes dead decomposition cleanup safe when shared branches are erased. Fixes llvm#4678
183474f to
071e930
Compare
| return nullptr; | ||
| } | ||
|
|
||
| static FailureOr<Value> matchZeroExtentRmsNormDecomposition(Value value) { |
There was a problem hiding this comment.
Thanks for the patch, but I am not following the motivation to match an already decomposed RMSNorm op here. torch-mlir doesn't decompose this op by default
ep = ep.run_decompositions() on the exported program and then feeds that to fx.export_and_import. Is there a real use-case to support this?
My concern is that it won't be scalable to add support for every decomposed op in this way. If there is truly a use-case, then I think you should raise the decomposed pattern back to rms_norm op in
rms_norm. That keeps the decomposition path free of raising operations/pattern-matching to keep the responsibilites (decompose vs recompose) separate in their own passes.
3 participants
Summary:
Validation: