Merge pull request #1131 from InfiniTensor/issue/1126

Issue/1126 - quickgelu, gelutanh

Author: wooway777

include/infinicore/adaptor/aten_adaptor.hpp

@@ -6,9 +6,9 @@
 #include <ATen/ATen.h>
 
 #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API)
-#include <c10/cuda/CUDAStream.h>
-#include <c10/cuda/CUDAGuard.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
 #endif
 
 namespace infinicore::adaptor {

include/infinicore/nn/embedding.hpp

@@ -1,7 +1,7 @@
 #pragma once
 
-#include "module.hpp"
 #include "../ops.hpp"
+#include "module.hpp"
 #include <optional>
 
 namespace infinicore::nn {
@@ -78,10 +78,10 @@ class Embedding : public Module {
     INFINICORE_NN_PARAMETER(weight);
 
 private:
-    size_t num_embeddings_;   // Vocabulary size
-    size_t embedding_dim_;    // Embedding dimension
-    std::optional<int64_t> padding_idx_;  // Optional padding index
-    DataType dtype_;           // Data type for embedding weights
+    size_t num_embeddings_;              // Vocabulary size
+    size_t embedding_dim_;               // Embedding dimension
+    std::optional<int64_t> padding_idx_; // Optional padding index
+    DataType dtype_;                     // Data type for embedding weights
 };
 
 } // namespace infinicore::nn

include/infinicore/nn/module.hpp

@@ -3,10 +3,10 @@
 #include "../tensor.hpp"
 #include "parameter.hpp"
 
+#include <spdlog/spdlog.h>
 #include <type_traits>
 #include <unordered_map>
 #include <vector>
-#include <spdlog/spdlog.h>
 
 namespace infinicore::nn {
 class Module {

include/infinicore/ops.hpp

@@ -14,27 +14,34 @@
 #include "ops/binary_cross_entropy_with_logits.hpp"
 #include "ops/causal_softmax.hpp"
 #include "ops/cdist.hpp"
+#include "ops/conv2d.hpp"
 #include "ops/cross_entropy.hpp"
 #include "ops/embedding.hpp"
 #include "ops/flash_attention.hpp"
 #include "ops/fmin.hpp"
 #include "ops/fmod.hpp"
+#include "ops/gelu.hpp"
+#include "ops/gelutanh.hpp"
 #include "ops/hardswish.hpp"
 #include "ops/hardtanh.hpp"
 #include "ops/kv_caching.hpp"
 #include "ops/layer_norm.hpp"
+#include "ops/linear.hpp"
 #include "ops/matmul.hpp"
 #include "ops/ones.hpp"
 #include "ops/paged_attention.hpp"
 #include "ops/paged_attention_prefill.hpp"
 #include "ops/paged_caching.hpp"
 #include "ops/per_tensor_dequant_i8.hpp"
 #include "ops/per_tensor_quant_i8.hpp"
+#include "ops/quickgelu.hpp"
 #include "ops/random_sample.hpp"
 #include "ops/rearrange.hpp"
 #include "ops/reciprocal.hpp"
+#include "ops/relu.hpp"
 #include "ops/rms_norm.hpp"
 #include "ops/rope.hpp"
 #include "ops/silu.hpp"
 #include "ops/silu_and_mul.hpp"
+#include "ops/softmax.hpp"
 #include "ops/swiglu.hpp"

include/infinicore/ops/conv2d.hpp

@@ -0,0 +1,38 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+#include <cstddef>
+#include <vector>
+
+namespace infinicore::op {
+class Conv2d {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor,
+                            const size_t *, const size_t *, const size_t *, size_t);
+    static void execute(Tensor output,
+                        Tensor input,
+                        Tensor weight,
+                        Tensor bias,
+                        const size_t *pads,
+                        const size_t *strides,
+                        const size_t *dilations,
+                        size_t n);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor conv2d(Tensor input,
+              Tensor weight,
+              Tensor bias,
+              const std::vector<size_t> &pads,
+              const std::vector<size_t> &strides,
+              const std::vector<size_t> &dilations);
+void conv2d_(Tensor output,
+             Tensor input,
+             Tensor weight,
+             Tensor bias,
+             const std::vector<size_t> &pads,
+             const std::vector<size_t> &strides,
+             const std::vector<size_t> &dilations);
+} // namespace infinicore::op

include/infinicore/ops/gelu.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class Gelu {
+public:
+    using schema = void (*)(Tensor, Tensor);
+    static void execute(Tensor output, Tensor input);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor gelu(Tensor input);
+void gelu_(Tensor output, Tensor input);
+} // namespace infinicore::op

include/infinicore/ops/gelutanh.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class GeluTanh {
+public:
+    using schema = void (*)(Tensor, Tensor);
+    static void execute(Tensor output, Tensor input);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor gelu_tanh(Tensor input);
+void gelu_tanh_(Tensor output, Tensor input);
+} // namespace infinicore::op

include/infinicore/ops/mha_kvcache.hpp

@@ -22,14 +22,14 @@ namespace infinicore::op {
 
 INFINICORE_GRAPH_OP_CLASS(
     MhaKVCache,
-    Tensor,           // out
-    const Tensor &,   // q
-    const Tensor &,   // k_cache
-    const Tensor &,   // v_cache
-    const Tensor &,   // seqlens_k
-    const Tensor &,   // block_table
-    std::optional<Tensor>,  // alibi_slopes
-    float);           // scale
+    Tensor,                // out
+    const Tensor &,        // q
+    const Tensor &,        // k_cache
+    const Tensor &,        // v_cache
+    const Tensor &,        // seqlens_k
+    const Tensor &,        // block_table
+    std::optional<Tensor>, // alibi_slopes
+    float);                // scale
 
 Tensor mha_kvcache(const Tensor &q,
                    const Tensor &k_cache,

include/infinicore/ops/quickgelu.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class QuickGelu {
+public:
+    using schema = void (*)(Tensor, Tensor);
+    static void execute(Tensor output, Tensor input);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor quick_gelu(Tensor input);
+void quick_gelu_(Tensor output, Tensor input);
+} // namespace infinicore::op

include/infinicore/ops/relu.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class Relu {
+public:
+    using schema = void (*)(Tensor, Tensor);
+    static void execute(Tensor output, Tensor input);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor relu(Tensor input);
+void relu_(Tensor output, Tensor input);
+} // namespace infinicore::op