Derive output_size of repeat_interleave when inputs are broadcast(fill(x)) (#109)

Author: mgehre-amd

e2e_testing/xfail_sets.py

@@ -18,6 +18,7 @@
     "Conv1dNoPaddingTransposeModule_basic",
     "Conv1dNoPaddingGroupModule_basic",
     "RepeatInterleaveStaticModule_basic",
+    "RepeatInterleaveFillModule_basic",
     # tm_tensor.scatter' op mismatch in shape of indices and update value at dim#0
     "IndexPutImpl2DNoneIndexBroadcastStaticModule_basic"
 }
@@ -277,6 +278,7 @@
     "ScatterValueIntModule_basic",
     # ERROR: Unsupported: dynamic shape operator: aten.repeat_interleave.Tensor
     "RepeatInterleaveModule_basic",
+    "RepeatInterleaveFillModule_basic",
 
     # failed to legalize operation 'torch.aten.convolution' that was explicitly marked illegal
     "Conv1dNoPaddingModule_basic",
@@ -1226,6 +1228,7 @@
     "ChunkListUnpack_Module_basic",
     "ChunkListUnpackUneven_Module_basic",
     "RepeatInterleaveStaticModule_basic",
+    "RepeatInterleaveFillModule_basic",
 }
 
 MAKE_FX_TOSA_PASS_SET = (TOSA_PASS_SET | {
@@ -1490,5 +1493,6 @@
     "ScatterValueFloatModule_basic",
     "ScatterValueIntModule_basic",
     "RepeatInterleaveModule_basic",
+    "RepeatInterleaveFillModule_basic",
     "Im2ColModule_basic",
 }

lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp

@@ -491,4 +491,11 @@ static void markDecomposedOpsAsIllegal(MLIRContext *context,
         auto opName = opOp->getAttr("name").cast<StringAttr>().getValue();
         return backendLegalOpsSet.contains(opName);
       });
+
+  // TODO: We need this for TOSA; other backends might be fine with this op
+  // having a dynamic sized output tensor.
+  target.addDynamicallyLegalOp<AtenRepeatInterleaveTensorOp>(
+    [](AtenRepeatInterleaveTensorOp op) {
+      return op.getOutputSize().getDefiningOp<ConstantIntOp>();
+  });
 }

lib/Dialect/Torch/Transforms/RecomposeComplexOps.cpp

@@ -393,6 +393,54 @@ class RecomposeChunkListUnpack : public OpRewritePattern<PrimListUnpackOp> {
     return success();
   }
 };
+class RecomposeRepeatInterleave : public OpRewritePattern<AtenRepeatInterleaveTensorOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenRepeatInterleaveTensorOp op,
+                                PatternRewriter &rewriter) const override {
+    if (!op.getOutputSize().getDefiningOp<ConstantNoneOp>())
+      return failure();
+
+    auto repeatsTy = dyn_cast<BaseTensorType>(op.getRepeats().getType());
+    if (!repeatsTy || !repeatsTy.areAllSizesKnown() || repeatsTy.getSizes().size() != 1) {
+      return rewriter.notifyMatchFailure(
+          op,
+          "Expected 1d tensor with static shape");
+    }
+    auto numElements = repeatsTy.getSizes()[0];
+
+    auto broadcast = op.getRepeats().getDefiningOp<AtenBroadcastToOp>();
+    if (!broadcast){
+      return rewriter.notifyMatchFailure(
+          op,
+          "Expected broadcast op defining repeat_interleave input");
+    }
+
+    auto fill = broadcast.getSelf().getDefiningOp<AtenFillScalarOp>();
+    if (!fill){
+      return rewriter.notifyMatchFailure(
+          op,
+          "Expected fill op defining broadcast/repeat_interleave input");
+    }
+
+    int64_t fillValue;
+    if (!matchPattern(fill.getValue(),
+                      m_TorchConstantInt(&fillValue))) {
+      return rewriter.notifyMatchFailure(
+          op,
+          "Expected fill value of fill.Scalar to be an integer constant");
+    }
+
+    auto outputSize = rewriter.create<Torch::ConstantIntOp>(
+          op->getLoc(), rewriter.getI64IntegerAttr(fillValue * numElements));
+    rewriter.replaceOpWithNewOp<AtenRepeatInterleaveTensorOp>(op, op.getType(), op.getRepeats(), outputSize);
+
+    if (op.getResult().use_empty())
+      rewriter.eraseOp(op);
+    return success();
+  }
+};
+
 } // namespace
 
 namespace {
@@ -412,6 +460,7 @@ class RecomposeComplexOpsPass
     patterns.add<RecomposeUnbindGetItem>(context);
     patterns.add<RecomposeSplitTensorPrimListUnpackOp>(context);
     patterns.add<RecomposeChunkListUnpack>(context);
+    patterns.add<RecomposeRepeatInterleave>(context);
 
     GreedyRewriteConfig config;
     config.useTopDownTraversal = true;

python/torch_mlir_e2e_test/test_suite/basic.py

@@ -1480,6 +1480,28 @@ def forward(self, x):
 def RepeatInterleaveModule_basic(module, tu: TestUtils):
     module.forward(torch.tensor([3, 1, 2, 4], dtype=torch.int))
 
+# ==============================================================================
+class RepeatInterleaveFillModule(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([1], torch.int, True),
+    ])
+    def forward(self, x):
+        x = torch.ops.aten.fill_(x, 2)
+        x = torch.ops.aten.expand(x, [16])
+        return torch.ops.aten.repeat_interleave(x)
+
+
+@register_test_case(module_factory=lambda: RepeatInterleaveFillModule())
+def RepeatInterleaveFillModule_basic(module, tu: TestUtils):
+    module.forward(torch.tensor([1], dtype=torch.int))
+
+
 # ==============================================================================
 
 class RepeatInterleaveStaticModule(torch.nn.Module):