BackendConv3D.h

#pragma once
 
#if _MSC_VER >= 1900 && defined(_M_X64)
 
#include "BackendBase.h"
#include "Parameters.h"
#include <immintrin.h>
 
namespace FLImaging
{
    namespace AI
    {
        template <typename T>
        class FL_EXPORT CTensor;
 
        template <typename T>
        class FL_EXPORT CBackendConv3D : public CBackendBase<T>
        {
        public:
            CBackendConv3D();
            CBackendConv3D(const CBackendConv3D<T>& bc);
            virtual ~CBackendConv3D();
 
            virtual const CResult SetConvolutionParams(const CConvolution3DParameters& convParams);
            virtual CConvolution3DParameters GetConvolutionParams();
 
            virtual const CResult Forward(CTensor<T>* pTsrX, CTensor<T>* pTsrW, CTensor<T>* pTsrY, const std::vector<int64_t>& vctYShape);
            virtual const CResult DerivativeImage(CTensor<T>* pTsrDy, CTensor<T>* pTsrW, CTensor<T>* pTsrDx, const std::vector<int64_t>& vctDxShape, bool bAddGradient);
            virtual const CResult DerivativeKernel(CTensor<T>* pTsrDy, CTensor<T>* pTsrX, CTensor<T>* pTsrDw, const std::vector<int64_t>& vctDwShape, bool bAddGradient);
 
            DeclareGetClassType();
            SupportToDuplicateObjectWithoutCreateNewObject(CBackendConv3D<T>, *this);
 
        protected:
            virtual const CResult Forward_Direct(CTensor<T>* pTsrOperand, CTensor<T>* pTsrKernel, CTensor<T>* pTsrResult);
            virtual const CResult DerivativeImage_Direct(CTensor<T>* pTsrDy, CTensor<T>* pTsrW, CTensor<T>* pTsrDx, bool bAddGradient);
            virtual const CResult DerivativeKernel_Direct(CTensor<T>* pTsrDy, CTensor<T>* pTsrX, CTensor<T>* pTsrDw, bool bAddGradient);
 
        protected:
            CConvolution3DParameters m_conv3DParams;
 
        private:
            template <typename T>
            struct AvxTemplate;
 
            template <>
            struct AvxTemplate<float>
            {
                using Vec = __m256;
                static constexpr int32_t LANES = 8;
 
                static inline Vec zero()
                {
                    return _mm256_setzero_ps();
                }
 
                static inline Vec set1(float f32V)
                {
                    return _mm256_set1_ps(f32V);
                }
 
                static inline Vec loadu(const float* pF32)
                {
                    return _mm256_loadu_ps(pF32);
                }
 
                static inline void storeu(float* pF32, Vec vec)
                {
                    _mm256_storeu_ps(pF32, vec);
                }
 
                static inline Vec fmadd(Vec vecA, Vec vecB, Vec vecC)
                {
                    return _mm256_add_ps(vecC, _mm256_mul_ps(vecA, vecB));
                }
 
                static inline float hsum(Vec v)
                {
                    alignas(32) float arrF32Buf[LANES];
                    _mm256_store_ps(arrF32Buf, v);
                    float f32Ret = 0.0f;
 
                    for(int i = 0; i < LANES; ++i)
                        f32Ret += arrF32Buf[i];
 
                    return f32Ret;
                }
            };
 
            template <>
            struct AvxTemplate<double>
            {
                using Vec = __m256d;
                static constexpr int32_t LANES = 4;
 
                static inline Vec zero()
                {
                    return _mm256_setzero_pd();
                }
 
                static inline Vec set1(double f64V)
                {
                    return _mm256_set1_pd(f64V);
                }
 
                static inline Vec loadu(const double* pF64)
                {
                    return _mm256_loadu_pd(pF64);
                }
 
                static inline void storeu(double* pF64, Vec vec)
                {
                    _mm256_storeu_pd(pF64, vec);
                }
 
                static inline Vec fmadd(Vec vecA, Vec vecB, Vec vecC)
                {
                    return _mm256_add_pd(vecC, _mm256_mul_pd(vecA, vecB));
                }
 
                static inline double hsum(Vec v)
                {
                    alignas(32) double arrF64Buf[LANES];
                    _mm256_store_pd(arrF64Buf, v);
                    double f64Ret = 0.0;
 
                    for(int i = 0; i < LANES; ++i)
                        f64Ret += arrF64Buf[i];
 
                    return f64Ret;
                }
            };
        };
    }
}
 
#endif