[Mlir-commits] [mlir] [mlir][Tensor][Func] Rewrite single-dim Tensor return to Scalar (PR #194404)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Apr 27 08:46:46 PDT 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: ioana ghiban (ioghiban)
<details>
<summary>Changes</summary>
Add a Tensor-dialect rewrite that scalarizes function results when a function returns exactly one statically-shaped single-element tensor. The rewrite updates the `func.func` result type from `tensor<...>` to the element type and rewrites each `func.return` by inserting a `tensor.extract` before returning the scalar.
Assisted-by: Codex (refine implementation + tests). I reviewed all code and tests before submission.
## Example
Before:
```mlir
func.func @<!-- -->rank1(%arg0: tensor<1xi64>) -> tensor<1xi64> {
return %arg0 : tensor<1xi64>
}
```
After:
```mlir
func.func @<!-- -->rank1(%arg0: tensor<1xi64>) -> i64 {
%c0 = arith.constant 0 : index
%0 = tensor.extract %arg0[%c0] : tensor<1xi64>
return %0 : i64
}
```
Rank-0 tensors are handled similarly, with `tensor.extract %t[]`.
## Motivation
The `tosa-converter-for-flite` may produce functions that return tensor-wrapped scalars. These are semantically scalar values, but they remain wrapped in `tensor<1xT>` or `tensor<T>` function results.
This becomes a problem later in lowering, in particular for EmitC, where such results eventually bufferize to MemRef arrays and trigger:
error: 'emitc.func' op cannot return array type
## Scope
This rewrite intentionally handles only a narrow case:
the function must return exactly one value,
that value must be a ranked, static shaped, one-element tensor
The rewrite only changes the function result type and the corresponding func.return operations. It does not attempt broader tensor-to-scalar canonicalization across arbitrary ops or block arguments.
## Why ?
EmitC does not handle array return types because it models C/C++ function signatures, and in C/C++ a function cannot return an array by value. Arrays are storage objects, not first-class function return values. In practice, such cases are represented via:
- an out parameter,
- a pointer return,
- or a wrapper struct.
EmitC preserves C/C++ ABI-valid function signatures, so tensor-wrapped scalar returns are addressed before lowering reaches EmitC.
This patch handles this issue at the highest level where the intent is still explicit. At the Tensor dialect level, a single-element tensor result can still be recognized as a scalar value with tensor packaging, so scalarizing it here is simpler and more semantic than deferring the fix to MemRef or EmitC lowering.
## Correctness
The rewrite is semantics-preserving for statically-shaped single-element tensors:
- a tensor with exactly one element can only contain one readable value
- `tensor.extract` retrieves that unique element,
- replacing the function result with the extracted element preserves the meaning of the returned value.
The transformation is restricted to verified `func.func` IR, so the rewritten `func.return` operations remain consistent with the updated function signature. Other tensor types and multi-result functions are left unchanged.
---
Full diff: https://github.com/llvm/llvm-project/pull/194404.diff
6 Files Affected:
- (modified) mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h (+7)
- (modified) mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt (+2)
- (added) mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResultPatterns.cpp (+93)
- (added) mlir/test/Dialect/Tensor/transform-single-dim-to-scalar.mlir (+62)
- (modified) mlir/test/lib/Dialect/Tensor/CMakeLists.txt (+1)
- (modified) mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp (+16-2)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
index 093393eca7436..b522d9429ca51 100644
--- a/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
@@ -103,6 +103,13 @@ using ControlFoldFn = std::function<bool(OpOperand *)>;
void populateRewriteAsConstantPatterns(RewritePatternSet &patterns,
const ControlFoldFn &controlFn);
+/// Populates `patterns` with patterns that rewrite a function returning a
+/// single statically shaped tensor with exactly one element into a function
+/// that returns the tensor element type. The function body is updated by
+/// inserting a `tensor.extract` before each `func.return`.
+void populateScalarizeSingleElementTensorReturnPatterns(
+ RewritePatternSet &patterns);
+
//===----------------------------------------------------------------------===//
// Transform helpers
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
index 33d32c592a844..2a9a5f913233d 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
@@ -9,6 +9,7 @@ add_mlir_dialect_library(MLIRTensorTransforms
ReshapePatterns.cpp
RewriteAsConstant.cpp
RuntimeOpVerification.cpp
+ ScalarizeFunctionResultPatterns.cpp
SwapExtractSliceWithProducerPatterns.cpp
SubsetInsertionOpInterfaceImpl.cpp
@@ -25,6 +26,7 @@ add_mlir_dialect_library(MLIRTensorTransforms
MLIRArithDialect
MLIRArithUtils
MLIRDialectUtils
+ MLIRFuncDialect
MLIRIR
MLIRLinalgDialect
MLIRMemRefDialect
diff --git a/mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResultPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResultPatterns.cpp
new file mode 100644
index 0000000000000..fc54327a24db4
--- /dev/null
+++ b/mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResultPatterns.cpp
@@ -0,0 +1,93 @@
+//===- ScalarizeFunctionResultPatterns.cpp - Scalarize tensor returns -----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/IR/SymbolTable.h"
+
+using namespace mlir;
+using namespace mlir::tensor;
+
+namespace {
+
+struct ScalarizeSingleElementTensorReturnPattern
+ : public OpRewritePattern<func::FuncOp> {
+ using OpRewritePattern<func::FuncOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(func::FuncOp funcOp,
+ PatternRewriter &rewriter) const override {
+ if (funcOp.isDeclaration())
+ return rewriter.notifyMatchFailure(funcOp, "function has no body");
+
+ FunctionType functionType = funcOp.getFunctionType();
+ if (functionType.getNumResults() != 1)
+ return rewriter.notifyMatchFailure(
+ funcOp, "function does not return exactly one value");
+
+ auto tensorType = dyn_cast<RankedTensorType>(functionType.getResult(0));
+ if (!tensorType)
+ return rewriter.notifyMatchFailure(
+ funcOp, "function result is not a ranked tensor");
+ if (!tensorType.hasStaticShape())
+ return rewriter.notifyMatchFailure(
+ funcOp, "function result tensor does not have a static shape");
+ if (tensorType.getNumElements() != 1)
+ return rewriter.notifyMatchFailure(
+ funcOp, "function result tensor does not have exactly one element");
+
+ Operation *symbolTable =
+ SymbolTable::getNearestSymbolTable(funcOp->getParentOp());
+ if (symbolTable && !SymbolTable::symbolKnownUseEmpty(funcOp, symbolTable))
+ return rewriter.notifyMatchFailure(funcOp, "function has symbol users");
+
+ SmallVector<func::ReturnOp> returnOps;
+ for (Block &block : funcOp.getBody()) {
+ auto returnOp = dyn_cast<func::ReturnOp>(block.getTerminator());
+ if (!returnOp)
+ return rewriter.notifyMatchFailure(
+ funcOp, "function has a non-func.return terminator");
+ if (returnOp.getNumOperands() != 1)
+ return rewriter.notifyMatchFailure(
+ returnOp, "return does not have exactly one operand");
+ assert(returnOp.getOperand(0).getType() == tensorType &&
+ "return operand type must match function result type");
+ returnOps.push_back(returnOp);
+ }
+
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPointToStart(&funcOp.getBody().front());
+ Value zeroIndex =
+ arith::ConstantIndexOp::create(rewriter, funcOp.getLoc(), 0);
+
+ SmallVector<Value> zeroIndices(tensorType.getRank(), zeroIndex);
+
+ Type scalarType = tensorType.getElementType();
+ for (func::ReturnOp returnOp : returnOps) {
+ rewriter.setInsertionPoint(returnOp);
+ Value extracted = tensor::ExtractOp::create(
+ rewriter, returnOp.getLoc(), returnOp.getOperand(0), zeroIndices);
+ rewriter.replaceOpWithNewOp<func::ReturnOp>(returnOp, extracted);
+ }
+
+ SmallVector<Type> newResults{scalarType};
+ FunctionType newFunctionType = FunctionType::get(
+ funcOp.getContext(), functionType.getInputs(), newResults);
+ rewriter.modifyOpInPlace(funcOp, [&] { funcOp.setType(newFunctionType); });
+ return success();
+ }
+};
+
+} // namespace
+
+void mlir::tensor::populateScalarizeSingleElementTensorReturnPatterns(
+ RewritePatternSet &patterns) {
+ patterns.add<ScalarizeSingleElementTensorReturnPattern>(
+ patterns.getContext());
+}
diff --git a/mlir/test/Dialect/Tensor/transform-single-dim-to-scalar.mlir b/mlir/test/Dialect/Tensor/transform-single-dim-to-scalar.mlir
new file mode 100644
index 0000000000000..c4b67cb92d8bb
--- /dev/null
+++ b/mlir/test/Dialect/Tensor/transform-single-dim-to-scalar.mlir
@@ -0,0 +1,62 @@
+// RUN: mlir-opt -test-tensor-transform-patterns=test-scalarize-single-element-tensor-return %s | FileCheck %s
+
+// Inserted ExtractOp gets constant folded for rank-0 tensors
+// i.e. no accessed indices
+func.func @rank0() -> tensor<i64> {
+ %0 = arith.constant dense<-1> : tensor<i64>
+ return %0 : tensor<i64>
+}
+// CHECK-LABEL: func.func @rank0
+// CHECK-SAME: -> i64
+// CHECK-NEXT: %[[CST:.*]] = arith.constant -1 : i64
+// CHECK-NEXT: return %[[CST]] : i64
+
+func.func @rank1(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func @rank1
+// CHECK-SAME: %[[SRC:.*]]: tensor<1xi64>) -> i64 {
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[EXT:.*]] = tensor.extract %[[SRC]][%[[C0]]] : tensor<1xi64>
+// CHECK: return %[[EXT]] : i64
+
+func.func @rank2_single_element(%arg0: tensor<1x1xi64>) -> tensor<1x1xi64> {
+ return %arg0 : tensor<1x1xi64>
+}
+// CHECK-LABEL: func.func @rank2_single_element
+// CHECK-SAME: %[[SRC:.*]]: tensor<1x1xi64>) -> i64
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[EXT:.*]] = tensor.extract %[[SRC]][%[[C0]], %[[C0]]] : tensor<1x1xi64>
+// CHECK: return %[[EXT]] : i64
+
+func.func @caller(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ %0 = func.call @callee(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+ return %0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func @caller
+// CHECK-SAME: -> i64
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[CALL:.*]] = call @callee(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+// CHECK: %[[EXT:.*]] = tensor.extract %[[CALL]][%[[C0]]] : tensor<1xi64>
+// CHECK-NEXT: return %[[EXT]] : i64
+
+//===----------------------------------------------------------------------===//
+// Negative tests (must NOT rewrite)
+//===----------------------------------------------------------------------===//
+
+func.func @multiple_elements(%arg0: tensor<2xi64>) -> tensor<2xi64> {
+ return %arg0 : tensor<2xi64>
+}
+// CHECK-LABEL: func.func @multiple_elements
+// CHECK-SAME: -> tensor<2xi64>
+// CHECK-NOT: tensor.extract
+// CHECK-NEXT: return %arg0 : tensor<2xi64>
+
+/// If no @caller, then rewrite applied to @callee
+func.func @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func @callee
+// CHECK-SAME: -> tensor<1xi64>
+// CHECK-NOT: tensor.extract
+// CHECK-NEXT: return %arg0 : tensor<1xi64>
\ No newline at end of file
diff --git a/mlir/test/lib/Dialect/Tensor/CMakeLists.txt b/mlir/test/lib/Dialect/Tensor/CMakeLists.txt
index 28eae8ffb670f..39d3d912eecd4 100644
--- a/mlir/test/lib/Dialect/Tensor/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/Tensor/CMakeLists.txt
@@ -6,6 +6,7 @@ add_mlir_library(MLIRTensorTestPasses
)
mlir_target_link_libraries(MLIRTensorTestPasses PUBLIC
MLIRArithDialect
+ MLIRFuncDialect
MLIRLinalgDialect
MLIRPass
MLIRSCFDialect
diff --git a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
index 687473ebe6d60..050f91f16afc4 100644
--- a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
@@ -11,6 +11,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -32,8 +33,8 @@ struct TestTensorTransforms
TestTensorTransforms(const TestTensorTransforms &pass) : PassWrapper(pass) {}
void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<arith::ArithDialect, scf::SCFDialect, linalg::LinalgDialect,
- transform::TransformDialect>();
+ registry.insert<arith::ArithDialect, func::FuncDialect, scf::SCFDialect,
+ linalg::LinalgDialect, transform::TransformDialect>();
}
StringRef getArgument() const final {
@@ -93,6 +94,11 @@ struct TestTensorTransforms
*this, "test-tracking-listener",
llvm::cl::desc("Test tensor TrackingListener for the transform dialect"),
llvm::cl::init(false)};
+
+ Option<bool> testScalarizeSingleElementTensorReturn{
+ *this, "test-scalarize-single-element-tensor-return",
+ llvm::cl::desc("Test scalarization of single-element tensor returns"),
+ llvm::cl::init(false)};
};
} // namespace
@@ -143,6 +149,12 @@ static void applyFoldExtractFromCollapseShapePatterns(Operation *rootOp) {
(void)applyPatternsGreedily(rootOp, std::move(patterns));
}
+static void applyScalarizeSingleElementTensorReturnPatterns(Operation *rootOp) {
+ RewritePatternSet patterns(rootOp->getContext());
+ tensor::populateScalarizeSingleElementTensorReturnPatterns(patterns);
+ (void)applyPatternsGreedily(rootOp, std::move(patterns));
+}
+
namespace {
/// Base pattern to rewrite a `tensor.collapse_shape -> tensor.extract_slice`.
/// The `tensor.extract_slice` is replaced by a loop or gather operation that
@@ -394,6 +406,8 @@ void TestTensorTransforms::runOnOperation() {
}
if (testFoldExtractFromCollapseShape)
applyFoldExtractFromCollapseShapePatterns(rootOp);
+ if (testScalarizeSingleElementTensorReturn)
+ applyScalarizeSingleElementTensorReturnPatterns(rootOp);
if (testTrackingListener)
if (failed(testTrackingListenerReplacements(rootOp)))
return signalPassFailure();
``````````
</details>
https://github.com/llvm/llvm-project/pull/194404
More information about the Mlir-commits
mailing list