[Fix][Relay][TOPI] Bug fix in relay.sum and topi.sum functions when w…

include/tvm/topi/reduction.h

@@ -325,7 +325,11 @@ inline PrimExpr ProdOp(PrimExpr source, Array<IterVar> axis, Array<PrimExpr> ini
  */
 inline Tensor sum(const Tensor& data, const Array<Integer>& axis, bool keepdims = false,
                   bool atleast1d = false) {
-  return CommReduce(data, axis, tvm::sum, keepdims, atleast1d);
+  if (data->dtype.is_bool()) {
+    return CommReduce(data, axis, tvm::any, keepdims, atleast1d);
+  } else {
+    return CommReduce(data, axis, tvm::sum, keepdims, atleast1d);
+  }
 }
 
 inline Tensor collapse_sum(const Tensor& data, Array<PrimExpr> target_shape) {

tests/python/relay/test_op_level4.py

@@ -369,6 +369,41 @@ def test_reduce(
         tvm.testing.assert_allclose(op_res1.numpy(), ref_res, rtol=1e-5)
 
 
+@tvm.testing.uses_gpu
+def test_sum_with_bool_input():
+    def verify(dshape, axis, keepdims, exclude):
+        x = relay.var("x", relay.TensorType(dshape, "bool"))
+
+        y = relay.sum(x, axis, keepdims, exclude)
+
+        func = relay.Function([x], y)
+        func = run_infer_type(func)
+
+        text = func.astext()
+        assert "sum" in text
+
+        data = np.random.choice([False, True], size=dshape)
+
+        if exclude and axis is not None:
+            axis = tuple(set(range(len(dshape))) - set(axis))
+
+        ref_res = np.sum(data, axis, keepdims=keepdims, dtype="bool")
+        for target, dev in tvm.testing.enabled_targets():
+            op_res = relay.create_executor("graph", device=dev, target=target).evaluate(func)(data)
+            tvm.testing.assert_allclose(op_res.numpy(), ref_res)
+
+    verify((3, 5, 7, 9), None, False, False)
+    verify((3, 5, 7, 9), None, True, False)
+    verify((3, 5, 7, 9), (0,), False, False)
+    verify((3, 5, 7, 9), (1,), True, False)
+    verify((3, 5, 7, 9), (2, 3), False, True)
+    verify((3, 5, 7, 9), (0, 2), True, True)
+    verify((3, 5, 7, 9), (0, 1, 2, 3), False, False)
+    verify((3, 5, 7, 9), (0, 1, 2, 3), False, True)
+    verify((3, 5, 7, 9), (0, 1, 2, 3), True, False)
+    verify((3, 5, 7, 9), (0, 1, 2, 3), True, True)
+
+
 @tvm.testing.uses_gpu
 def test_argmin_argmax_get_last_elements():
     def get_test_case(shape, gt_func, test_argmin=False):
@@ -638,7 +673,6 @@ def verify(dshape, begin, end, strides, vshape, test_ref=True):
         func = run_infer_type(func)
         text = func.astext()
         assert "strided_set" in text
-        print(text)
         assert func.body.checked_type == relay.ty.TensorType(dshape, "float32")
         if not test_ref:
             return

tests/python/topi/python/test_topi_reduce.py

@@ -44,6 +44,8 @@
     ((32, 128, 24), None, True, "any", "bool"),
     ((1, 4, 7), 1, True, "any", "bool"),
     ((128, 24, 128, 24), 2, False, "any", "bool"),
+    ((128, 24, 128, 24), 2, False, "sum", "bool"),
+    ((128, 24, 128, 24), 0, True, "sum", "bool"),
 )
 
 
@@ -57,7 +59,10 @@ def ref_data(in_shape, axis, keepdims, reduce_type, dtype):
         in_npy_map = np.sqrt(np.exp(in_npy)).astype(dtype)
 
     if reduce_type == "sum":
-        out_npy = in_npy_map.sum(axis=axis, keepdims=keepdims)
+        if dtype == "bool":
+            out_npy = in_npy_map.sum(axis=axis, keepdims=keepdims, dtype="bool")
+        else:
+            out_npy = in_npy_map.sum(axis=axis, keepdims=keepdims)
     elif reduce_type == "all" and dtype == "bool":
         out_npy = in_npy_map.all(axis=axis, keepdims=keepdims)
     elif reduce_type == "any" and dtype == "bool":
@@ -113,7 +118,10 @@ def test_reduce_map(target, dev, ref_data, in_shape, axis, keepdims, reduce_type
     A1 = topi.sqrt(topi.exp(A))
     out_dtype = dtype
     if reduce_type == "sum":
-        B = topi.sum(A1, axis=axis, keepdims=keepdims)
+        if dtype == "bool":
+            B = topi.sum(A, axis=axis, keepdims=keepdims)
+        else:
+            B = topi.sum(A1, axis=axis, keepdims=keepdims)
     elif reduce_type == "all":
         B = topi.all(A, axis=axis, keepdims=keepdims)
     elif reduce_type == "any":