Add E2E support for aten.as_strided

e2e_testing/xfail_sets.py

@@ -607,6 +607,7 @@
     "ViewCollapseDynamicWithAtenSizeIntModule_basic",
     "AtenEmbeddingBagSumExample_basic",
     "Aten_EmbeddingBagExample_basic",
+    "AsStridedModule_basic",
     "ElementwiseRemainderScalarModule_Int_Float_basic",
     "ElementwiseRemainderScalarModule_Float_basic",
     "ElementwiseRemainderScalarModule_Int_basic",
@@ -627,5 +628,5 @@
     "ConvolutionBackwardModule2DPadded_basic",
     "ConvolutionBackwardModule3D_basic",
     "VarMeanCorrectionModule_basic",
-    "VarMeanCorrectionNoneModule_basic"
+    "VarMeanCorrectionNoneModule_basic",
 }

include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td

@@ -7466,6 +7466,32 @@ def Torch_AtenAliasCopyOp : Torch_Op<"aten.alias_copy", [
   }];
 }
 
+def Torch_AtenAsStridedOp : Torch_Op<"aten.as_strided", [
+    AllowsTypeRefinement,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::as_strided : (Tensor, int[], int[], int?) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    AnyTorchListOfTorchIntType:$size,
+    AnyTorchListOfTorchIntType:$stride,
+    AnyTorchOptionalIntType:$storage_offset
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenAsStridedOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 4, 1);
+    }
+    void AtenAsStridedOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 4, 1);
+    }
+  }];
+  let hasFolder = 1;
+}
+
 def Torch_AtenAsStridedCopyOp : Torch_Op<"aten.as_strided_copy", [
     AllowsTypeRefinement,
     HasValueSemantics,

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -2132,6 +2132,18 @@ OpFoldResult AtenSqrtIntOp::fold(ArrayRef<Attribute> operands) {
   return nullptr;
 }
 
+//===----------------------------------------------------------------------===//
+// AtenAsStridedOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult AtenAsStridedOp::fold(ArrayRef<Attribute> operands) {
+  if (auto tensorType = getOperand(0).getType().dyn_cast<BaseTensorType>()) {
+    if (tensorType.hasSizes() && tensorType.getSizes().size() == 0)
+      return getOperand(0);
+  }
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // PrimDtypeOp
 //===----------------------------------------------------------------------===//

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -2948,6 +2948,34 @@ class DecomposeAtenSelectScatterOp
 };
 } // namespace
 
