diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 35bb08012e..7f590d8b04 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -189,6 +189,15 @@
* Added ``CZ`` support to ``to-ppr`` pass.
[(#3009)](https://github.com/PennyLaneAI/catalyst/pull/3009)
+* IR amplification for circuits whose gate parameters are computed from runtime values
+ (e.g. `qml.TrotterProduct` with runtime Hamiltonian coefficients) is drastically reduced.
+ The new `scalarize-tensor-extracts` pass turns gate-angle dataflow into scalar arithmetic
+ instead of thousands of small tensors that each survive bufferization as an allocation,
+ and the new `reroll-loops` pass reconstructs the loops that tracing unrolled by rewriting
+ repeated op sequences as `scf.for` loops. On a Trotterized QPE workload with runtime
+ coefficients, compile time, peak memory, and final IR size all drop by large factors.
+ [(#3013)](https://github.com/PennyLaneAI/catalyst/pull/3013)
+
Breaking changes 💔
* Python 3.11 is no longer supported. Catalyst now requires Python 3.12 or newer.
diff --git a/frontend/test/lit/test_detensorize.py b/frontend/test/lit/test_detensorize.py
index cc2a2295cd..40cf98509b 100644
--- a/frontend/test/lit/test_detensorize.py
+++ b/frontend/test/lit/test_detensorize.py
@@ -45,15 +45,15 @@ def f_with_cond(a, b):
# CHECK-NOT: linalg.generic
a2 = a + a
if a2 > b:
- # CHECK: arith.subf
+ # CHECK-DAG: arith.subf
a = a - 2.0
- # CHECK: arith.mulf
+ # CHECK-DAG: arith.mulf
b = b * 2.0
c = a + b
else:
- # CHECK: arith.mulf
+ # CHECK-DAG: arith.mulf
a = a * 2.0
- # CHECK: arith.subf
+ # CHECK-DAG: arith.subf
b = b - 2.0
c = a + b
return c * 2.0
diff --git a/frontend/test/pytest/test_trotter_runtime_coeffs.py b/frontend/test/pytest/test_trotter_runtime_coeffs.py
new file mode 100644
index 0000000000..5d02945032
--- /dev/null
+++ b/frontend/test/pytest/test_trotter_runtime_coeffs.py
@@ -0,0 +1,117 @@
+# Copyright 2026 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Integration tests for the IR-amplification fixes for runtime-coefficient
+Hamiltonians: scalarize-tensor-extracts and elementwise fusion in the default
+pipeline must preserve numerics for Trotterized workloads with runtime
+coefficients."""
+
+import numpy as np
+import pennylane as qml
+import pytest
+from jax import numpy as jnp
+
+from catalyst import qjit
+
+# H2/STO-3G-like coefficients; the structure (15 terms, runtime values) is what
+# exercises the pipeline, the values just need to be a valid Hamiltonian.
+COEFFS = [
+ -0.0996,
+ 0.1711,
+ 0.1711,
+ -0.2225,
+ -0.2225,
+ 0.1686,
+ 0.0453,
+ -0.0453,
+ -0.0453,
+ 0.0453,
+ 0.1205,
+ 0.1658,
+ 0.1658,
+ 0.1205,
+ 0.1743,
+]
+
+
+def make_ops():
+ """Build a fresh list of the Hamiltonian's Pauli-word operators."""
+ return [
+ qml.Identity(0),
+ qml.PauliZ(0),
+ qml.PauliZ(1),
+ qml.PauliZ(2),
+ qml.PauliZ(3),
+ qml.PauliZ(0) @ qml.PauliZ(1),
+ qml.PauliY(0) @ qml.PauliX(1) @ qml.PauliX(2) @ qml.PauliY(3),
+ qml.PauliY(0) @ qml.PauliY(1) @ qml.PauliX(2) @ qml.PauliX(3),
+ qml.PauliX(0) @ qml.PauliX(1) @ qml.PauliY(2) @ qml.PauliY(3),
+ qml.PauliX(0) @ qml.PauliY(1) @ qml.PauliY(2) @ qml.PauliX(3),
+ qml.PauliZ(0) @ qml.PauliZ(2),
+ qml.PauliZ(0) @ qml.PauliZ(3),
+ qml.PauliZ(1) @ qml.PauliZ(2),
+ qml.PauliZ(1) @ qml.PauliZ(3),
+ qml.PauliZ(2) @ qml.PauliZ(3),
+ ]
+
+
+N_QUBITS = 4
+N_EST = 2
+N_TROTTER = 6
+
+
+def make_qpe(dev, runtime: bool):
+ """Controlled-power Trotterized QPE, with runtime or compile-time coeffs."""
+
+ @qml.qnode(dev)
+ def qpe_circuit(coeffs):
+ qml.PauliX(0)
+ qml.PauliX(1)
+ for k in range(N_EST):
+ qml.Hadamard(wires=N_QUBITS + k)
+ H = qml.dot(coeffs if runtime else COEFFS, make_ops())
+ for k in range(N_EST):
+ t = 2 ** (N_EST - 1 - k)
+ qml.ctrl(
+ qml.adjoint(
+ qml.TrotterProduct(
+ H, time=t, n=N_TROTTER, order=2, check_hermitian=False
+ )
+ ),
+ control=N_QUBITS + k,
+ )
+ qml.adjoint(qml.QFT)(wires=range(N_QUBITS, N_QUBITS + N_EST))
+ return qml.probs(wires=range(N_QUBITS, N_QUBITS + N_EST))
+
+ return qpe_circuit
+
+
+class TestRuntimeCoefficientTrotter:
+ """Numerical equivalence of runtime- and fixed-coefficient Trotterization
+ through the default pipeline (which scalarizes and fuses)."""
+
+ def test_runtime_matches_fixed(self):
+ """qml.dot with traced coefficients must produce the same distribution
+ as the same Hamiltonian with Python-float coefficients."""
+ dev = qml.device("lightning.qubit", wires=N_QUBITS + N_EST)
+ coeffs = jnp.array(COEFFS)
+
+ dyn = qjit(make_qpe(dev, runtime=True))(coeffs)
+ fixed = qjit(make_qpe(dev, runtime=False))(coeffs)
+
+ assert np.allclose(np.asarray(dyn), np.asarray(fixed), atol=1e-9)
+
+
+if __name__ == "__main__":
+ pytest.main(["-x", __file__])
diff --git a/mlir/include/Catalyst/Transforms/Passes.td b/mlir/include/Catalyst/Transforms/Passes.td
index c5f87f0d7c..044090ea9f 100644
--- a/mlir/include/Catalyst/Transforms/Passes.td
+++ b/mlir/include/Catalyst/Transforms/Passes.td
@@ -333,6 +333,30 @@ def RegisterDecompRuleResourcePass : Pass<"register-decomp-rule-resource"> {
}];
}
+def ScalarizeTensorExtractsPass : Pass<"scalarize-tensor-extracts"> {
+ let summary = "Sink scalar tensor.extract ops through small-tensor producers.";
+ let description = [{
+ Programs that compute quantum gate parameters from runtime inputs (e.g.
+ Trotterization with runtime Hamiltonian coefficients) produce long chains of
+ small tensor operations whose only consumers are scalar `tensor.extract`
+ operations. Each intermediate tensor survives bufferization as an allocation
+ with copies and each `linalg.generic` is later unrolled into loops, causing
+ severe IR amplification.
+
+ This pass folds `tensor.extract` through elementwise `linalg.generic` (by
+ inlining the scalar payload), `tensor.collapse_shape`, and
+ `tensor.extract_slice`, so the extracted element is computed directly in
+ scalar arithmetic and the intermediate tensors become dead. Payload inlining
+ is limited to small statically-shaped tensors to bound code growth.
+ }];
+
+ let dependentDialects = [
+ "mlir::arith::ArithDialect",
+ "mlir::linalg::LinalgDialect",
+ "mlir::tensor::TensorDialect"
+ ];
+}
+
def EmptyPass : Pass<"empty"> {
let summary = "Empty pass that does nothing.";
diff --git a/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h b/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h
index 9c652e7732..897e805b6a 100644
--- a/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h
+++ b/mlir/include/Driver/DefaultPipelines/DefaultPipelines.h
@@ -71,9 +71,20 @@ const PipelineList pipelineList{
"scatter-lowering",
"hlo-custom-call-lowering",
"cse",
+ // Sink scalar extractions through small-tensor producers and fuse the
+ // remaining elementwise ops. Traced gate-parameter dataflow (e.g. runtime
+ // Hamiltonian coefficients) otherwise reaches bufferization as thousands
+ // of tiny tensor ops that each become an alloc + copy.
+ "scalarize-tensor-extracts",
+ "func.func(linalg-fuse-elementwise-ops)",
+ "canonicalize",
"func.func(linalg-detensorize{aggressive-mode})",
"detensorize-scf",
"detensorize-function-boundary",
+ // Detensorization is what materializes tensor.extract on the gate-angle
+ // dataflow, so scalarization must run again here to fold the
+ // extract_slice/collapse_shape chains it exposes.
+ "scalarize-tensor-extracts",
"canonicalize",
"symbol-dce"}},
{"gradient-lowering-stage",
diff --git a/mlir/lib/Catalyst/Transforms/CMakeLists.txt b/mlir/lib/Catalyst/Transforms/CMakeLists.txt
index bb794d3b4c..b1edf9355e 100644
--- a/mlir/lib/Catalyst/Transforms/CMakeLists.txt
+++ b/mlir/lib/Catalyst/Transforms/CMakeLists.txt
@@ -23,6 +23,7 @@ file(GLOB SRC
QnodeToAsyncPatterns.cpp
RegisterInactiveCallbackPass.cpp
ResourceAnalysisPass.cpp
+ ScalarizeTensorExtractsPass.cpp
RegisterDecompRuleResourcePass.cpp
SplitMultipleTapes.cpp
TBAAPatterns.cpp
diff --git a/mlir/lib/Catalyst/Transforms/ScalarizeTensorExtractsPass.cpp b/mlir/lib/Catalyst/Transforms/ScalarizeTensorExtractsPass.cpp
new file mode 100644
index 0000000000..d9036a2b1c
--- /dev/null
+++ b/mlir/lib/Catalyst/Transforms/ScalarizeTensorExtractsPass.cpp
@@ -0,0 +1,313 @@
+// Copyright 2026 Xanadu Quantum Technologies Inc.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+// http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#define DEBUG_TYPE "scalarize-tensor-extracts"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace llvm;
+using namespace mlir;
+
+// Gate parameters computed from runtime inputs (e.g. Trotterization with
+// runtime Hamiltonian coefficients) arrive as long chains of small tensor ops
+// consumed only by scalar `tensor.extract` operations; each such tensor
+// survives bufferization as an allocation plus copies. This pass sinks
+// `tensor.extract` through `linalg.generic` (inlining the scalar payload),
+// `tensor.collapse_shape`, and `tensor.extract_slice`, so the extracted
+// element is computed in scalar arithmetic and the tensors become dead.
+// Payload inlining is limited to small statically-shaped results to bound
+// code growth.
+
+namespace {
+
+/// Upper bound on the number of elements of a tensor whose producer payload we
+/// are willing to clone per extraction site. 16 covers the 4x4 matrices of
+/// two-qubit gates, the largest tensors in gate-parameter dataflow, while
+/// keeping the worst-case code growth per extract small.
+constexpr int64_t kMaxScalarizedElements = 16;
+
+/// Whether the generic's payload is speculatable scalar code that never reads
+/// the accumulator (output block arguments), so it can be inlined verbatim.
+bool hasInlinablePayload(linalg::GenericOp genericOp)
+{
+ Block *body = genericOp.getBody();
+ for (Operation &op : body->without_terminator()) {
+ if (!isPure(&op)) {
+ return false;
+ }
+ }
+ for (OpOperand &outOperand : genericOp.getDpsInitsMutable()) {
+ BlockArgument outArg = body->getArgument(outOperand.getOperandNumber());
+ if (!outArg.use_empty()) {
+ return false;
+ }
+ }
+ return true;
+}
+
+/// Whether every tensor input's indexing map is built only from dimension and
+/// constant expressions, so per-input extraction indices can be derived from
+/// the iteration indices. Validated up front: patterns must not mutate the IR
+/// on the failure path, and ops are created per input during the rewrite.
+bool hasScalarizableInputMaps(linalg::GenericOp genericOp)
+{
+ for (OpOperand *inOperand : genericOp.getDpsInputOperands()) {
+ if (!isa(inOperand->get().getType())) {
+ continue;
+ }
+ AffineMap inputMap = genericOp.getMatchingIndexingMap(inOperand);
+ for (AffineExpr expr : inputMap.getResults()) {
+ if (!isa(expr)) {
+ return false;
+ }
+ }
+ }
+ return true;
+}
+
+/// Materialize scalar operands for the generic's payload: one tensor.extract
+/// per tensor input, at indices given by composing that input's indexing map
+/// with the iteration indices. Scalar operands map through unchanged. Returns
+/// the block-argument-to-scalar mapping used to clone the payload.
+IRMapping materializeScalarOperands(PatternRewriter &rewriter, Location loc,
+ linalg::GenericOp genericOp, ArrayRef iterIndices)
+{
+ Block *body = genericOp.getBody();
+ IRMapping mapping;
+ for (OpOperand *inOperand : genericOp.getDpsInputOperands()) {
+ BlockArgument blockArg = body->getArgument(inOperand->getOperandNumber());
+ Value input = inOperand->get();
+ if (!isa(input.getType())) {
+ mapping.map(blockArg, input);
+ continue;
+ }
+ AffineMap inputMap = genericOp.getMatchingIndexingMap(inOperand);
+ SmallVector inputIndices;
+ for (AffineExpr expr : inputMap.getResults()) {
+ if (auto dimExpr = dyn_cast(expr)) {
+ inputIndices.push_back(iterIndices[dimExpr.getPosition()]);
+ }
+ else {
+ auto constExpr = cast(expr);
+ inputIndices.push_back(
+ arith::ConstantIndexOp::create(rewriter, loc, constExpr.getValue()));
+ }
+ }
+ Value scalar = tensor::ExtractOp::create(rewriter, loc, input, inputIndices);
+ mapping.map(blockArg, scalar);
+ }
+ return mapping;
+}
+
+/// Fold tensor.extract(linalg.generic) by inlining the generic's scalar payload
+/// at the extraction point, for elementwise (all-parallel) generics.
+struct ExtractOfGeneric : public OpRewritePattern {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tensor::ExtractOp extractOp,
+ PatternRewriter &rewriter) const override
+ {
+ auto genericOp = extractOp.getTensor().getDefiningOp();
+ if (!genericOp) {
+ return failure();
+ }
+
+ // Only elementwise generics: every iterator is parallel.
+ if (genericOp.getNumParallelLoops() != genericOp.getNumLoops()) {
+ return failure();
+ }
+
+ // Restrict to small statically shaped results to bound code growth.
+ auto resultType = dyn_cast(extractOp.getTensor().getType());
+ if (!resultType || !resultType.hasStaticShape() ||
+ resultType.getNumElements() > kMaxScalarizedElements) {
+ return failure();
+ }
+
+ // Identify which result of the generic is being extracted and require an
+ // identity indexing map for it, so the iteration indices equal the
+ // extraction indices.
+ auto resultNumber = cast(extractOp.getTensor()).getResultNumber();
+ OpOperand *initOperand = genericOp.getDpsInitOperand(resultNumber);
+ AffineMap outputMap = genericOp.getMatchingIndexingMap(initOperand);
+ if (!outputMap.isIdentity()) {
+ return failure();
+ }
+
+ if (!hasInlinablePayload(genericOp) || !hasScalarizableInputMaps(genericOp)) {
+ return failure();
+ }
+
+ Location loc = extractOp.getLoc();
+ SmallVector iterIndices(extractOp.getIndices());
+ IRMapping mapping = materializeScalarOperands(rewriter, loc, genericOp, iterIndices);
+
+ // Clone the payload, resolving linalg.index to the extraction indices.
+ Block *body = genericOp.getBody();
+ for (Operation &op : body->without_terminator()) {
+ if (auto indexOp = dyn_cast(op)) {
+ mapping.map(indexOp.getResult(), iterIndices[indexOp.getDim()]);
+ continue;
+ }
+ rewriter.clone(op, mapping);
+ }
+
+ auto yieldOp = cast(body->getTerminator());
+ Value result = mapping.lookupOrDefault(yieldOp.getOperand(resultNumber));
+ rewriter.replaceOp(extractOp, result);
+ return success();
+ }
+};
+
+/// Fold tensor.extract(tensor.collapse_shape) for collapses that only drop or
+/// merge unit dimensions (at most one non-unit dimension per reassociation
+/// group), by extracting directly from the source.
+struct ExtractOfCollapseShape : public OpRewritePattern {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tensor::ExtractOp extractOp,
+ PatternRewriter &rewriter) const override
+ {
+ auto collapseOp = extractOp.getTensor().getDefiningOp();
+ if (!collapseOp) {
+ return failure();
+ }
+ auto srcType = collapseOp.getSrcType();
+ if (!srcType.hasStaticShape()) {
+ return failure();
+ }
+
+ // Validate every reassociation group up front (at most one non-unit
+ // source dimension per group): patterns must not mutate the IR on the
+ // failure path, and index constants are created per group below.
+ SmallVector groups = collapseOp.getReassociationIndices();
+ SmallVector nonUnitDims(groups.size(), -1);
+ for (const auto &[groupIdx, group] : llvm::enumerate(groups)) {
+ for (int64_t srcDim : group) {
+ if (srcType.getDimSize(srcDim) != 1) {
+ if (nonUnitDims[groupIdx] != -1) {
+ return failure(); // true merge of two non-unit dims
+ }
+ nonUnitDims[groupIdx] = srcDim;
+ }
+ }
+ }
+
+ Location loc = extractOp.getLoc();
+ SmallVector srcIndices(srcType.getRank());
+
+ Value zero;
+ auto getZero = [&]() {
+ if (!zero) {
+ zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ }
+ return zero;
+ };
+
+ // A rank-0 result means every source dimension is a unit dimension.
+ if (groups.empty()) {
+ for (int64_t dim = 0; dim < srcType.getRank(); ++dim) {
+ srcIndices[dim] = getZero();
+ }
+ }
+ for (const auto &[groupIdx, group] : llvm::enumerate(groups)) {
+ for (int64_t srcDim : group) {
+ srcIndices[srcDim] = (srcDim == nonUnitDims[groupIdx])
+ ? extractOp.getIndices()[groupIdx]
+ : getZero();
+ }
+ }
+
+ rewriter.replaceOpWithNewOp(extractOp, collapseOp.getSrc(), srcIndices);
+ return success();
+ }
+};
+
+/// Fold tensor.extract(tensor.extract_slice) by extracting from the source at
+/// offset + index * stride.
+struct ExtractOfExtractSlice : public OpRewritePattern {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tensor::ExtractOp extractOp,
+ PatternRewriter &rewriter) const override
+ {
+ auto sliceOp = extractOp.getTensor().getDefiningOp();
+ if (!sliceOp) {
+ return failure();
+ }
+
+ Location loc = extractOp.getLoc();
+ int64_t srcRank = sliceOp.getSourceType().getRank();
+
+ // The slice may be rank-reducing: map each source dim to its position in
+ // the result (or none if the dim was dropped).
+ llvm::SmallBitVector droppedDims = sliceOp.getDroppedDims();
+
+ auto materialize = [&](OpFoldResult ofr) -> Value {
+ if (auto val = dyn_cast(ofr)) {
+ return val;
+ }
+ return arith::ConstantIndexOp::create(rewriter, loc,
+ cast(cast(ofr)).getInt());
+ };
+
+ SmallVector srcIndices;
+ unsigned resultDim = 0;
+ for (int64_t dim = 0; dim < srcRank; ++dim) {
+ Value offset = materialize(sliceOp.getMixedOffsets()[dim]);
+ if (droppedDims.test(dim)) {
+ srcIndices.push_back(offset);
+ continue;
+ }
+ Value index = extractOp.getIndices()[resultDim++];
+ Value stride = materialize(sliceOp.getMixedStrides()[dim]);
+ Value scaled = arith::MulIOp::create(rewriter, loc, index, stride);
+ srcIndices.push_back(arith::AddIOp::create(rewriter, loc, offset, scaled));
+ }
+
+ rewriter.replaceOpWithNewOp(extractOp, sliceOp.getSource(), srcIndices);
+ return success();
+ }
+};
+
+} // namespace
+
+namespace catalyst {
+
+#define GEN_PASS_DEF_SCALARIZETENSOREXTRACTSPASS
+#include "Catalyst/Transforms/Passes.h.inc"
+
+struct ScalarizeTensorExtractsPass
+ : public impl::ScalarizeTensorExtractsPassBase {
+ using impl::ScalarizeTensorExtractsPassBase<
+ ScalarizeTensorExtractsPass>::ScalarizeTensorExtractsPassBase;
+
+ void runOnOperation() override
+ {
+ MLIRContext *context = &getContext();
+ RewritePatternSet patterns(context);
+ patterns.add(context);
+ if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
+ signalPassFailure();
+ }
+ }
+};
+
+} // namespace catalyst
diff --git a/mlir/test/Catalyst/ScalarizeTensorExtractsTest.mlir b/mlir/test/Catalyst/ScalarizeTensorExtractsTest.mlir
new file mode 100644
index 0000000000..d41df225da
--- /dev/null
+++ b/mlir/test/Catalyst/ScalarizeTensorExtractsTest.mlir
@@ -0,0 +1,146 @@
+// Copyright 2026 Xanadu Quantum Technologies Inc.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+// http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// RUN: quantum-opt --scalarize-tensor-extracts --canonicalize --split-input-file %s | FileCheck %s
+
+// The extract_slice + collapse_shape + extract chain produced when tracing
+// indexes a 1-D runtime tensor collapses to a single extract.
+
+// CHECK-LABEL: @slice_collapse_extract
+// CHECK-SAME: (%[[ARG:.+]]: tensor<15xf64>)
+// CHECK: %[[C14:.+]] = arith.constant 14 : index
+// CHECK: %[[RES:.+]] = tensor.extract %[[ARG]][%[[C14]]] : tensor<15xf64>
+// CHECK-NOT: tensor.extract_slice
+// CHECK-NOT: tensor.collapse_shape
+// CHECK: return %[[RES]]
+func.func @slice_collapse_extract(%arg0: tensor<15xf64>) -> f64 {
+ %s = tensor.extract_slice %arg0[14] [1] [1] : tensor<15xf64> to tensor<1xf64>
+ %c = tensor.collapse_shape %s [] : tensor<1xf64> into tensor
+ %e = tensor.extract %c[] : tensor
+ return %e : f64
+}
+
+// -----
+
+// Extracting one element of an elementwise linalg.generic inlines the scalar
+// payload; the generic and its tensor.empty become dead.
+
+// CHECK-LABEL: @extract_of_generic
+// CHECK-SAME: (%[[A:.+]]: tensor<2x2xf64>, %[[B:.+]]: tensor<2x2xf64>)
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[EA:.+]] = tensor.extract %[[A]][%[[C0]], %[[C1]]]
+// CHECK-DAG: %[[EB:.+]] = tensor.extract %[[B]][%[[C0]], %[[C1]]]
+// CHECK: %[[RES:.+]] = arith.mulf %[[EA]], %[[EB]] : f64
+// CHECK-NOT: linalg.generic
+// CHECK: return %[[RES]]
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @extract_of_generic(%arg0: tensor<2x2xf64>, %arg1: tensor<2x2xf64>) -> f64 {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %empty = tensor.empty() : tensor<2x2xf64>
+ %prod = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]}
+ ins(%arg0, %arg1 : tensor<2x2xf64>, tensor<2x2xf64>) outs(%empty : tensor<2x2xf64>) {
+ ^bb0(%in0: f64, %in1: f64, %out: f64):
+ %m = arith.mulf %in0, %in1 : f64
+ linalg.yield %m : f64
+ } -> tensor<2x2xf64>
+ %res = tensor.extract %prod[%c0, %c1] : tensor<2x2xf64>
+ return %res : f64
+}
+
+// -----
+
+// Broadcast (rank-0 to 2x2) generics fold to an extract of the rank-0 source.
+
+// CHECK-LABEL: @extract_of_broadcast
+// CHECK-SAME: (%[[A:.+]]: tensor)
+// CHECK: %[[E:.+]] = tensor.extract %[[A]][] : tensor
+// CHECK-NOT: linalg.generic
+// CHECK: return %[[E]]
+#map0 = affine_map<(d0, d1) -> ()>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+func.func @extract_of_broadcast(%arg0: tensor) -> f64 {
+ %c1 = arith.constant 1 : index
+ %empty = tensor.empty() : tensor<2x2xf64>
+ %bcast = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor) outs(%empty : tensor<2x2xf64>) {
+ ^bb0(%in: f64, %out: f64):
+ linalg.yield %in : f64
+ } -> tensor<2x2xf64>
+ %res = tensor.extract %bcast[%c1, %c1] : tensor<2x2xf64>
+ return %res : f64
+}
+
+// -----
+
+// Reductions must not be scalarized: the payload reads the accumulator.
+
+// CHECK-LABEL: @reduction_untouched
+// CHECK: linalg.generic
+// CHECK: tensor.extract
+#map_in = affine_map<(d0) -> (d0)>
+#map_out = affine_map<(d0) -> ()>
+func.func @reduction_untouched(%arg0: tensor<8xf64>) -> f64 {
+ %cst = arith.constant 0.0 : f64
+ %empty = tensor.empty() : tensor
+ %fill = linalg.fill ins(%cst : f64) outs(%empty : tensor) -> tensor
+ %sum = linalg.generic {indexing_maps = [#map_in, #map_out], iterator_types = ["reduction"]}
+ ins(%arg0 : tensor<8xf64>) outs(%fill : tensor) {
+ ^bb0(%in: f64, %acc: f64):
+ %a = arith.addf %in, %acc : f64
+ linalg.yield %a : f64
+ } -> tensor
+ %res = tensor.extract %sum[] : tensor
+ return %res : f64
+}
+
+// -----
+
+// Large tensors must not be scalarized (payload cloning is capped).
+
+// CHECK-LABEL: @large_tensor_untouched
+// CHECK: linalg.generic
+// CHECK: tensor.extract
+#map2 = affine_map<(d0) -> (d0)>
+func.func @large_tensor_untouched(%arg0: tensor<100xf64>) -> f64 {
+ %c5 = arith.constant 5 : index
+ %empty = tensor.empty() : tensor<100xf64>
+ %sq = linalg.generic {indexing_maps = [#map2, #map2], iterator_types = ["parallel"]}
+ ins(%arg0 : tensor<100xf64>) outs(%empty : tensor<100xf64>) {
+ ^bb0(%in: f64, %out: f64):
+ %m = arith.mulf %in, %in : f64
+ linalg.yield %m : f64
+ } -> tensor<100xf64>
+ %res = tensor.extract %sq[%c5] : tensor<100xf64>
+ return %res : f64
+}
+
+// -----
+
+// Rank-reducing extract_slice: index arithmetic offset + i * stride.
+
+// CHECK-LABEL: @strided_slice_extract
+// CHECK-SAME: (%[[ARG:.+]]: tensor<4x6xf64>, %[[I:.+]]: index)
+// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[SCALED:.+]] = arith.muli %[[I]], %[[C2]]
+// CHECK-DAG: %[[COL:.+]] = arith.addi %[[SCALED]], %[[C1]]
+// CHECK: %[[RES:.+]] = tensor.extract %[[ARG]][%[[C2]], %[[COL]]] : tensor<4x6xf64>
+// CHECK: return %[[RES]]
+func.func @strided_slice_extract(%arg0: tensor<4x6xf64>, %i: index) -> f64 {
+ %s = tensor.extract_slice %arg0[2, 1] [1, 3] [1, 2] : tensor<4x6xf64> to tensor<3xf64>
+ %e = tensor.extract %s[%i] : tensor<3xf64>
+ return %e : f64
+}