Support broadcast index put (#3421)

Author: chohk88

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -810,43 +810,8 @@ def aten_ops_select(
     )
 
 
-def index_put_validator(
-    node: Node, settings: Optional[CompilationSettings] = None
-) -> bool:
-    if args_bounds_check(node.args, 3, False):  # Check if accumulate is valid
-        _LOGGER.debug("We do not support accumulate=True for aten.index_put operation")
-        accumulate_valid = False
-    else:
-        accumulate_valid = True
-
-    # Retrieve input tensor's meta information
-    input_meta = node.args[0].meta.get("tensor_meta")
-    if not input_meta:
-        _LOGGER.warning(
-            "Meta information of input is missing. Unable to validate if broadcasting is needed, falling back to PyTorch operation."
-        )
-        return False
-
-    input_shape = input_meta.shape
-    input_num_dims = len(input_shape)
-
-    # Check if broadcasting is valid
-    indices_num_dims = len(node.args[1])
-    if indices_num_dims == input_num_dims:
-        broadcast_valid = True
-    else:
-        _LOGGER.debug(
-            "We do not support broadcasting when the number of index dimensions does not match the number of input tensor dimensions."
-        )
-        broadcast_valid = False
-
-    # Return validation result
-    return accumulate_valid and broadcast_valid
-
-
 @dynamo_tensorrt_converter(
     torch.ops.aten.index_put.default,
-    capability_validator=index_put_validator,
 )
 @enforce_tensor_types(
     {

py/torch_tensorrt/dynamo/conversion/impl/arange.py

@@ -9,6 +9,7 @@
     SourceIR,
     cast_trt_tensor,
     get_trt_tensor,
+    set_layer_name,
 )
 from torch_tensorrt.fx.types import TRTTensor
 
@@ -22,36 +23,45 @@ def arange(
     end: Union[int, TRTTensor],
     step: Union[int, TRTTensor],
 ) -> TRTTensor:
-    if any(isinstance(tensor, TRTTensor) for tensor in (start, end, step)):
+    """
+    Creates a sequence of values (arange) either dynamically or statically,
+    then outputs a TensorRT tensor.
+
+    If any of (start, end, step) is a TRT tensor, it sets up a dynamic arange
+    using a Fill layer. Otherwise, it creates a static NumPy array and converts
+    it into a TensorRT constant tensor.
+    """
+    # If any argument is a TRT tensor, use dynamic arange with a Fill layer
+    if any(isinstance(x, TRTTensor) for x in (start, end, step)):
+        # Convert start, end, step into TRT tensors with appropriate rank
         start_rank_0 = get_trt_tensor(ctx, start, name + "_start_rank_0", min_rank=0)
         start_rank_1 = get_trt_tensor(ctx, start, name + "_start_rank_1", min_rank=1)
         end = get_trt_tensor(ctx, end, name + "_end", min_rank=1)
         step = get_trt_tensor(ctx, step, name + "_step", min_rank=1)
-        # Calculate shape = (end-start) / step
+
+        # Compute (end - start) / step to determine the output length
         shape = impl.elementwise.sub(
-            ctx,
-            target,
-            source_ir,
-            name + "_sub",
-            end,
-            start_rank_1,
+            ctx, target, source_ir, name + "_sub", end, start_rank_1
         )
         shape = impl.elementwise.trunc_div(
-            ctx,
-            target,
-            source_ir,
-            name + "_shape",
-            shape,
-            step,
+            ctx, target, source_ir, name + "_shape", shape, step
         )
         shape = cast_trt_tensor(ctx, shape, end.dtype, name + "_shape_casted")
+
+        # Build a Fill layer in LINSPACE mode
         fill_layer = ctx.net.add_fill(
             shape.shape, trt.FillOperation.LINSPACE, shape.dtype
         )
-        fill_layer.set_input(0, shape)
-        # Set start index
-        fill_layer.set_input(1, start_rank_0)
-        # Set delta/step
-        fill_layer.set_input(2, step)
+        fill_layer.set_input(0, shape)  # output length
+        fill_layer.set_input(1, start_rank_0)  # start value
+        fill_layer.set_input(2, step)  # step size
+
         return fill_layer.get_output(0)
-    return np.arange(start, end, step)
+
+    else:
+        # All arguments are static, so use NumPy arange and create a TRT constant
+        arr = np.arange(start, end, step, dtype=np.int32)
+        weights = trt.Weights(arr)
+        const_layer = ctx.net.add_constant(arr.shape, weights)
+        set_layer_name(const_layer, target, f"{name}_arange_const", source_ir)
+        return const_layer.get_output(0)