Repository: alibaba/BladeDISC (language: python)

Issue:
fix lowering failed if `aten.split`'s operand is dynamic shape
IR example lowering failed on DiscDecomposeComplexOps pass:

``` mlir
  %9 = torch.aten.reshape %7, %8 : !torch.vtensor<[?,?],f16>, !torch.list<int> -> !torch.vtensor<[?,?,?],f16>
  %10 = torch.operator "aten.split.Tensor"(%9, %int1280, %int-1) : (!torch.vtensor<[?,?,?],f16>, !torch.int, !torch.int) -> !torch.list<vtensor>
  %11 = torch.aten.__getitem__.t %10, %int0 : !torch.list<vtensor>, !torch.int -> !torch.vtensor<[?,?,?],f16>
  %12 = torch.aten.__getitem__.t %10, %int1 : !torch.list<vtensor>, !torch.int -> !torch.vtensor<[?,?,?],f16>
  %13 = torch.aten.to.dtype %12, %int6, %false, %false, %none : !torch.vtensor<[?,?,?],f16>, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor
```


Discussion / hints:

Fix patch:
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/BUILD b/pytorch_blade/pytorch_blade/torch-mlir/BUILD
index ebd78979bd0..2e8b077fc31 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/BUILD
+++ b/pytorch_blade/pytorch_blade/torch-mlir/BUILD
@@ -58,6 +58,7 @@ cc_library(
         "lib/Dialect/TorchConversion/Transforms/DiscConvertTorchToDiscMhlo.cpp",
         "lib/Dialect/TorchConversion/Transforms/UpgradeLegacyOps.cpp",
         "lib/Dialect/TorchConversion/Transforms/VerifyMhloBackendContract.cpp",
+        "lib/Dialect/TorchConversion/Transforms/DiscSimplifyPatterns.cpp",
     ],
     hdrs = [
         "include/torch-mlir/Conversion/MhloPasses.h",
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/include/torch-mlir/Conversion/MhloPasses.h b/pytorch_blade/pytorch_blade/torch-mlir/include/torch-mlir/Conversion/MhloPasses.h
index 4a16f71e247..ccb3dd6f7d8 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/include/torch-mlir/Conversion/MhloPasses.h
+++ b/pytorch_blade/pytorch_blade/torch-mlir/include/torch-mlir/Conversion/MhloPasses.h
@@ -68,6 +68,7 @@ std::unique_ptr<OperationPass<func::FuncOp>> createApplyDiscPdlPatternsPass(
     const std::string& pdll_files = "",
     const std::string& pdll_include_dirs = "");
 std::unique_ptr<OperationPass<func::FuncOp>> createDiscConvertTorchToMhloPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createDiscSimplifyPatternsPass();
 std::unique_ptr<OperationPass<func::FuncOp>> createDiscDecomposeComplexOpsPass();
 std::unique_ptr<OperationPass<func::FuncOp>> createDiscConvertTorchToDiscMhlo();
 std::unique_ptr<OperationPass<func::FuncOp>> createDiscUpgradeLegacyOpsPass();
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/include/torch-mlir/Conversion/MhloPasses.td b/pytorch_blade/pytorch_blade/torch-mlir/include/torch-mlir/Conversion/MhloPasses.td
index d20d62e87ce..0c3efc91257 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/include/torch-mlir/Conversion/MhloPasses.td
+++ b/pytorch_blade/pytorch_blade/torch-mlir/include/torch-mlir/Conversion/MhloPasses.td
@@ -98,4 +98,12 @@ def DiscConvertTorchToDiscMhlo : Pass<"convert-torch-to-disc-mhlo", "func::FuncO
  let constructor = "mlir::torch::TorchConversion::createDiscConvertTorchToDiscMhlo()";
 }
 
+def DiscSimplifyPatternsPass : Pass<"disc-simplify-patterns", "func::FuncOp"> {
+  let summary = "simplify Torch Dialect patterns";
+  let description = [{
+    simplify patterns based on Torch Dialect.
+  }];
+  let constructor = "mlir::torch::TorchConversion::createDiscSimplifyPatternsPass()";
+}
+
 #endif // TORCHMLIR_CONVERSION_DISC_MHLO_PASSES
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp
index 0a29e003d47..642a2afed33 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp
@@ -58,6 +58,10 @@ void mlir::torch::createDiscTorchBackendToMhloBackendPipeline(
   // Add decompose passes
   pm.addNestedPass<func::FuncOp>(createDiscDecomposeComplexOpsPass());
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+  // Add simplify patterns pass
+  pm.addNestedPass<func::FuncOp>(createDiscSimplifyPatternsPass());
+  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+
   pm.addNestedPass<func::FuncOp>(
       Torch::createDecomposeComplexOpsPass(/*legalOps*/ {}));
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp
index 4e80752fa5b..fb1ba009ce4 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp
@@ -84,6 +84,22 @@ class ConvertAtenOp : public OpConversionPattern<AtenOpT> {
       ConversionPatternRewriter& rewriter) const override;
 };
 
