Implement lowering of torch.aten.triu_indices (#3451)

Closes
[nod-ai/SHARK-ModelDev/issues/709](nod-ai/SHARK-ModelDev#709)

---------

Co-authored-by: Branko Trifkovic <[email protected]>

Author: BaneTrifa

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

@@ -15517,6 +15517,36 @@ def Torch_AtenScalarImplicitOp : Torch_Op<"aten.ScalarImplicit", [
   let hasCanonicalizer = 1;
 }
 
+def Torch_AtenTriuIndicesOp : Torch_Op<"aten.triu_indices", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::triu_indices : (int, int, int, int?, int?, Device?, bool?) -> (Tensor)`";
+  let arguments = (ins
+    Torch_IntType:$row,
+    Torch_IntType:$col,
+    Torch_IntType:$offset,
+    AnyTorchOptionalIntType:$dtype,
+    AnyTorchOptionalIntType:$layout,
+    AnyTorchOptionalDeviceType:$device,
+    AnyTorchOptionalBoolType:$pin_memory
+  );
+  let results = (outs
+    AnyTorchOptionalTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenTriuIndicesOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 7, 1);
+    }
+    void AtenTriuIndicesOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 7, 1);
+    }
+  }];
+  let hasVerifier = 1;
+}
+
 def Torch_Aten_SoftmaxBackwardDataOp : Torch_Op<"aten._softmax_backward_data", [
     AllowsTypeRefinement,
     HasValueSemantics,

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -5212,3 +5212,39 @@ LogicalResult BindSymbolicShapeOp::verify() {
 
   return success();
 }
+// AtenTriuIndicesOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult AtenTriuIndicesOp::verify() {
+
+  // Check if row, col and offset are constant ints
+  int64_t row;
+  if (!matchPattern(getRow(), m_TorchConstantInt(&row)))
+    return success();
+
+  int64_t col;
+  if (!matchPattern(getCol(), m_TorchConstantInt(&col)))
+    return success();
+
+  int64_t offset;
+  if (!matchPattern(getOffset(), m_TorchConstantInt(&offset)))
+    return success();
+
+  // Check if values of row, and col are valid
+  if (row < 0)
+    return emitOpError("row must be non-negative, got ") << row;
+
+  if (col < 0)
+    return emitOpError("col must be non-negative, got ") << col;
+
+  // Check if dtype is valid
+  int64_t dtype;
+  if (!matchPattern(getDtype(), m_TorchConstantInt(&dtype)))
+    return success();
+  if (dtype != (int)torch_upstream::ScalarType::Int &&
+      dtype != (int)torch_upstream::ScalarType::Long)
+    return emitOpError(
+        "'triu_indices' implemented only for torch.int32 and torch.int64");
+
+  return success();
+}

lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp

@@ -9729,6 +9729,68 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    %1 = call @__torch__._embedding_bag_helper(%arg0, %arg1, %arg2, %arg7, %arg4, %arg6, %0) : (!torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool, !torch.int, !torch.optional<list<int>>, !torch.optional<int>) -> !torch.tuple<list<int>, list<int>, list<int>, list<int>>\n"
 "    return %1 : !torch.tuple<list<int>, list<int>, list<int>, list<int>>\n"
 "  }\n"
+"  func.func @\"__torch_mlir_shape_fn.aten.triu_indices\"(%arg0: !torch.int, %arg1: !torch.int, %arg2: !torch.int, %arg3: !torch.optional<int>, %arg4: !torch.optional<int>, %arg5: !torch.optional<Device>, %arg6: !torch.optional<bool>) -> !torch.list<int> {\n"
+"    %true = torch.constant.bool true\n"
+"    %int0 = torch.constant.int 0\n"
+"    %int2 = torch.constant.int 2\n"
+"    %int1 = torch.constant.int 1\n"
+"    %0 = torch.aten.eq.int %arg0, %int0 : !torch.int, !torch.int -> !torch.bool\n"
+"    %1 = torch.prim.If %0 -> (!torch.bool) {\n"
+"      torch.prim.If.yield %true : !torch.bool\n"
+"    } else {\n"
+"      %3 = torch.aten.eq.int %arg1, %int0 : !torch.int, !torch.int -> !torch.bool\n"
+"      torch.prim.If.yield %3 : !torch.bool\n"
+"    }\n"
+"    %2 = torch.prim.If %1 -> (!torch.list<int>) {\n"
+"      %3 = torch.prim.ListConstruct %int2, %int0 : (!torch.int, !torch.int) -> !torch.list<int>\n"
+"      torch.prim.If.yield %3 : !torch.list<int>\n"
+"    } else {\n"
+"      %3 = torch.aten.sub.int %arg2, %int1 : !torch.int, !torch.int -> !torch.int\n"
+"      %4 = torch.aten.eq.int %arg0, %int0 : !torch.int, !torch.int -> !torch.bool\n"
+"      %5 = torch.prim.If %4 -> (!torch.bool) {\n"
+"        torch.prim.If.yield %true : !torch.bool\n"
+"      } else {\n"
+"        %11 = torch.aten.eq.int %arg1, %int0 : !torch.int, !torch.int -> !torch.bool\n"
+"        torch.prim.If.yield %11 : !torch.bool\n"
+"      }\n"
+"      %6:2 = torch.prim.If %5 -> (!torch.int, !torch.int) {\n"
+"        torch.prim.If.yield %int0, %int0 : !torch.int, !torch.int\n"
+"      } else {\n"
+"        %11 = torch.aten.gt.int %3, %int0 : !torch.int, !torch.int -> !torch.bool\n"
+"        %12 = torch.prim.If %11 -> (!torch.int) {\n"
+"          %27 = torch.aten.add.int %int1, %3 : !torch.int, !torch.int -> !torch.int\n"
+"          %28 = torch.prim.min.int %arg1, %27 : !torch.int, !torch.int -> !torch.int\n"
+"          torch.prim.If.yield %28 : !torch.int\n"
+"        } else {\n"
+"          %27 = torch.aten.add.int %arg0, %3 : !torch.int, !torch.int -> !torch.int\n"
+"          %28 = torch.aten.gt.int %27, %int0 : !torch.int, !torch.int -> !torch.bool\n"
+"          %29 = torch.aten.Int.bool %28 : !torch.bool -> !torch.int\n"
+"          torch.prim.If.yield %29 : !torch.int\n"
+"        }\n"
+"        %13 = torch.aten.add.int %arg0, %3 : !torch.int, !torch.int -> !torch.int\n"
+"        %14 = torch.prim.min.int %arg1, %13 : !torch.int, !torch.int -> !torch.int\n"
+"        %15 = torch.prim.max.int %int0, %14 : !torch.int, !torch.int -> !torch.int\n"
+"        %16 = torch.aten.add.int %arg0, %3 : !torch.int, !torch.int -> !torch.int\n"
+"        %17 = torch.prim.min.int %arg0, %16 : !torch.int, !torch.int -> !torch.int\n"
+"        %18 = torch.prim.max.int %int0, %17 : !torch.int, !torch.int -> !torch.int\n"
+"        %19 = torch.aten.sub.int %15, %12 : !torch.int, !torch.int -> !torch.int\n"
+"        %20 = torch.aten.add.int %19, %int1 : !torch.int, !torch.int -> !torch.int\n"
+"        %21 = torch.aten.add.int %12, %15 : !torch.int, !torch.int -> !torch.int\n"
+"        %22 = torch.aten.mul.int %21, %20 : !torch.int, !torch.int -> !torch.int\n"
+"        %23 = torch.aten.floordiv.int %22, %int2 : !torch.int, !torch.int -> !torch.int\n"
+"        %24 = torch.aten.sub.int %18, %20 : !torch.int, !torch.int -> !torch.int\n"
+"        %25 = torch.aten.mul.int %24, %arg1 : !torch.int, !torch.int -> !torch.int\n"
+"        %26 = torch.prim.max.int %int0, %25 : !torch.int, !torch.int -> !torch.int\n"
+"        torch.prim.If.yield %23, %26 : !torch.int, !torch.int\n"
+"      }\n"
+"      %7 = torch.aten.mul.int %arg0, %arg1 : !torch.int, !torch.int -> !torch.int\n"
+"      %8 = torch.aten.add.int %6#0, %6#1 : !torch.int, !torch.int -> !torch.int\n"
+"      %9 = torch.aten.sub.int %7, %8 : !torch.int, !torch.int -> !torch.int\n"
+"      %10 = torch.prim.ListConstruct %int2, %9 : (!torch.int, !torch.int) -> !torch.list<int>\n"
+"      torch.prim.If.yield %10 : !torch.list<int>\n"
+"    }\n"
+"    return %2 : !torch.list<int>\n"
+"  }\n"
 "  func.func @\"__torch_mlir_shape_fn.aten.nll_loss_forward\"(%arg0: !torch.list<int>, %arg1: !torch.list<int>, %arg2: !torch.optional<list<int>>, %arg3: !torch.int, %arg4: !torch.int) -> !torch.tuple<list<int>, list<int>> {\n"
 "    %0 = call @__torch__.torch.jit._shape_functions.nll_loss_forward(%arg0, %arg1, %arg2, %arg3) : (!torch.list<int>, !torch.list<int>, !torch.optional<list<int>>, !torch.int) -> !torch.tuple<list<int>, list<int>>\n"
 "    return %0 : !torch.tuple<list<int>, list<int>>\n"
@@ -14023,6 +14085,18 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    %int6 = torch.constant.int 6\n"
 "    return %int6 : !torch.int\n"
 "  }\n"
+"  func.func @\"__torch_mlir_dtype_fn.aten.triu_indices\"(%arg0: !torch.int, %arg1: !torch.int, %arg2: !torch.int, %arg3: !torch.optional<int>, %arg4: !torch.optional<int>, %arg5: !torch.optional<Device>, %arg6: !torch.optional<bool>) -> !torch.int {\n"
+"    %int4 = torch.constant.int 4\n"
+"    %none = torch.constant.none\n"
+"    %0 = torch.aten.__is__ %arg3, %none : !torch.optional<int>, !torch.none -> !torch.bool\n"
+"    %1 = torch.prim.If %0 -> (!torch.int) {\n"
+"      torch.prim.If.yield %int4 : !torch.int\n"
+"    } else {\n"
+"      %2 = torch.prim.unchecked_cast %arg3 : !torch.optional<int> -> !torch.int\n"
+"      torch.prim.If.yield %2 : !torch.int\n"
+"    }\n"
+"    return %1 : !torch.int\n"
+"  }\n"
 "  func.func @\"__torch_mlir_dtype_fn.aten.int_repr\"(%arg0: !torch.tuple<int, int>) -> !torch.int {\n"
 "    %int3 = torch.constant.int 3\n"
 "    %int1 = torch.constant.int 1\n"