[Mlir-commits] [mlir] [MLGO][EmitC] Scalarize single-element tensor returns (PR #199686)
ioana ghiban
llvmlistbot at llvm.org
Tue May 26 07:07:37 PDT 2026
https://github.com/ioghiban created https://github.com/llvm/llvm-project/pull/199686
Add an EmitC-owned interprocedural preparation pass, `mlgo-scalarize-single-element-tensor-return`, that scalarizes private function results when a function returns exactly one statically-shaped, single-element ranked tensor, e.g. `tensor<T>`, `tensor<1xT>`, `tensor<1x1xT>` etc.
The pass 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 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 can 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 rewrite is intentionally conservative:
- the function must be a private definition
- it must return exactly one value
- that value must be a statically-shaped ranked tensor with exactly one element
The transform is interprocedural, but it only rewrites:
- the `func.func` result type
- the `func.return` operations in the rewritten function
- direct `func.call` users that must be adapted to the new result type
It does not attempt general tensor-to-scalar canonicalization across arbitrary ops, block arguments, or indirect call targets.
## Analysis
This rewrite cannot be decided locally on a single `func.func`. Whether a callee can be scalarized depends on all of its transitive users.
The implementation therefore runs as a module-level pass and performs analysis in two stages:
1. A linear outer loop scans all direct `func.func` ops in the module and caches `ScalarizableFunctionInfo` for locally-eligible candidates.
2. For each candidate, a memoized DFS walks its transitive private call users to determine whether the function is globally rewritable.
The DFS uses per-function states such as `unknown`, `visiting`, `blocked`, and `rewritable`, so each function’s transitive user graph is traversed at most once by the DFS. If a function has already been visited through an earlier DFS, the outer loop reuses the memoized result instead of re-traversing it.
The DFS is also truncated conservatively:
- if a public caller is found, scalarization is blocked
- if a non-call symbol user is found, scalarization is blocked
- if a caller cannot be updated, scalarization is blocked
- if a private caller is not itself a scalarization candidate, the DFS stops there, but the callee may still be rewritten
To avoid repeated symbol-table walks, the pass builds a `SymbolUserMap` once per module and uses that snapshot for the entire analysis/rewrite. This is preferable to repeatedly querying symbol tables or rediscovering users per function. The snapshot is taken before any rewriting starts, so analysis is performed against a stable view of the original symbol-use graph.
During DFS, the pass also caches the original direct `func.call` users for each rewritable function. Those cached call sites are then reused during rewriting.
## Rewrite
After analysis, only functions proven `rewritable` are transformed.
The pass does not invoke `applyPatternsGreedily` over the module. Instead, it rewrites functions directly from the analyzed result. The DFS memoization order is recorded, and the final rewrite walks that order in reverse so callees are rewritten before candidate callers. This keeps the cached original call sites valid while signatures are updated.
For each rewritten function, the transform:
- inserts `tensor.extract` before each `func.return`
- changes the function result type from the single-element tensor type to its element type
- updates cached direct `func.call` users to the new scalar result type
If a caller is not rewritten but still expects the old tensor result type, the call result is reboxed with `tensor.from_elements` so existing tensor-typed uses remain valid.
## Why `ModuleOp`
Although the rewrite is applied to `func.func`, the legality decision is module-global.
We need module-level state for:
- the `SymbolUserMap` snapshot
- cached `ScalarizableFunctionInfo`
- memoized DFS state
- cached original direct call users
- the set of functions proven rewritable
- the final rewrite order
That information is shared across functions, so the natural root operation is `ModuleOp`. A local greedy rewrite on individual `func.func` ops is not a good fit here because the transform depends on interprocedural information and coordinated updates to related call users.
## Reboxing Strategy
When a callee is rewritten from:
```mlir
-> tensor<...1 element...>
```
to
```mlir
-> element-type
```
some private callers may remain unchanged. In those cases, existing uses in the caller may still require the original tensor type. The transform handles this by rebuilding the single-element tensor with `tensor.from_elements`.
Conceptually, this is the inverse of the `tensor.extract` inserted in the callee:
- callee rewrite: tensor -> scalar
- unchanged caller adaptation: scalar -> tensor
This allows scalarization to propagate through private call chains without requiring every intermediate caller to be rewritten.
## Function Type Update
The function type is updated with `funcOp.setType(...)`.
This is preferred over `function_interface_impl::eraseFunctionResults` plus `function_interface_impl::insertFunctionResults` because this transform is constrained to a single-result-to-single-result replacement. In that case, there is no result reordering, so preserving attribute order explicitly is unnecessary, and `setType(...)` is simpler and clearer.
## Limitations
This transform is intentionally conservative.
In particular, when the DFS detects a cycle, none of the cycle members are rewritten. Recursive SCCs are therefore left unchanged for now.
Similarly, functions with public users, non-call symbol users, unsupported callers, or non-eligible result types are left untouched.
The implementation is a direct module pass, not a `matchAndRewrite`-driven greedy pattern, so it does not provide a transform-dialect pattern-testing path.
## Correctness
The rewrite is semantics-preserving for statically-shaped single-element ranked tensors:
- such a tensor contains exactly one value
- `tensor.extract` retrieves that unique value
- replacing the function result with that scalar preserves the meaning of the returned result
For unchanged callers that still require the tensor type, `tensor.from_elements` reconstructs the original single-element tensor value, preserving type correctness and behavior at the call boundary.
>From fcf2b97d761ba6f5fcab84effa58c7f33598e917 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Tue, 26 May 2026 14:45:41 +0200
Subject: [PATCH] [MLGO][EmitC] Scalarize single-element tensor returns
---
.../mlir/Dialect/EmitC/Transforms/Passes.h | 1 +
.../mlir/Dialect/EmitC/Transforms/Passes.td | 14 +
.../Dialect/EmitC/Transforms/CMakeLists.txt | 4 +
.../MLGOScalarizeFunctionResult.cpp | 309 ++++++++++++++++++
.../mlgo-scalarize-single-elem-return.mlir | 256 +++++++++++++++
5 files changed, 584 insertions(+)
create mode 100644 mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp
create mode 100644 mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir
diff --git a/mlir/include/mlir/Dialect/EmitC/Transforms/Passes.h b/mlir/include/mlir/Dialect/EmitC/Transforms/Passes.h
index 1af4aa06fa811..cfa081b41177d 100644
--- a/mlir/include/mlir/Dialect/EmitC/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/EmitC/Transforms/Passes.h
@@ -15,6 +15,7 @@ namespace mlir {
namespace emitc {
#define GEN_PASS_DECL_FORMEXPRESSIONSPASS
+#define GEN_PASS_DECL_MLGOSCALARIZESINGLEELEMENTTENSORRETURNPASS
#define GEN_PASS_DECL_WRAPFUNCINCLASSPASS
#include "mlir/Dialect/EmitC/Transforms/Passes.h.inc"
diff --git a/mlir/include/mlir/Dialect/EmitC/Transforms/Passes.td b/mlir/include/mlir/Dialect/EmitC/Transforms/Passes.td
index 40ecef33448d7..08b994c6cdfaf 100644
--- a/mlir/include/mlir/Dialect/EmitC/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/EmitC/Transforms/Passes.td
@@ -20,6 +20,20 @@ def FormExpressionsPass : Pass<"form-expressions"> {
let dependentDialects = ["emitc::EmitCDialect"];
}
+def MLGOScalarizeSingleElementTensorReturnPass
+ : Pass<"mlgo-scalarize-single-element-tensor-return", "ModuleOp"> {
+ let summary = "Scalarize private tensor-returning functions for EmitC lowering of MLGO models";
+ 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", "emitc::EmitCDialect", "func::FuncDialect",
+ "tensor::TensorDialect"
+ ];
+}
+
def WrapFuncInClassPass : Pass<"wrap-emitc-func-in-class"> {
let summary = "Wrap functions in classes, using arguments as fields.";
let description = [{
diff --git a/mlir/lib/Dialect/EmitC/Transforms/CMakeLists.txt b/mlir/lib/Dialect/EmitC/Transforms/CMakeLists.txt
index baf67afc30072..eba28c626b106 100644
--- a/mlir/lib/Dialect/EmitC/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/EmitC/Transforms/CMakeLists.txt
@@ -1,6 +1,7 @@
add_mlir_dialect_library(MLIREmitCTransforms
Transforms.cpp
FormExpressions.cpp
+ MLGOScalarizeFunctionResult.cpp
TypeConversions.cpp
WrapFuncInClass.cpp
@@ -11,8 +12,11 @@ add_mlir_dialect_library(MLIREmitCTransforms
MLIREmitCTransformsIncGen
LINK_LIBS PUBLIC
+ MLIRArithDialect
+ MLIRFuncDialect
MLIRIR
MLIRPass
MLIREmitCDialect
+ MLIRTensorDialect
MLIRTransforms
)
diff --git a/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp b/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp
new file mode 100644
index 0000000000000..1b58ce4aa7181
--- /dev/null
+++ b/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp
@@ -0,0 +1,309 @@
+//===- MLGOScalarizeFunctionResult.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/EmitC/IR/EmitC.h"
+#include "mlir/Dialect/EmitC/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/SymbolTable.h"
+#include "llvm/ADT/DenseMap.h"
+
+namespace mlir {
+namespace emitc {
+#define GEN_PASS_DEF_MLGOSCALARIZESINGLEELEMENTTENSORRETURNPASS
+#include "mlir/Dialect/EmitC/Transforms/Passes.h.inc"
+} // namespace emitc
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+enum class ScalarizationState { unknown, visiting, blocked, rewritable };
+
+struct ScalarizableFunctionInfo {
+ RankedTensorType tensorType;
+ SmallVector<func::ReturnOp> returnOps;
+};
+
+static FailureOr<ScalarizableFunctionInfo>
+getScalarizableFunctionInfo(func::FuncOp funcOp) {
+ // Only private function definitions with one single-element ranked tensor
+ // result are locally eligible for scalarization.
+ if (funcOp.isDeclaration() || !funcOp.isPrivate())
+ return failure();
+
+ FunctionType functionType = funcOp.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 info{tensorType, {}};
+ for (Block &block : funcOp.getBody()) {
+ auto returnOp = dyn_cast<func::ReturnOp>(block.getTerminator());
+ if (!returnOp)
+ return failure();
+ info.returnOps.push_back(returnOp);
+ }
+
+ if (info.returnOps.empty())
+ return failure();
+ return info;
+}
+
+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;
+ // Functions are appended here when proven rewritable by the DFS. The final
+ // rewrite walks this list in reverse so callees are rewritten before private
+ // callers that may need their updated call signatures.
+ SmallVector<func::FuncOp> rewriteOrder;
+};
+
+static bool isDirectFunctionInModule(func::FuncOp funcOp,
+ const ScalarizationAnalysis &analysis) {
+ return analysis.moduleFunctions.contains(funcOp);
+}
+
+static ScalarizationState
+computeScalarizationState(func::FuncOp funcOp,
+ ScalarizationAnalysis &analysis) {
+ auto [it, inserted] =
+ analysis.states.try_emplace(funcOp, 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[funcOp] = ScalarizationState::blocked;
+ return ScalarizationState::blocked;
+ };
+ auto setRewritable = [&] {
+ analysis.states[funcOp] = ScalarizationState::rewritable;
+ analysis.rewriteOrder.push_back(funcOp);
+ return ScalarizationState::rewritable;
+ };
+
+ if (!analysis.candidateInfos.contains(funcOp))
+ return setBlocked();
+ analysis.states[funcOp] = ScalarizationState::visiting;
+
+ SmallVector<func::CallOp> callUsers;
+ for (Operation *user : analysis.userMap.getUsers(funcOp.getOperation())) {
+ auto callOp = 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 (!callOp)
+ return setBlocked();
+ callUsers.push_back(callOp);
+
+ func::FuncOp caller = callOp->getParentOfType<func::FuncOp>();
+ // Only direct callers in the same module participate in this analysis.
+ if (!caller)
+ return setBlocked();
+ // computeScalarizationState assumes getScalarizableFunctionInfo has already
+ // categorized every direct func.func in the current module, so any direct
+ // caller in this module must already appear in the analysis tables.
+ if (!isDirectFunctionInModule(caller, analysis))
+ return setBlocked();
+ // Public and external callers keep the current function blocked because the
+ // pass cannot rewrite every visible call boundary.
+ if (caller.isPublic())
+ return setBlocked();
+ if (caller.isExternal())
+ return setBlocked();
+
+ // 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(funcOp, std::move(callUsers));
+ return setRewritable();
+}
+
+static void analyzeModule(ModuleOp module, ScalarizationAnalysis &analysis) {
+ // First collect every direct function in the module and record the subset
+ // that is locally eligible. The second pass runs the memoized DFS only for
+ // candidates to determine the transitive blocked/rewritable state and build
+ // the rewrite order.
+ for (func::FuncOp funcOp : module.getOps<func::FuncOp>()) {
+ analysis.moduleFunctions.insert(funcOp);
+ FailureOr<ScalarizableFunctionInfo> info =
+ getScalarizableFunctionInfo(funcOp);
+ if (succeeded(info))
+ analysis.candidateInfos.try_emplace(funcOp, std::move(*info));
+ }
+
+ for (func::FuncOp funcOp : module.getOps<func::FuncOp>())
+ if (analysis.candidateInfos.contains(funcOp))
+ (void)computeScalarizationState(funcOp, analysis);
+}
+
+static void rewriteScalarizableFunction(func::FuncOp funcOp,
+ const ScalarizableFunctionInfo &info,
+ ArrayRef<func::CallOp> callOps,
+ RewriterBase &rewriter) {
+ // Scalarize eligible functions, as decided by analyzeModule: extract the
+ // unique element before each return, update the function type, and fix direct
+ // call users that were recorded during analysis.
+ RankedTensorType tensorType = info.tensorType;
+ SmallVector<Value> zeroIndices;
+ if (tensorType.getRank() != 0) {
+ rewriter.setInsertionPointToStart(&funcOp.getBody().front());
+ Value zero = arith::ConstantIndexOp::create(rewriter, funcOp.getLoc(), 0);
+ zeroIndices.assign(tensorType.getRank(), zero);
+ }
+
+ Type scalarType = tensorType.getElementType();
+ for (func::ReturnOp returnOp : info.returnOps) {
+ assert(returnOp.getNumOperands() == 1 &&
+ "func.return must have exactly one operand");
+ assert(returnOp.getOperand(0).getType() == tensorType &&
+ "func.return operand type must match the function result type");
+ rewriter.setInsertionPoint(returnOp);
+ Value scalar = rewriter.createOrFold<tensor::ExtractOp>(
+ returnOp.getLoc(), returnOp.getOperand(0), zeroIndices);
+ rewriter.replaceOpWithNewOp<func::ReturnOp>(returnOp, scalar);
+ }
+
+ FunctionType functionType = funcOp.getFunctionType();
+ // 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.
+ funcOp.setType(FunctionType::get(
+ funcOp.getContext(), functionType.getInputs(), TypeRange{scalarType}));
+
+ for (func::CallOp callOp : callOps) {
+ rewriter.setInsertionPoint(callOp);
+ func::CallOp newCallOp = func::CallOp::create(rewriter, callOp.getLoc(),
+ funcOp, callOp.getOperands());
+ newCallOp->setAttrs(callOp->getAttrs());
+
+ if (!callOp.getResult(0).use_empty()) {
+ Value wrappedResult =
+ tensor::FromElementsOp::create(rewriter, callOp.getLoc(), tensorType,
+ ValueRange{newCallOp.getResult(0)});
+ rewriter.replaceOp(callOp, wrappedResult);
+ } else {
+ rewriter.eraseOp(callOp);
+ }
+ }
+}
+
+/// Scalarizes private functions that return a statically-shaped ranked
+/// tensor with exactly one element.
+///
+/// The transform first analyzes module-wide symbol users and only rewrites
+/// functions whose transitive private call users can be updated safely.
+///
+/// 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
+/// }
+///
+/// Public callers and callees, non-call symbol users, and recursive cycles
+/// conservatively block scalarization.
+static LogicalResult
+MLGOScalarizeSingleElementTensorReturns(ModuleOp module,
+ RewriterBase &rewriter) {
+ // This pass is intentionally run as a direct module analysis + rewrite,
+ // rather than through matchAndRewrite on a ModuleOp pattern. The transform
+ // depends on module-scoped SymbolUserMap state and on a precomputed DFS
+ // result, so it does not expose a pattern-testing / transform-dialect pattern
+ // path.
+ SymbolTableCollection symbolTable;
+ // Take a snapshot of symbol users for the original module. This is
+ // safe because the snapshot is consulted only during analyzeModule, before
+ // any rewriting starts, and the rewrite phase relies exclusively on the
+ // cached analysis result instead of querying SymbolUserMap again. It is also
+ // safe with pass-manager multithreading: this pass invocation has exclusive
+ // access to the current ModuleOp, so no other pass mutates the same module
+ // concurrently and invalidates the snapshot underneath this analysis.
+ SymbolUserMap userMap(symbolTable, module);
+ ScalarizationAnalysis analysis(userMap);
+ analyzeModule(module, analysis);
+
+ for (func::FuncOp funcOp : llvm::reverse(analysis.rewriteOrder)) {
+ const ScalarizableFunctionInfo &info =
+ analysis.candidateInfos.find(funcOp)->second;
+ ArrayRef<func::CallOp> callOps = analysis.callUsers.find(funcOp)->second;
+ rewriteScalarizableFunction(funcOp, info, callOps, rewriter);
+ }
+
+ return success();
+}
+
+struct MLGOScalarizeSingleElementTensorReturnPass
+ : public emitc::impl::MLGOScalarizeSingleElementTensorReturnPassBase<
+ MLGOScalarizeSingleElementTensorReturnPass> {
+ using Base::Base;
+
+ void runOnOperation() override {
+ IRRewriter rewriter(&getContext());
+ if (failed(
+ MLGOScalarizeSingleElementTensorReturns(getOperation(), rewriter)))
+ signalPassFailure();
+ }
+};
+
+} // namespace
diff --git a/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir b/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir
new file mode 100644
index 0000000000000..52da8d1ab96e9
--- /dev/null
+++ b/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir
@@ -0,0 +1,256 @@
+// RUN: mlir-opt -mlgo-scalarize-single-element-tensor-return -split-input-file %s | FileCheck %s
+
+/// Both aller and callee 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]]] : tensor<1x1xi64>
+// CHECK: return %[[EXT]] : i64
+
+// -----
+
+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
+
+// -----
+/// Only callee updated
+
+func.func private @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+/// Private non-scalarizeable callers break the DFS loop but still allow rewritten
+/// callees. The caller's signature remains unchanged.
+// 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>
+
+// -----
+
+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
+
+// -----
+
+// 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
+
+// -----
+
+/// Only caller updated
+
+func.func private @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ return %arg0 : tensor<1xi64>
+}
+/// `func.constant` is a non-call symbol user of the target, so it blocks
+/// scalarization.
+// 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
+
+// -----
+
+/// 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>
+
+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: %[[ARG:.*]]: tensor<1xi64>,
+// CHECK-SAME: %[[COND:.*]]: i1
+// CHECK-SAME: -> tensor<1xi64>
+// CHECK: cf.cond_br %[[COND]], ^bb1, ^bb2
+// CHECK: ^bb1:
+// CHECK: cf.br ^bb3(%[[ARG]] : tensor<1xi64>)
+// CHECK: ^bb2:
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1xi64>
+// CHECK: cf.br ^bb3(%[[EMPTY]] : tensor<1xi64>)
+// CHECK: ^bb3(%[[RESULT:.*]]: tensor<1xi64>):
+// CHECK: return %[[RESULT]] : tensor<1xi64>
+// CHECK-NOT: tensor.extract
+
+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