[Mlir-commits] [mlir] 9ed8b1d - [MLGO][EmitC] Scalarize single-element tensor returns (#199686)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jun 16 02:03:15 PDT 2026
Author: ioana ghiban
Date: 2026-06-16T11:03:10+02:00
New Revision: 9ed8b1dff880f4000523e6cab3c726c9fb5a5b7d
URL: https://github.com/llvm/llvm-project/commit/9ed8b1dff880f4000523e6cab3c726c9fb5a5b7d
DIFF: https://github.com/llvm/llvm-project/commit/9ed8b1dff880f4000523e6cab3c726c9fb5a5b7d.diff
LOG: [MLGO][EmitC] Scalarize single-element tensor returns (#199686)
Add an EmitC-owned preparation pass,
`mlgo-scalarize-single-element-tensor-return`, that rewrites private
functions returning a statically-shaped ranked tensor with exactly one
element into functions returning the element type directly.
Assisted-by: Codex (refine implementation + tests). I reviewed all code
and tests before submission.
## Example
Before:
```mlir
func.func private @rank1(%arg0: tensor<1xi64>) -> tensor<1xi64> {
return %arg0 : tensor<1xi64>
}
```
After:
```mlir
func.func private @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[]`, which
canonicalization/constant-folding can simplify further.
## Motivation
`tosa-converter-for-flite` may produce functions that return
tensor-wrapped scalars. These values are semantically scalars, but
remain packaged as single-element tensors in function signatures. This
becomes problematic later in lowering, especially for EmitC, where such
returns eventually bufferize to MemRef arrays and fail with:
```bash
error: 'emitc.func' op cannot return array type
```
This patch fixes the issue while the IR is still in `func`/`tensor`
form, where the scalar intent remains explicit and the transform can
stay semantics-driven.
## Scope
The transform is intentionally conservative: it only handles private
definitions with one single-element ranked tensor result, and it blocks
on public callers, non-call symbol users, unsupported callers, and
recursive cycles.
## Analysis
The pass runs at `ModuleOp` scope because legality depends on transitive
private call users. It first computes module-level scalarization
analysis from a `SymbolUserMap` snapshot, then rewrites only the
precomputed rewritable functions.
## Rewrite
Each rewritten function gets tensor.extract inserted before
`func.return`, its result type is updated with `funcOp.setType(...)`,
and direct `func.call` users are adjusted.
Unchanged private callers that still need the original tensor type are
adapted with `tensor.from_elements`.
Added:
mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResult.cpp
mlir/test/Dialect/Tensor/scalarize-single-elem-return.mlir
Modified:
mlir/include/mlir/Dialect/Tensor/Transforms/Passes.h
mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td
mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.h
index 03cb30bc07357..c13294635610a 100644
--- a/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.h
@@ -18,7 +18,8 @@ namespace tensor {
// Passes
//===----------------------------------------------------------------------===//
-/// Creates an instance of the `tensor` subset folding pass.
+/// Create instances of `tensor` transformation passes.
+#define GEN_PASS_DECL_SCALARIZESINGLEELEMENTTENSORRETURNPASS
#define GEN_PASS_DECL_FOLDTENSORSUBSETOPSPASS
#include "mlir/Dialect/Tensor/Transforms/Passes.h.inc"
diff --git a/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td
index 62f467f2513ac..4dc5ab91b4a41 100644
--- a/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td
@@ -26,4 +26,18 @@ def FoldTensorSubsetOpsPass : Pass<"fold-tensor-subset-ops"> {
];
}
+def ScalarizeSingleElementTensorReturnPass
+ : Pass<"scalarize-single-element-tensor-return", "ModuleOp"> {
+ let summary = "Scalarize private functions that return single-element tensor";
+ let description = [{
+ Rewrites private functions returning exactly one statically shaped ranked
+ tensor of exactly one element into functions returning the tensor element
+ type, provided all transitive users can be updated safely.
+ }];
+ let dependentDialects = [
+ "arith::ArithDialect", "func::FuncDialect",
+ "tensor::TensorDialect"
+ ];
+}
+
#endif // MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES
diff --git a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
index 33d32c592a844..78e64c82544e0 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
@@ -6,6 +6,7 @@ add_mlir_dialect_library(MLIRTensorTransforms
ExtractSliceFromReshapeUtils.cpp
FoldTensorSubsetOps.cpp
IndependenceTransforms.cpp
+ ScalarizeFunctionResult.cpp
ReshapePatterns.cpp
RewriteAsConstant.cpp
RuntimeOpVerification.cpp
@@ -25,6 +26,7 @@ add_mlir_dialect_library(MLIRTensorTransforms
MLIRArithDialect
MLIRArithUtils
MLIRDialectUtils
+ MLIRFuncDialect # Dependency of scalarize-single-element-tensor-return
MLIRIR
MLIRLinalgDialect
MLIRMemRefDialect
diff --git a/mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResult.cpp b/mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResult.cpp
new file mode 100644
index 0000000000000..60385d0c8b50d
--- /dev/null
+++ b/mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResult.cpp
@@ -0,0 +1,318 @@
+//===- ScalarizeFunctionResult.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/Passes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/SymbolTable.h"
+#include "llvm/ADT/DenseMap.h"
+
+namespace mlir {
+namespace tensor {
+#define GEN_PASS_DEF_SCALARIZESINGLEELEMENTTENSORRETURNPASS
+#include "mlir/Dialect/Tensor/Transforms/Passes.h.inc"
+} // namespace tensor
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+// Analysis state used by the memoized DFS below. It classifies whether a
+// candidate private function can be scalarized. The rewrite only consumes the
+// functions classified as `rewritable`.
+enum class ScalarizationState {
+ // No DFS classification has been computed for this function yet.
+ unknown,
+ // The function is currently on the recursive DFS stack. Re-entering this
+ // state means a cycle was found, which this pass conservatively blocks.
+ visiting,
+ // The function cannot be scalarized safely because either it is not locally
+ // eligible or one of its transitive users blocks the rewrite.
+ blocked,
+ // The function is locally eligible and all transitive users considered by
+ // this pass can be updated consistently.
+ rewritable
+};
+
+// Info analyzed to decide scalarizing a locally eligible function: a private
+// definition with exactly one statically-shaped ranked tensor result containing
+// one element. Functions not meeting these criteria are not represented here,
+// although they may still appear in the broader module analysis as callers or
+// blockers.
+struct ScalarizableFunctionInfo {
+ RankedTensorType tensorType;
+ SmallVector<func::ReturnOp> returnOps;
+};
+
+// Returns per-function scalarization info when this function is locally
+// eligible, i.e. it is a private definition with one statically-shaped ranked
+// tensor result containing exactly one element.
+static FailureOr<ScalarizableFunctionInfo>
+getScalarizableFunctionInfoIfEligible(func::FuncOp func) {
+ if (func.isDeclaration() || !func.isPrivate())
+ return failure();
+
+ FunctionType functionType = func.getFunctionType();
+ if (functionType.getNumResults() != 1)
+ return failure();
+
+ auto tensorType = dyn_cast<RankedTensorType>(functionType.getResult(0));
+ if (!tensorType || !tensorType.hasStaticShape() ||
+ tensorType.getNumElements() != 1)
+ return failure();
+
+ ScalarizableFunctionInfo sfi{tensorType, {}};
+ for (Block &block : func.getBody()) {
+ auto returnOp = dyn_cast<func::ReturnOp>(block.getTerminator());
+ if (returnOp)
+ sfi.returnOps.push_back(returnOp);
+ }
+
+ // While FuncOp is guaranteed to contain terminator ops, there is no guarantee
+ // that it will contain ReturnOp(s). Hence the check.
+ if (sfi.returnOps.empty())
+ return failure();
+ return sfi;
+}
+
+struct ScalarizationAnalysis {
+ explicit ScalarizationAnalysis(SymbolUserMap &userMap) : userMap(userMap) {}
+
+ SymbolUserMap &userMap;
+ DenseSet<func::FuncOp> moduleFunctions;
+ DenseMap<func::FuncOp, ScalarizableFunctionInfo> candidateInfos;
+ DenseMap<func::FuncOp, SmallVector<func::CallOp>> callUsers;
+ DenseMap<func::FuncOp, ScalarizationState> states;
+ // DFS completion order for rewritable functions. The rewrite phase walks
+ // this list in reverse so callees are rewritten before their call sites.
+ SmallVector<func::FuncOp> rewriteOrder;
+};
+
+// Runs the memoized DFS that classifies one candidate function by walking its
+// transitive private call users.
+static ScalarizationState
+computeScalarizationState(func::FuncOp func, ScalarizationAnalysis &analysis) {
+ auto [it, inserted] =
+ analysis.states.try_emplace(func, ScalarizationState::unknown);
+ if (!inserted) {
+ // Conservatively reject recursive cycles instead of reasoning about SCCs.
+ if (it->second == ScalarizationState::visiting)
+ return ScalarizationState::blocked;
+ return it->second;
+ }
+
+ // Starting from one locally-eligible private function, walk its symbol users
+ // upward through private callers. Memoization avoids re-traversing the same
+ // subgraph from the outer linear scan, and early blocking truncates the DFS
+ // as soon as a public/unsupported user is found.
+ auto setBlocked = [&] {
+ analysis.states[func] = ScalarizationState::blocked;
+ return ScalarizationState::blocked;
+ };
+ auto setRewritable = [&] {
+ analysis.states[func] = ScalarizationState::rewritable;
+ analysis.rewriteOrder.push_back(func);
+ return ScalarizationState::rewritable;
+ };
+
+ if (!analysis.candidateInfos.contains(func))
+ return setBlocked();
+ analysis.states[func] = ScalarizationState::visiting;
+
+ SmallVector<func::CallOp> directCallUsers;
+ for (Operation *user : analysis.userMap.getUsers(func.getOperation())) {
+ auto directCall = dyn_cast<func::CallOp>(user);
+ // Non-call symbol uses, such as func.constant, prevent updating all users
+ // consistently, so the current function stays blocked.
+ if (!directCall)
+ return setBlocked();
+ directCallUsers.push_back(directCall);
+
+ func::FuncOp caller = directCall->getParentOfType<func::FuncOp>();
+ // A direct call user outside any func.func can still be updated in place,
+ // but it terminates the DFS because there is no caller signature to
+ // analyze or rewrite transitively.
+ if (!caller)
+ continue;
+ // Since `getScalarizableFunctionInfoIfEligible` has already categorized
+ // every direct func.func in the current module, any direct caller must
+ // already appear in the analysis tables.
+ assert(analysis.moduleFunctions.contains(caller) &&
+ "Caller of private function is not a direct function in the module");
+
+ // Public and external callers keep the current function blocked because the
+ // pass cannot rewrite every visible call boundary.
+ if (caller.isPublic())
+ return setBlocked();
+
+ assert(!caller.isExternal() && "Caller of private function is external.");
+
+ // A private non-candidate caller can absorb the scalarized call by
+ // reboxing the scalar result back into a single-element tensor.
+ if (!analysis.candidateInfos.contains(caller))
+ continue;
+
+ if (computeScalarizationState(caller, analysis) !=
+ ScalarizationState::rewritable)
+ return setBlocked();
+ }
+
+ analysis.callUsers.try_emplace(func, std::move(directCallUsers));
+ return setRewritable();
+}
+
+// Builds the module-level analysis state used by the rewrite phase.
+static void computeScalarizationAnalysis(ModuleOp module,
+ ScalarizationAnalysis &analysis) {
+ // First collect all direct functions and the subset that is locally
+ // eligible.
+ for (func::FuncOp func : module.getOps<func::FuncOp>()) {
+ analysis.moduleFunctions.insert(func);
+ FailureOr<ScalarizableFunctionInfo> sfi =
+ getScalarizableFunctionInfoIfEligible(func);
+ if (succeeded(sfi))
+ analysis.candidateInfos.try_emplace(func, std::move(*sfi));
+ }
+ // Then run the memoized DFS for candidate roots.
+ for (func::FuncOp func : module.getOps<func::FuncOp>())
+ if (analysis.candidateInfos.contains(func))
+ (void)computeScalarizationState(func, analysis);
+}
+
+// Rewrites one function that has already been proven rewritable and updates
+// the direct call users cached before any IR mutation started.
+static void rewriteScalarizableFunction(func::FuncOp func,
+ const ScalarizableFunctionInfo &sfi,
+ ArrayRef<func::CallOp> directCalls,
+ RewriterBase &rewriter) {
+ OpBuilder::InsertionGuard guard(rewriter);
+ // Scalarize the unique element before each return.
+ RankedTensorType tensorType = sfi.tensorType;
+ SmallVector<Value> zeroIndices;
+ if (tensorType.getRank() != 0) {
+ rewriter.setInsertionPointToStart(&func.getBody().front());
+ Value zero = arith::ConstantIndexOp::create(rewriter, func.getLoc(), 0);
+ zeroIndices.assign(tensorType.getRank(), zero);
+ }
+
+ Type scalarType = tensorType.getElementType();
+ for (func::ReturnOp funcReturn : sfi.returnOps) {
+ assert(funcReturn.getNumOperands() == 1 &&
+ "func.return must have exactly one operand");
+ assert(funcReturn.getOperand(0).getType() == tensorType &&
+ "func.return operand type must match the function result type");
+ rewriter.setInsertionPoint(funcReturn);
+ Value scalar = rewriter.createOrFold<tensor::ExtractOp>(
+ funcReturn.getLoc(), funcReturn.getOperand(0), zeroIndices);
+ rewriter.replaceOpWithNewOp<func::ReturnOp>(funcReturn, scalar);
+ }
+
+ FunctionType functionType = func.getFunctionType();
+ // Update the function type:
+ // This is a 1-result to 1-result type replacement, so the existing result
+ // attribute dictionary remains attached to result #0 without reordering,
+ // hence the rewrite is done directly without function_interface methods.
+ func.setType(FunctionType::get(func.getContext(), functionType.getInputs(),
+ TypeRange{scalarType}));
+ // Fix direct call users that were recorded during analysis.
+ for (func::CallOp directCall : directCalls) {
+ rewriter.setInsertionPoint(directCall);
+ func::CallOp newDirectCall = func::CallOp::create(
+ rewriter, directCall.getLoc(), func, directCall.getOperands());
+ newDirectCall->setAttrs(directCall->getAttrs());
+
+ if (!directCall.getResult(0).use_empty()) {
+ Value wrappedResult = tensor::FromElementsOp::create(
+ rewriter, directCall.getLoc(), tensorType,
+ ValueRange{newDirectCall.getResult(0)});
+ rewriter.replaceOp(directCall, wrappedResult);
+ } else {
+ rewriter.eraseOp(directCall);
+ }
+ }
+}
+
+/// Drives the complete module-level transform: analyze the original module,
+/// determine which private functions can be scalarized safely, then
+/// rewrite only that precomputed set.
+///
+/// BEFORE (both callee and private caller rewritten)
+/// private callee(x : tensor<1xT>) -> tensor<1xT> { return x }
+/// private caller(x : tensor<1xT>) -> tensor<1xT> {
+/// y = call callee(x)
+/// return y
+/// }
+///
+/// AFTER
+/// private callee(x : tensor<1xT>) -> T {
+/// return tensor.extract x[0]
+/// }
+/// private caller(x : tensor<1xT>) -> T {
+/// y = call callee(x)
+/// return y
+/// }
+///
+/// BEFORE (callee rewritten, unchanged private caller reboxes)
+/// private callee(x : tensor<1xT>) -> tensor<1xT> { return x }
+/// private caller(x : tensor<1xT>, z : tensor<1xT>) -> tensor<1xT> {
+/// y = call callee(x)
+/// r = tensor_op(y, z)
+/// return r
+/// }
+///
+/// AFTER
+/// private callee(x : tensor<1xT>) -> T {
+/// return tensor.extract x[0]
+/// }
+/// private caller(x : tensor<1xT>, z : tensor<1xT>) -> tensor<1xT> {
+/// y = call callee(x)
+/// y_boxed = tensor.from_elements y
+/// r = tensor_op(y_boxed, z)
+/// return r
+/// }
+static LogicalResult
+ScalarizeSingleElementTensorReturns(ModuleOp module, RewriterBase &rewriter) {
+ // The transform depends on module-scoped `SymbolUserMap` state and on a
+ // precomputed DFS result, so it runs as a direct module analysis + rewrite
+ // instead of exposing a pattern-testing / transform-dialect pattern path.
+ SymbolTableCollection symbolTable;
+ // Take a snapshot of symbol users for the original module. This is safe
+ // because it is consulted only during `computeScalarizationAnalysis`, before
+ // rewriting starts, and this pass invocation has exclusive access to the
+ // current module.
+ SymbolUserMap userMap(symbolTable, module);
+ ScalarizationAnalysis analysis(userMap);
+ computeScalarizationAnalysis(module, analysis);
+
+ for (func::FuncOp func : llvm::reverse(analysis.rewriteOrder)) {
+ const ScalarizableFunctionInfo &sfi =
+ analysis.candidateInfos.find(func)->second;
+ ArrayRef<func::CallOp> directCalls = analysis.callUsers.find(func)->second;
+ rewriteScalarizableFunction(func, sfi, directCalls, rewriter);
+ }
+
+ return success();
+}
+
+struct ScalarizeSingleElementTensorReturnPass
+ : public tensor::impl::ScalarizeSingleElementTensorReturnPassBase<
+ ScalarizeSingleElementTensorReturnPass> {
+ using Base::Base;
+
+ void runOnOperation() override {
+ IRRewriter rewriter(&getContext());
+ if (failed(ScalarizeSingleElementTensorReturns(getOperation(), rewriter)))
+ signalPassFailure();
+ }
+};
+
+} // namespace
diff --git a/mlir/test/Dialect/Tensor/scalarize-single-elem-return.mlir b/mlir/test/Dialect/Tensor/scalarize-single-elem-return.mlir
new file mode 100644
index 0000000000000..0a80155a8ae6c
--- /dev/null
+++ b/mlir/test/Dialect/Tensor/scalarize-single-elem-return.mlir
@@ -0,0 +1,294 @@
+// RUN: mlir-opt -scalarize-single-element-tensor-return -split-input-file %s | FileCheck %s
+
+/// Positive tests: functions with no users updated.
+
+func.func private @rank0() -> tensor<i64> {
+ %0 = arith.constant dense<-1> : tensor<i64>
+ return %0 : tensor<i64>
+}
+/// Inserted ExtractOp gets folded for rank-0 tensors.
+// CHECK-LABEL: func.func private @rank0
+// CHECK-SAME: -> i64
+// CHECK: %[[CST:.*]] = arith.constant -1 : i64
+// CHECK: return %[[CST]] : i64
+
+func.func private @rank1(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @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 private @rank2_single_element(%arg0: tensor<1x1xi64>)
+ -> tensor<1x1xi64> {
+ return %arg0 : tensor<1x1xi64>
+}
+// CHECK-LABEL: func.func private @rank2_single_element
+// CHECK-SAME: %[[SRC:.*]]: tensor<1x1xi64>) -> i64
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[EXT:.*]] = tensor.extract %[[SRC]][%[[C0]], %[[C0]]]
+// CHECK-SAME: : tensor<1x1xi64>
+// CHECK: return %[[EXT]] : i64
+
+func.func private @non_return_terminator(%arg0: tensor<1xi64>, %cond: i1)
+ -> tensor<1xi64> {
+ cf.cond_br %cond, ^bb1, ^bb2
+^bb1:
+ cf.br ^bb3(%arg0 : tensor<1xi64>)
+^bb2:
+ %0 = tensor.empty() : tensor<1xi64>
+ cf.br ^bb3(%0 : tensor<1xi64>)
+^bb3(%result: tensor<1xi64>):
+ return %result : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @non_return_terminator(
+// CHECK-SAME: %[[ARG0:.*]]: tensor<1xi64>,
+// CHECK-SAME: %[[ARG1:.*]]: i1) -> i64 {
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: cf.cond_br %[[ARG1]], ^bb1, ^bb2
+// CHECK: ^bb1:
+// CHECK: cf.br ^bb3(%[[ARG0]] : tensor<1xi64>)
+// CHECK: ^bb2:
+// CHECK: %[[EMPTY_0:.*]] = tensor.empty() : tensor<1xi64>
+// CHECK: cf.br ^bb3(%[[EMPTY_0]] : tensor<1xi64>)
+// CHECK: ^bb3(%[[VAL_0:.*]]: tensor<1xi64>):
+// CHECK: %[[EXT:.*]] = tensor.extract %[[VAL_0]][%[[C0]]]
+// CHECK-SAME: : tensor<1xi64>
+// CHECK: return %[[EXT]] : i64
+
+/// Positive tests: both caller and callee updated.
+
+func.func private @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @callee
+// 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 private @caller(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ %0 = call @callee(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+ return %0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @caller
+// CHECK-SAME: -> i64
+// CHECK: %[[CALL:.*]] = call @callee(%arg0) : (tensor<1xi64>) -> i64
+// CHECK: return %[[CALL]] : i64
+
+// -----
+
+/// Positive tests: only callee updated.
+
+/// Private non-scalarizeable callers still allow rewritten callees.
+/// The caller's signature remains unchanged.
+func.func private @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @callee
+// 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 private @non_scalarizable_caller_multiple_dimensions_return(
+ %arg0: tensor<1xi64>) -> tensor<2xi64> {
+ %0 = call @callee(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+ %1 = tensor.empty() : tensor<2xi64>
+ return %1 : tensor<2xi64>
+}
+// CHECK-LABEL: func.func private @non_scalarizable_caller_multiple_dimensions_return
+// CHECK-SAME: -> tensor<2xi64>
+// CHECK: %[[CALL:.*]] = call @callee(%arg0) : (tensor<1xi64>) -> i64
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<2xi64>
+// CHECK: return %[[EMPTY]] : tensor<2xi64>
+
+// -----
+
+/// Positive tests: only callee updated, with callsite reboxing.
+///
+/// The private caller is not scalarizable because it has multiple results,
+/// so the rewritten callee result is reboxed at the callsite.
+func.func private @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @callee
+// 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 private @non_scalarizeable_caller_with_multiple_returns(
+ %arg0: tensor<1xi64>) -> (tensor<1xi64>, i64) {
+ %0 = call @callee(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+ %1 = arith.constant 7 : i64
+ return %0, %1 : tensor<1xi64>, i64
+}
+
+// CHECK-LABEL: func.func private @non_scalarizeable_caller_with_multiple_returns
+// CHECK-SAME: -> (tensor<1xi64>, i64)
+// CHECK-DAG: %[[CST:.*]] = arith.constant 7 : i64
+// CHECK-DAG: %[[CALL:.*]] = call @callee(%arg0) : (tensor<1xi64>) -> i64
+// CHECK-DAG: %[[BOX:.*]] = tensor.from_elements %[[CALL]] : tensor<1xi64>
+// CHECK: return %[[BOX]], %[[CST]] : tensor<1xi64>, i64
+
+// -----
+
+/// Positive tests: only callee updated, with callsite reboxing.
+///
+/// Some existing uses in the caller still expect the old tensor type.
+func.func private @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @callee
+// 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 private @non_scalarizeable_caller_expecting_tensor_type(
+ %arg0: tensor<1xi64>, %arg1: tensor<1xi64>) -> (tensor<1xi64>) {
+ %0 = call @callee(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+ %init = tensor.empty() : tensor<1xi64>
+ %1 = linalg.map { arith.addi } ins(%0, %arg1 : tensor<1xi64>, tensor<1xi64>)
+ outs(%init : tensor<1xi64>)
+ return %1 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @non_scalarizeable_caller_expecting_tensor_type
+// CHECK-SAME: %[[ARG0:.*]]: tensor<1xi64>,
+// CHECK-SAME: %[[ARG1:.*]]: tensor<1xi64>) -> i64
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[CALL:.*]] = call @callee(%[[ARG0]]) : (tensor<1xi64>) -> i64
+// CHECK: %[[BOX:.*]] = tensor.from_elements %[[CALL]] : tensor<1xi64>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1xi64>
+// CHECK: %[[MAP:.*]] = linalg.map { arith.addi
+// CHECK: ins(%[[BOX]], %[[ARG1]] : tensor<1xi64>, tensor<1xi64>)
+// CHECK: outs(%[[EMPTY]] : tensor<1xi64>)
+// CHECK: %[[EXT:.*]] = tensor.extract %[[MAP]][%[[C0]]] : tensor<1xi64>
+// CHECK: return %[[EXT]] : i64
+
+// -----
+
+/// Positive tests: callee updated for a direct call user that is not enclosed
+/// by a func.func.
+
+func.func private @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @callee
+// 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
+
+%0 = tensor.empty() : tensor<1xi64>
+%1 = func.call @callee(%0) : (tensor<1xi64>) -> tensor<1xi64>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1xi64>
+// CHECK: %[[CALL:.*]] = func.call @callee(%[[EMPTY]]) : (tensor<1xi64>) -> i64
+
+// -----
+
+/// Positive tests: only caller updated.
+
+/// `func.constant` is a non-call symbol user of the target, so it blocks
+/// scalarization.
+func.func private @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @callee
+// CHECK-SAME: -> tensor<1xi64>
+// CHECK-NOT: tensor.extract
+// CHECK: return %arg0 : tensor<1xi64>
+
+func.func private @caller_using_constant(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ %fn = func.constant @callee : (tensor<1xi64>) -> tensor<1xi64>
+ %result = func.call_indirect %fn(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+ return %result : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @caller_using_constant
+// CHECK-SAME: %[[SRC:.*]]: tensor<1xi64>) -> i64
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[CONST:.*]] = constant @callee : (tensor<1xi64>) -> tensor<1xi64>
+// CHECK: %[[CALL:.*]] = call_indirect %[[CONST]](%[[SRC]])
+// CHECK-SAME: : (tensor<1xi64>) -> tensor<1xi64>
+// CHECK: %[[EXT:.*]] = tensor.extract %[[CALL]][%[[C0]]] : tensor<1xi64>
+// CHECK: return %[[EXT]] : i64
+
+// -----
+
+/// Negative tests: neither callee nor caller updated.
+
+func.func private @multiple_elements(%arg0: tensor<2xi64>) -> tensor<2xi64> {
+ return %arg0 : tensor<2xi64>
+}
+// CHECK-LABEL: func.func private @multiple_elements
+// CHECK-SAME: -> tensor<2xi64>
+// CHECK-NOT: tensor.extract
+// CHECK: return %arg0 : tensor<2xi64>
+
+func.func private @dynamic_shape(%arg0: tensor<?xi64>) -> tensor<?xi64> {
+ return %arg0 : tensor<?xi64>
+}
+// CHECK-LABEL: func.func private @dynamic_shape
+// CHECK-SAME: -> tensor<?xi64>
+// CHECK-NOT: tensor.extract
+// CHECK: return %arg0 : tensor<?xi64>
+
+func.func private @unranked(%arg0: tensor<*xi64>) -> tensor<*xi64> {
+ return %arg0 : tensor<*xi64>
+}
+// CHECK-LABEL: func.func private @unranked
+// CHECK-SAME: -> tensor<*xi64>
+// CHECK-NOT: tensor.extract
+// CHECK: return %arg0 : tensor<*xi64>
+
+func.func @public_function(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func @public_function
+// CHECK-SAME: -> tensor<1xi64>
+// CHECK-NOT: tensor.extract
+// CHECK: return %arg0 : tensor<1xi64>
+
+func.func private @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @callee
+// CHECK-SAME: -> tensor<1xi64>
+// CHECK-NOT: tensor.extract
+// CHECK: return %arg0 : tensor<1xi64>
+
+func.func @public_caller(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ %0 = call @callee(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+ return %0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func @public_caller
+// CHECK-SAME: -> tensor<1xi64>
+// CHECK: %[[CALL:.*]] = call @callee(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+// CHECK-NOT: tensor.extract
+// CHECK: return %[[CALL]] : tensor<1xi64>
+
+// -----
+
+/// Recursive cycle - the DFS marks the cycle conservatively as blocked,
+/// so neither function is scalarized.
+
+func.func private @recursive_a(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ %0 = call @recursive_b(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+ return %0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @recursive_a
+// CHECK-SAME: -> tensor<1xi64>
+// CHECK: %[[CALL:.*]] = call @recursive_b(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+// CHECK: return %[[CALL]] : tensor<1xi64>
+
+func.func private @recursive_b(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ %0 = call @recursive_a(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+ return %0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @recursive_b
+// CHECK-SAME: -> tensor<1xi64>
+// CHECK: %[[CALL:.*]] = call @recursive_a(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+// CHECK: return %[[CALL]] : tensor<1xi64>
More information about the Mlir-commits
mailing list