+llvm::Optional<int64_t> getMaxIndexFromItemOps(OperatorOp op) {
+  int64_t maxIndex = -1;
+  for (Operation* user : op.getResult(0).getUsers()) {
+    if (mlir::isa<Aten__Getitem__TOp>(user)) {
+      int64_t indexInt;
+      auto indexValue = user->getOperand(1);
+      if (!matchPattern(indexValue, m_TorchConstantInt(&indexInt)))
+        return llvm::None;
+      maxIndex = std::max(maxIndex, indexInt);
+    } else {
+      return llvm::None;
+    }
+  }
+  return maxIndex;
+}
+
 LogicalResult decomposeSplits(
     ConversionPatternRewriter& rewriter,
     OperatorOp op,
@@ -91,9 +107,8 @@ LogicalResult decomposeSplits(
     Value dim,
     int64_t chunks,
     bool keepDim = true) {
-  if (chunks < 0) {
+  if (chunks < 0)
     return failure();
-  }
   int64_t dimInt;
   if (!matchPattern(dim, m_TorchConstantInt(&dimInt)))
     return rewriter.notifyMatchFailure(op, "unknown dim");
@@ -220,6 +235,12 @@ LogicalResult ConvertAtenOp<OperatorOp>::matchAndRewrite(
         chunksInt = inputShape[dimInt] / chunkSizeInt;
       }
     }
+    if (chunksInt < 0) {
+      // inference result number according to the max index of aten.item ops.
+      auto maxItemIndex = getMaxIndexFromItemOps(op);
+      if (maxItemIndex)
+        chunksInt = maxItemIndex.value() + 1;
+    }
     return decomposeSplits(rewriter, op, chunkSize, dim, chunksInt);
   } else if ("aten.chunk" == name) {
     int64_t chunksInt = -1;
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscSimplifyPatterns.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscSimplifyPatterns.cpp
new file mode 100644
index 00000000000..0f85a5be626
--- /dev/null
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscSimplifyPatterns.cpp
@@ -0,0 +1,83 @@
+// Copyright 2023 The BladeDISC Authors. All rights reserved.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "torch-mlir/Conversion/MhloPasses.h"
+#include "torch-mlir/Conversion/TorchToMhlo/TorchToMhlo.h"
+#include "torch-mlir/Conversion/Utils/Utils.h"
+#include "torch-mlir/Dialect/Torch/IR/TorchDialect.h"
+#include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
+#include "torch-mlir/Dialect/Torch/IR/TorchTypes.h"
+#include "torch-mlir/Dialect/Torch/Utils/Utils.h"
+#include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionOps.h"
+
+using namespace mlir;
+using namespace mlir::torch;
+using namespace mlir::torch::Torch;
+using namespace mlir::torch::TorchConversion;
+namespace {
+
+class SimplifyGetitemOp : public OpConversionPattern<Aten__Getitem__TOp> {
+ public:
+  using OpConversionPattern<Aten__Getitem__TOp>::OpConversionPattern;
+  using OpAdaptor = typename Aten__Getitem__TOp::Adaptor;
+  LogicalResult matchAndRewrite(
+      Aten__Getitem__TOp op,
+      OpAdaptor adaptor,
+      ConversionPatternRewriter& rewriter) const override {
+    auto input = op.getOperand(0);
+    auto index = op.getOperand(1);
+    auto result = op.getResult();
+
+    int64_t indexInt;
+    if (!matchPattern(op.idx(), m_TorchConstantInt(&indexInt)))
+      return failure();
+
+    auto listOp = input.getDefiningOp<PrimListConstructOp>();
+    if (!listOp)
+      return failure();
+
+    SmallVector<Value> torchTensors;
+    if (!getListConstructElements(listOp, torchTensors)) {
+      return failure();
+    }
+    for (Operation* user : result.getUsers()) {
+      user->replaceUsesOfWith(result, torchTensors[indexInt]);
+    }
+    op.erase();
+    return success();
+  }
+};
+
+class DiscSimplifyPatternsPass
+    : public DiscSimplifyPatternsPassBase<DiscSimplifyPatternsPass> {
+  void runOnOperation() override {
+    MLIRContext* context = &getContext();
+    ConversionTarget target(*context);
+    target.addLegalDialect<Torch::TorchDialect>();
+
+    RewritePatternSet patterns(context);
+    patterns.add<SimplifyGetitemOp>(context);
+
+    if (failed(applyPartialConversion(
+            getOperation(), target, std::move(patterns)))) {
+      return signalPassFailure();
+    }
+  }
+};
+
+} // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> mlir::torch::TorchConversion::
+    createDiscSimplifyPatternsPass() {
+  return std::make_unique<DiscSimplifyPatternsPass>();
+}


Test patch:
diff --git a/pytorch_blade/tests/mhlo/disc_simlifier.mlir b/pytorch_blade/tests/mhlo/disc_simlifier.mlir
new file mode 100644
index 00000000000..c65561475b6
--- /dev/null
+++ b/pytorch_blade/tests/mhlo/disc_simlifier.mlir
@@ -0,0 +1,11 @@
+// RUN: torch-mlir-opt <%s --disc-simplify-patterns --canonicalize -split-input-file -verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL:  func.func @disc_simplifier_pass(
+// CHECK-SAME:         %[[ARG0:.*]]: !torch.vtensor<[?,?,?],f16>) -> !torch.vtensor<[?,?,?],f16>
+// CHECK: return %[[ARG0]] : !torch.vtensor<[?,?,?],f16>
+func.func @disc_simplifier_pass(%arg0: !torch.vtensor<[?,?,?],f16>) -> !torch.vtensor<[?,?,?],f16> {
+  %int0 = torch.constant.int 0
+  %1 = torch.prim.ListConstruct %arg0 : (!torch.vtensor<[?,?,?],f16>) -> !torch.list<vtensor>
+  %2 = torch.aten.__getitem__.t %1, %int0 : !torch.list<vtensor>, !torch.int -> !torch.vtensor<[?,?,?],f16>
+  return %2 : !torch.vtensor<[?,?,?],f16>
+}
diff --git a/pytorch_blade/tests/mhlo/split.mlir b/pytorch_blade/tests/mhlo/split.mlir
new file mode 100644
index 00000000000..62fe714c6ea
--- /dev/null
+++ b/pytorch_blade/tests/mhlo/split.mlir
@@ -0,0 +1,19 @@
+// RUN: torch-mlir-opt <%s --torch-disc-decompose-complex-ops -split-input-file -verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL:  func.func @decompose_split(
+// CHECK-SAME:         %[[ARG0:.*]]: !torch.vtensor<[?,?,?],f16>) -> (!torch.vtensor<[?,?,?],f16>, !torch.vtensor<[?,?,?],f16>) {
+// CHECK:           %[[T0:.*]] = torch.aten.add.int %int0_1, %int1280 : !torch.int, !torch.int -> !torch.int
+// CHECK:           %[[T1:.*]] = torch.aten.slice.Tensor %arg0, %int-1, %int0_1, %0, %int1_0 : !torch.vtensor<[?,?,?],f16>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[?,?,?],f16>
+// CHECK:           %[[T2:.*]] = torch.aten.add.int %0, %int1280 : !torch.int, !torch.int -> !torch.int
+// CHECK:           %[[T3:.*]] = torch.aten.slice.Tensor %arg0, %int-1, %0, %2, %int1_0 : !torch.vtensor<[?,?,?],f16>, !torch.int, !torch.int, !torch.int, !torch.int -> !torch.vtensor<[?,?,?],f16>
+// CHECK:           %[[L0:.*]] = torch.prim.ListConstruct %[[T1]], %[[T3]] : (!torch.vtensor<[?,?,?],f16>, !torch.vtensor<[?,?,?],f16>) -> !torch.list<vtensor>
+func.func @decompose_split(%arg0: !torch.vtensor<[?,?,?],f16>) -> (!torch.vtensor<[?,?,?],f16>, !torch.vtensor<[?,?,?],f16>) {
+  %int1280 = torch.constant.int 1280
+  %int-1 = torch.constant.int -1
+  %int0 = torch.constant.int 0
+  %int1 = torch.constant.int 1
+  %10 = torch.operator "aten.split.Tensor"(%arg0, %int1280, %int-1) : (!torch.vtensor<[?,?,?],f16>, !torch.int, !torch.int) -> !torch.list<vtensor>
+  %11 = torch.aten.__getitem__.t %10, %int0 : !torch.list<vtensor>, !torch.int -> !torch.vtensor<[?,?,?],f16>
+  %12 = torch.aten.__getitem__.t %10, %int1 : !torch.list<vtensor>, !torch.int -> !torch.vtensor<[?,?,?],f16>
+  return %11, %12 : !torch.vtensor<[?,?,?],f16>, !torch.vtensor<[?,?,?],f16>
+}
\ No newline at end of file