gemm.h

#ifndef GEMM_H_
#define GEMM_H_
 
#include <adf.h>
#include <assert.h>
 
 
template <int M, int K, int N, int IS_RELU>
class GemmReluScalarMKNKStream {
  private:
    alignas(32) float (&bias)[N];
    alignas(32) float in_row[K];
 
  public:
    GemmReluScalarMKNKStream (
      float (&b)[N]
    ): bias(b) {};
 
    void filter(
      input_stream<float>* in,      // MxK
      input_stream<float>* weight,  // NxK
      output_window<float>* out     // MxN
    );
    static void registerKernelClass() {
      REGISTER_FUNCTION(GemmReluScalarMKNKStream::filter);
      REGISTER_PARAMETER(bias);
    };
};
 
 
// xA^T + b as per torch,nn.Linear
template <int M, int K, int N, int IS_RELU>
class GemmReluScalarMKNK {
  
  private:
    alignas(32) float (&weights)[N*K]; // NxK
    alignas(32) float (&bias)[N];      // N
  
  public:
    GemmReluScalarMKNK (
      float (&w)[N*K],
      float (&b)[N]
    ): weights(w), bias(b) {};
 
    void filter(
      input_window<float>* in,  // MxK
      output_window<float>* out // MxN
    );
    
    static void registerKernelClass() {
      REGISTER_FUNCTION(GemmReluScalarMKNK::filter);
      REGISTER_PARAMETER(weights);
      REGISTER_PARAMETER(bias);
    };
};
 
 
template <int M, int K, int N, int IS_RELU>
class GemmReluScalarMKKN {
  
  private:
    alignas(32) float (&weights)[K*N]; // KxN
    alignas(32) float (&bias)[N];      // N
  
  public:
    GemmReluScalarMKKN (
      float (&w)[N*K],
      float (&b)[N]
    ): weights(w), bias(b) {};
 
    void filter(
      input_window<float>* in,  // MxK
      output_window<float>* out // MxN
    );
    
    static void registerKernelClass() {
      REGISTER_FUNCTION(GemmReluScalarMKKN::filter);
      REGISTER_PARAMETER(weights);
      REGISTER_PARAMETER(bias);
    };
};
 
 
template <int M, int K, int N, int IS_RELU>
class GemmReluMKKNTwoAccsStream {
  private:
    alignas(32) float (&bias)[N];
    alignas(32) float in_row[2*K];
    alignas(32) float out_row[N];
 
  public:
    GemmReluMKKNTwoAccsStream (
      float (&b)[N]
    ): bias(b) {};
 
    void filter(
      input_stream<float>* in,      // MxK
      input_stream<float>* weight,  // NxK
      output_window<float>* out     // MxN
    );
    static void registerKernelClass() {
      static_assert((2*K + N)*4 <= 24576);
      static_assert(N%8 == 0);
      REGISTER_FUNCTION(GemmReluMKKNTwoAccsStream::filter);
      REGISTER_PARAMETER(bias);
    };
};
 
 
template <int M, int K, int N, int IS_RELU>
class GemmReluMKKNStream {
  private:
    alignas(32) float (&bias)[N];
    alignas(32) float in_row[4*K];
    alignas(32) float out_row[3*N];
 
  public:
    GemmReluMKKNStream (
      float (&b)[N]
    ): bias(b) {};
 
    void filter(
      input_stream<float>* in,      // MxK
      input_stream<float>* weight,  // NxK
      output_window<float>* out     // MxN
    );
    static void registerKernelClass() {
      static_assert((4*K + 3*N)*4 <= 24576);
      static_assert(N%8 == 0);
      REGISTER_FUNCTION(GemmReluMKKNStream::filter);
      REGISTER_PARAMETER(bias);
    };
};
 
 
template <int M, int K, int N, int IS_RELU>
class GemmReluMKKN {
  
  private:
    alignas(32) float (&weights)[K*N]; // KxN
    alignas(32) float (&bias)[N];      // N
 
    static constexpr int K_REM8 = K%8;
    static constexpr int RUN_LASTCHUNK = K_REM8 > 0;
  
  public:
    GemmReluMKKN (
      float (&w)[N*K],
      float (&b)[N]
    ): weights(w), bias(b) {};
 
    void filter(
      input_window<float>* in,  // MxK
      output_window<float>* out // MxN
    );
    
    static void registerKernelClass() {
      static_assert(K%4==0 && N%8==0);
      REGISTER_FUNCTION(GemmReluMKKN::filter);
      REGISTER_PARAMETER(weights);
      REGISTER_PARAMETER(bias);
    };
};
#endif // GEMM_H_