Add lowering for aten.unfold.

PiperOrigin-RevId: 822295986

Author: junjiang-lab

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -325,6 +325,85 @@ def _aten_cat(lctx: LoweringContext, tensors, dim=0):
   return stablehlo.concatenate(non_empty_tensors, dim)
 
 
+# Schema:
+#   - aten::unfold(Tensor self, int dim, int size, int step) -> Tensor
+# Torch Reference:
+#   - https://pytorch.org/docs/stable/generated/torch.Tensor.unfold.html
+@lower(torch.ops.aten.unfold.default)
+def _aten_unfold(lctx, x: ir.Value, dim: int, size: int, step: int):
+  x_shape = x.type.shape
+  rank = len(x_shape)
+  if dim < 0:
+    dim += rank
+
+  num_windows = (x_shape[dim] - size) // step + 1
+  batch_shape = list(x_shape[:dim]) + [num_windows] + list(x_shape[dim + 1 :])
+
+  # Create start_indices for gather.
+  # The shape of start_indices will be batch_shape + [rank].
+  # start_indices[b_0,...,b_{rank-1}] will be [p_0,...,p_{rank-1}] where
+  # p_j = b_j for j != dim and p_dim = b_dim * step.
+  indices_parts = []
+  i64 = ir.IntegerType.get_signless(64)
+  for i in range(rank):
+    bshape = [1] * rank
+    bshape[i] = batch_shape[i]
+    dim_len = batch_shape[i]
+
+    iota = stablehlo.IotaOp(
+        ir.RankedTensorType.get([dim_len], i64),
+        iota_dimension=ir.IntegerAttr.get(i64, 0),
+    ).result
+    if i == dim:
+      iota = stablehlo.multiply(iota, utils.splat(step, i64, [dim_len]))
+
+    iota_reshaped = stablehlo.reshape(
+        ir.RankedTensorType.get(bshape, i64), iota
+    )
+    indices_parts.append(
+        stablehlo.broadcast_in_dim(
+            ir.RankedTensorType.get(batch_shape, i64),
+            iota_reshaped,
+            ir.DenseI64ArrayAttr.get(list(range(rank))),
+        )
+    )
+
+  # For each dimension i, indices_parts[i] contains the i-th coordinate
+  # of start_indices. We unsqueeze each part to shape batch_shape + [1]
+  # and concatenate along the new dimension to produce start_indices of
+  # shape batch_shape + [rank].
+  unsqueezed_parts = [
+      stablehlo.reshape(ir.RankedTensorType.get(batch_shape + [1], i64), part)
+      for part in indices_parts
+  ]
+  start_indices = stablehlo.concatenate(
+      unsqueezed_parts, ir.IntegerAttr.get(i64, rank)
+  )
+
+  slice_sizes_list = [1] * rank
+  slice_sizes_list[dim] = size
+  slice_sizes = ir.DenseI64ArrayAttr.get(slice_sizes_list)
+
+  collapsed_slice_dims_list = [i for i in range(rank) if i != dim]
+
+  dnums = stablehlo.GatherDimensionNumbers.get(
+      offset_dims=[rank],
+      collapsed_slice_dims=collapsed_slice_dims_list,
+      operand_batching_dims=[],
+      start_indices_batching_dims=[],
+      start_index_map=list(range(rank)),
+      index_vector_dim=rank,
+  )
+
+  return stablehlo.gather(
+      x,
+      start_indices,
+      dnums,
+      slice_sizes,
+      indices_are_sorted=ir.BoolAttr.get(False),
+  )
+
+
 # Schema:
 #   - aten::slice_scatter(Tensor self, Tensor src, int dim=0, SymInt?
 #       start=None, SymInt? end=None, SymInt step=1) -> Tensor

ai_edge_torch/odml_torch/test/test_core_aten_ops.py

@@ -393,6 +393,10 @@ def _run_export_and_compare(
       ("aten_tril_0", torch.ops.aten.tril, (rnd(torch.float32, (10, 10)),), dict()),
       ("aten_trunc_0", torch.ops.aten.trunc, (rnd(torch.float32, (10, 10)),), dict()),
       ("aten_unbind_copy_int_0", torch.ops.aten.unbind_copy.int, (rnd(torch.float32, (10, 10)),), dict()),
+      ("aten_unfold_copy_0", torch.ops.aten.unfold, (rnd(torch.float32, (4, 4)), 0, 2, 1), dict()),
+      ("aten_unfold_copy_1", torch.ops.aten.unfold, (rnd(torch.float32, (4, 4)), 1, 2, 1), dict()),
+      ("aten_unfold_copy_2", torch.ops.aten.unfold, (rnd(torch.float32, (5, 5)), 1, 1, 2), dict()),
+      ("aten_unfold_copy_3", torch.ops.aten.unfold, (rnd(torch.float32, (2, 4, 4)), -1, 2, 1), dict()),
       ("aten_unsqueeze_copy_0", torch.ops.aten.unsqueeze_copy, (rnd(torch.float32, (2, 0, 2)), 1,), dict()),
       ("aten_upsample_bilinear2d_0", torch.ops.aten.upsample_bilinear2d, (rnd(torch.float32, (1, 3, 2, 10)), [3, 20], False,), dict()),
       ("aten_upsample_nearest2d_0", torch.ops.aten.upsample_nearest2d, (rnd(torch.float32, (1, 3, 2, 10)), [3, 20],), dict()),