[TorchToLinalg] Lower count_nonzero to Linalg

Author: um-urosmarkovic

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

@@ -9902,6 +9902,31 @@ def Torch_AtenRot90Op : Torch_Op<"aten.rot90", [
   let hasVerifier = 1;
 }
 
+def Torch_AtenCountNonzeroOp : Torch_Op<"aten.count_nonzero", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::count_nonzero : (Tensor, int?) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    AnyTorchOptionalIntType:$dim
+  );
+  let results = (outs
+    AnyTorchOptionalTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenCountNonzeroOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 2, 1);
+    }
+    void AtenCountNonzeroOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 2, 1);
+    }
+  }];
+  let hasVerifier = 1;
+}
+
 def Torch_AtenConstantPadNdOp : Torch_Op<"aten.constant_pad_nd", [
     AllowsTypeRefinement,
     HasValueSemantics,

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -6033,6 +6033,33 @@ LogicalResult AtenRot90Op::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// AtenCountNonzero
+//===----------------------------------------------------------------------===//
+
+LogicalResult AtenCountNonzeroOp::verify() {
+
+  auto selfType = cast<BaseTensorType>(getSelf().getType());
+
+  if (!selfType.hasDtype() || !selfType.hasSizes())
+    return success();
+
+  if (!isa<Torch::IntType>(getDim().getType()) &&
+      !isa<Torch::NoneType>(getDim().getType()))
+    return emitOpError("parameter dim must be none or int type");
+
+  int64_t dim;
+  if (!matchPattern(getDim(), m_TorchConstantInt(&dim)))
+    return success();
+
+  int selfRank = selfType.getSizes().size();
+  if (dim >= selfRank || dim < -selfRank)
+    return emitOpError("expected dim to be in [ ")
+           << -selfRank << ", " << selfRank - 1 << " ], but got dim = " << dim;
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // OnnxVariantRotaryEmbeddingOp
 //===----------------------------------------------------------------------===//

lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp

@@ -9191,6 +9191,39 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    }\n"
 "    return %arg0 : !torch.list<int>\n"
 "  }\n"
+"  func.func @\"__torch_mlir_shape_fn.aten.count_nonzero\"(%arg0: !torch.list<int>, %arg1: !torch.optional<int>) -> !torch.list<int> {\n"
+"    %false = torch.constant.bool false\n"
+"    %str = torch.constant.str \"AssertionError: \"\n"
+"    %true = torch.constant.bool true\n"
+"    %none = torch.constant.none\n"
+"    %0 = torch.aten.__is__ %arg1, %none : !torch.optional<int>, !torch.none -> !torch.bool\n"
+"    %1 = torch.prim.If %0 -> (!torch.list<int>) {\n"
+"      %2 = torch.prim.ListConstruct  : () -> !torch.list<int>\n"
+"      torch.prim.If.yield %2 : !torch.list<int>\n"
+"    } else {\n"
+"      %2 = torch.prim.unchecked_cast %arg1 : !torch.optional<int> -> !torch.int\n"
+"      %3 = torch.aten.len.t %arg0 : !torch.list<int> -> !torch.int\n"
+"      %4 = torch.aten.neg.int %3 : !torch.int -> !torch.int\n"
+"      %5 = torch.aten.lt.int %2, %4 : !torch.int, !torch.int -> !torch.bool\n"
+"      %6 = torch.prim.If %5 -> (!torch.bool) {\n"
+"        torch.prim.If.yield %true : !torch.bool\n"
+"      } else {\n"
+"        %9 = torch.aten.len.t %arg0 : !torch.list<int> -> !torch.int\n"
+"        %10 = torch.aten.ge.int %2, %9 : !torch.int, !torch.int -> !torch.bool\n"
+"        torch.prim.If.yield %10 : !torch.bool\n"
+"      }\n"
+"      %7 = torch.aten.__not__ %6 : !torch.bool -> !torch.bool\n"
+"      torch.prim.If %7 -> () {\n"
+"        torch.prim.If.yield\n"
+"      } else {\n"
+"        torch.prim.RaiseException %str, %none : !torch.str, !torch.none\n"
+"        torch.prim.If.yield\n"
+"      }\n"
+"      %8 = func.call @__torch__.torch.jit._shape_functions.argmax(%arg0, %arg1, %false) : (!torch.list<int>, !torch.optional<int>, !torch.bool) -> !torch.list<int>\n"
+"      torch.prim.If.yield %8 : !torch.list<int>\n"
+"    }\n"
+"    return %1 : !torch.list<int>\n"
+"  }\n"
 "  func.func @\"__torch_mlir_shape_fn.aten._to_copy\"(%arg0: !torch.list<int>, %arg1: !torch.optional<int>, %arg2: !torch.optional<int>, %arg3: !torch.optional<Device>, %arg4: !torch.optional<bool>, %arg5: !torch.bool, %arg6: !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"
@@ -15782,6 +15815,10 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    %0:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple<int, int> -> !torch.int, !torch.int\n"
 "    return %0#1 : !torch.int\n"
 "  }\n"
+"  func.func @\"__torch_mlir_dtype_fn.aten.count_nonzero\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.optional<int>) -> !torch.int {\n"
+"    %int4 = torch.constant.int 4\n"
+"    return %int4 : !torch.int\n"
+"  }\n"
 "  func.func @\"__torch_mlir_dtype_fn.aten.rot90\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.int, %arg2: !torch.list<int>) -> !torch.int {\n"
 "    %0:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple<int, int> -> !torch.int, !torch.int\n"
 "    return %0#1 : !torch.int\n"

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -6661,6 +6661,47 @@ class DecomposeAtenRot90Op : public OpRewritePattern<AtenRot90Op> {
 };
 } // namespace
 
