Ttir Decomposition For Ttmetal

TTIR Composite Op Decomposition for TTMetal

Decompose a high-level fused TTIR op into primitive TTIR ops so the D2M/TTMetal backend can lower them individually. The TTNN backend keeps native fused ops; these decomposition patterns only run in the TTMetal pipeline via the unified TTIRDecomposeComposites pass.

When to Use

• The op has no D2M/TTMetal lowering but can be expressed as a sequence of ops that do (matmul, add, multiply, reduce, reshape, permute, etc.).
• The TTNN backend already has native support, so it skips the decomposition.

Architecture

All composite decompositions live in a single pass (ttir-decompose-composites) that uses MLIR's greedy pattern rewriter. Each op decomposition is an OpRewritePattern<T> with a configurable benefit level that controls application order. For example, SDPA has higher benefit than softmax so it runs first — the softmax ops it produces are then caught by the softmax pattern on subsequent rewriter iterations.

Files to Modify

You should NOT need to touch CMakeLists.txt or Passes.td pass registration in the common case. However, when adding a new composite decomposition, verify that ttir-decompose-composites is scheduled in the relevant D2M pipeline in D2MPipelines.cpp, and update that pipeline if necessary.

Step-by-Step

1. Understand the Op

Read the op definition in include/ttmlir/Dialect/TTIR/IR/TTIROps.td. Note tensor shapes, attributes (optional mask, scale, etc.), and the mathematical decomposition into primitives.

2. Add a Pattern to DecomposeComposites.cpp

Open lib/Dialect/TTIR/Transforms/DecomposeComposites.cpp and add a new OpRewritePattern<YourOp> struct. Follow the existing patterns as examples.

Pattern template:

struct DecomposeYourOpPattern : public OpRewritePattern<YourOp> {
  using OpRewritePattern<YourOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(YourOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    // ... decomposition logic using rewriter.create<T>(...) ...
    rewriter.replaceOp(op, result);
    return success();
  }
};

Then register the pattern in TTIRDecomposeComposites::runOnOperation():

void runOnOperation() final {
  RewritePatternSet patterns(&getContext());
  patterns.add<DecomposeSDPAPattern>(&getContext(), /*benefit=*/2);
  patterns.add<DecomposeRMSNormPattern>(&getContext(), /*benefit=*/1);
  patterns.add<DecomposeSoftmaxPattern>(&getContext(), /*benefit=*/0);
  patterns.add<DecomposeYourOpPattern>(&getContext(), /*benefit=*/N);  // NEW

  if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
    signalPassFailure();
  }
}

Benefit ordering: If your decomposition produces ops that another pattern needs to decompose further (e.g. SDPA produces softmax), give your pattern a higher benefit number than the downstream pattern.

Key conventions:

• Use OpRewritePattern<T> and PatternRewriter, not IRRewriter.
• Always return success() after rewriter.replaceOp(op, result).
• Call rewriter.replaceOp(op, result) at the end to replace the original.

3. Creating TTIR Ops in Decomposition Code

Common op creation patterns (use rewriter.create<T>(...)):

// Elementwise binary (add, multiply, subtract, etc.)
auto add = rewriter.create<AddOp>(loc, resultType, lhs, rhs);

// MatmulOp
auto mm = rewriter.create<MatmulOp>(loc, resultType, a, b);

// SoftmaxOp
auto sm = rewriter.create<SoftmaxOp>(loc, resultType, input,
    rewriter.getSI32IntegerAttr(dim),
    rewriter.getBoolAttr(false));

// FullOp (scalar constant broadcast to shape)
auto full = rewriter.create<FullOp>(loc, resultType,
    rewriter.getF32FloatAttr(value));

// ReshapeOp
SmallVector<int32_t> shapeI32(newShape.begin(), newShape.end());
auto reshape = rewriter.create<ReshapeOp>(loc, newType, input,
    rewriter.getI32ArrayAttr(shapeI32));

// PermuteOp
auto permute = rewriter.create<PermuteOp>(loc, permutedType, input,
    rewriter.getDenseI64ArrayAttr(permutation));

