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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: zishu0131-png

<details>
<summary>Changes</summary>

The bufferization interface for `scf.yield` only recognized `scf::IfOp` and `scf::ExecuteRegionOp` as parents. A value yielded from `scf.index_switch` therefore had no alias relation with the switch result, so one-shot bufferization could write through a non-writable input in place.

The same model's `bufferize()` already lists `scf::IndexSwitchOp`, so this adds it on the aliasing side as well. `isWritable` and `bufferizesToMemoryWrite` need no change: the switch result is an `OpResult`, and the switch itself does not write memory.

Tests: alias-set coverage for single- and multi-result switches in `one-shot-bufferize-analysis.mlir`, plus an end-to-end case in `one-shot-bufferize.mlir` that now materializes a `memref.alloc` before the write instead of writing into the input buffer.

Fixes #<!-- -->236567.

Assisted-by: Claude Code


---
Full diff: https://github.com/llvm/llvm-project/pull/236644.diff


3 Files Affected:

- (modified) mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp (+1-2) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir (+41) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir (+24) 


``````````diff
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>
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/236644