feat(dnn): enable eye to support bool

GitOrigin-RevId: 76d874d5b7

dnn/src/cuda/eye/eye.cu

@@ -39,6 +39,7 @@ void exec_internal(T* dst, size_t m, size_t n, int k, cudaStream_t stream) {
 #define INST(T)   template void exec_internal<T>(T*, size_t, size_t, int, cudaStream_t);
 #define cb(DType) INST(typename DTypeTrait<DType>::ctype)
 MEGDNN_FOREACH_COMPUTING_DTYPE(cb)
+cb(::megdnn::dtype::Bool)
 
 }  // namespace eye
 }  // namespace cuda

dnn/src/cuda/eye/opr_impl.cpp

@@ -26,6 +26,7 @@ void EyeImpl::exec(_megdnn_tensor_out dst, _megdnn_workspace workspace) {
                 cuda_stream(handle()));                                                \
     }
     MEGDNN_FOREACH_COMPUTING_DTYPE(cb)
+    cb(::megdnn::dtype::Bool)
 #undef cb
 }
 

dnn/src/naive/eye/opr_impl.cpp

@@ -31,6 +31,7 @@ void EyeImpl::exec(_megdnn_tensor_out dst, _megdnn_workspace workspace) {
         MEGDNN_DISPATCH_CPU_KERN_OPR(exec_internal<ctype>(dst.ptr<ctype>(), m, n)); \
     }
     MEGDNN_FOREACH_COMPUTING_DTYPE(cb)
+    cb(::megdnn::dtype::Bool)
 #undef cb
 }
 

dnn/src/rocm/eye/eye.cpp.hip

@@ -44,7 +44,7 @@ void exec_internal(T* dst, size_t m, size_t n, int k, hipStream_t stream) {
     template void exec_internal<T>(T*, size_t, size_t, int, hipStream_t);
 #define cb(DType) INST(typename DTypeTrait<DType>::ctype)
 MEGDNN_FOREACH_COMPUTING_DTYPE(cb)
-
+cb(::megdnn::dtype::Bool)
 }  // namespace eye
 }  // namespace rocm
 }  // namespace megdnn

dnn/src/rocm/eye/opr_impl.cpp

@@ -27,6 +27,7 @@ void EyeImpl::exec(_megdnn_tensor_out dst, _megdnn_workspace workspace) {
                 hip_stream(handle()));                                                 \
     }
     MEGDNN_FOREACH_COMPUTING_DTYPE(cb)
+    cb(::megdnn::dtype::Bool)
 #undef cb
 }
 

imperative/python/test/unit/functional/test_tensor.py

@@ -24,21 +24,22 @@ from megengine.utils.network_node import VarNode
 
 
 def test_eye():
-    dtype = np.float32
+    dtypes = [np.float32, np.bool]
     cases = [{"input": [10, 20]}, {"input": [30]}]
-    for case in cases:
-        np.testing.assert_allclose(
-            F.eye(case["input"], dtype=dtype).numpy(),
-            np.eye(*case["input"]).astype(dtype),
-        )
-        np.testing.assert_allclose(
-            F.eye(*case["input"], dtype=dtype).numpy(),
-            np.eye(*case["input"]).astype(dtype),
-        )
-        np.testing.assert_allclose(
-            F.eye(tensor(case["input"]), dtype=dtype).numpy(),
-            np.eye(*case["input"]).astype(dtype),
-        )
+    for dtype in dtypes:
+        for case in cases:
+            np.testing.assert_allclose(
+                F.eye(case["input"], dtype=dtype).numpy(),
+                np.eye(*case["input"]).astype(dtype),
+            )
+            np.testing.assert_allclose(
+                F.eye(*case["input"], dtype=dtype).numpy(),
+                np.eye(*case["input"]).astype(dtype),
+            )
+            np.testing.assert_allclose(
+                F.eye(tensor(case["input"]), dtype=dtype).numpy(),
+                np.eye(*case["input"]).astype(dtype),
+            )
 
 
 def test_full():