[torchlib] Reimplement as_strided without an ONNX loop

onnxscript/function_libs/tools/torch_lib/deduce_type_constraints_test.py

@@ -28,7 +28,6 @@ def torch_lib_onnx_functions_from_registry() -> Generator[onnxscript.OnnxFunctio
 
 class TestDeduceTypeConstraints(unittest.TestCase):
     _SKIP_FUNCTIONS_WITH_LOOP_OR_SCAN = (
-        "_aten_as_strided_onnx",
         "_aten_unfold_onnx",
         "_aten_embedding_bag_onnx",
         "_aten_embedding_bag_1d_padding_idx_onnx",

onnxscript/function_libs/torch_lib/ops/core.py

@@ -57,6 +57,7 @@
 _INT64_MAX = 9223372036854775807
 _INT64_MIN = -9223372036854775808
 _MATH_PI = math.pi
+_AS_STRIDED_STATIC_INDEX_SIZE_LIMIT = 512 * 512
 
 
 @functools.lru_cache(maxsize=None)
@@ -900,76 +901,92 @@ def aten_argwhere(self: TensorType) -> TensorType:
 
 @torch_op("aten::as_strided", trace_only=True)
 def aten_as_strided(
-    self: TTensor, size: INT64, stride: Sequence[int], storage_offset: int = 0
+    self: TTensor,
+    size: Sequence[INT64],
+    stride: Sequence[INT64],
+    storage_offset: Optional[INT64] = None,
 ) -> TTensor:
     """as_strided(Tensor(a) self, SymInt[] size, SymInt[] stride, SymInt? storage_offset=None) -> Tensor(a)"""
 
-    rank = len(stride)
-    return _aten_as_strided_onnx(self, size, stride, storage_offset, rank)
-
-
-@torch_op("aten::as_strided", private=True)
-def _aten_as_strided_onnx(
-    self: TTensor, size: INT64, stride: INT64, storage_offset: int = 0, rank: int = 0
-) -> TTensor:
-    # e.g. when size=[2,3,4], stride=[2,1,3], indices=[0]
-    # i = 0
-    # indices=[0], add_value=[0,3,6,9]
-    # expand(shape=[4]) to [0,0,0,0]
-    # then + add_value = [0,3,6,9]
-    # i = 1
-    # indices=[0,3,6,9], add_value=[0,1,2]
-    # expand(shape=[3,4] to [[0,3,6,9],[0,3,6,9],[0,3,6,9]]
-    # indices + add_value = [[0,3,6,9],[1,3,7,10],[2,5,8,11]]
-    # i = 2
-    # indices = [[0,3,6,9],[1,3,7,10],[2,5,8,11]], add_value=[0,2]
-    # expand(shape=[2,3,4]) to [[[0,3,6,9],[1,3,7,10],[2,5,8,11]]],[[0,3,6,9],[1,3,7,10],[2,5,8,11]]]
-    # indices + add_value = [[[0,3,6,9],[1,3,7,10],[2,5,8,11]]],[[2,5,8,11],[3,5,9,12],[4,7,10,13]]]
-    neg_1 = op.Constant(value_ints=[-1])
-    rank_tensor = op.Reshape(rank, neg_1)  # should be 3
-    # The final indices for op.Gather(data, indices), will be continually changed during the loop
-    indices = op.Constant(value_int=0)
-    one_seq = op.SequenceEmpty()
-    for i in range(rank):
-        # Get the index from back to front, should be 2,1,0 when to i=0,1,2
-        j = rank - i - 1
-        j_tensor = op.Reshape(j, neg_1)
-        # Get size according to index_j, should be 4,3,2 when i=0,1,2
-        size_dim_j = op.Gather(size, j_tensor, axis=0)
-        # Get right size according to index_j, should be [4],[3,4],[2,3,4] when i=0,1,2
-        size_after_j = op.Slice(size, j_tensor, rank_tensor)
-        # Get stride according to index_j, should be 3,1,2 when i=0,1,2
-        stride_dim_j = op.Gather(stride, j_tensor, axis=0)
-        indices = op.Expand(indices, size_after_j)
-        # When size[j]=4, stride[j]=3, then add_value = [0,1,2,3] * 3 = [0,3,6,9]
-        # When size[j]=3, stride[j]=1, then add_value = [0,1,2] * 1 = [0,1,2]
-        # When size[j]=2, stride[j]=2, then add_value = [0,1] * 2 = [0,2]
-        add_value = op.Range(0, size_dim_j, 1) * stride_dim_j
-        # Compute the shape for add_value for correct broadcasting
-        if i == 0:
-            # shape = [dim_size]
-            shape = size_dim_j
-        else:
-            # shape = [dim_size, 1, 1, ...], the count of 1 euqal to i
-            ones = op.ConcatFromSequence(one_seq, axis=0)
-            shape = op.Concat(op.Cast(size_dim_j, to=FLOAT.dtype), ones, axis=0)
-            shape = op.Cast(shape, to=INT64.dtype)
-
-        add_value = op.Reshape(add_value, shape)
-        # Broadcasting add value to indices according to size and stride value
-        indices = indices + add_value
-        # Dims after dim_size to reshape(add_value), should be [1],[1,1],[1,1,1] when i=0,1,2
-        one_seq = op.SequenceInsert(one_seq, op.Constant(value_floats=[1.0]))
-
+    # torch.as_strided produces a view of `self`'s underlying linear storage
+    # with the requested `size` (the output shape) and `stride` (the step, in
+    # elements of storage, taken along each output dimension), starting at
+    # `storage_offset` elements into the storage. For an output element at
+    # position (i_0, ..., i_{n-1}) the element read from storage lives at the flat
+    # index storage_offset + sum_d i_d * stride[d]. So if we flatten `self` to 1-D
+    # and gather it with a tensor of those flat indices shaped like the output, we
+    # reproduce the view as a single Gather. This avoids the hard-to-fold loop of
+    # the previous implementation.
+    rank = len(size)
+    # ONNX exposes logical values, not backing storage. Flattening is correct only
+    # when logical row-major order matches storage order; non-contiguous views remain
+    # unsupported, as they were in the previous lowering.
     self_flatten = op.Reshape(self, op.Constant(value_ints=[-1]))
-    indices = op.Add(indices, storage_offset)
-    result = op.Gather(self_flatten, indices)
 
-    return result
+    # A missing storage_offset means "start at the beginning of the storage".
+    if storage_offset is None:
+        storage_offset = 0
+
+    static_inputs = (
+        all(isinstance(s, int) for s in size)
+        and all(isinstance(s, int) for s in stride)
+        and isinstance(storage_offset, int)
+    )
+    static_index_count = math.prod(size) if static_inputs else None
+    if (
+        static_index_count is not None
+        and static_index_count <= _AS_STRIDED_STATIC_INDEX_SIZE_LIMIT
+    ):
+        # Static fast path: every size/stride/offset is known at trace time, so we
+        # compute the full index tensor with NumPy and emit it as a single
+        # constant that downstream passes can fold trivially.
+        # Start from the storage_offset; the per-dimension contributions are added in.
+        indices = np.array(storage_offset, dtype=np.int64)
+        for dim, (dim_size, dim_stride) in enumerate(zip(size, stride)):
+            # Contribution of dimension `dim`: index i_dim contributes i_dim * stride[dim].
+            add_value = np.arange(dim_size, dtype=np.int64) * dim_stride
+            # Reshape that 1-D contribution so it broadcasts along `dim` only
+            # (length dim_size at position `dim`, length 1 everywhere else), which
+            # lets the running sum build the full n-D index grid.
+            broadcast_shape = [1] * rank
+            broadcast_shape[dim] = dim_size
+            indices = indices + add_value.reshape(broadcast_shape)
+        indices = op.Constant(value=ir.tensor(indices))
+    else:
+        # Build runtime indices when a SymInt is dynamic or a static index tensor
+        # would exceed the optimizer's default folded-output size limit. The loop
+        # is unrolled at trace time because rank is static, so no Loop/Scan is emitted.
+        zero = op.Constant(value_int=0)
+        one = op.Constant(value_int=1)
+        # `scalar_shape` reshapes a value to a 0-D scalar (shape []).
+        scalar_shape = op.Constant(value=ir.tensor(np.array([], dtype=np.int64)))
+        # Start the running index from storage_offset as an INT64 scalar; SymInt
+        # runtime values are assumed to be INT64.
+        indices = op.Reshape(storage_offset, scalar_shape)
+        for dim in range(rank):
+            # Reshape this dimension's size and stride to INT64 scalars.
+            dim_size = op.Reshape(size[dim], scalar_shape)
+            dim_stride = op.Reshape(stride[dim], scalar_shape)
+            # add_value = arange(dim_size) * dim_stride, a 1-D tensor of length dim_size
+            # holding the storage offsets contributed by index 0..dim_size-1 along `dim`.
+            add_value = op.Mul(op.Range(zero, dim_size, one), dim_stride)
+            # Insert singleton axes everywhere except `dim` so this 1-D contribution
+            # broadcasts along dimension `dim` only when added to the running index,
+            # matching the NumPy `reshape(broadcast_shape)` in the static path.
+            unsqueeze_axes = [axis for axis in range(rank) if axis != dim]
+            if unsqueeze_axes:
+                add_value = op.Unsqueeze(add_value, op.Constant(value_ints=unsqueeze_axes))
+            indices = op.Add(indices, add_value)
+
+    # `indices` now has shape `size`; gathering yields the strided view.
+    return op.Gather(self_flatten, indices)
 
 
 def aten_as_strided_copy(
-    self: TensorType, size: INT64, stride: INT64, storage_offset: Optional[INT64] = None
+    self: TensorType,
+    size: Sequence[INT64],
+    stride: Sequence[INT64],
+    storage_offset: Optional[INT64] = None,
 ) -> TensorType:
     """as_strided_copy(Tensor self, SymInt[] size, SymInt[] stride, SymInt? storage_offset=None) -> Tensor"""
 
@@ -979,8 +996,8 @@ def aten_as_strided_copy(
 def aten_as_strided_scatter(
     self: TensorType,
     src: TensorType,
-    size: INT64,
-    stride: INT64,
+    size: Sequence[INT64],
+    stride: Sequence[INT64],
     storage_offset: Optional[INT64] = None,
 ) -> TensorType:
     """as_strided_scatter(Tensor self, Tensor src, SymInt[] size, SymInt[] stride, SymInt? storage_offset=None) -> Tensor"""
@@ -3724,7 +3741,7 @@ def aten_empty_quantized(
 @torch_op("aten::empty_strided", trace_only=True)
 def aten_empty_strided(
     size: Sequence[INT64],
-    stride: INT64,
+    stride: Sequence[INT64],
     layout: str = "",
     dtype: int = FLOAT.dtype,
     device: str = "",
@@ -7581,8 +7598,8 @@ def aten_new_empty(
 @torch_op("aten::new_empty_strided", trace_only=True)
 def aten_new_empty_strided(
     self: TTensor,
-    size: INT64,
-    stride: INT64,
+    size: Sequence[INT64],
+    stride: Sequence[INT64],
     dtype: int = -1,
     layout: str = "",
     device: str = "",
@@ -7591,6 +7608,7 @@ def aten_new_empty_strided(
     """new_empty_strided(Tensor self, SymInt[] size, SymInt[] stride, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor"""
 
     # using zero to simulate empty array
+    size = common_ops.merge_dims(size)
     zero = op.ConstantOfShape(size)
     if dtype == -1:
         return op.CastLike(zero, self)

tests/function_libs/torch_lib/e2e_ops_tests.py

@@ -1580,6 +1580,124 @@ def forward(self, x):
                 got = onnx_program.call_reference({"x": inputs[0]})
                 torch.testing.assert_close(expected, got[0])
 
+    def test_aten_as_strided_static_multi_dim(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.as_strided(x, (2, 3), (4, 1), 2)
+
+        model = Model()
+        x = torch.arange(24, dtype=torch.float32).reshape(4, 6)
+        onnx_program = torch.onnx.export(
+            model, (x,), dynamo=True, optimize=False, verbose=False
+        )
+        op_types = [node.op_type for node in onnx_program.model.graph]
+        self.assertNotIn("Range", op_types)
+        index_constants = [
+            node.attributes["value"].as_tensor()
+            for node in onnx_program.model.graph
+            if node.op_type == "Constant"
+            and "value" in node.attributes
+            and node.attributes["value"].as_tensor().shape == [2, 3]
+        ]
+        self.assertEqual(len(index_constants), 1)
+        _testing.assert_onnx_program(onnx_program)
+
+    def test_aten_as_strided_large_static_shape_avoids_large_constant(self):
+        output_shape = (512, 513)
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.as_strided(x, output_shape, (0, 0))
+
+        onnx_program = torch.onnx.export(
+            Model(),
+            (torch.tensor([3.0]),),
+            dynamo=True,
+            optimize=True,
+            verbose=False,
+        )
+        self.assertIn("Add", [node.op_type for node in onnx_program.model.graph])
+
+        constant_sizes = [
+            node.attributes["value"].as_tensor().size
+            for node in onnx_program.model.graph
+            if node.op_type == "Constant" and "value" in node.attributes
+        ]
+        initializer_sizes = [
+            initializer.const_value.size
+            for initializer in onnx_program.model.graph.initializers.values()
+            if initializer.const_value is not None
+        ]
+        self.assertLess(
+            max([*constant_sizes, *initializer_sizes], default=0),
+            math.prod(output_shape),
+        )
+        _testing.assert_onnx_program(onnx_program)
+
+    def test_aten_as_strided_static_single_dim(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.as_strided(x, (4,), (2,))
+
+        model = Model()
+        x = torch.arange(12, dtype=torch.float32)
+        onnx_program = torch.onnx.export(model, (x,), dynamo=True, verbose=False)
+        _testing.assert_onnx_program(onnx_program)
+
+    def test_aten_as_strided_static_overlapping(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.as_strided(x, (3, 3), (1, 1))
+
+        model = Model()
+        x = torch.arange(10, dtype=torch.float32)
+        onnx_program = torch.onnx.export(model, (x,), dynamo=True, verbose=False)
+        _testing.assert_onnx_program(onnx_program)
+
+    def test_aten_as_strided_static_scalar(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.as_strided(x, (), (), 3)
+
+        model = Model()
+        x = torch.arange(12, dtype=torch.float32)
+        onnx_program = torch.onnx.export(model, (x,), dynamo=True, verbose=False)
+        _testing.assert_onnx_program(onnx_program)
+
+    def test_aten_as_strided_dynamic_size(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                n = x.shape[0] - 1
+                return torch.as_strided(x, (n, 2), (1, 1))
+
+        model = Model()
+        x = torch.arange(12, dtype=torch.float32)
+        onnx_program = torch.onnx.export(
+            model,
+            (x,),
+            dynamic_shapes=({0: "length"},),
+            dynamo=True,
+            verbose=False,
+        )
+        _testing.assert_onnx_program(onnx_program)
+
+    def test_aten_as_strided_dynamic_size_with_offset(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                n = x.shape[0] - 2
+                return torch.as_strided(x, (n,), (1,), 1)
+
+        model = Model()
+        x = torch.arange(12, dtype=torch.float32)
+        onnx_program = torch.onnx.export(
+            model,
+            (x,),
+            dynamic_shapes=({0: "length"},),
+            dynamo=True,
+            verbose=False,
+        )
+        _testing.assert_onnx_program(onnx_program)
+
     @parameterized.parameterized.expand(
         [
             ("float32", torch.float32),