[CANN]Support COUNT_EQUAL && STEP && SGN

Signed-off-by: noemotiovon <[email protected]>

Author: noemotiovon

ggml/src/ggml-cann/acl_tensor.cpp

@@ -41,6 +41,8 @@ aclDataType ggml_cann_type_mapping(ggml_type type) {
             return ACL_INT4;
         case GGML_TYPE_Q8_0:
             return ACL_INT8;
+        case GGML_TYPE_I64:
+            return ACL_INT64;
         default:
             return ACL_DT_UNDEFINED;
     }

ggml/src/ggml-cann/aclnn_ops.cpp

@@ -62,6 +62,8 @@
 #include <aclnnop/aclnn_log.h>
 #include <aclnnop/aclnn_mean.h>
 #include <aclnnop/aclnn_reflection_pad1d.h>
+#include <aclnnop/aclnn_eq_tensor.h>
+#include <aclnnop/aclnn_gt_scalar.h>
 #include <float.h>
 
 #include <cmath>
@@ -2697,3 +2699,35 @@ void ggml_cann_pad_reflect_1d(ggml_backend_cann_context& ctx, ggml_tensor* dst){
     }
     ACL_CHECK(aclDestroyIntArray(paddings));
 }
+
+void ggml_cann_count_equal(ggml_backend_cann_context& ctx, ggml_tensor* dst){
+    ggml_tensor * src0 = dst->src[0];
+    ggml_tensor * src1 = dst->src[1];
+
+    aclTensor* acl_self = ggml_cann_create_tensor(src0);
+    aclTensor* acl_other = ggml_cann_create_tensor(src1);
+    
+    GGML_CANN_CALL_ACLNN_OP(InplaceEqTensor, acl_self, acl_other);
+
+    ggml_cann_sum(ctx, dst);
+
+    ACL_CHECK(aclDestroyTensor(acl_self));
+    ACL_CHECK(aclDestroyTensor(acl_other));
+}
+
+void ggml_cann_step(ggml_backend_cann_context& ctx, ggml_tensor* dst){
+    ggml_tensor * src0 = dst->src[0];
+
+    aclTensor* acl_src = ggml_cann_create_tensor(src0);
+    aclTensor* acl_dst = ggml_cann_create_tensor(dst);
+
+    float alphaValue = 0.0f;
+    aclScalar* alpha = nullptr;
+    alpha = aclCreateScalar(&alphaValue, aclDataType::ACL_FLOAT);
+    
+    GGML_CANN_CALL_ACLNN_OP(GtScalar, acl_src, alpha, acl_dst);
+
+    ACL_CHECK(aclDestroyTensor(acl_src));
+    ACL_CHECK(aclDestroyTensor(acl_dst));
+    ACL_CHECK(aclDestroyScalar(alpha));
+}

ggml/src/ggml-cann/aclnn_ops.h

@@ -43,6 +43,7 @@
 #include <aclnnop/aclnn_sin.h>
 #include <aclnnop/aclnn_cos.h>
 #include <aclnnop/aclnn_log.h>
+#include <aclnnop/aclnn_sign.h>
 #include "acl_tensor.h"
 #include "common.h"
 
@@ -681,6 +682,37 @@ void ggml_cann_mean(ggml_backend_cann_context& ctx, ggml_tensor* dst);
  */
 void ggml_cann_pad_reflect_1d(ggml_backend_cann_context& ctx, ggml_tensor* dst);
 
+/**
+ * @brief   Counts the number of equal elements in two ggml tensors using the CANN backend.
+ * 
+ * @details This function performs an element-wise comparison between two input tensors,
+ *          and counts the number of positions where the elements are equal. The result is 
+ *          stored in the destination tensor `dst` as a scalar.
+ * 
+ *          The operation is optimized using the CANN backend, making it suitable for
+ *          high-performance inference or training scenarios.
+ * 
+ * @param ctx The CANN context used for operations.
+ * @param dst The destination tensor where the result will be stored.
+ *            dst->op is expected to be `GGML_OP_COUNT_EQUAL`.
+ */
+void ggml_cann_count_equal(ggml_backend_cann_context& ctx, ggml_tensor* dst);
+
+/**
+ * @brief   Applies the Step activation function to a ggml tensor using the CANN backend.
+ * 
+ * @details This function applies a step function element-wise to the input tensor, where
+ *          each element is transformed to 1.0 if it is greater than 0, and 0.0 otherwise.
+ *          The result is stored in the destination tensor `dst`.
+ * 
+ *          This operation is accelerated using the CANN backend to improve runtime performance.
+ * 
+ * @param ctx The CANN context used for operations.
+ * @param dst The destination tensor where the result will be stored.
+ *            dst->op is expected to be `GGML_OP_STEP`.
+ */
+void ggml_cann_step(ggml_backend_cann_context& ctx, ggml_tensor* dst);
+
 /**
  * @brief Applies a element-wise operation to two input tensors using the CANN
  * backend.

ggml/src/ggml-cann/ggml-cann.cpp

@@ -1358,6 +1358,12 @@ static bool ggml_cann_compute_forward(ggml_backend_cann_context& ctx,
                 case GGML_UNARY_OP_ELU:
                     ggml_cann_elu(ctx, dst);
                     break;
+                case GGML_UNARY_OP_SGN:
+                    GGML_CANN_CALL_UNARY_OP(Sign);
+                    break;
+                case GGML_UNARY_OP_STEP:
+                    ggml_cann_step(ctx, dst);
+                    break;
                 default:
                     return false;
             }
@@ -1465,6 +1471,9 @@ static bool ggml_cann_compute_forward(ggml_backend_cann_context& ctx,
         case GGML_OP_PAD_REFLECT_1D:
             ggml_cann_pad_reflect_1d(ctx, dst);
             break;
+        case GGML_OP_COUNT_EQUAL:
+            ggml_cann_count_equal(ctx, dst);
+            break;
         default:
             return false;
     }
@@ -1727,6 +1736,8 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev,
                 case GGML_UNARY_OP_TANH:
                 case GGML_UNARY_OP_EXP:
                 case GGML_UNARY_OP_ELU:
+                case GGML_UNARY_OP_SGN:
+                case GGML_UNARY_OP_STEP:
                     return true;
                 default:
                     return false;
@@ -1863,6 +1874,7 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev,
         case GGML_OP_LOG:
         case GGML_OP_MEAN:
         case GGML_OP_PAD_REFLECT_1D:
+        case GGML_OP_COUNT_EQUAL:
             return true;
         default:
             return false;