Integration of I-ViT operators

Context

Implementation of the 3 I-ViT operators: ShiftMax, ShiftGELU and ILayerNorm, which are quantifiable arithmetic approximations of SoftMax, GELU and LayerNorm. The forward and backward operators are implemented to enable quantization-aware-training of networks containing these layers (notably Transformers). Unit tests are also implemented to illustrate how these operators work, and to measure the difference between these approximators and the original operators.

Modified files

• include/aidge/backend/cuda.hpp: Import of ShiftMax, ShiftGELU and ILayerNorm;

Added files

• src/operator/ShiftMaxImpl.cpp and include/aidge/backend/cuda/operator/ShiftMaxImpl.hpp: Definition of the forward and backward for the ShiftMax method (highly inspired by ReLU implementation);

• src/operator/ShiftMaxImpl_CUDA_kernels.cu and include/aidge/backend/cuda/operator/ShiftMaxImpl_CUDA_kernels.hpp: Implementation of the CUDA kernels which compute the ShiftMax operation. This CUDA code is conform to the Python implementation of I-ViT;

• unit_tests/Test_ShiftMaxImpl.cpp: Unit test that illustrate how the forward and backward function of ShiftMax work and compare it with SoftMax;

• src/operator/ShiftGELUImpl.cpp and include/aidge/backend/cuda/operator/ShiftGELUImpl.hpp: Definition of the forward and backward for the ShiftGELU method (highly inspired by ReLU implementation);

• src/operator/ShiftGELUImpl_CUDA_kernels.cu and include/aidge/backend/cuda/operator/ShiftGELUImpl_CUDA_kernels.hpp: Implementation of the CUDA kernels which compute the ShiftGELU operation. This CUDA code is conform to the Python implementation of I-ViT;

• unit_tests/Test_ShiftGELUImpl.cpp: Unit test that illustrate how the forward and backward function of ShiftGELU work and compare it with SoftMax;

• src/operator/ILayerNormImpl.cpp and include/aidge/backend/cuda/operator/ILayerNormImpl.hpp: Definition of the forward and backward for the ILayerNorm method (highly inspired by FC implementation);

• src/operator/ILayerNormImpl_CUDA_kernels.cu and include/aidge/backend/cuda/operator/ILayerNormImpl_CUDA_kernels.hpp: Implementation of the CUDA kernels which compute the ILayerNorm operation. This CUDA code is conform to the Python implementation of I-ViT;

• unit_tests/Test_IlayerNormImpl.cpp: Unit test that illustrate how the forward and backward function of ILayerNorm work;

TODO

Merge into main if pipeline successful

• NOT DONE
• DONE
• TO DO

include/aidge/backend/cuda/operator/ILayerNormImpl.hpp

+/********************************************************************************
+ * Copyright (c) 2024 Thales
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ * Author: Lucas RAKOTOARIVONY, Thales Research & Technology France
+ * Date: 10.09.2024
+ *
+ ********************************************************************************/
+
+#ifndef AIDGE_BACKEND_CUDA_OPERATOR_ILAYERNORMIMPL_H_
+#define AIDGE_BACKEND_CUDA_OPERATOR_ILAYERNORMIMPL_H_
+
+#include <array>
+#include <memory>
+#include <tuple>
+#include <vector>
+
+#include <cudnn.h>
+
+#include "aidge/backend/OperatorImpl.hpp"
+#include "aidge/operator/ILayerNorm.hpp"
+#include "aidge/utils/Registrar.hpp"
+#include "aidge/utils/Types.h"
+
+#include "aidge/backend/cuda/utils/CudaUtils.hpp"
+
+namespace Aidge {
+class ILayerNormImpl_cuda : public OperatorImpl {
+public:
+    ILayerNormImpl_cuda(const ILayerNorm_Op &op) : OperatorImpl(op, "cuda") {}
+
+    static std::unique_ptr<ILayerNormImpl_cuda> create(const ILayerNorm_Op &op) {
+        return std::make_unique<ILayerNormImpl_cuda>(op);
+    }
+
+    virtual std::set<ImplSpec> getAvailableImplSpecs() const override {
+        return {
+            {DataType::Float64},
+            {DataType::Float32},
+            {DataType::Float16},
+        };
+    }
+
+    void forward() override;
+    void backward() override;
+
+private:
+    std::shared_ptr<Tensor> mInput0Fallback;
+    std::shared_ptr<Tensor> mInput1Fallback;
+    std::shared_ptr<Tensor> mInput2Fallback;
+    std::shared_ptr<Tensor> mOutputGradFallback;
+
+    template <class T> void forward_(const Tensor& input0, const Tensor& input1, const Tensor& input2);
+    template <class T> void backward_(const Tensor& output_grad);
+};
+
+// Implementation entry point registration to Operator
+REGISTRAR(ILayerNorm_Op, "cuda", Aidge::ILayerNormImpl_cuda::create);
+}  // namespace Aidge
+
+#endif /* AIDGE_BACKEND_CUDA_OPERATOR_ILAYERNORMIMPL_H_ */