+// Decompose aten.count_nonzero to aten.ne.Scalar and
+// aten.sum/aten.sum.dim_IntList
+namespace {
+class DecomposeAtenCountNonzeroOp
+    : public OpRewritePattern<AtenCountNonzeroOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenCountNonzeroOp op,
+                                PatternRewriter &rewriter) const override {
+    auto dim = op.getDim();
+    if (!isa<Torch::NoneType>(dim.getType()) &&
+        !isa<Torch::IntType>(dim.getType())) {
+      return rewriter.notifyMatchFailure(
+          op, "expected `dim` to be `None` or `int`");
+    }
+    Location loc = op.getLoc();
+    auto self = op.getSelf();
+    auto inputType = cast<BaseTensorType>(self.getType());
+    auto inpBoolTy = inputType.getWithSizesAndDtype(inputType.getSizes(),
+                                                    rewriter.getI1Type());
+    auto cstZero =
+        rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(0));
+    auto nonZeroMask =
+        rewriter.create<AtenNeScalarOp>(loc, inpBoolTy, self, cstZero);
+    auto none = rewriter.create<ConstantNoneOp>(loc);
+    if (isa<Torch::NoneType>(dim.getType())) {
+      rewriter.replaceOpWithNewOp<AtenSumOp>(op, op.getResult().getType(),
+                                             nonZeroMask, none);
+    } else {
+      auto cstFalse = rewriter.create<ConstantBoolOp>(loc, false);
+      auto dimIntList = rewriter.create<PrimListConstructOp>(
+          loc, ListType::get(IntType::get(op.getContext())),
+          SmallVector<Value>{dim});
+      rewriter.replaceOpWithNewOp<AtenSumDimIntListOp>(
+          op, op.getResult().getType(), nonZeroMask, dimIntList, cstFalse,
+          none);
+    }
+    return success();
+  }
+};
+} // namespace
+
 // Decompose aten.std.correction to sqrt(var.correction(x))
 namespace {
 class DecomposeAtenStdCorrectionOp
@@ -12018,6 +12059,7 @@ class DecomposeComplexOpsPass
     addPatternIfTargetOpIsIllegal<DecomposeAtenVarCorrectionOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenStdDimOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenRot90Op>(patterns);
+    addPatternIfTargetOpIsIllegal<DecomposeAtenCountNonzeroOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenStdCorrectionOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenSplitWithSizesOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenNarrowOp>(patterns);

lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp

@@ -413,6 +413,7 @@ static void markDecomposedOpsAsIllegal(MLIRContext *context,
   target.addIllegalOp<AtenMvOp>();
   target.addIllegalOp<AtenRenormOp>();
   target.addIllegalOp<AtenRot90Op>();
+  target.addIllegalOp<AtenCountNonzeroOp>();
   target.addIllegalOp<AtenLinalgCrossOp>();
   target.addIllegalOp<Aten_LinalgDetOp>();
   target.addIllegalOp<AtenLinalgSlogdetOp>();

projects/pt1/e2e_testing/xfail_sets.py

@@ -2846,6 +2846,10 @@
     "ConvolutionModule2DTransposeNonUnitOutputPadding_basic",
     "ConvolutionModule2DTransposeStrided_basic",
     "ConvolutionModule2DTranspose_basic",
+    # Error: onnx lowering,
+    "CountNonzeroModuleBool_Basic",
+    "CountNonzeroModuleF32_basic",
+    "CountNonzeroModuleI64_basic",
     "Deg2radModule_basic",
     "DivFloatModule_basic",
     "DivIntModule_basic",

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py

@@ -1436,6 +1436,21 @@ def aten〇rot90〡shape(self: List[int], k: int = 1, dims: List[int] = (0, 1,))
 
     return self
 
+@check_shape_function([
+    Invocation(TensorOfShape(2, 3, 4)), # Basic case.
+    Invocation(TensorOfShape(2, 3, 4), dim = 1), # Test explicit dim.
+    Invocation(TensorOfShape(2, 3, 4), dim = -1), # Test explicit dim(negative).
+    Invocation(TensorOfShape(2, 3, 4), dim = -3), # Test explicit dim(negative).
+    Invocation(TensorOfShape(2, 3, 4), dim = 0), # Test explicit dim.
+    Invocation(TensorOfShape(2, 3, 4), dim = 2), # Test explicit maximum valid dim.
+    ErrorInvocation(TensorOfShape(2, 3, 4), dim = -4), # Test dim out of bound.
+    ErrorInvocation(TensorOfShape(2, 3, 4), dim = 3), # Test dim out of bound.
+])
+def aten〇count_nonzero〡shape(self: List[int], dim: Optional[int] = None) -> List[int]:
+    if dim is None: return []
+    assert not (dim < -len(self) or dim >= len(self))
+    return upstream_shape_functions.argmax(self, dim)
+
 def aten〇_to_copy〡shape(self: List[int], dtype: Optional[int] = None, layout: Optional[int] = None, device: Optional[device] = None, pin_memory: Optional[bool] = None, non_blocking: bool = False, memory_format: Optional[int] = None) -> List[int]:
     return upstream_shape_functions.unary(self)
 
@@ -5514,6 +5529,9 @@ def aten〇diag_embed〡dtype(self_rank_dtype: Tuple[int, int], offset: int = 0,
     self_rank, self_dtype = self_rank_dtype
     return self_dtype
 
+def aten〇count_nonzero〡dtype(self_rank_dtype: Tuple[int, int], dim: Optional[int] = None) -> int:
+    return torch.int64
+
 def aten〇rot90〡dtype(self_rank_dtype: Tuple[int, int], k: int = 1, dims: List[int] = (0, 1,)) -> int:
     self_rank, self_dtype = self_rank_dtype
     return self_dtype

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py

@@ -787,6 +787,7 @@ def emit_with_mutating_variants(key, **kwargs):
     emit("aten::diag_embed : (Tensor, int, int, int) -> (Tensor)")
     emit("aten::_weight_norm_interface : (Tensor, Tensor, int) -> (Tensor, Tensor)")
     emit("aten::rot90 : (Tensor, int, int[]) -> (Tensor)", has_verifier=True)
+    emit("aten::count_nonzero : (Tensor, int?) -> (Tensor)", has_verifier=True)
 
     # Misc tensor ops.
     emit("aten::constant_pad_nd : (Tensor, int[], Scalar) -> (Tensor)")

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

@@ -2496,3 +2496,53 @@ def TraceUnsignedIntModule_basic(module, tu: TestUtils):
 @register_test_case(module_factory=lambda: TraceIntModule())
 def TraceUnsignedIntModule_empty(module, tu: TestUtils):
     module.forward(tu.randint(0, 0, low=0, high=10))
+
+
+# ==============================================================================
+
+
+class CountNonzeroModuleI64(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.dim = 2
+
+    @export
+    @annotate_args([None, ([2, 3, 4], torch.int64, True)])
+    def forward(self, x):
+        return torch.ops.aten.count_nonzero(x, self.dim)
+
+
+@register_test_case(module_factory=lambda: CountNonzeroModuleI64())
+def CountNonzeroModuleI64_basic(module, tu: TestUtils):
+    module.forward(tu.randint(2, 3, 4, low=-2, high=2))
+
+
+class CountNonzeroModuleF32(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.dim = -3
+
+    @export
+    @annotate_args([None, ([2, 3, 4], torch.float32, True)])
+    def forward(self, x):
+        return torch.ops.aten.count_nonzero(x, self.dim)
+
+
+@register_test_case(module_factory=lambda: CountNonzeroModuleF32())
+def CountNonzeroModuleF32_basic(module, tu: TestUtils):
+    module.forward(tu.rand(2, 3, 4))
+
+
+class CountNonzeroModuleBool(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([None, ([2, 3, 4], torch.bool, True)])
+    def forward(self, x):
+        return torch.ops.aten.count_nonzero(x)
+
+
+@register_test_case(module_factory=lambda: CountNonzeroModuleBool())
+def CountNonzeroModuleBool_Basic(module, tu: TestUtils):
+    module.forward(tu.randint(2, 3, 4, low=0, high=2).to(torch.bool))