[Torch] handle duplicate symbolic shape bindings (#24875)

torch might emit multiple bindings for a single value, this breaks a
invariant the surrounding code expects which makes this pass error out.

This fixes this by only accepting one binding per value and discarding
the rest.

Assisted by Claude

Signed-off-by: default <ziereis@roofline.ai>

compiler/plugins/input/Torch/InputConversion/BindSymbolicShapes.cpp

diff --git a/compiler/plugins/input/Torch/InputConversion/BindSymbolicShapes.cpp b/compiler/plugins/input/Torch/InputConversion/BindSymbolicShapes.cpp
index 2291372..692b60c 100644
--- a/compiler/plugins/input/Torch/InputConversion/BindSymbolicShapes.cpp
+++ b/compiler/plugins/input/Torch/InputConversion/BindSymbolicShapes.cpp
@@ -413,6 +413,7 @@
 
     llvm::SmallVector<Operation *> cleanupOpList;
     llvm::SmallVector<TensorBinding> bindings;
+    llvm::DenseMap<Value, Torch::BindSymbolicShapeOp> boundValues;
     // Mapping of SSA value for a torch.symbolic_int (or related op) to its
     // info.
     llvm::DenseMap<Value, SymbolInfo> symbolInfos;
@@ -428,6 +429,22 @@
         if (!isEligibleBinding(bindOp)) {
           return;
         }
+        auto [it, inserted] =
+            boundValues.try_emplace(bindOp.getOperand(), bindOp);
+        if (!inserted) {
+          auto previousBindOp = it->second;
+          if (bindOp.getShapeExpressions() !=
+                  previousBindOp.getShapeExpressions() ||
+              !llvm::equal(bindOp.getShapeSymbols(),
+                           previousBindOp.getShapeSymbols())) {
+            auto diagnostic = bindOp.emitError(
+                "conflicting symbolic shape bindings for the same value");
+            diagnostic.attachNote(previousBindOp.getLoc())
+                << "previous binding is here";
+            return signalPassFailure();
+          }
+          return;
+        }
         auto torchType =
             cast<Torch::ValueTensorType>(bindOp.getOperand().getType());
         auto builtinType = dyn_cast_if_present<RankedTensorType>(

compiler/plugins/input/Torch/InputConversion/test/bind_symbolic_shapes.mlir

diff --git a/compiler/plugins/input/Torch/InputConversion/test/bind_symbolic_shapes.mlir b/compiler/plugins/input/Torch/InputConversion/test/bind_symbolic_shapes.mlir
index 335ca30..998d34a 100644
--- a/compiler/plugins/input/Torch/InputConversion/test/bind_symbolic_shapes.mlir
+++ b/compiler/plugins/input/Torch/InputConversion/test/bind_symbolic_shapes.mlir
@@ -96,6 +96,36 @@
 }
 
 // -----
+// Verify duplicate bindings are only materialized once.
+// CHECK-LABEL: @duplicate_binding
+module @duplicate_binding {
+  func.func @main(%arg0: !torch.vtensor<[1,?],si64>) -> !torch.vtensor<[1,?],si64> attributes {torch.assume_strict_symbolic_shapes} {
+    // CHECK-COUNT-1: torch_c.to_builtin_tensor %arg0
+    // CHECK-COUNT-1: flow.tensor.tie_shape
+    // CHECK-COUNT-1: torch_c.from_builtin_tensor
+    // CHECK-NOT: torch.bind_symbolic_shape
+    %0 = torch.symbolic_int "s0" {min_val = 2, max_val = 65} : !torch.int
+    torch.bind_symbolic_shape %arg0, [%0], affine_map<()[s0] -> (1, s0 * 32)> : !torch.vtensor<[1,?],si64>
+    %int1 = torch.constant.int 1
+    %1 = torch.aten.add.Scalar %arg0, %int1, %int1 : !torch.vtensor<[1,?],si64>, !torch.int, !torch.int -> !torch.vtensor<[1,?],si64>
+    torch.bind_symbolic_shape %arg0, [%0], affine_map<()[s0] -> (1, s0 * 32)> : !torch.vtensor<[1,?],si64>
+    return %1 : !torch.vtensor<[1,?],si64>
+  }
+}
+
+// -----
+module @conflicting_duplicate_binding {
+  func.func @main(%arg0: !torch.vtensor<[1,?],si64>) {
+    %0 = torch.symbolic_int "s0" {min_val = 2, max_val = 65} : !torch.int
+    // expected-note@+1 {{previous binding is here}}
+    torch.bind_symbolic_shape %arg0, [%0], affine_map<()[s0] -> (1, s0 * 32)> : !torch.vtensor<[1,?],si64>
+    // expected-error@+1 {{conflicting symbolic shape bindings for the same value}}
+    torch.bind_symbolic_shape %arg0, [%0], affine_map<()[s0] -> (1, s0 * 16)> : !torch.vtensor<[1,?],si64>
+    return
+  }
+}
+
+// -----
 // CHECK-LABEL: @add_expr
 module @add_expr {
   func.func @main(%arg0: !torch.vtensor<[?,?],f32>, %arg1: !torch.vtensor<[?,?],f32>) {