netkit/include/xnnpack_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
70
71
72
73
74
75
76
77
78
#pragma once

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

// XNNPACK LayerFast backend (float32). Used on cpu/MPU when NETKIT_XNNPACK=1.
// MCU builds leave this off (NETKIT_XNNPACK_ALLOWED=0).
struct XnnpackKernel
{
    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 TryFullyConnectedWithBias(const Tensor& input,
                                          const Tensor& weights,
                                          const Tensor& bias,
                                          NetkitKernelActivation fuse_activation,
                                          Tensor& output);

    // Unsupported by this backend — return false so ComposedKernel falls back.
    static bool TryBatchNorm2dForward(const Tensor& /*input*/,
                                      const float* /*scale*/,
                                      const float* /*bias*/,
                                      int /*channels*/,
                                      Tensor& /*output*/)
    {
        return false;
    }

    static bool TryActivationForward(const Tensor& /*input*/,
                                     Tensor& /*output*/,
                                     NetkitKernelActivation /*activation*/,
                                     float /*leaky_alpha*/)
    {
        return false;
    }

    static bool TrySoftmaxForward(const Tensor& /*input*/, Tensor& /*output*/) { return false; }

    static bool TryGeluForward(const Tensor& /*input*/, Tensor& /*output*/) { return false; }

    static bool TryMatAdd(const Tensor& /*a*/, const Tensor& /*b*/, Tensor& /*c*/) { return false; }
};