Add sigmoid op wrapper

Made-with: Cursor

Author: Ceng23333

include/infinicore/ops.hpp

@@ -34,6 +34,7 @@
 #include "ops/reciprocal.hpp"
 #include "ops/rms_norm.hpp"
 #include "ops/rope.hpp"
+#include "ops/sigmoid.hpp"
 #include "ops/silu.hpp"
 #include "ops/silu_and_mul.hpp"
 #include "ops/swiglu.hpp"

include/infinicore/ops/sigmoid.hpp

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

src/infinicore/ops/sigmoid/sigmoid.cc

@@ -0,0 +1,39 @@
+#include "infinicore/ops/sigmoid.hpp"
+
+#include "../../utils.hpp"
+
+#include <stdexcept>
+
+namespace infinicore::op {
+
+common::OpDispatcher<Sigmoid::schema> &Sigmoid::dispatcher() {
+    static common::OpDispatcher<Sigmoid::schema> dispatcher_;
+    return dispatcher_;
+};
+
+void Sigmoid::execute(Tensor output, Tensor input) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input);
+    infinicore::context::setDevice(output->device());
+    auto device_type = output->device().getType();
+    auto func = dispatcher().lookup(device_type);
+
+    if (func == nullptr) {
+        throw std::runtime_error("No Sigmoid implementation found for device type: " + std::to_string(static_cast<int>(device_type)));
+    }
+
+    func(output, input);
+}
+
+Tensor sigmoid(Tensor input) {
+    Shape shape = input->shape();
+    auto output = Tensor::empty(shape, input->dtype(), input->device());
+    sigmoid_(output, input);
+    return output;
+}
+
+void sigmoid_(Tensor output, Tensor input) {
+    Sigmoid::execute(output, input);
+}
+
+} // namespace infinicore::op
+

src/infinicore/ops/sigmoid/sigmoid_infiniop.cc

@@ -0,0 +1,52 @@
+#include "../../utils.hpp"
+#include "infinicore/common/hash.hpp"
+#include "infinicore/ops/common/cache.hpp"
+#include "infinicore/ops/sigmoid.hpp"
+
+#include <infiniop.h>
+
+namespace infinicore::op::sigmoid_impl::infiniop {
+
+thread_local common::OpCache<size_t, infiniopSigmoidDescriptor_t> caches(
+    100, // capacity
+    [](infiniopSigmoidDescriptor_t &desc) {
+        if (desc != nullptr) {
+            INFINICORE_CHECK_ERROR(infiniopDestroySigmoidDescriptor(desc));
+            desc = nullptr;
+        }
+    });
+
+void calculate(Tensor output, Tensor input) {
+    size_t seed = hash_combine(output, input);
+
+    auto device = context::getDevice();
+    auto &cache = caches.getCache(device);
+
+    auto desc_opt = cache.get(seed);
+    infiniopSigmoidDescriptor_t desc = nullptr;
+
+    if (!desc_opt) {
+        INFINICORE_CHECK_ERROR(infiniopCreateSigmoidDescriptor(
+            context::getInfiniopHandle(device), &desc,
+            output->desc(), input->desc()));
+        cache.put(seed, desc);
+    } else {
+        desc = *desc_opt;
+    }
+
+    size_t workspace_size = 0;
+    INFINICORE_CHECK_ERROR(infiniopGetSigmoidWorkspaceSize(desc, &workspace_size));
+    std::shared_ptr<Memory> workspace = context::allocateMemory(workspace_size);
+
+    INFINICORE_CHECK_ERROR(infiniopSigmoid(
+        desc, workspace->data(), workspace_size,
+        output->data(), input->data(), context::getStream()));
+}
+
+static bool registered = []() {
+    Sigmoid::dispatcher().registerAll(&calculate, false);
+    return true;
+}();
+
+} // namespace infinicore::op::sigmoid_impl::infiniop
+