[Mlir-commits] [mlir] e650426 - [MLIR][Tensor] Allow tensor-result scalarization with public callers (#231447)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 16 06:52:42 PDT 2026
Author: ioana ghiban
Date: 2026-10-16T15:52:34+02:00
New Revision: e650426926c8d8b802a978b65f6204667c82deef
URL: https://github.com/llvm/llvm-project/commit/e650426926c8d8b802a978b65f6204667c82deef
DIFF: https://github.com/llvm/llvm-project/commit/e650426926c8d8b802a978b65f6204667c82deef.diff
LOG: [MLIR][Tensor] Allow tensor-result scalarization with public callers (#231447)
Private callees should be scalarized even when called by public
functions. `tensor.from_elements` preserves existing tensor uses without
changing public signatures.
- Simplify DFS to recurse only into candidate callers, removing
redundant checks and the `moduleFunctions` set.
- Add coverage for private call chains with public callers and unchanged
public call boundaries.
- Rename the tensor-consumer test to reflect that its caller is
scalarized.
Validation: C++ syntax and diff checks passed. Runtime verification
requires rebuilding mlir-opt.
I leveraged AI to write this PR.
Added:
Modified:
mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResult.cpp
mlir/test/Dialect/Tensor/scalarize-single-elem-return.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResult.cpp b/mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResult.cpp
index c0c92db18db80c..f29f9384cbff47 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResult.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/ScalarizeFunctionResult.cpp
@@ -35,8 +35,7 @@ enum class ScalarizationState {
// 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.
+ // A symbol user or a candidate caller blocks the rewrite.
blocked,
// The function is locally eligible and all transitive users considered by
// this pass can be updated consistently.
@@ -88,17 +87,16 @@ struct ScalarizationAnalysis {
explicit ScalarizationAnalysis(SymbolUserMap &userMap) : userMap(userMap) {}
SymbolUserMap &userMap;
- DenseSet<func::FuncOp> moduleFunctions;
DenseMap<func::FuncOp, ScalarizableFunctionInfo> candidateInfos;
DenseMap<func::FuncOp, SmallVector<func::CallOp>> callUsers;
DenseMap<func::FuncOp, ScalarizationState> states;
// DFS completion order for rewritable functions. The rewrite phase walks
- // this list in reverse so callees are rewritten before their call sites.
+ // this list in reverse so callees are rewritten before candidate callers.
SmallVector<func::FuncOp> rewriteOrder;
};
-// Runs the memoized DFS that classifies one candidate function by walking its
-// transitive private call users.
+// Classifies a locally eligible candidate using memoized DFS through candidate
+// callers. Local eligibility must be cached for the entire module first.
static ScalarizationState
computeScalarizationState(func::FuncOp func, ScalarizationAnalysis &analysis) {
auto [it, inserted] =
@@ -113,7 +111,9 @@ computeScalarizationState(func::FuncOp func, ScalarizationAnalysis &analysis) {
// 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.
+ // when an unsupported symbol user or a blocked candidate caller is found.
+ // Non-candidate callers, including public functions, stop recursion without
+ // blocking callsite adaptation.
auto setBlocked = [&] {
analysis.states[func] = ScalarizationState::blocked;
return ScalarizationState::blocked;
@@ -124,8 +124,6 @@ computeScalarizationState(func::FuncOp func, ScalarizationAnalysis &analysis) {
return ScalarizationState::rewritable;
};
- if (!analysis.candidateInfos.contains(func))
- return setBlocked();
analysis.states[func] = ScalarizationState::visiting;
SmallVector<func::CallOp> directCallUsers;
@@ -138,27 +136,10 @@ computeScalarizationState(func::FuncOp func, ScalarizationAnalysis &analysis) {
directCallUsers.push_back(directCall);
func::FuncOp caller = directCall->getParentOfType<func::FuncOp>();
- // A direct call user outside any func.func can still be updated in place,
- // but it terminates the DFS because there is no caller signature to
- // analyze or rewrite transitively.
- if (!caller)
- continue;
- // Since `getScalarizableFunctionInfoIfEligible` has already categorized
- // every direct func.func in the current module, any direct caller must
- // already appear in the analysis tables.
- assert(analysis.moduleFunctions.contains(caller) &&
- "Caller of private function is not a direct function in the module");
-
- // Public and external callers keep the current function blocked because the
- // pass cannot rewrite every visible call boundary.
- if (caller.isPublic())
- return setBlocked();
-
- assert(!caller.isExternal() && "Caller of private function is external.");
-
- // A private non-candidate caller can absorb the scalarized call by
- // reboxing the scalar result back into a single-element tensor.
- if (!analysis.candidateInfos.contains(caller))
+ // Cache every direct call, but recurse only into candidate callers. Other
+ // callers, including public functions, keep their signatures and existing
+ // tensor uses through reboxing. Calls outside any func.func stop here too.
+ if (!caller || !analysis.candidateInfos.contains(caller))
continue;
if (computeScalarizationState(caller, analysis) !=
@@ -173,10 +154,8 @@ computeScalarizationState(func::FuncOp func, ScalarizationAnalysis &analysis) {
// 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.
+ // Cache local eligibility for all direct functions before starting DFS.
for (func::FuncOp func : module.getOps<func::FuncOp>()) {
- analysis.moduleFunctions.insert(func);
FailureOr<ScalarizableFunctionInfo> sfi =
getScalarizableFunctionInfoIfEligible(func);
if (succeeded(sfi))
diff --git a/mlir/test/Dialect/Tensor/scalarize-single-elem-return.mlir b/mlir/test/Dialect/Tensor/scalarize-single-elem-return.mlir
index 0a80155a8ae6c7..1fd83999416929 100644
--- a/mlir/test/Dialect/Tensor/scalarize-single-elem-return.mlir
+++ b/mlir/test/Dialect/Tensor/scalarize-single-elem-return.mlir
@@ -136,7 +136,7 @@ func.func private @non_scalarizeable_caller_with_multiple_returns(
// -----
-/// Positive tests: only callee updated, with callsite reboxing.
+/// Positive tests: both caller and 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> {
@@ -148,7 +148,7 @@ 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_expecting_tensor_type(
+func.func private @caller_with_tensor_consumer(
%arg0: tensor<1xi64>, %arg1: tensor<1xi64>) -> (tensor<1xi64>) {
%0 = call @callee(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
%init = tensor.empty() : tensor<1xi64>
@@ -156,7 +156,7 @@ func.func private @non_scalarizeable_caller_expecting_tensor_type(
outs(%init : tensor<1xi64>)
return %1 : tensor<1xi64>
}
-// CHECK-LABEL: func.func private @non_scalarizeable_caller_expecting_tensor_type
+// CHECK-LABEL: func.func private @caller_with_tensor_consumer
// CHECK-SAME: %[[ARG0:.*]]: tensor<1xi64>,
// CHECK-SAME: %[[ARG1:.*]]: tensor<1xi64>) -> i64
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
@@ -252,21 +252,47 @@ func.func @public_function(%arg0: tensor<1xi64>) -> tensor<1xi64> {
// CHECK-NOT: tensor.extract
// CHECK: return %arg0 : tensor<1xi64>
+// -----
+
+/// Positive tests: private callees updated while public caller signatures and
+/// calls to public functions remain unchanged through reboxing.
+
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 @private_caller(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ %0 = call @callee(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+ return %0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func private @private_caller
+// CHECK-SAME: -> i64
+// CHECK: %[[CALL:.*]] = call @callee(%arg0) : (tensor<1xi64>) -> i64
+// CHECK: return %[[CALL]] : i64
+
+func.func @public_caller_of_private_function(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ %0 = call @private_caller(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+ return %0 : tensor<1xi64>
+}
+// CHECK-LABEL: func.func @public_caller_of_private_function
// CHECK-SAME: -> tensor<1xi64>
+// CHECK: %[[CALL:.*]] = call @private_caller(%arg0) : (tensor<1xi64>) -> i64
+// CHECK: %[[BOX:.*]] = tensor.from_elements %[[CALL]] : tensor<1xi64>
// CHECK-NOT: tensor.extract
-// CHECK: return %arg0 : tensor<1xi64>
+// CHECK: return %[[BOX]] : tensor<1xi64>
-func.func @public_caller(%arg0: tensor<1xi64>) -> tensor<1xi64> {
- %0 = call @callee(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+func.func @public_caller_of_public_function(%arg0: tensor<1xi64>) -> tensor<1xi64> {
+ %0 = call @public_caller_of_private_function(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
return %0 : tensor<1xi64>
}
-// CHECK-LABEL: func.func @public_caller
+// CHECK-LABEL: func.func @public_caller_of_public_function
// CHECK-SAME: -> tensor<1xi64>
-// CHECK: %[[CALL:.*]] = call @callee(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
+// CHECK: %[[CALL:.*]] = call @public_caller_of_private_function(%arg0) : (tensor<1xi64>) -> tensor<1xi64>
// CHECK-NOT: tensor.extract
// CHECK: return %[[CALL]] : tensor<1xi64>
More information about the Mlir-commits
mailing list