+namespace {
+class DecomposeAtenAsStridedOp
+    : public OpRewritePattern<AtenAsStridedOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenAsStridedOp op,
+                                PatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    Value self = op.self();
+    Value size = op.size();
+    Value stride = op.stride();
+    Value storage_offset = op.storage_offset();
+
+    //Value result = 
+
+
+    //Value src = rewriter.create<AtenAsStridedOp>(loc, op.self().getType(), size, stride, storage_offset);
+    rewriter.replaceOpWithNewOp<AtenAsStridedOp>(
+        op, op.getType(), self, size, stride, storage_offset);
+
+
+
+
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 class DecomposeAten_EmbeddingBagOp
     : public OpRewritePattern<Aten_EmbeddingBagOp> {
@@ -3328,6 +3356,8 @@ class DecomposeComplexOpsPass
     target.addIllegalOp<AtenStdDimOp>();
     patterns.add<DecomposeAtenNarrowOp>(context);
     target.addIllegalOp<AtenNarrowOp>();
+    patterns.add<DecomposeAtenAsStridedOp>(context);
+    target.addIllegalOp<AtenAsStridedOp>();
     patterns.add<DecomposeAten_EmbeddingBagOp>(context);
     target.addIllegalOp<Aten_EmbeddingBagOp>();
     patterns.add<DecomposeAtenLiftFreshCopyOp>(context);

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -699,7 +699,7 @@ void TypeAnalysis::visitOperation(Operation *op,
           AtenCopyOp, AtenZeroOp, AtenIndexPutHackedTwinOp,
           AtenMaskedFillScalarOp, AtenFlipOp, PrimAbsScalarOp, AtenNumpyTOp,
           AtenTriuOp, AtenMaskedFillTensorOp, AtenRollOp, AtenPowTensorTensorOp,
-          AtenLiftFreshCopyOp, AtenIndexTensorHackedTwinOp,
+          AtenLiftFreshCopyOp, AtenIndexTensorHackedTwinOp, AtenAsStridedOp,
           AtenUpsampleNearest2dVecOp, AtenMishOp, AtenRoundOp, AtenFillTensorOp,
           AtenUpsampleNearest2dBackwardOp>(op)) {
     return incorporateKnowledge(op->getResult(0), operands[0]->getValue());

lib/Dialect/Torch/Transforms/ShapeLibrary.cpp

@@ -6601,6 +6601,10 @@ StringRef mlir::torch::Torch::getShapeLibrary() {
 "    %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list<int>) -> !torch.list<int>\n"
 "    return %0 : !torch.list<int>\n"
 "  }\n"
+"  func.func @\"__torch_mlir_shape_fn.aten.as_strided\"(%arg0: !torch.list<int>, %arg1: !torch.list<int>, %arg2: !torch.list<int>, %arg3: !torch.optional<int>) -> !torch.list<int> {\n"
+"    %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list<int>) -> !torch.list<int>\n"
+"    return %0 : !torch.list<int>\n"
+"  }\n"
 "  func.func @\"__torch_mlir_shape_fn.aten.embedding_bag.padding_idx\"(%arg0: !torch.list<int>, %arg1: !torch.list<int>, %arg2: !torch.list<int>, %arg3: !torch.bool, %arg4: !torch.int, %arg5: !torch.bool, %arg6: !torch.optional<list<int>>, %arg7: !torch.bool, %arg8: !torch.optional<int>) -> !torch.tuple<list<int>, list<int>, list<int>, list<int>> {\n"
 "    %0 = call @__torch__._embedding_bag_helper(%arg0, %arg1, %arg2, %arg7, %arg4) : (!torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.int) -> !torch.tuple<list<int>, list<int>, list<int>, list<int>>\n"
 "    return %0 : !torch.tuple<list<int>, list<int>, list<int>, list<int>>\n"

python/torch_mlir/__init__.py

@@ -344,6 +344,7 @@ def compile(model: torch.nn.Module,
     mb = ModuleBuilder()
     import_options = ImportOptions()
     import_options.ignoreExistingTensorShapesAndDtypes = ignore_traced_shapes
+
     try:
         original_stderr = sys.stderr
         sys.stderr = StringIO()
@@ -363,12 +364,15 @@ def compile(model: torch.nn.Module,
         return mb.module
 
     option_string = "{backend-legal-ops=" + ",".join(backend_legal_ops) + "}"
+    #mb.module.dump()
     run_pipeline_with_repro_report(
         mb.module,
         f"builtin.module(torchscript-module-to-torch-backend-pipeline{option_string})",
         "Lowering TorchScript IR -> Torch Backend IR",
     )
 
+
+
     if verbose:
         print("\n====================")
         print("Torch Backend IR")

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/shape_lib_gen.py

@@ -1027,6 +1027,9 @@ def aten〇index_put(self: List[int], indices: List[Optional[List[int]]], values
 def aten〇index_put〇hacked_twin(self: List[int], indices: List[List[int]], values: List[int], accumulate: bool = False) -> List[int]:
     return upstream_shape_functions.unary(self)
 
+def aten〇as_strided(self: List[int], size: List[int], stride: List[int], storage_offset: Optional[int] = None) -> List[int]:
+    return upstream_shape_functions.unary(self)
+
 def aten〇embedding(weight: List[int], indices: List[int], padding_idx: int = -1, scale_grad_by_freq: bool = False, sparse: bool = False) -> List[int]:
     return upstream_shape_functions.embedding(weight, indices, padding_idx, scale_grad_by_freq, sparse)
 

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py

@@ -508,6 +508,7 @@ def emit_with_mutating_variants(key, **kwargs):
 
     # Functionalization ops
     emit("aten::alias_copy : (Tensor) -> (Tensor)")
+    emit("aten::as_strided : (Tensor, int[], int[], int?) -> (Tensor)", has_folder=True)
     emit("aten::as_strided_copy : (Tensor, int[], int[], int?) -> (Tensor)")
     emit("aten::diagonal_copy : (Tensor, int, int, int) -> (Tensor)")
     emit("aten::expand_copy : (Tensor, int[], bool) -> (Tensor)")

python/torch_mlir_e2e_test/test_suite/basic.py

@@ -2954,6 +2954,26 @@ def Aten_EmbeddingBagExample_basic(module, tu: TestUtils):
 
 # ==============================================================================
 
+class AsStridedModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+    ])
+
+    def forward(self):
+        x = torch.randn(5, 5)
+        print (x)
+        return torch.ops.aten.as_strided(x, (2, 2), (5, 3))
+
+@register_test_case(module_factory=lambda: AsStridedModule())
+def AsStridedModule_basic(module, tu: TestUtils):
+    module.forward()
+
+# ==============================================================================
+
 class CumsumModule(torch.nn.Module):
 
     def __init__(self):