[TorchToStablehlo] [TorchToTosa] Add lowering support for torch.aten.triu (#4637)

## Description

Add support for lowering `torch.aten.triu` in both `TorchToStablehlo`
and `TorchToTosa` conversion passes.

To prevent code duplication, the implementation refactors and unifies
the core logic of `AtenTrilOp` and `AtenTriuOp` under shared helper
templates (`convertTrilOrTriu` and `convertTrilOrTriuTosa`),
parameterizing only the comparison direction (GE for triu, LE for tril)
and mask generation.

## Tests

1. see unit tests added in the commit.
2. run e2e test via
https://gist.github.com/hsqStephenZhang/66e0ed68cc43acdc18a4f4d9959c07a7.
all backends generate code as efficient as code generated for `tril`

Fixes #4604

---------

Signed-off-by: hsqStephenZhang <stephenzhang666666@gmail.com>

Author: hsqStephenZhang

lib/Conversion/TorchToStablehlo/Basic.cpp

@@ -2109,15 +2109,13 @@ LogicalResult ConvertAtenOp<AtenBitwiseRightShiftTensorOp>::matchAndRewrite(
   return success();
 }
 
-template <>
-LogicalResult ConvertAtenOp<AtenTrilOp>::matchAndRewrite(
-    AtenTrilOp op, OpAdaptor adaptor,
-    ConversionPatternRewriter &rewriter) const {
-
+template <typename AtenOpT>
+LogicalResult convertTrilOrTriu(AtenOpT op, Value self, Value diagonal,
+                                stablehlo::ComparisonDirection dir,
+                                const TypeConverter *typeConverter,
+                                ConversionPatternRewriter &rewriter) {
   Location loc = op.getLoc();
 
-  Value self = adaptor.getSelf();
-
   auto selfTy = cast<RankedTensorType>(self.getType());
   if (!selfTy.hasStaticShape()) {
     return op->emitError("dynamic shaped input is not supported");
@@ -2133,25 +2131,23 @@ LogicalResult ConvertAtenOp<AtenTrilOp>::matchAndRewrite(
   Value rowIdxTensor =
       stablehlo::IotaOp::create(rewriter, loc, iotaTy, 0).getResult();
 
-  Value diagonal = adaptor.getDiagonal();
   Value diagonalTensor =
       tensor::FromElementsOp::create(rewriter, loc, diagonal).getResult();
 
   auto bcastDimensions = rewriter.getDenseI64ArrayAttr({1});
   Value shiftedRowIdxTensor = chlo::BroadcastAddOp::create(
       rewriter, loc, rowIdxTensor, diagonalTensor, bcastDimensions);
 
-  auto cmpDirectionAttr = stablehlo::ComparisonDirectionAttr::get(
-      rewriter.getContext(), stablehlo::ComparisonDirection::LE);
+  auto cmpDirectionAttr =
+      stablehlo::ComparisonDirectionAttr::get(rewriter.getContext(), dir);
   auto cmpTypeAttr = stablehlo::ComparisonTypeAttr::get(
       rewriter.getContext(), stablehlo::ComparisonType::SIGNED);
   auto cmpTy = iotaTy.clone(rewriter.getI1Type());
   Value cmpRes = stablehlo::CompareOp::create(rewriter, loc, cmpTy,
                                               colIdxTensor, shiftedRowIdxTensor,
                                               cmpDirectionAttr, cmpTypeAttr);
 
-  auto resTy =
-      cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));
+  auto resTy = cast<RankedTensorType>(typeConverter->convertType(op.getType()));
 
   auto bcastTy = resTy.clone(rewriter.getI1Type());
   auto bcastAttr = rewriter.getDenseI64ArrayAttr({selfRank - 2, selfRank - 1});
@@ -2162,16 +2158,17 @@ LogicalResult ConvertAtenOp<AtenTrilOp>::matchAndRewrite(
   Value zeroTensor;
   if (isa<mlir::FloatType>(resElemTy)) {
     auto constAttr = SplatElementsAttr::get(
-        resTy, llvm::APFloat::getZero(
-                   cast<FloatType>(resElemTy).getFloatSemantics(), false));
+        resTy,
+        llvm::APFloat::getZero(
+            cast<mlir::FloatType>(resElemTy).getFloatSemantics(), false));
     zeroTensor = stablehlo::ConstantOp::create(rewriter, loc, resTy, constAttr);
   } else if (isa<mlir::IntegerType>(resElemTy)) {
     auto constAttr = SplatElementsAttr::get(
         resTy,
         llvm::APInt::getZero(cast<mlir::IntegerType>(resElemTy).getWidth()));
     zeroTensor = stablehlo::ConstantOp::create(rewriter, loc, resTy, constAttr);
   } else {
-    return op.emitError("element type is not float or integer");
+    return op->emitError("element type is not float or integer");
   }
 
   rewriter.replaceOpWithNewOp<stablehlo::SelectOp>(
@@ -2180,6 +2177,24 @@ LogicalResult ConvertAtenOp<AtenTrilOp>::matchAndRewrite(
   return success();
 }
 
+template <>
+LogicalResult ConvertAtenOp<AtenTrilOp>::matchAndRewrite(
+    AtenTrilOp op, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  return convertTrilOrTriu(op, adaptor.getSelf(), adaptor.getDiagonal(),
+                           stablehlo::ComparisonDirection::LE,
+                           getTypeConverter(), rewriter);
+}
+
+template <>
+LogicalResult ConvertAtenOp<AtenTriuOp>::matchAndRewrite(
+    AtenTriuOp op, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  return convertTrilOrTriu(op, adaptor.getSelf(), adaptor.getDiagonal(),
+                           stablehlo::ComparisonDirection::GE,
+                           getTypeConverter(), rewriter);
+}
+
 template <>
 LogicalResult ConvertAtenOp<AtenIsfiniteOp>::matchAndRewrite(
     AtenIsfiniteOp op, OpAdaptor adaptor,
@@ -2484,6 +2499,7 @@ void mlir::torch::torch_to_stablehlo::populateBasicOpPatternsAndLegality(
   INSERT_ATENOP_PATTERN(AtenBitwiseRightShiftTensorOp);
 
   INSERT_ATENOP_PATTERN(AtenTrilOp);
+  INSERT_ATENOP_PATTERN(AtenTriuOp);
   INSERT_ATENOP_PATTERN(AtenIsfiniteOp);
   INSERT_ATENOP_PATTERN(AtenSortOp);
 

lib/Conversion/TorchToTosa/TorchToTosa.cpp

@@ -8795,19 +8795,20 @@ ConvertAtenOp<Aten__InterpolateSizeListScaleListOp>::matchAndRewriteImpl(
   return success();
 }
 
-// Template to create supporting tril mask tensor for aten.tril
+// Template to create supporting mask tensor for aten.tril/triu
 template <typename T>
-Value createTrilMask(PatternRewriter &rewriter, Operation *op,
-                     ArrayRef<int64_t> shape, int64_t h, int64_t w,
-                     int64_t diagonal) {
+Value createTrilOrTriuMask(PatternRewriter &rewriter, Operation *op,
+                           ArrayRef<int64_t> shape, int64_t h, int64_t w,
+                           int64_t diagonal, bool isTril) {
   SmallVector<T> vec;
 
   for (int64_t i = 0; i < h; i++) {
     for (int64_t j = 0; j < w; j++) {
       // Positive diagonal value includes as many diagonals above the main
       // diagonal, while negative diagonal value excludes as many diagonals
       // below the main diagonal.
-      if (i >= j - diagonal) {
+      auto cmp = isTril ? i >= j - diagonal : i <= j - diagonal;
+      if (cmp) {
         vec.push_back(static_cast<T>(1));
       } else {
         vec.push_back(static_cast<T>(0));
@@ -8818,13 +8819,10 @@ Value createTrilMask(PatternRewriter &rewriter, Operation *op,
   return tosa::getConstTensor<T>(rewriter, op, vec, shape).value();
 }
 
-// Legalization for aten.tril
-template <>
-LogicalResult ConvertAtenOp<AtenTrilOp>::matchAndRewriteImpl(
-    AtenTrilOp op, OpAdaptor adaptor,
-    ConversionPatternRewriter &rewriter) const {
-  auto self = adaptor.getSelf();
-
+template <typename AtenOpT>
+LogicalResult convertTrilOrTriu(AtenOpT op, Value self, Value diagonalVal,
+                                bool isTril, const TypeConverter *typeConverter,
+                                ConversionPatternRewriter &rewriter) {
   // Not a ranked tensor type
   auto selfType = dyn_cast<RankedTensorType>(self.getType());
   if (!selfType)
@@ -8841,27 +8839,33 @@ LogicalResult ConvertAtenOp<AtenTrilOp>::matchAndRewriteImpl(
     return rewriter.notifyMatchFailure(
         op, "Currently only static shapes are supported");
 
-  const TypeConverter *typeConverter = this->getTypeConverter();
   RankedTensorType resultType = cast<RankedTensorType>(
       typeConverter->convertType(op->getResult(0).getType()));
   if (!resultType)
     return rewriter.notifyMatchFailure(op, "Result type cannot be empty");
 
+  // clang-format off
   // Get height, width of input tensor, and diagonal arg to create
   // a const mask tensor to multiply with input.
   // This mask tensor has the same height and width of input tensor
-  // and consists of 1's for the lower triangle part and 0's for the rest.
-  // For example, with h=4, w=6, diagonal=1:
+  // and consists of 1's for the lower/higher triangle part and 0's for the rest.
+  // For tril with h=4, w=6, diagonal=1:
   // tensor([[1, 1, 0, 0, 0, 0],
   //         [1, 1, 1, 0, 0, 0],
   //         [1, 1, 1, 1, 0, 0],
   //         [1, 1, 1, 1, 1, 0]])
+  // For triu with h=4, w=6, diagonal=1:
+  // tensor([[0, 1, 1, 1, 1, 1],
+  //         [0, 0, 1, 1, 1, 1],
+  //         [0, 0, 0, 1, 1, 1],
+  //         [0, 0, 0, 0, 1, 1]])
+  // clang-format on
   auto selfShape = selfType.getShape();
   int64_t h = selfShape[selfRank - 2];
   int64_t w = selfShape[selfRank - 1];
   int64_t diagonal;
 
-  if (!matchPattern(op.getDiagonal(), m_TorchConstantInt(&diagonal)))
+  if (!matchPattern(diagonalVal, m_TorchConstantInt(&diagonal)))
     return rewriter.notifyMatchFailure(op, "Diagonal value is not an integer");
 
   // Define shape for mask tensor based on rank
@@ -8871,39 +8875,56 @@ LogicalResult ConvertAtenOp<AtenTrilOp>::matchAndRewriteImpl(
   maskShape.push_back(h);
   maskShape.push_back(w);
 
-  Value trilMask = TypeSwitch<Type, Value>(resultType.getElementType())
-                       .Case<mlir::FloatType>([&](auto) {
-                         return createTrilMask<float>(rewriter, op, maskShape,
-                                                      h, w, diagonal);
-                       })
-                       .Case<mlir::IntegerType>([&](auto intType) {
-                         switch (intType.getWidth()) {
-                         case 1:
-                           return createTrilMask<bool>(rewriter, op, maskShape,
-                                                       h, w, diagonal);
-                         case 32:
-                           return createTrilMask<int32_t>(
-                               rewriter, op, maskShape, h, w, diagonal);
-                         case 64:
-                           return createTrilMask<int64_t>(
-                               rewriter, op, maskShape, h, w, diagonal);
-                         }
-                         llvm_unreachable("Invalid integer width");
-                       });
-
-  if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), self, trilMask)
-          .failed())
+  Value mask =
+      TypeSwitch<Type, Value>(resultType.getElementType())
+          .Case<mlir::FloatType>([&](auto) {
+            return createTrilOrTriuMask<float>(rewriter, op, maskShape, h, w,
+                                               diagonal, isTril);
+          })
+          .template Case<mlir::IntegerType>([&](auto intType) {
+            switch (intType.getWidth()) {
+            case 1:
+              return createTrilOrTriuMask<bool>(rewriter, op, maskShape, h, w,
+                                                diagonal, isTril);
+            case 32:
+              return createTrilOrTriuMask<int32_t>(rewriter, op, maskShape, h,
+                                                   w, diagonal, isTril);
+            case 64:
+              return createTrilOrTriuMask<int64_t>(rewriter, op, maskShape, h,
+                                                   w, diagonal, isTril);
+            }
+            llvm_unreachable("Invalid integer width");
+          });
+
+  if (mlir::tosa::EqualizeRanks(rewriter, op->getLoc(), self, mask).failed())
     return rewriter.notifyMatchFailure(
         op, "Failed to equalize ranks among operands and result");
 
-  auto result =
-      tosa::createMulOpAndCast(rewriter, op, resultType, self, trilMask,
-                               /*shift=*/0);
+  auto result = tosa::createMulOpAndCast(rewriter, op, resultType, self, mask,
+                                         /*shift=*/0);
   rewriter.replaceOp(op, result.getResult());
 
   return success();
 }
 
+// Legalization for aten.triu
+template <>
+LogicalResult ConvertAtenOp<AtenTriuOp>::matchAndRewriteImpl(
+    AtenTriuOp op, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  return convertTrilOrTriu(op, adaptor.getSelf(), op.getDiagonal(),
+                           /*isTril=*/false, getTypeConverter(), rewriter);
+}
+
+// Legalization for aten.tril
+template <>
+LogicalResult ConvertAtenOp<AtenTrilOp>::matchAndRewriteImpl(
+    AtenTrilOp op, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  return convertTrilOrTriu(op, adaptor.getSelf(), op.getDiagonal(),
+                           /*isTril=*/true, getTypeConverter(), rewriter);
+}
+
 // Legalization for aten.flip
 template <>
 LogicalResult ConvertAtenOp<AtenFlipOp>::matchAndRewriteImpl(
@@ -11756,6 +11777,7 @@ std::set<StringRef> populateTorchToTosaConversionPatternsAndIllegalOps(
   INSERT_ATENOP_PATTERN(AtenIscloseOp);
   INSERT_ATENOP_PATTERN(Aten__InterpolateSizeListScaleListOp);
   INSERT_ATENOP_PATTERN(AtenTrilOp);
+  INSERT_ATENOP_PATTERN(AtenTriuOp);
   INSERT_ATENOP_PATTERN(AtenDiagonalOp);
   INSERT_ATENOP_PATTERN(AtenIndexSelectOp);
   INSERT_ATENOP_PATTERN(AtenFlipOp);

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -769,75 +769,6 @@ static Value performLastReduceAndPermute(PatternRewriter &rewriter,
   return out;
 }
 
-namespace {
-class DecomposeAtenTriuOp : public OpRewritePattern<AtenTriuOp> {
-public:
-  using OpRewritePattern::OpRewritePattern;
-  LogicalResult matchAndRewrite(AtenTriuOp op,
-                                PatternRewriter &rewriter) const override {
-    Location loc = op.getLoc();
-    Value input = op.getSelf();
-    auto inputType = cast<BaseTensorType>(input.getType());
-    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 cstZero =
-        ConstantIntOp::create(rewriter, loc, rewriter.getI64IntegerAttr(0));
-    Value cstOne =
-        ConstantIntOp::create(rewriter, loc, rewriter.getI64IntegerAttr(1));
-    Value none = ConstantNoneOp::create(rewriter, loc);
-
-    Value rowSize = getTensorDimSize(rewriter, input, -2);
-    Value colSize = getTensorDimSize(rewriter, input, -1);
-
-    auto si64Type = rewriter.getIntegerType(/*width=*/64, /*isSigned*/ true);
-    auto int64DtypeInt = getDtypeIntValueForType(rewriter, loc, si64Type);
-    auto rowArrangeType = getTensorTypeFromShapeValues({rowSize}, si64Type);
-    auto colArrangeType = getTensorTypeFromShapeValues({colSize}, si64Type);
-
-    Value rowArange =
-        AtenArangeOp::create(rewriter, loc, rowArrangeType, rowSize,
-                             /*dtype=*/int64DtypeInt, /*layout=*/none,
-                             /*device=*/none, /*pin_memory=*/none);
-    Value colArange =
-        AtenArangeOp::create(rewriter, loc, colArrangeType, colSize,
-                             /*dtype=*/int64DtypeInt, /*layout=*/none,
-                             /*device=*/none, /*pin_memory=*/none);
-
-    auto unsqueezeRowArangeInfo =
-        unsqueezeTensor(rewriter, op, rowArange, cstOne);
-    auto unsqueezeColArangeInfo =
-        unsqueezeTensor(rewriter, op, colArange, cstZero);
-
-    if (failed(unsqueezeRowArangeInfo) || failed(unsqueezeColArangeInfo)) {
-      return rewriter.notifyMatchFailure(op,
-                                         "cannot generate unsqueeze tensor");
-    }
-
-    Value unsqueezeRowArange = unsqueezeRowArangeInfo.value();
-    Value unsqueezeColArange = unsqueezeColArangeInfo.value();
-
-    Value unsqueezeRowArangePlusDiagonal =
-        AtenAddScalarOp::create(rewriter, loc, unsqueezeRowArange.getType(),
-                                unsqueezeRowArange, op.getDiagonal(), cstOne);
-
-    auto boolType = rewriter.getI1Type();
-    auto condType = getTensorTypeFromShapeValues({rowSize, colSize}, boolType);
-    Value condTensor =
-        AtenGeTensorOp::create(rewriter, loc, condType, unsqueezeColArange,
-                               unsqueezeRowArangePlusDiagonal);
-
-    rewriter.replaceOpWithNewOp<AtenWhereScalarOtherOp>(
-        op, op.getResult().getType(), condTensor, input, cstZero);
-    return success();
-  }
-};
-} // namespace
-
 /*
  This function calculates the number of elements in the lower triangle (below
  the main diagonal) of a tensor with dimensions [row, col]. The main diagonal
@@ -13583,7 +13514,6 @@ class DecomposeComplexOpsPass
     addPatternIfTargetOpIsIllegal<DecomposeAtenTypeAsOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenTileOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenReshapeAsOp>(patterns);
-    addPatternIfTargetOpIsIllegal<DecomposeAtenTriuOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenTriuIndicesOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenTrilIndicesOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenDeg2radOp>(patterns);

lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp

@@ -582,7 +582,6 @@ static void markDecomposedOpsAsIllegal(MLIRContext *context,
   target.addIllegalOp<AtenTypeAsOp>();
   target.addIllegalOp<AtenTileOp>();
   target.addIllegalOp<AtenReshapeAsOp>();
-  target.addIllegalOp<AtenTriuOp>();
   target.addIllegalOp<AtenTriuIndicesOp>();
   target.addIllegalOp<AtenTrilIndicesOp>();
   target.addIllegalOp<AtenDeg2radOp>();

test/Conversion/TorchToStablehlo/basic.mlir

@@ -360,3 +360,25 @@ func.func @torch.aten.sort(%arg0: !torch.vtensor<[2,3],f32>) -> (!torch.vtensor<
   %values, %indices = torch.aten.sort %arg0, %int-1, %true : !torch.vtensor<[2,3],f32>, !torch.int, !torch.bool -> !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],si64>
   return %values, %indices : !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],si64>
 }
+
+// -----
+
+// CHECK-LABEL:   func.func @torch.aten.triu(
+// CHECK-SAME:                   %[[ARG_0:.*]]: !torch.vtensor<[2,3,5],f32>,
+// CHECK-SAME:                   %[[ARG_1:.*]]: !torch.int) -> !torch.vtensor<[2,3,5],f32>
+// CHECK-DAG:       %[[VAL_0:.*]] = torch_c.to_builtin_tensor %[[ARG_0]] : !torch.vtensor<[2,3,5],f32> -> tensor<2x3x5xf32>
+// CHECK-DAG:       %[[VAL_1:.*]] = torch_c.to_i64 %[[ARG_1]]
+// CHECK:           %[[VAL_2:.*]] = stablehlo.iota dim = 1 : tensor<3x5xi64>
+// CHECK:           %[[VAL_3:.*]] = stablehlo.iota dim = 0 : tensor<3x5xi64>
+// CHECK:           %[[VAL_4:.*]] = tensor.from_elements %[[VAL_1]] : tensor<1xi64>
+// CHECK:           %[[VAL_5:.*]] = chlo.broadcast_add %[[VAL_3]], %[[VAL_4]] {broadcast_dimensions = array<i64: 1>} : (tensor<3x5xi64>, tensor<1xi64>) -> tensor<3x5xi64>
+// CHECK:           %[[VAL_6:.*]] = stablehlo.compare  GE, %[[VAL_2]], %[[VAL_5]],  SIGNED : (tensor<3x5xi64>, tensor<3x5xi64>) -> tensor<3x5xi1>
+// CHECK:           %[[VAL_7:.*]] = stablehlo.broadcast_in_dim %[[VAL_6]], dims = [1, 2] : (tensor<3x5xi1>) -> tensor<2x3x5xi1>
+// CHECK:           %[[VAL_8:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<2x3x5xf32>
+// CHECK:           %[[VAL_9:.*]] = stablehlo.select %[[VAL_7]], %[[VAL_0]], %[[VAL_8]] : tensor<2x3x5xi1>, tensor<2x3x5xf32>
+// CHECK:           %[[VAL_10:.*]] = torch_c.from_builtin_tensor %[[VAL_9]] : tensor<2x3x5xf32> -> !torch.vtensor<[2,3,5],f32>
+// CHECK:           return %[[VAL_10:.*]] : !torch.vtensor<[2,3,5],f32>
+func.func @torch.aten.triu(%arg0: !torch.vtensor<[2,3,5],f32>, %arg1: !torch.int) -> !torch.vtensor<[2,3,5],f32> {
+  %0 = torch.aten.triu %arg0, %arg1:!torch.vtensor<[2,3,5],f32>, !torch.int -> !torch.vtensor<[2,3,5],f32>
+  return %0 : !torch.vtensor<[2,3,5],f32>
+}