[FIX][TOPI] Broadcast values in index_put

python/tvm/topi/index_put.py

@@ -17,12 +17,15 @@
 # pylint: disable=invalid-name
 """IndexPut operator"""
 
+import tvm_ffi
+
 from tvm import te
 from tvm_tirx import tirx
 from tvm_tirx.script.ir_builder import IRBuilder
 from tvm_tirx.tirx.script import ir_builder as T
 
 from . import utils
+from .broadcast import broadcast_to
 
 
 def index_put(data, indices, values, accumulate=False):
@@ -38,7 +41,8 @@ def index_put(data, indices, values, accumulate=False):
         Index tensors are broadcast together following NumPy broadcasting rules.
 
     values : tvm.te.Tensor
-        The values to place at the specified indices.
+        The values to place at the specified indices.  ``values`` is
+        broadcast to the shape of the indexed region following NumPy rules.
 
     accumulate : bool, optional
         Whether to accumulate (add) values rather than replace.
@@ -84,6 +88,16 @@ def index_put(data, indices, values, accumulate=False):
                         raise ValueError(f"Cannot broadcast index shapes: {index_shapes}")
         broadcast_shape.append(max_dim)
 
+    # The indexed region has the broadcast index shape.  Broadcast ``values``
+    # to that shape so that the flat index used below is also valid when
+    # ``values`` needs broadcasting (e.g. x[:, [0, 2]] = v with v of shape
+    # (2, 1)).  Without this, the flat index reads past the end of values.
+    if len(values.shape) != len(broadcast_shape) or any(
+        not tvm_ffi.structural_equal(values_dim, region_dim)
+        for values_dim, region_dim in zip(values.shape, broadcast_shape)
+    ):
+        values = broadcast_to(values, broadcast_shape)
+
     # Compute total number of elements after broadcasting
     index_len = 1
     for dim in broadcast_shape:

tests/python/relax/test_op_manipulate.py

@@ -15,8 +15,10 @@
 # specific language governing permissions and limitations
 # under the License.
 # ruff: noqa: E731, F841
+import numpy as np
 import pytest
 
+import tvm
 import tvm_tirx
 import tvm_tirx.testing
 from tvm import relax
@@ -3532,5 +3534,43 @@ def test_one_hot_infer_ty():
         )
 
 
+@pytest.mark.parametrize("accumulate", [False, True])
+def test_index_put_broadcast_values(accumulate):
+    bb = relax.BlockBuilder()
+    data = relax.Var("data", R.Tensor((2, 3), "float32"))
+    indices_0 = relax.Var("indices_0", R.Tensor((2, 1), "int32"))
+    indices_1 = relax.Var("indices_1", R.Tensor((2,), "int32"))
+    values = relax.Var("values", R.Tensor((2, 1), "float32"))
+
+    with bb.function("main", [data, indices_0, indices_1, values]):
+        with bb.dataflow():
+            out = bb.emit(
+                relax.op.index_put(data, relax.Tuple([indices_0, indices_1]), values, accumulate)
+            )
+            gv = bb.emit_output(out)
+        bb.emit_func_output(gv)
+
+    exe = tvm.compile(bb.get(), target="llvm")
+    vm = relax.VirtualMachine(exe, tvm_tirx.cpu())
+
+    data_np = np.arange(6, dtype="float32").reshape(2, 3)
+    indices_0_np = np.array([[0], [1]], dtype="int32")
+    indices_1_np = np.array([0, 2], dtype="int32")
+    values_np = np.array([[10.0], [20.0]], dtype="float32")
+    expected = data_np.copy()
+    if accumulate:
+        np.add.at(expected, (indices_0_np, indices_1_np), np.broadcast_to(values_np, (2, 2)))
+    else:
+        expected[indices_0_np, indices_1_np] = values_np
+
+    out = vm["main"](
+        *(
+            tvm_tirx.runtime.tensor(value)
+            for value in (data_np, indices_0_np, indices_1_np, values_np)
+        )
+    )
+    tvm_tirx.testing.assert_allclose(out.numpy(), expected)
+
+
 if __name__ == "__main__":
     tvm_tirx.testing.main()

tests/python/topi/test_topi_ir_builder.py

@@ -172,6 +172,52 @@ def test_gpu_imperative_tensors(kind):
 
 
 @pytest.mark.skipif(not tvm_tirx.runtime.enabled("llvm"), reason="LLVM is not enabled")
+@pytest.mark.parametrize(
+    "values_shape, accumulate",
+    [
+        ((2, 1), False),
+        ((2,), False),
+        ((), False),
+        ((2, 1), True),
+        ((2,), True),
+    ],
+)
+def test_index_put_broadcast_values(values_shape, accumulate):
+    data = np.arange(6, dtype="float32").reshape(2, 3)
+    indices_0 = np.array([[0], [1]], dtype="int32")
+    indices_1 = np.array([0, 2], dtype="int32")
+    if values_shape == (2,):
+        values = np.array([10.0, 20.0], dtype="float32")
+    elif values_shape == (2, 1):
+        values = np.array([[10.0], [20.0]], dtype="float32")
+    else:
+        values = np.array(7.0, dtype="float32")
+
+    indexed = np.broadcast_to(values, (2, 2))
+    expected = data.copy()
+    if accumulate:
+        np.add.at(expected, (indices_0, indices_1), indexed)
+    else:
+        expected[indices_0, indices_1] = indexed
+
+    x = te.placeholder(data.shape, "float32")
+    i = te.placeholder(indices_0.shape, "int32")
+    j = te.placeholder(indices_1.shape, "int32")
+    v = te.placeholder(values.shape, "float32")
+    out = topi.index_put(x, [i, j], v, accumulate=accumulate)
+
+    executable = tvm_tirx.tirx.compile(te.create_function([x, i, j, v, out]), target="llvm")
+    actual = tvm_tirx.runtime.empty(data.shape, "float32")
+    executable(
+        tvm_tirx.runtime.tensor(data),
+        tvm_tirx.runtime.tensor(indices_0),
+        tvm_tirx.runtime.tensor(indices_1),
+        tvm_tirx.runtime.tensor(values),
+        actual,
+    )
+    tvm_tirx.testing.assert_allclose(actual.numpy(), expected, atol=1e-5, rtol=1e-5)
+
+
 @pytest.mark.parametrize(
     "kind",
     [