netkit/include/cmsis_nn_kernel.hpp at main · NetKit-Labs/netkit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
#pragma once

#include "kernel_activation.hpp"
#include "tensor.hpp"

struct CmsisNnKernel
{
    static bool TryConv2dForward(const Tensor& input,
                                 float* weights,
                                 float* bias,
                                 int kernel_size,
                                 int stride,
                                 int pad_h,
                                 int pad_w,
                                 int in_channels,
                                 int out_channels,
                                 NetkitKernelActivation fuse_activation,
                                 Tensor& output);

    static bool TryDepthwiseConv2dForward(const Tensor& input,
                                          float* weights,
                                          float* bias,
                                          int kernel_h,
                                          int kernel_w,
                                          int stride,
                                          int pad_h,
                                          int pad_w,
                                          int channels,
                                          NetkitKernelActivation fuse_activation,
                                          Tensor& output);

    static bool TryMaxPool2dForward(const Tensor& input,
                                    int pool_size,
                                    int stride,
                                    int pad_h,
                                    int pad_w,
                                    NetkitKernelActivation fuse_activation,
                                    Tensor& output);

    static bool TryAvgPool2dForward(const Tensor& input,
                                    int pool_size,
                                    int stride,
                                    int pad_h,
                                    int pad_w,
                                    Tensor& output);

    static bool TryBatchNorm2dForward(const Tensor& input,
                                      const float* scale,
                                      const float* bias,
                                      int channels,
                                      Tensor& output);

    static bool TryFullyConnectedWithBias(const Tensor& input,
                                          const Tensor& weights,
                                          const Tensor& bias,
                                          NetkitKernelActivation fuse_activation,
                                          Tensor& output);

    static bool TryActivationForward(const Tensor& input,
                                     Tensor& output,
                                     NetkitKernelActivation activation,
                                     float leaky_alpha);

    static bool TrySoftmaxForward(const Tensor& input, Tensor& output);

    static bool TryGeluForward(const Tensor& input, Tensor& output);

    static bool TryMatAdd(const Tensor& a, const Tensor& b, Tensor& c);
};