Lower replication pad 1d to linalg backend

Signed-off-by: Zahid Wakeel <[email protected]>

Author: zahidwx

lib/Conversion/TorchToLinalg/TensorConstructors.cpp

@@ -116,6 +116,96 @@ class ConvertAtenConstantPadNdOp
 
 namespace {
 
+class ConvertAtenReplicationPad1dOp
+    : public OpConversionPattern<AtenReplicationPad1dOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(AtenReplicationPad1dOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
+      return failure();
+
+    Location loc = op.getLoc();
+    Value input = adaptor.getSelf();
+    auto inputType = llvm::cast<RankedTensorType>(input.getType());
+    int64_t inputRank = inputType.getRank();
+
+    if (inputRank < 2)
+      return rewriter.notifyMatchFailure(op, "input rank must be at least 2");
+
+    SmallVector<int64_t> padInts;
+    if (!matchPattern(op.getPadding(), m_TorchListOfConstantInts(padInts)))
+      return rewriter.notifyMatchFailure(
+          op, "only support constant int pad ranges");
+
+    if (padInts.size() != 2)
+      return rewriter.notifyMatchFailure(
+          op, "pad range must have exactly two values");
+
+    int64_t leftPad = padInts[0];
+    int64_t rightPad = padInts[1];
+
+    int64_t widthDim = inputRank - 1;
+    Type indexType = rewriter.getIndexType();
+    Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
+    Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+
+    SmallVector<Value> inputShape = getTensorSizes(rewriter, loc, input);
+    Value widthSize = inputShape[widthDim];
+    Value widthMinusOne = rewriter.create<arith::SubIOp>(loc, widthSize, one);
+
+    // Build offset and size arrays for slicing
+    SmallVector<OpFoldResult> allOneStrides(inputRank,
+                                            rewriter.getIndexAttr(1));
+    SmallVector<OpFoldResult> leftOffsets(inputRank, rewriter.getIndexAttr(0));
+    SmallVector<OpFoldResult> rightOffsets(inputRank, rewriter.getIndexAttr(0));
+    SmallVector<OpFoldResult> sizes(inputRank, rewriter.getIndexAttr(0));
+    for (int i = 0; i < inputRank; ++i)
+      sizes[i] = (i == widthDim) ? rewriter.getIndexAttr(1)
+                                 : getAsOpFoldResult(inputShape[i]);
+
+    rightOffsets[widthDim] = getAsOpFoldResult(widthMinusOne);
+
+    // Extract leftmost and rightmost slices
+    Value leftSlice = rewriter.create<tensor::ExtractSliceOp>(
+        loc, input, leftOffsets, sizes, allOneStrides);
+    Value rightSlice = rewriter.create<tensor::ExtractSliceOp>(
+        loc, input, rightOffsets, sizes, allOneStrides);
+
+    // Create repeated tiles
+    SmallVector<Value> resultParts;
+
+    if (leftPad > 0) {
+      SmallVector<Value> leftTiles(leftPad, leftSlice);
+      Value leftConcat =
+          rewriter.create<tensor::ConcatOp>(loc, widthDim, leftTiles);
+      resultParts.push_back(leftConcat);
+    }
+
+    resultParts.push_back(input);
+
+    if (rightPad > 0) {
+      SmallVector<Value> rightTiles(rightPad, rightSlice);
+      Value rightConcat =
+          rewriter.create<tensor::ConcatOp>(loc, widthDim, rightTiles);
+      resultParts.push_back(rightConcat);
+    }
+
+    Value result =
+        rewriter.create<tensor::ConcatOp>(loc, widthDim, resultParts);
+    Type resultType = getTypeConverter()->convertType(op.getType());
+    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultType, result);
+
+    return success();
+  }
+};
+
+} // namespace
+
+namespace {
+
 // Lower aten.replication_pad2d operator into a sequence of
 // tensor.extract_slice and tensor.concat operations.
 
@@ -621,6 +711,8 @@ void mlir::torch::torch_to_linalg::
   MLIRContext *context = patterns.getContext();
   target.addIllegalOp<AtenReplicationPad2dOp>();
   patterns.add<ConvertAtenReplicationPad2dOp>(typeConverter, context);
+  target.addIllegalOp<AtenReplicationPad1dOp>();
+  patterns.add<ConvertAtenReplicationPad1dOp>(typeConverter, context);
   target.addIllegalOp<AtenConstantPadNdOp>();
   patterns.add<ConvertAtenConstantPadNdOp>(typeConverter, context);
   target.addIllegalOp<AtenZerosOp, AtenOnesOp>();

projects/pt1/python/torch_mlir_e2e_test/test_suite/basic.py

@@ -685,6 +685,26 @@ def ReplicationPad2dModule_left0(module, tu: TestUtils):
 # ==============================================================================
 
 
+class ReplicationPad1dModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1, -1, -1], torch.float32, True),
+        ]
+    )
+    def forward(self, x):
+        return torch.ops.aten.replication_pad1d(x, [3, 5])
+
+
+@register_test_case(module_factory=lambda: ReplicationPad1dModule())
+def ReplicationPad1dModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(1, 15, 20, low=-1))
+
+
 class ReplicationPad2dModule_right0_module(torch.nn.Module):
     def __init__(self):
         super().__init__()