[Torch Dialect] Decompose AtenTriuOp

Author: Mi-Jiazhi

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -246,6 +246,66 @@ class DecomposeAtenAmaxOp : public OpRewritePattern<AtenAmaxOp> {
 };
 } // end namespace
 
+namespace {
+class DecomposeAtenTriuOp : public OpRewritePattern<AtenTriuOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenTriuOp op,
+                                PatternRewriter &rewriter) const override {
+    MLIRContext *context = op.getContext();
+    Location loc = op.getLoc();
+    Value input = op.getSelf();
+    auto inputType = input.getType().cast<BaseTensorType>();
+    if (!inputType.hasSizes() || !inputType.hasDtype()) {
+      return rewriter.notifyMatchFailure(op, "should have shape and dtype");
+    }
+    if (inputType.getSizes().size() < 2) {
+      return rewriter.notifyMatchFailure(op, "the rank of tensor should >= 2");
+    }
+
+    Value rowDim = rewriter.create<Torch::ConstantIntOp>(
+        loc, rewriter.getI64IntegerAttr(-2));
+    Value colDim = rewriter.create<Torch::ConstantIntOp>(
+        loc, rewriter.getI64IntegerAttr(-1));
+    Value rowSize = rewriter.create<AtenSizeIntOp>(loc, input, rowDim);
+    Value colSize = rewriter.create<AtenSizeIntOp>(loc, input, colDim);
+
+    auto baseType = ValueTensorType::getWithLeastStaticInformation(context);
+
+    Value none = rewriter.create<ConstantNoneOp>(loc);
+    Value rowArange = rewriter.create<AtenArangeOp>(
+        loc, baseType, rowSize, /*dtype=*/none, /*layout=*/none,
+        /*device=*/none, /*pin_memory=*/none);
+    Value colArange = rewriter.create<AtenArangeOp>(
+        loc, baseType, colSize, /*dtype=*/none, /*layout=*/none,
+        /*device=*/none, /*pin_memory=*/none);
+
+    Value unsqueezeRowArange = rewriter.create<AtenUnsqueezeOp>(
+        loc, baseType, rowArange,
+        rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(1)));
+
+    Value unsqueezeColArange = rewriter.create<AtenUnsqueezeOp>(
+        loc, baseType, colArange,
+        rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(0)));
+
+    Value unsqueezeRowArangePlusDiagonal = rewriter.create<AtenAddScalarOp>(
+        loc, baseType, unsqueezeRowArange,
+        rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(1)),
+        op.getDiagonal());
+
+    Value condTensor = rewriter.create<AtenGeTensorOp>(
+        loc, baseType, unsqueezeColArange, unsqueezeRowArangePlusDiagonal);
+
+    auto others =
+        rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(0));
+
+    rewriter.replaceOpWithNewOp<AtenWhereScalarOtherOp>(
+        op, op.getResult().getType(), condTensor, input, others);
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 class DecomposeAtenSizeOp : public OpRewritePattern<AtenSizeOp> {
 public:
@@ -5819,6 +5879,7 @@ class DecomposeComplexOpsPass
     addPatternIfTargetOpIsIllegal<DecomposeAtenTileOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenReshapeAsOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenIndexTensorOp>(patterns);
+    addPatternIfTargetOpIsIllegal<DecomposeAtenTriuOp>(patterns);
 
     GreedyRewriteConfig config;
     config.useTopDownTraversal = true;

lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp

@@ -500,6 +500,7 @@ static void markDecomposedOpsAsIllegal(MLIRContext *context,
   target.addIllegalOp<AtenTypeAsOp>();
   target.addIllegalOp<AtenTileOp>();
   target.addIllegalOp<AtenReshapeAsOp>();
+  target.addIllegalOp<AtenTriuOp>();
   for (auto &opName : backendLegalOpsSet) {
     target.addLegalOp(
         OperationName(kTorchOpPrefix + opName.first().str(), context));

projects/pt1/e2e_testing/xfail_sets.py

@@ -1308,6 +1308,8 @@
     "MeanModule_basic",
     "ArangeStartOutModule_basic",
     "ArangeStartOutViewModule_basic",
+    "TriuModule_basic",
+    "TriuBroadcastModule_basic",
 }
 
 MAKE_FX_TOSA_PASS_SET = (TOSA_PASS_SET | {

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

@@ -3251,6 +3251,53 @@ def AtenTriuWithPosDiagonalModule_basic(module, tu: TestUtils):
 # ==============================================================================
 
 
+class TriuModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([4,5], torch.float32, True),
+    ])
+    def forward(self, x):
+        return torch.ops.aten.triu(x, 1)
+
+
+@register_test_case(module_factory=lambda: TriuModule())
+def TriuModule_basic(module, tu: TestUtils):
+    x=torch.tensor([[ 0.5876, -0.0794, -1.8373,  0.6654, 0.2],
+        [-0.2447,  0.9556, -1.2919,  1.3378, 0.3],
+        [ 0.4333,  0.3146,  0.6576, -1.0432, 0.4],
+        [-0.9888,  torch.nan, torch.inf, -torch.inf, 0.5]])
+    module.forward(x)
+
+
+# ==============================================================================
+
+
+class TriuBroadcastModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([3,4,5,6], torch.float32, True),
+        ([], torch.int32, True),
+    ])
+    def forward(self, x, diagonal):
+        return torch.ops.aten.triu(x, diagonal)
+
+
+@register_test_case(module_factory=lambda: TriuBroadcastModule())
+def TriuBroadcastModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3,4,5,6), tu.randint())
+
+
+# ==============================================================================
+
+
 class AtenTriuWithNegDiagonalModule(torch.nn.Module):
 
     def __init__(self):