[TorchToLinalg] : Lower count_nonzero.dim_IntList to Linalg-on-Tensors (#4146)

- Lowered count_nonzero.dim_IntList op to Linalg-on-Tensors
- Decomposed aten.count_nonzero using aten.ne.Scalar and
aten.sum.dim_IntList
- Added tests to
torch-mlir/projects/pt1/python/torch_mlir_e2e_test/test_suite/reduction.py

Author: um-urosmarkovic

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

@@ -9927,6 +9927,31 @@ def Torch_AtenCountNonzeroOp : Torch_Op<"aten.count_nonzero", [
   let hasVerifier = 1;
 }
 
+def Torch_AtenCountNonzeroDimIntListOp : Torch_Op<"aten.count_nonzero.dim_IntList", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::count_nonzero.dim_IntList : (Tensor, int[]) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    AnyTorchListOfTorchIntType:$dim
+  );
+  let results = (outs
+    AnyTorchOptionalTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenCountNonzeroDimIntListOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 2, 1);
+    }
+    void AtenCountNonzeroDimIntListOp::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

@@ -6060,6 +6060,32 @@ LogicalResult AtenCountNonzeroOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// AtenCountNonzeroDimIntListOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult AtenCountNonzeroDimIntListOp::verify() {
+
+  auto selfType = cast<BaseTensorType>(getSelf().getType());
+
+  if (!selfType.hasDtype() || !selfType.hasSizes())
+    return success();
+
+  SmallVector<int64_t> dims;
+  if (!matchPattern(getDim(), m_TorchListOfConstantInts(dims)))
+    return emitOpError("expected dim to be constructed from list construct");
+
+  int64_t selfRank = selfType.getSizes().size();
+
+  for (auto d : dims) {
+    if (d >= selfRank || d < -selfRank)
+      return emitOpError("expected to be in [ ")
+             << -selfRank << " , " << selfRank - 1 << " ], but got dim = " << d;
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // OnnxVariantRotaryEmbeddingOp
 //===----------------------------------------------------------------------===//

lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp

@@ -9234,6 +9234,40 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    }\n"
 "    return %1 : !torch.list<int>\n"
 "  }\n"
+"  func.func @\"__torch_mlir_shape_fn.aten.count_nonzero.dim_IntList\"(%arg0: !torch.list<int>, %arg1: !torch.list<int>) -> !torch.list<int> {\n"
+"    %false = torch.constant.bool false\n"
+"    %none = torch.constant.none\n"
+"    %str = torch.constant.str \"AssertionError: \"\n"
+"    %true = torch.constant.bool true\n"
+"    %int0 = torch.constant.int 0\n"
+"    %0 = torch.aten.len.t %arg1 : !torch.list<int> -> !torch.int\n"
+"    torch.prim.Loop %0, %true, init() {\n"
+"    ^bb0(%arg2: !torch.int):\n"
+"      %4 = torch.aten.__getitem__.t %arg1, %arg2 : !torch.list<int>, !torch.int -> !torch.int\n"
+"      %5 = torch.aten.len.t %arg0 : !torch.list<int> -> !torch.int\n"
+"      %6 = torch.aten.neg.int %5 : !torch.int -> !torch.int\n"
+"      %7 = torch.aten.lt.int %4, %6 : !torch.int, !torch.int -> !torch.bool\n"
+"      %8 = torch.prim.If %7 -> (!torch.bool) {\n"
+"        torch.prim.If.yield %true : !torch.bool\n"
+"      } else {\n"
+"        %10 = torch.aten.len.t %arg0 : !torch.list<int> -> !torch.int\n"
+"        %11 = torch.aten.ge.int %4, %10 : !torch.int, !torch.int -> !torch.bool\n"
+"        torch.prim.If.yield %11 : !torch.bool\n"
+"      }\n"
+"      %9 = torch.aten.__not__ %8 : !torch.bool -> !torch.bool\n"
+"      torch.prim.If %9 -> () {\n"
+"        torch.prim.If.yield\n"
+"      } else {\n"
+"        torch.prim.RaiseException %str, %none : !torch.str, !torch.none\n"
+"        torch.prim.If.yield\n"
+"      }\n"
+"      torch.prim.Loop.condition %true, iter()\n"
+"    } : (!torch.int, !torch.bool) -> ()\n"
+"    %1 = torch.derefine %arg1 : !torch.list<int> to !torch.optional<list<int>>\n"
+"    %2 = torch.derefine %int0 : !torch.int to !torch.any\n"
+"    %3 = call @__torch__.torch.jit._shape_functions.sum_mean_dim(%arg0, %1, %false, %2) : (!torch.list<int>, !torch.optional<list<int>>, !torch.bool, !torch.any) -> !torch.list<int>\n"
+"    return %3 : !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"

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -6702,6 +6702,42 @@ class DecomposeAtenCountNonzeroOp
 };
 } // namespace
 
