[Mlir-commits] [mlir] [MLIR][SCF] Include IndexSwitchOp in scf.yield aliasing (PR #236644)

[llvmlistbot at llvm.org]

https://github.com/zishu0131-png updated https://github.com/llvm/llvm-project/pull/236644

>From 6469d0518363ef4d60d7c9046f26e8d268c4f169 Mon Sep 17 00:00:00 2001
From: zishu0131-png <zishu0131 at gmail.com>
Date: Sun, 1 Nov 2026 01:52:36 +0800
Subject: [PATCH] [MLIR][SCF] Include IndexSwitchOp in scf.yield aliasing

The bufferization interface for scf.yield only recognized scf.if and
scf.execute_region parents, so values yielded from scf.index_switch lost
their alias relation with the switch result. One-shot bufferization could
then write through a non-writable input. The same model's bufferize() already
handles scf::IndexSwitchOp, so this adds it to the aliasing side as well.

Fixes #236567.

Assisted-by: Claude Code
---
 .../BufferizableOpInterfaceImpl.cpp           |  3 +-
 .../one-shot-bufferize-analysis.mlir          | 41 +++++++++++++++++++
 .../Transforms/one-shot-bufferize.mlir        | 24 +++++++++++
 3 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index 35357b150e60f5..26f9088e0b87f4 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -1139,10 +1139,9 @@ struct YieldOpInterface
 
   AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                       const AnalysisState &state) const {
-    if (auto ifOp = dyn_cast<scf::IfOp>(op->getParentOp())) {
+    if (isa<scf::IfOp, scf::IndexSwitchOp>(op->getParentOp()))
       return {{op->getParentOp()->getResult(opOperand.getOperandNumber()),
                BufferRelation::Equivalent, /*isDefinite=*/false}};
-    }
     if (isa<scf::ExecuteRegionOp>(op->getParentOp()))
       return {{op->getParentOp()->getResult(opOperand.getOperandNumber()),
                BufferRelation::Equivalent}};
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
index 7d7b1c8eb9bee1..ebbe562d6e3d81 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
@@ -201,3 +201,44 @@ func.func @shorter_inplace_attr(%arg0: tensor<4xi32>) -> (tensor<4xi32>, tensor<
   // tensor operands.  The analysis must expand the annotation without crashing.
   return {__inplace_operands_attr__ = ["false"]} %arg0, %arg0 : tensor<4xi32>, tensor<4xi32>
 }
+
+// -----
+
+// A tensor yielded from `scf.index_switch` aliases with the switch's result,
+// just like `scf.if`. Writing to that result must therefore not modify a
+// non-writable input in place.
+// CHECK-LABEL: func @index_switch_non_writable(
+// CHECK-ALIAS-SETS-LABEL: func @index_switch_non_writable(
+func.func @index_switch_non_writable(%t: tensor<4xf32> {bufferization.writable = false},
+                                     %v: vector<4xf32>, %i: index) -> tensor<4xf32> {
+  %c0 = arith.constant 0 : index
+  // CHECK-ALIAS-SETS: %[[switch:.*]] = scf.index_switch
+  // CHECK-ALIAS-SETS-SAME: __opresult_alias_set_attr__ = {{.*}}%[[switch]]{{.*}}%arg0
+  %0 = scf.index_switch %i -> tensor<4xf32>
+  default {
+    scf.yield %t : tensor<4xf32>
+  }
+  // CHECK: vector.transfer_write {{.*}} {__inplace_operands_attr__ = ["none", "false", "none"]
+  %1 = vector.transfer_write %v, %0[%c0] {in_bounds = [true]} : vector<4xf32>, tensor<4xf32>
+  return %1 : tensor<4xf32>
+}
+
+// -----
+
+// Each `scf.index_switch` result aliases the correspondingly-indexed operand
+// yielded by every case and by the default region.
+// CHECK-ALIAS-SETS-LABEL: func @index_switch_multi_result(
+func.func @index_switch_multi_result(%a: tensor<4xf32> {bufferization.writable = false},
+                                     %b: tensor<4xf32>, %i: index)
+    -> (tensor<4xf32>, tensor<4xf32>) {
+  // CHECK-ALIAS-SETS: %[[sw:.*]]:2 = scf.index_switch
+  // CHECK-ALIAS-SETS-SAME: __opresult_alias_set_attr__ = {{.*}}"%[[sw]]#0", "%arg0"{{.*}}"%[[sw]]#1", "%arg1"
+  %0, %1 = scf.index_switch %i -> tensor<4xf32>, tensor<4xf32>
+  case 1 {
+    scf.yield %a, %b : tensor<4xf32>, tensor<4xf32>
+  }
+  default {
+    scf.yield %a, %b : tensor<4xf32>, tensor<4xf32>
+  }
+  return %0, %1 : tensor<4xf32>, tensor<4xf32>
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
index 2e834c4f158c6c..f07c1a29b3af29 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
@@ -284,3 +284,27 @@ func.func @materialize_in_dest_raw(%f: f32, %f2: f32, %idx: index) -> (tensor<5x
 
   return %0, %r : tensor<5xf32>, f32
 }
+
+// -----
+
+// Writing through a value yielded by `scf.index_switch` must not modify a
+// non-writable input in place: the write goes to a fresh buffer.
+// CHECK-LABEL: func @index_switch_non_writable
+func.func @index_switch_non_writable(%t: tensor<4xf32> {bufferization.writable = false},
+                                     %v: vector<4xf32>, %i: index) -> tensor<4xf32> {
+  %c0 = arith.constant 0 : index
+  // CHECK: scf.index_switch
+  %0 = scf.index_switch %i -> tensor<4xf32>
+  case 1 {
+    %e = tensor.empty() : tensor<4xf32>
+    scf.yield %e : tensor<4xf32>
+  }
+  // CHECK: default
+  default {
+    scf.yield %t : tensor<4xf32>
+  }
+  // CHECK: %[[ALLOC:.*]] = memref.alloc
+  // CHECK-NEXT: vector.transfer_write {{.*}}, %[[ALLOC]]
+  %1 = vector.transfer_write %v, %0[%c0] {in_bounds = [true]} : vector<4xf32>, tensor<4xf32>
+  return %1 : tensor<4xf32>
+}