// MeanOp (reduction)
auto mean = rewriter.create<MeanOp>(loc, reducedType, input,
    rewriter.getBoolAttr(/*keep_dim=*/true),
    rewriter.getI32ArrayAttr(reduceDims));

// RsqrtOp (unary)
auto rsqrt = rewriter.create<RsqrtOp>(loc, type, input);

4. Add MLIR Lit Tests

Add test functions and FileCheck assertions to test/ttmlir/Dialect/TTIR/Transforms/metal_composite_decompositions.mlir.

The file uses a single pass (--ttir-decompose-composites) with multiple check prefixes. Add a new check prefix for your op and a new RUN line:

// RUN: ttmlir-opt --ttir-decompose-composites %s | FileCheck %s --check-prefix=YOUROP

Then add test functions:

// YOUROP-LABEL: func.func @your_op_basic
// YOUROP-NOT: ttir.your_op
// YOUROP: "ttir.multiply"
// YOUROP: "ttir.add"
// YOUROP: return
func.func @your_op_basic(%input: tensor<...xbf16>) -> tensor<...xbf16> {
  %0 = "ttir.your_op"(%input) <{...}> : (...) -> ...
  return %0 : ...
}

5. Add Python Builder Tests

Add tests to test/python/golden/d2m/test_composite_ops.py. This file contains all composite decomposition tests for the TTMetal pipeline.

Follow the existing patterns (SDPA, RMSNorm, softmax) as examples:

@pytest.mark.parametrize("shape", [...])
@pytest.mark.parametrize("target", ["ttmetal"])
def test_your_op_decomposition(
    shape: Shape,
    target: str,
    request,
    device,
):
    """Test your_op decomposition for the TTMetal pipeline."""

    def module(builder: TTIRBuilder):
        @builder.func([shape], [torch.float32])
        def your_op(
            in0: Operand,
            builder: TTIRBuilder,
            unit_attrs: Optional[List[str]] = None,
        ):
            return builder.your_op(in0, ..., unit_attrs=unit_attrs)

    compile_and_execute_ttir(
        module,
        target=target,
        **get_request_kwargs(request),
        device=device,
    )

Key points:

• Always use target="ttmetal" — decompositions only run in the TTMetal pipeline.
• Use compile_and_execute_ttir from builder.base.builder_apis.
• Parametrize over the attributes that affect decomposition logic (e.g. causal vs non-causal, with/without weight, different dims).

6. Build and Iterate

Run ./build_and_test.sh or:

source env/activate
cmake --build build

Fix compilation errors, then test with the lit test:

build/bin/ttmlir-opt --ttir-decompose-composites test/ttmlir/Dialect/TTIR/Transforms/metal_composite_decompositions.mlir

And the Python tests:

pytest -svv test/python/golden/d2m/test_composite_ops.py

Reference Implementations

All live in lib/Dialect/TTIR/Transforms/DecomposeComposites.cpp:

• DecomposeRMSNormPattern (benefit 1): Decomposes rms_norm(x, w, b, eps) into x^2 -> mean -> +eps -> rsqrt -> *x -> *w -> +b.

• DecomposeSDPAPattern (benefit 2): Decomposes scaled_dot_product_attention(Q, K, V, mask) into Q @ K^T -> *scale -> +mask -> softmax -> @ V, with GQA head expansion via reshape. Produces ttir.softmax ops that the softmax pattern then decomposes.

• DecomposeSoftmaxPattern (benefit 0): Decomposes softmax(x, dim) into max -> subtract -> exp -> sum -> div (uses ttir.div rather than reciprocal -> multiply to work around a broadcast-multiply bug). When numericStable=false, the max-subtract step is skipped.

• Inner-min decomposition (in TTIRToD2M): D2MInnerMinDecompositionRewriter in lib/Conversion/TTIRToD2M/TTIRToD2M.cpp rewrites inner-dim ttir.min into neg(max(neg(x))) during conversion (there is no tile_reduce_min kernel). Outer-dim min reductions use the accumulation path with d2m.tile_minimum.