+// Decompose aten.count_nonzero to aten.ne.Scalar and aten.sum.dim_IntList
+namespace {
+class DecomposeAtenCountNonzeroDimIntListOp
+    : public OpRewritePattern<AtenCountNonzeroDimIntListOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenCountNonzeroDimIntListOp op,
+                                PatternRewriter &rewriter) const override {
+    Value dimList = op.getDim();
+    if (isa<Torch::NoneType>(dimList.getType()))
+      return rewriter.notifyMatchFailure(
+          op, "expected `dim` to be constructed from list");
+
+    SmallVector<int64_t> dimIntListElem;
+    if (!matchPattern(dimList, m_TorchListOfConstantInts(dimIntListElem)))
+      return rewriter.notifyMatchFailure(
+          op, "expected `dim` to be constructed from list of integers");
+
+    Location loc = op.getLoc();
+    Value self = op.getSelf();
+    BaseTensorType inputType = cast<BaseTensorType>(self.getType());
+    auto inpBoolTy = inputType.getWithSizesAndDtype(inputType.getSizes(),
+                                                    rewriter.getI1Type());
+    Value cstZero =
+        rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(0));
+    Value nonZeroMask =
+        rewriter.create<AtenNeScalarOp>(loc, inpBoolTy, self, cstZero);
+    Value none = rewriter.create<ConstantNoneOp>(loc);
+    Value cstFalse = rewriter.create<ConstantBoolOp>(loc, false);
+    rewriter.replaceOpWithNewOp<AtenSumDimIntListOp>(
+        op, op.getResult().getType(), nonZeroMask, dimList, cstFalse, none);
+    return success();
+  }
+};
+} // namespace
+
 // Decompose aten.std.correction to sqrt(var.correction(x))
 namespace {
 class DecomposeAtenStdCorrectionOp
@@ -12064,6 +12100,8 @@ class DecomposeComplexOpsPass
     addPatternIfTargetOpIsIllegal<DecomposeAtenStdDimOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenRot90Op>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenCountNonzeroOp>(patterns);
+    addPatternIfTargetOpIsIllegal<DecomposeAtenCountNonzeroDimIntListOp>(
+        patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenStdCorrectionOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenSplitWithSizesOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenNarrowOp>(patterns);

lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp

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

projects/pt1/e2e_testing/xfail_sets.py

@@ -2862,6 +2862,9 @@
     "CountNonzeroModuleBool_Basic",
     "CountNonzeroModuleF32_basic",
     "CountNonzeroModuleI64_basic",
+    "CountNonzeroDimIntListModuleBool_Basic",
+    "CountNonzeroDimIntListModuleI64_basic",
+    "CountNonzeroDimIntListModuleF32_basic",
     "Deg2radModule_basic",
     "DivFloatModule_basic",
     "DivIntModule_basic",

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

@@ -1453,6 +1453,23 @@ def aten〇count_nonzero〡shape(self: List[int], dim: Optional[int] = None) ->
     assert not (dim < -len(self) or dim >= len(self))
     return upstream_shape_functions.argmax(self, dim)
 
+@check_shape_function([
+    Invocation(TensorOfShape(2, 3, 4), dim = []), # Basic case.
+    Invocation(TensorOfShape(2, 3, 4), dim=[1]), # Test explicit dim.
+    Invocation(TensorOfShape(2, 3, 4), dim=[-1]), # Test explicit dim.
+    Invocation(TensorOfShape(2, 3, 4), dim=[0,1]), # Test explicit dim.
+    Invocation(TensorOfShape(2, 3, 4), dim=[-3]), # Test explicit dim(negative).
+    Invocation(TensorOfShape(2, 3, 4), dim=[2]), # Maximum valid dim.
+    ErrorInvocation(TensorOfShape(2, 3, 4), dim=[-4]), # Test dim out of bound.
+     ErrorInvocation(TensorOfShape(2, 3, 4), dim=[1,-4]), # Test dim out of bound.
+    ErrorInvocation(TensorOfShape(2, 3, 4), dim=[3]), # Test dim out of bound.
+])
+def aten〇count_nonzero〇dim_IntList〡shape(self: List[int], dim: List[int]) -> List[int]:
+    if dim is None: return []
+    for d in dim:
+        assert not (d < -len(self) or d >= len(self))
+    return upstream_shape_functions.sum_mean_dim(self, dim, keep_dim=False, dt=0)
+
 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)
 

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

@@ -788,6 +788,10 @@ def emit_with_mutating_variants(key, **kwargs):
     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)
+    emit(
+        "aten::count_nonzero.dim_IntList : (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

@@ -2546,3 +2546,54 @@ def forward(self, 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))
+
+
+# ==============================================================================
+
+
+class CountNonzeroDimIntListModuleF32(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.dim = [-1, 1]
+
+    @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: CountNonzeroDimIntListModuleF32())
+def CountNonzeroDimIntListModuleF32_basic(module, tu: TestUtils):
+    module.forward(tu.rand(2, 3, 4))
+
+
+class CountNonzeroDimIntListModuleI64(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.dim = [-2, -0, -1]
+
+    @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: CountNonzeroDimIntListModuleI64())
+def CountNonzeroDimIntListModuleI64_basic(module, tu: TestUtils):
+    module.forward(tu.randint(2, 3, 4))
+
+
+class CountNonzeroDimIntListModuleBool(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.dim = [1]
+
+    @export
+    @annotate_args([None, ([2, 3, 4], torch.bool, True)])
+    def forward(self, x):
+        return torch.ops.aten.count_nonzero(x, self.dim)
+
+
+@register_test_case(module_factory=lambda: CountNonzeroDimIntListModuleBool())
+def CountNonzeroDimIntListModuleBool_Basic(module, tu: TestUtils):
+    module.forward(tu.randint(2, 3, 4, low=0, high=2).to(torch.bool))