Merge pull request #71 from InfiniTensor/dev-erf

feat: add Erf cpu/cuda kernel

Author: YdrMaster

src/04kernel/src/kernels/simple_unary/cpu_kernel.cc

@@ -18,6 +18,7 @@ namespace refactor::kernel {
             Op::Sigmoid,
             Op::Tanh,
             Op::Neg,
+            Op::Erf,
         };
         return supportedOp.contains(op) && a.dataType.isCpuNumberic()
                    ? std::make_unique<K>(op, a.dataType, a.elementsSize())
@@ -155,6 +156,21 @@ namespace refactor::kernel {
                     default:
                         UNREACHABLE();
                 }
+            case Op::Erf:
+                switch (dataType) {
+                    CASE(std::erf, F32);
+                    CASE(std::erf, F64);
+                    CASE(std::erf, I8);
+                    CASE(std::erf, I16);
+                    CASE(std::erf, I32);
+                    CASE(std::erf, I64);
+                    CASE(std::erf, U8);
+                    CASE(std::erf, U16);
+                    CASE(std::erf, U32);
+                    CASE(std::erf, U64);
+                    default:
+                        UNREACHABLE();
+                }
             default:
                 UNREACHABLE();
         }

src/04kernel/src/kernels/simple_unary/cuda_kernel.cc

@@ -18,7 +18,8 @@ namespace refactor::kernel {
     auto K::build(Op op, Tensor const &a) noexcept -> KernelBox {
         static const std::unordered_set<Op>
             supportedOp{Op::Abs, Op::Relu, Op::Sqrt,
-                        Op::Sigmoid, Op::Tanh, Op::Neg};
+                        Op::Sigmoid, Op::Tanh, Op::Neg,
+                        Op::Erf};
 #ifndef USE_CUDA
         return nullptr;
 #endif
@@ -140,6 +141,19 @@ extern "C" __global__ void kernel(
             {__(Op::Neg, DT::BF16), "-x"},
             {__(Op::Neg, DT::F32 ), "-x"},
             {__(Op::Neg, DT::F64 ), "-x"},
+
+            {__(Op::Erf, DT::F32 ), "erff(x)"},
+            {__(Op::Erf, DT::F64 ), "erf(x)"},
+            {__(Op::Erf, DT::U8  ), "erff(static_cast<float>(x))"},
+            {__(Op::Erf, DT::I8  ), "erff(static_cast<float>(x))"},
+            {__(Op::Erf, DT::U16 ), "erff(static_cast<float>(x))"},
+            {__(Op::Erf, DT::I16 ), "erff(static_cast<float>(x))"},
+            {__(Op::Erf, DT::U32 ), "erf(static_cast<double>(x))"},
+            {__(Op::Erf, DT::I32 ), "erf(static_cast<double>(x))"},
+            {__(Op::Erf, DT::U64 ), "erf(static_cast<double>(x))"},
+            {__(Op::Erf, DT::I64 ), "erf(static_cast<double>(x))"},
+            {__(Op::Erf, DT::FP16), "__float2half(erff(__half2float(x)))"},
+            {__(Op::Erf, DT::BF16), "__float2bfloat16(erff(__bfloat162float(x)))"},
         };
         // clang-format on
 

src/04kernel/test/kernels/simple_unary/test_cpu.cpp

@@ -31,4 +31,5 @@ TEST(kernel, SimpleUnaryCpu) {
     testOp(SimpleUnaryType::Abs, std::abs);
     testOp(SimpleUnaryType::Sqrt, std::sqrt);
     testOp(SimpleUnaryType::Tanh, std::tanh);
+    testOp(SimpleUnaryType::Erf, std::erf);
 }

src/04kernel/test/kernels/simple_unary/test_cuda.cpp

@@ -51,6 +51,7 @@ TEST(kernel, SimpleUnaryCuda) {
     testOp(SimpleUnaryType::Sqrt);
     testOp(SimpleUnaryType::Sigmoid);
     testOp(SimpleUnaryType::Tanh);
+    testOp(SimpleUnaryType::Erf);
 }
 
 #endif

src/07onnx/test/test_simple_unary.cpp

@@ -0,0 +1,25 @@
+#include "../src/operators/simple_unary.hh"
+#include "onnx/operators.h"
+#include <gtest/gtest.h>
+
+using namespace refactor;
+using namespace onnx;
+
+TEST(infer, SimpleUnary) {
+    onnx::register_();
+
+    {
+        // Erf Test
+        auto edges = Edges{
+            {Tensor::share(DataType::F32, Shape{DimExpr(2), DimExpr(3)}, {}), ""},
+        };
+        count_t inputs[]{0};
+        auto infered = SimpleUnary(SimpleUnaryType::Erf).infer(TensorRefs(edges, inputs), {true});
+        ASSERT_TRUE(infered.isOk());
+        auto outputs = std::move(infered.unwrap());
+        ASSERT_EQ(outputs.size(), 1);
+        auto y = std::move(outputs[0]);
+        ASSERT_EQ(y->dataType, DataType::F32);
+        ASSERT_EQ(y->shape, (Shape{DimExpr(2), DimExpr(3)}));
+    }
+}