[Mlir-commits] [mlir] [MLGO][EmitC] Scalarize single-element tensor returns (PR #199686)
ioana ghiban
llvmlistbot at llvm.org
Thu Jun 4 03:14:31 PDT 2026
https://github.com/ioghiban updated https://github.com/llvm/llvm-project/pull/199686
>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 1/4] [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>
>From 58a99529e29cbd3be399a17201d0055d73183d01 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Mon, 1 Jun 2026 16:24:14 +0200
Subject: [PATCH 2/4] Address first round of comments
---
.../MLGOScalarizeFunctionResult.cpp | 157 +++++++++---------
.../mlgo-scalarize-single-elem-return.mlir | 96 ++++++-----
2 files changed, 135 insertions(+), 118 deletions(-)
diff --git a/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp b/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp
index 1b58ce4aa7181..7015708da3fd6 100644
--- a/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp
+++ b/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp
@@ -27,21 +27,34 @@ using namespace mlir;
namespace {
-enum class ScalarizationState { unknown, visiting, blocked, 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
+};
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>
-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())
+getScalarizableFunctionInfoIfEligible(func::FuncOp func) {
+ if (func.isDeclaration() || !func.isPrivate())
return failure();
- FunctionType functionType = funcOp.getFunctionType();
+ FunctionType functionType = func.getFunctionType();
if (functionType.getNumResults() != 1)
return failure();
@@ -51,11 +64,10 @@ getScalarizableFunctionInfo(func::FuncOp funcOp) {
return failure();
ScalarizableFunctionInfo info{tensorType, {}};
- for (Block &block : funcOp.getBody()) {
+ for (Block &block : func.getBody()) {
auto returnOp = dyn_cast<func::ReturnOp>(block.getTerminator());
- if (!returnOp)
- return failure();
- info.returnOps.push_back(returnOp);
+ if (returnOp)
+ info.returnOps.push_back(returnOp);
}
if (info.returnOps.empty())
@@ -71,22 +83,24 @@ struct ScalarizationAnalysis {
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.
+ // 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;
};
-static bool isDirectFunctionInModule(func::FuncOp funcOp,
+// Returns whether `func` is one of the direct function definitions in the
+// current module-level analysis scope.
+static bool isDirectFunctionInModule(func::FuncOp func,
const ScalarizationAnalysis &analysis) {
- return analysis.moduleFunctions.contains(funcOp);
+ return analysis.moduleFunctions.contains(func);
}
+// Runs the memoized DFS that classifies one candidate function by walking its
+// transitive private call users.
static ScalarizationState
-computeScalarizationState(func::FuncOp funcOp,
- ScalarizationAnalysis &analysis) {
+computeScalarizationState(func::FuncOp func, ScalarizationAnalysis &analysis) {
auto [it, inserted] =
- analysis.states.try_emplace(funcOp, ScalarizationState::unknown);
+ analysis.states.try_emplace(func, ScalarizationState::unknown);
if (!inserted) {
// Conservatively reject recursive cycles instead of reasoning about SCCs.
if (it->second == ScalarizationState::visiting)
@@ -99,21 +113,21 @@ computeScalarizationState(func::FuncOp funcOp,
// 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;
+ analysis.states[func] = ScalarizationState::blocked;
return ScalarizationState::blocked;
};
auto setRewritable = [&] {
- analysis.states[funcOp] = ScalarizationState::rewritable;
- analysis.rewriteOrder.push_back(funcOp);
+ analysis.states[func] = ScalarizationState::rewritable;
+ analysis.rewriteOrder.push_back(func);
return ScalarizationState::rewritable;
};
- if (!analysis.candidateInfos.contains(funcOp))
+ if (!analysis.candidateInfos.contains(func))
return setBlocked();
- analysis.states[funcOp] = ScalarizationState::visiting;
+ analysis.states[func] = ScalarizationState::visiting;
SmallVector<func::CallOp> callUsers;
- for (Operation *user : analysis.userMap.getUsers(funcOp.getOperation())) {
+ for (Operation *user : analysis.userMap.getUsers(func.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.
@@ -125,9 +139,9 @@ computeScalarizationState(func::FuncOp 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.
+ // `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
@@ -147,40 +161,41 @@ computeScalarizationState(func::FuncOp funcOp,
return setBlocked();
}
- analysis.callUsers.try_emplace(funcOp, std::move(callUsers));
+ analysis.callUsers.try_emplace(func, 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);
+// 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> info =
- getScalarizableFunctionInfo(funcOp);
+ getScalarizableFunctionInfoIfEligible(func);
if (succeeded(info))
- analysis.candidateInfos.try_emplace(funcOp, std::move(*info));
+ analysis.candidateInfos.try_emplace(func, std::move(*info));
}
-
- for (func::FuncOp funcOp : module.getOps<func::FuncOp>())
- if (analysis.candidateInfos.contains(funcOp))
- (void)computeScalarizationState(funcOp, analysis);
+ // 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);
}
-static void rewriteScalarizableFunction(func::FuncOp funcOp,
+// 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 &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.
+ // 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);
+ rewriter.setInsertionPointToStart(&func.getBody().front());
+ Value zero = arith::ConstantIndexOp::create(rewriter, func.getLoc(), 0);
zeroIndices.assign(tensorType.getRank(), zero);
}
@@ -196,17 +211,17 @@ static void rewriteScalarizableFunction(func::FuncOp funcOp,
rewriter.replaceOpWithNewOp<func::ReturnOp>(returnOp, scalar);
}
- FunctionType functionType = funcOp.getFunctionType();
+ FunctionType functionType = func.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}));
+ func.setType(FunctionType::get(func.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());
+ func, callOp.getOperands());
newCallOp->setAttrs(callOp->getAttrs());
if (!callOp.getResult(0).use_empty()) {
@@ -220,11 +235,9 @@ static void rewriteScalarizableFunction(func::FuncOp funcOp,
}
}
-/// 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.
+/// 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 }
@@ -260,34 +273,26 @@ static void rewriteScalarizableFunction(func::FuncOp funcOp,
/// 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.
+ // 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 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.
+ // 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);
- analyzeModule(module, analysis);
+ computeScalarizationAnalysis(module, analysis);
- for (func::FuncOp funcOp : llvm::reverse(analysis.rewriteOrder)) {
+ for (func::FuncOp func : 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);
+ analysis.candidateInfos.find(func)->second;
+ ArrayRef<func::CallOp> callOps = analysis.callUsers.find(func)->second;
+ rewriteScalarizableFunction(func, info, callOps, rewriter);
}
return success();
diff --git a/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir b/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir
index 52da8d1ab96e9..a6022967a7b9b 100644
--- a/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir
+++ b/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir
@@ -1,6 +1,6 @@
// RUN: mlir-opt -mlgo-scalarize-single-element-tensor-return -split-input-file %s | FileCheck %s
-/// Both aller and callee updated
+/// Positive tests: both caller and callee updated.
func.func private @rank0() -> tensor<i64> {
%0 = arith.constant dense<-1> : tensor<i64>
@@ -21,16 +21,42 @@ func.func private @rank1(%arg0: tensor<1xi64>) -> tensor<1xi64> {
// CHECK: %[[EXT:.*]] = tensor.extract %[[SRC]][%[[C0]]] : tensor<1xi64>
// CHECK: return %[[EXT]] : i64
-func.func private @rank2_single_element(%arg0: tensor<1x1xi64>) -> tensor<1x1xi64> {
+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: %[[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
func.func private @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
return %arg0 : tensor<1xi64>
@@ -51,20 +77,22 @@ func.func private @caller(%arg0: tensor<1xi64>) -> tensor<1xi64> {
// CHECK: return %[[CALL]] : i64
// -----
-/// Only callee updated
+/// Positive tests: only callee updated.
+
+/// Private non-scalarizeable callers break the DFS loop but still allow
+/// rewritten callees. The caller's signature remains unchanged.
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> {
+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>
@@ -77,6 +105,10 @@ func.func private @non_scalarizable_caller_multiple_dimensions_return(%arg0: ten
// -----
+/// 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>
}
@@ -86,7 +118,8 @@ func.func private @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
// 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) {
+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
@@ -101,7 +134,9 @@ func.func private @non_scalarizeable_caller_with_multiple_returns(%arg0: tensor<
// -----
-// Some existing uses in the caller still expect the old tensor type
+/// 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>
}
@@ -110,11 +145,13 @@ func.func private @callee(%arg0: tensor<1xi64>) -> tensor<1xi64> {
// 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>) {
+
+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>)
+ %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
@@ -132,13 +169,13 @@ func.func private @non_scalarizeable_caller_expecting_tensor_type(%arg0: tensor<
// -----
-/// Only caller updated
+/// 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>
}
-/// `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
@@ -160,7 +197,7 @@ func.func private @caller_using_constant(%arg0: tensor<1xi64>) -> tensor<1xi64>
// -----
-/// Neither callee nor caller updated
+/// Negative tests: neither callee nor caller updated.
func.func private @multiple_elements(%arg0: tensor<2xi64>) -> tensor<2xi64> {
return %arg0 : tensor<2xi64>
@@ -212,31 +249,6 @@ func.func @public_caller(%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>
>From 6c6e363cce8cdc98cd17b410503994d635cddf7b Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Wed, 3 Jun 2026 17:03:12 +0200
Subject: [PATCH 3/4] Address second round of comments
---
.../Dialect/EmitC/Transforms/CMakeLists.txt | 3 +
.../MLGOScalarizeFunctionResult.cpp | 105 +++++++++---------
.../mlgo-scalarize-single-elem-return.mlir | 26 ++++-
3 files changed, 78 insertions(+), 56 deletions(-)
diff --git a/mlir/lib/Dialect/EmitC/Transforms/CMakeLists.txt b/mlir/lib/Dialect/EmitC/Transforms/CMakeLists.txt
index eba28c626b106..61a4bebb0c185 100644
--- a/mlir/lib/Dialect/EmitC/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/EmitC/Transforms/CMakeLists.txt
@@ -12,11 +12,14 @@ add_mlir_dialect_library(MLIREmitCTransforms
MLIREmitCTransformsIncGen
LINK_LIBS PUBLIC
+ # MLGO related
MLIRArithDialect
+ # MLGO related
MLIRFuncDialect
MLIRIR
MLIRPass
MLIREmitCDialect
+ # MLGO related
MLIRTensorDialect
MLIRTransforms
)
diff --git a/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp b/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp
index 7015708da3fd6..2697169542593 100644
--- a/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp
+++ b/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp
@@ -63,16 +63,16 @@ getScalarizableFunctionInfoIfEligible(func::FuncOp func) {
tensorType.getNumElements() != 1)
return failure();
- ScalarizableFunctionInfo info{tensorType, {}};
+ ScalarizableFunctionInfo sfi{tensorType, {}};
for (Block &block : func.getBody()) {
auto returnOp = dyn_cast<func::ReturnOp>(block.getTerminator());
if (returnOp)
- info.returnOps.push_back(returnOp);
+ sfi.returnOps.push_back(returnOp);
}
- if (info.returnOps.empty())
+ if (sfi.returnOps.empty())
return failure();
- return info;
+ return sfi;
}
struct ScalarizationAnalysis {
@@ -88,13 +88,6 @@ struct ScalarizationAnalysis {
SmallVector<func::FuncOp> rewriteOrder;
};
-// Returns whether `func` is one of the direct function definitions in the
-// current module-level analysis scope.
-static bool isDirectFunctionInModule(func::FuncOp func,
- const ScalarizationAnalysis &analysis) {
- return analysis.moduleFunctions.contains(func);
-}
-
// Runs the memoized DFS that classifies one candidate function by walking its
// transitive private call users.
static ScalarizationState
@@ -126,30 +119,31 @@ computeScalarizationState(func::FuncOp func, ScalarizationAnalysis &analysis) {
return setBlocked();
analysis.states[func] = ScalarizationState::visiting;
- SmallVector<func::CallOp> callUsers;
+ SmallVector<func::CallOp> directCallUsers;
for (Operation *user : analysis.userMap.getUsers(func.getOperation())) {
- auto callOp = dyn_cast<func::CallOp>(user);
+ 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 (!callOp)
+ if (!directCall)
return setBlocked();
- callUsers.push_back(callOp);
+ directCallUsers.push_back(directCall);
- func::FuncOp caller = callOp->getParentOfType<func::FuncOp>();
- // Only direct callers in the same module participate in this analysis.
+ func::FuncOp caller = directCall->getParentOfType<func::FuncOp>();
+ // The symbol user must be enclosed by a func.func.
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();
+ // 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();
- if (caller.isExternal())
- 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.
@@ -161,7 +155,7 @@ computeScalarizationState(func::FuncOp func, ScalarizationAnalysis &analysis) {
return setBlocked();
}
- analysis.callUsers.try_emplace(func, std::move(callUsers));
+ analysis.callUsers.try_emplace(func, std::move(directCallUsers));
return setRewritable();
}
@@ -172,10 +166,10 @@ static void computeScalarizationAnalysis(ModuleOp module,
// eligible.
for (func::FuncOp func : module.getOps<func::FuncOp>()) {
analysis.moduleFunctions.insert(func);
- FailureOr<ScalarizableFunctionInfo> info =
+ FailureOr<ScalarizableFunctionInfo> sfi =
getScalarizableFunctionInfoIfEligible(func);
- if (succeeded(info))
- analysis.candidateInfos.try_emplace(func, std::move(*info));
+ 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>())
@@ -186,12 +180,11 @@ static void computeScalarizationAnalysis(ModuleOp module,
// 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 &info,
- ArrayRef<func::CallOp> callOps,
+ const ScalarizableFunctionInfo &sfi,
+ ArrayRef<func::CallOp> directCalls,
RewriterBase &rewriter) {
- // 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;
+ // Scalarize the unique element before each return.
+ RankedTensorType tensorType = sfi.tensorType;
SmallVector<Value> zeroIndices;
if (tensorType.getRank() != 0) {
rewriter.setInsertionPointToStart(&func.getBody().front());
@@ -200,37 +193,39 @@ static void rewriteScalarizableFunction(func::FuncOp func,
}
Type scalarType = tensorType.getElementType();
- for (func::ReturnOp returnOp : info.returnOps) {
- assert(returnOp.getNumOperands() == 1 &&
+ for (func::ReturnOp funcReturn : sfi.returnOps) {
+ assert(funcReturn.getNumOperands() == 1 &&
"func.return must have exactly one operand");
- assert(returnOp.getOperand(0).getType() == tensorType &&
+ assert(funcReturn.getOperand(0).getType() == tensorType &&
"func.return operand type must match the function result type");
- rewriter.setInsertionPoint(returnOp);
+ rewriter.setInsertionPoint(funcReturn);
Value scalar = rewriter.createOrFold<tensor::ExtractOp>(
- returnOp.getLoc(), returnOp.getOperand(0), zeroIndices);
- rewriter.replaceOpWithNewOp<func::ReturnOp>(returnOp, scalar);
+ 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}));
-
- for (func::CallOp callOp : callOps) {
- rewriter.setInsertionPoint(callOp);
- func::CallOp newCallOp = func::CallOp::create(rewriter, callOp.getLoc(),
- func, 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);
+ // Fix direct call users that were recorded during analysis.
+ for (func::CallOp directCall : directCalls) {
+ rewriter.setInsertionPoint(directCall);
+ OpBuilder::InsertionGuard guard(rewriter);
+ 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(callOp);
+ rewriter.eraseOp(directCall);
}
}
}
@@ -289,10 +284,10 @@ MLGOScalarizeSingleElementTensorReturns(ModuleOp module,
computeScalarizationAnalysis(module, analysis);
for (func::FuncOp func : llvm::reverse(analysis.rewriteOrder)) {
- const ScalarizableFunctionInfo &info =
+ const ScalarizableFunctionInfo &sfi =
analysis.candidateInfos.find(func)->second;
- ArrayRef<func::CallOp> callOps = analysis.callUsers.find(func)->second;
- rewriteScalarizableFunction(func, info, callOps, rewriter);
+ ArrayRef<func::CallOp> directCalls = analysis.callUsers.find(func)->second;
+ rewriteScalarizableFunction(func, sfi, directCalls, rewriter);
}
return success();
diff --git a/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir b/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir
index a6022967a7b9b..61ce510a4eeb5 100644
--- a/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir
+++ b/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir
@@ -1,6 +1,6 @@
// RUN: mlir-opt -mlgo-scalarize-single-element-tensor-return -split-input-file %s | FileCheck %s
-/// Positive tests: both caller and callee updated.
+/// Positive tests: functions with no users updated.
func.func private @rank0() -> tensor<i64> {
%0 = arith.constant dense<-1> : tensor<i64>
@@ -58,6 +58,8 @@ func.func private @non_return_terminator(%arg0: tensor<1xi64>, %cond: i1)
// 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>
}
@@ -249,6 +251,28 @@ func.func @public_caller(%arg0: tensor<1xi64>) -> tensor<1xi64> {
// CHECK-NOT: tensor.extract
// CHECK: return %[[CALL]] : tensor<1xi64>
+// -----
+
+/// Callee blocked by a symbol 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>) -> tensor<1xi64>
+// CHECK-NOT: tensor.extract
+// CHECK: return %[[SRC]] : tensor<1xi64>
+
+%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>) -> 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>
>From 6c962f6630d8c36f22e4a30313bc47b4e8836f20 Mon Sep 17 00:00:00 2001
From: Ioana Ghiban <ioana.ghiban at arm.com>
Date: Thu, 4 Jun 2026 12:05:26 +0200
Subject: [PATCH 4/4] Address third round of comments
---
.../MLGOScalarizeFunctionResult.cpp | 16 ++++++--
.../mlgo-scalarize-single-elem-return.mlir | 40 ++++++++++---------
2 files changed, 34 insertions(+), 22 deletions(-)
diff --git a/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp b/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp
index 2697169542593..cea73cb066fa7 100644
--- a/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp
+++ b/mlir/lib/Dialect/EmitC/Transforms/MLGOScalarizeFunctionResult.cpp
@@ -27,6 +27,9 @@ 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,
@@ -41,6 +44,11 @@ enum class ScalarizationState {
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;
@@ -129,9 +137,11 @@ computeScalarizationState(func::FuncOp func, ScalarizationAnalysis &analysis) {
directCallUsers.push_back(directCall);
func::FuncOp caller = directCall->getParentOfType<func::FuncOp>();
- // The symbol user must be enclosed by a func.func.
+ // 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)
- return setBlocked();
+ continue;
// Since `getScalarizableFunctionInfoIfEligible` has already categorized
// every direct func.func in the current module, any direct caller must
// already appear in the analysis tables.
@@ -183,6 +193,7 @@ 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;
@@ -214,7 +225,6 @@ static void rewriteScalarizableFunction(func::FuncOp func,
// Fix direct call users that were recorded during analysis.
for (func::CallOp directCall : directCalls) {
rewriter.setInsertionPoint(directCall);
- OpBuilder::InsertionGuard guard(rewriter);
func::CallOp newDirectCall = func::CallOp::create(
rewriter, directCall.getLoc(), func, directCall.getOperands());
newDirectCall->setAttrs(directCall->getAttrs());
diff --git a/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir b/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir
index 61ce510a4eeb5..8ca0f758eb901 100644
--- a/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir
+++ b/mlir/test/Dialect/EmitC/mlgo-scalarize-single-elem-return.mlir
@@ -82,8 +82,8 @@ func.func private @caller(%arg0: tensor<1xi64>) -> tensor<1xi64> {
/// Positive tests: only callee updated.
-/// Private non-scalarizeable callers break the DFS loop but still allow
-/// rewritten callees. The caller's signature remains unchanged.
+/// 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>
}
@@ -171,6 +171,25 @@ func.func private @non_scalarizeable_caller_expecting_tensor_type(
// -----
+/// 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
@@ -253,23 +272,6 @@ func.func @public_caller(%arg0: tensor<1xi64>) -> tensor<1xi64> {
// -----
-/// Callee blocked by a symbol 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>) -> tensor<1xi64>
-// CHECK-NOT: tensor.extract
-// CHECK: return %[[SRC]] : tensor<1xi64>
-
-%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>) -> tensor<1xi64>
-
-// -----
-
/// Recursive cycle - the DFS marks the cycle conservatively as blocked,
/// so neither function is scalarized.
More information about the Mlir-commits
mailing list