[Torch] optimize repeat singletons (#4572)

Optimize aten.repeat decomposition for singleton input dimensions.

When an existing input dim has size 1 and its repeat value is greater
than 1, the repeated dim can be represented directly in the
aten.broadcast_to shape. This avoids the pattern of inserting an extra
dimension, broadcasting, and then folding it back with flatten.
Non-singleton repeated dims keep the existing decomposition path.
Added a lit test covering mixed repeated dims, where singleton dims use
direct broadcast and a non-singleton repeated dim still uses the
existing reshape path.

Author: ahagro01

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -5096,6 +5096,17 @@ class DecomposeAtenRepeatOp : public OpRewritePattern<AtenRepeatOp> {
       repeatInts.push_back(repeat);
     }
 
+    // Track repeated singleton dims that can be materialized with broadcast.
+    llvm::SmallVector<int64_t> selfSizes(selfTy.getSizes().begin(),
+                                         selfTy.getSizes().end());
+    llvm::SmallVector<bool> broadcastRepeatedSingletonDims(repeats.size(),
+                                                           false);
+    for (int i = batch, s = repeats.size(); i < s; ++i) {
+      int64_t inputDim = i - batch;
+      broadcastRepeatedSingletonDims[i] =
+          selfSizes[inputDim] == 1 && repeatInts[i] > 1;
+    }
+
     // Unsqueeze all newly created dims
     llvm::SmallVector<int> unsqueezeDims;
     for (int i = 0; i < batch; ++i) {
@@ -5106,9 +5117,9 @@ class DecomposeAtenRepeatOp : public OpRewritePattern<AtenRepeatOp> {
       unsqueezeDims.push_back(i);
     }
 
-    // Unsqueeze any non-unary repeats for existing dims
+    // Unsqueeze non-unary repeats, except singleton dims handled by broadcast.
     for (int i = batch, s = repeats.size(); i < s; ++i) {
-      if (repeatInts[i] == 1)
+      if (repeatInts[i] == 1 || broadcastRepeatedSingletonDims[i])
         continue;
       int64_t dim = i + unsqueezeDims.size() - batch;
       Value iv =
@@ -5127,6 +5138,12 @@ class DecomposeAtenRepeatOp : public OpRewritePattern<AtenRepeatOp> {
     }
 
     for (int i = batch, s = repeats.size(); i < s; ++i) {
+      if (broadcastRepeatedSingletonDims[i]) {
+        lengths.push_back(repeats[i]);
+        expandShape.push_back(repeatInts[i]);
+        continue;
+      }
+
       if (repeatInts[i] != 1) {
         lengths.push_back(repeats[i]);
         expandShape.push_back(repeatInts[i]);
@@ -5149,7 +5166,7 @@ class DecomposeAtenRepeatOp : public OpRewritePattern<AtenRepeatOp> {
 
     auto outShape = cast<ValueTensorType>(op.getResult().getType()).getSizes();
     for (int i = batch, s = repeats.size(); i < s; ++i) {
-      if (repeatInts[i] == 1)
+      if (repeatInts[i] == 1 || broadcastRepeatedSingletonDims[i])
         continue;
 
       auto selfShape = selfTy.getSizes();

test/Dialect/Torch/decompose-complex-ops.mlir

@@ -1090,3 +1090,49 @@ func.func @torch.aten.mish$f8E8M0FNU(%arg0: !torch.vtensor<[2,3],f8E8M0FNU>) ->
   %0 = torch.aten.mish %arg0 : !torch.vtensor<[2,3],f8E8M0FNU> -> !torch.vtensor<[2,3],f8E8M0FNU>
   return %0 : !torch.vtensor<[2,3],f8E8M0FNU>
 }
+
+// -----
+
+// CHECK-LABEL: func.func @repeat_mixed_dims_broadcast_singletons
+// CHECK-SAME: (%[[ARG0:.*]]: !torch.vtensor<[1,2,1],f32>) -> !torch.vtensor<[3,8,5],f32>
+// CHECK-DAG: %[[CNEG1:.*]] = torch.constant.int -1
+// CHECK-DAG: %[[C1:.*]] = torch.constant.int 1
+// CHECK-DAG: %[[C2:.*]] = torch.constant.int 2
+// CHECK-DAG: %[[C3:.*]] = torch.constant.int 3
+// CHECK-DAG: %[[C4:.*]] = torch.constant.int 4
+// CHECK-DAG: %[[C5:.*]] = torch.constant.int 5
+// CHECK: %[[UNSQUEEZE:.*]] = torch.aten.unsqueeze %[[ARG0]], %[[C1]] : !torch.vtensor<[1,2,1],f32>, !torch.int -> !torch.vtensor<[1,1,2,1],f32>
+// CHECK: %[[SHAPE:.*]] = torch.prim.ListConstruct %[[C3]], %[[C4]], %[[C2]], %[[C5]] : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
+// CHECK: %[[BROADCAST:.*]] = torch.aten.broadcast_to %[[UNSQUEEZE]], %[[SHAPE]] : !torch.vtensor<[1,1,2,1],f32>, !torch.list<int> -> !torch.vtensor<[3,4,2,5],f32>
+// CHECK: %[[VIEW_SHAPE:.*]] = torch.prim.ListConstruct %[[C3]], %[[CNEG1]], %[[C5]] : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
+// CHECK: %[[VIEW:.*]] = torch.aten.view %[[BROADCAST]], %[[VIEW_SHAPE]] : !torch.vtensor<[3,4,2,5],f32>, !torch.list<int> -> !torch.vtensor<[3,8,5],f32>
+// CHECK: return %[[VIEW]] : !torch.vtensor<[3,8,5],f32>
+func.func @repeat_mixed_dims_broadcast_singletons(%arg0: !torch.vtensor<[1,2,1],f32>) -> !torch.vtensor<[3,8,5],f32> {
+  %int3 = torch.constant.int 3
+  %int4 = torch.constant.int 4
+  %int5 = torch.constant.int 5
+  %0 = torch.prim.ListConstruct %int3, %int4, %int5 : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
+  %1 = torch.aten.repeat %arg0, %0 : !torch.vtensor<[1,2,1],f32>, !torch.list<int> -> !torch.vtensor<[3,8,5],f32>
+  return %1 : !torch.vtensor<[3,8,5],f32>
+}
+
+// -----
+
+// CHECK-LABEL: func @repeat_broadcasts_static_singleton_dims
+func.func @repeat_broadcasts_static_singleton_dims(%arg0: !torch.vtensor<[1,1,6,1,4,4],f32>) -> !torch.vtensor<[4,1,6,2500,4,4],f32> {
+  %int4 = torch.constant.int 4
+  %int1 = torch.constant.int 1
+  %int2500 = torch.constant.int 2500
+  // CHECK-DAG: %[[C4:.*]] = torch.constant.int 4
+  // CHECK-DAG: %[[C2500:.*]] = torch.constant.int 2500
+  // CHECK-DAG: %[[C1:.*]] = torch.constant.int 1
+  // CHECK-DAG: %[[C6:.*]] = torch.constant.int 6
+  // CHECK: %[[SHAPE:.*]] = torch.prim.ListConstruct %[[C4]], %[[C1]], %[[C6]], %[[C2500]], %[[C4]], %[[C4]] : (!torch.int, !torch.int, !torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
+  // CHECK-NOT: torch.aten.unsqueeze
+  // CHECK-NOT: torch.aten.flatten
+  // CHECK: %[[BROADCAST:.*]] = torch.aten.broadcast_to %arg0, %[[SHAPE]] : !torch.vtensor<[1,1,6,1,4,4],f32>, !torch.list<int> -> !torch.vtensor<[4,1,6,2500,4,4],f32>
+  // CHECK: return %[[BROADCAST]] : !torch.vtensor<[4,1,6,2500,4,4],f32>
+  %repeats = torch.prim.ListConstruct %int4, %int1, %int1, %int2500, %int1, %int1 : (!torch.int, !torch.int, !torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
+  %0 = torch.aten.repeat %arg0, %repeats : !torch.vtensor<[1,1,6,1,4,4],f32>, !torch.list<int> -> !torch.vtensor<[4,1,6,2500,4,4],f32>
+  return %0 : !torch.vtensor<[4,1,6,2500,4,4],f32>
+}