issue/427 - the sigmoid, topksoftmax, and topkrouter ops

include/infiniop.h

@@ -19,6 +19,8 @@
 #include "infiniop/ops/sub.h"
 #include "infiniop/ops/swiglu.h"
 #include "infiniop/ops/topkrouter.h"
+#include "infiniop/ops/topksoftmax.h"
+#include "infiniop/ops/sigmoid.h"
 #include "infiniop/tensor_descriptor.h"
 
 #endif // __INFINIOP_API_H__

include/infiniop/ops/sigmoid.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_SIGMOID_API_H__
+#define __INFINIOP_SIGMOID_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopSigmoidDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateSigmoidDescriptor(infiniopHandle_t handle,
+                                                            infiniopSigmoidDescriptor_t *desc_ptr,
+                                                            infiniopTensorDescriptor_t y,
+                                                            infiniopTensorDescriptor_t x);
+
+__C __export infiniStatus_t infiniopGetSigmoidWorkspaceSize(infiniopSigmoidDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopSigmoid(infiniopSigmoidDescriptor_t desc,
+                                            void *workspace,
+                                            size_t workspace_size,
+                                            void *y,
+                                            const void *x,
+                                            void *stream);
+
+__C __export infiniStatus_t infiniopDestroySigmoidDescriptor(infiniopSigmoidDescriptor_t desc);
+
+#endif

include/infiniop/ops/topkrouter.h

@@ -5,16 +5,23 @@
 
 typedef struct InfiniopDescriptor *infiniopTopkrouterDescriptor_t;
 
-__C __export infiniStatus_t infiniopCreateTopkrouterDescriptor(
-    infiniopHandle_t handle,
-    infiniopTopkrouterDescriptor_t *desc_ptr,
-    infiniopTensorDescriptor_t x_desc,
-    infiniopTensorDescriptor_t correction_bias_desc);
+__C __export infiniStatus_t infiniopCreateTopkrouterDescriptor(infiniopHandle_t handle,
+                                                               infiniopTopkrouterDescriptor_t *desc_ptr,
+                                                               infiniopTensorDescriptor_t x_desc,
+                                                               infiniopTensorDescriptor_t correction_bias_desc);
 
 __C __export infiniStatus_t infiniopGetTopkrouterWorkspaceSize(infiniopTopkrouterDescriptor_t desc, size_t *size);
 
-__C __export infiniStatus_t infiniopTopkrouter(infiniopTopkrouterDescriptor_t desc, void *workspace, size_t workspace_size,
-                                               void *values, void *indices, void *x, void *correction_bias, float routed_scaling_factor, size_t topk, void *stream);
+__C __export infiniStatus_t infiniopTopkrouter(infiniopTopkrouterDescriptor_t desc,
+                                               void *workspace,
+                                               size_t workspace_size,
+                                               void *values,
+                                               void *indices,
+                                               const void *x,
+                                               const void *correction_bias,
+                                               const float routed_scaling_factor,
+                                               const size_t topk,
+                                               void *stream);
 
 __C __export infiniStatus_t infiniopDestroyTopkrouterDescriptor(infiniopTopkrouterDescriptor_t desc);
 

include/infiniop/ops/topksoftmax.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_TOPKSOFTMAX_API_H__
+#define __INFINIOP_TOPKSOFTMAX_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopTopksoftmaxDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateTopksoftmaxDescriptor(infiniopHandle_t handle,
+                                                                infiniopTopksoftmaxDescriptor_t *desc_ptr,
+                                                                infiniopTensorDescriptor_t x_desc);
+
+__C __export infiniStatus_t infiniopGetTopksoftmaxWorkspaceSize(infiniopTopksoftmaxDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopTopksoftmax(infiniopTopksoftmaxDescriptor_t desc,
+                                                void *workspace,
+                                                size_t workspace_size,
+                                                void *values,
+                                                void *indices,
+                                                const void *x,
+                                                const size_t topk,
+                                                const int norm,
+                                                void *stream);
+
+__C __export infiniStatus_t infiniopDestroyTopksoftmaxDescriptor(infiniopTopksoftmaxDescriptor_t desc);
+
+#endif

scripts/python_test.py

@@ -25,6 +25,9 @@ def run_tests(args):
         "sub.py",
         "swiglu.py",
         "softplus.py",
+        "sigmoid.py",
+        "topkrouter.py",
+        "topksoftmax.py",
     ]:
         result = subprocess.run(
             f"python {test} {args} --debug", text=True, encoding="utf-8", shell=True

src/infiniop-test/include/ops.hpp

@@ -16,7 +16,9 @@ DECLARE_INFINIOP_TEST(add)
 DECLARE_INFINIOP_TEST(causal_softmax)
 DECLARE_INFINIOP_TEST(rearrange)
 DECLARE_INFINIOP_TEST(sub)
-
+DECLARE_INFINIOP_TEST(sigmoid)
+DECLARE_INFINIOP_TEST(topkrouter)
+DECLARE_INFINIOP_TEST(topksoftmax)
 #define REGISTER_INFINIOP_TEST(name)                      \
     {                                                     \
         #name,                                            \
@@ -43,6 +45,9 @@ DECLARE_INFINIOP_TEST(sub)
         REGISTER_INFINIOP_TEST(causal_softmax) \
         REGISTER_INFINIOP_TEST(rearrange)      \
         REGISTER_INFINIOP_TEST(sub)            \
+        REGISTER_INFINIOP_TEST(sigmoid)        \
+        REGISTER_INFINIOP_TEST(topkrouter)     \
+        REGISTER_INFINIOP_TEST(topksoftmax)    \
     }
 
 namespace infiniop_test {

src/infiniop-test/src/ops/sigmoid.cpp

@@ -0,0 +1,103 @@
+#include "ops.hpp"
+#include "utils.hpp"
+#include <infinirt.h>
+#include <iomanip>
+#include <iostream>
+
+namespace infiniop_test::sigmoid {
+struct Test::Attributes {
+    std::shared_ptr<Tensor> x;
+    std::shared_ptr<Tensor> y;
+    std::shared_ptr<Tensor> ans;
+};
+
+std::shared_ptr<Test> Test::build(
+    std::unordered_map<std::string, std::vector<uint8_t>> attributes,
+    std::unordered_map<std::string, std::shared_ptr<Tensor>> tensors,
+    double rtol, double atol) {
+    auto test = std::shared_ptr<Test>(new Test(rtol, atol));
+    test->_attributes = new Attributes();
+    if (tensors.find("x") == tensors.end()
+        || tensors.find("y") == tensors.end()
+        || tensors.find("ans") == tensors.end()) {
+        throw std::runtime_error("Invalid Test");
+    }
+
+    test->_attributes->x = tensors["x"];
+    test->_attributes->y = tensors["y"];
+    test->_attributes->ans = tensors["ans"];
+
+    return test;
+}
+
+std::shared_ptr<infiniop_test::Result> Test::run(
+    infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) {
+    infiniopSigmoidDescriptor_t op_desc;
+    auto x = _attributes->x->to(device, device_id);
+    auto y = _attributes->y->to(device, device_id);
+    CHECK_OR(infiniopCreateSigmoidDescriptor(handle, &op_desc,
+                                             y->desc(),
+                                             x->desc()),
+             return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor."));
+    size_t workspace_size;
+    CHECK_OR(infiniopGetSigmoidWorkspaceSize(op_desc, &workspace_size),
+             return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size."));
+    void *workspace;
+    CHECK_OR(infinirtMalloc(&workspace, workspace_size),
+             return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace."));
+    CHECK_OR(infiniopSigmoid(op_desc, workspace, workspace_size,
+                             y->data(),
+                             x->data(),
+                             nullptr),
+             return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution."));
+
+    try {
+        allClose(y, _attributes->ans, _rtol, _atol);
+    } catch (const std::exception &e) {
+        return TEST_FAILED(RESULT_INCORRECT, e.what());
+    }
+
+    double elapsed_time = 0.;
+
+    elapsed_time = benchmark(
+        [=]() {
+            infiniopSigmoid(
+                op_desc, workspace, workspace_size,
+                y->data(),
+                x->data(),
+                nullptr);
+        },
+        warm_ups, iterations);
+
+    infiniopDestroySigmoidDescriptor(op_desc);
+    infinirtFree(workspace);
+    return TEST_PASSED(elapsed_time);
+}
+
+std::vector<std::string> Test::attribute_names() {
+    return {};
+}
+
+std::vector<std::string> Test::tensor_names() {
+    return {"x", "y", "ans"};
+}
+
+std::vector<std::string> Test::output_names() {
+    return {"y"};
+}
+
+std::string Test::toString() const {
+    std::ostringstream oss;
+    oss << op_name() << std::endl;
+    oss << "- x: " << _attributes->x->info() << std::endl;
+    oss << "- y: " << _attributes->y->info() << std::endl;
+    oss << std::scientific << std::setprecision(2);
+    oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl;
+    return oss.str();
+}
+
+Test::~Test() {
+    delete _attributes;
+}
+
+} // namespace infiniop_test::sigmoid

src/infiniop-test/src/ops/topkrouter.cpp

@@ -0,0 +1,130 @@
+#include "ops.hpp"
+#include "utils.hpp"
+#include <infinirt.h>
+#include <iomanip>
+#include <iostream>
+
+namespace infiniop_test::topkrouter {
+struct Test::Attributes {
+    std::shared_ptr<Tensor> values;
+    std::shared_ptr<Tensor> indices;
+    std::shared_ptr<Tensor> x;
+    std::shared_ptr<Tensor> correction_bias;
+    float routed_scaling_factor;
+    int topk;
+    std::shared_ptr<Tensor> lable_values;
+    std::shared_ptr<Tensor> lable_indices;
+};
+
+std::shared_ptr<Test> Test::build(std::unordered_map<std::string, std::vector<uint8_t>> attributes,
+                                  std::unordered_map<std::string, std::shared_ptr<Tensor>> tensors, double rtol,
+                                  double atol) {
+    auto test = std::shared_ptr<Test>(new Test(rtol, atol));
+    test->_attributes = new Attributes();
+
+    if (attributes.find("routed_scaling_factor") == attributes.end() || attributes.find("topk") == attributes.end() || tensors.find("values") == tensors.end() || tensors.find("indices") == tensors.end() || tensors.find("x") == tensors.end() || tensors.find("correction_bias") == tensors.end() || tensors.find("lable_values") == tensors.end() || tensors.find("lable_indices") == tensors.end()) {
+        throw std::runtime_error("Invalid Test: Missing attributes or tensors");
+    }
+
+    test->_attributes->values = tensors["values"];
+    test->_attributes->indices = tensors["indices"];
+    test->_attributes->x = tensors["x"];
+    test->_attributes->correction_bias = tensors["correction_bias"];
+
+    test->_attributes->routed_scaling_factor = *reinterpret_cast<float *>(attributes["routed_scaling_factor"].data());
+    test->_attributes->topk = *reinterpret_cast<int *>(attributes["topk"].data());
+
+    test->_attributes->lable_values = tensors["lable_values"];
+    test->_attributes->lable_indices = tensors["lable_indices"];
+
+    return test;
+}
+
+std::shared_ptr<infiniop_test::Result> Test::run(infiniopHandle_t handle, infiniDevice_t device, int device_id,
+                                                 size_t warm_ups, size_t iterations) {
+    infiniopTopkrouterDescriptor_t op_desc;
+    CHECK_OR(infiniopCreateTopkrouterDescriptor(handle, &op_desc, _attributes->x->desc(),
+                                                _attributes->correction_bias->desc()),
+             return TEST_FAILED(OP_CREATION_FAILED, "Failed to create Topkrouter descriptor"));
+
+    //
+    auto values = _attributes->values->to(device, device_id);
+    auto indices = _attributes->indices->to(device, device_id);
+    auto x = _attributes->x->to(device, device_id);
+    auto correction_bias = _attributes->correction_bias->to(device, device_id);
+
+    float routed_scaling_factor = _attributes->routed_scaling_factor;
+    int topk = _attributes->topk;
+
+    size_t workspace_size;
+    CHECK_OR(infiniopGetTopkrouterWorkspaceSize(op_desc, &workspace_size),
+             return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size"));
+    void *workspace = nullptr;
+    if (workspace_size > 0) {
+        CHECK_OR(infinirtMalloc(&workspace, workspace_size),
+                 return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace"));
+    }
+
+    CHECK_OR(infiniopTopkrouter(op_desc, workspace, workspace_size, values->data(), indices->data(), x->data(),
+                                correction_bias->data(), routed_scaling_factor, topk, nullptr),
+             return TEST_FAILED(OP_EXECUTION_FAILED, "Topkrouter execution failed"));
+
+    try {
+        allClose(values, _attributes->lable_values, _rtol, _atol);
+        allClose(indices, _attributes->lable_indices, _rtol, _atol);
+    } catch (const std::exception &e) {
+        return TEST_FAILED(RESULT_INCORRECT, e.what());
+    }
+
+    double elapsed_time = 0.;
+
+    elapsed_time = benchmark(
+        [=]() {
+            infiniopTopkrouter(op_desc, workspace, workspace_size, values->data(), indices->data(), x->data(),
+                               correction_bias->data(), routed_scaling_factor, topk, nullptr);
+        },
+        warm_ups, iterations);
+
+    if (workspace != nullptr) {
+        infinirtFree(workspace);
+    }
+
+    return TEST_PASSED(elapsed_time);
+}
+
+std::vector<std::string> Test::attribute_names() {
+    return {"routed_scaling_factor", "topk"};
+}
+
+std::vector<std::string> Test::tensor_names() {
+    return {"values", "indices", "x", "correction_bias", "lable_values", "lable_indices"};
+}
+
+std::vector<std::string> Test::output_names() {
+    return {"values", "indices"};
+}
+
+std::string Test::toString() const {
+    std::ostringstream oss;
+    oss << op_name() << std::endl;
+    oss << "- routed_scaling_factor=" << _attributes->routed_scaling_factor << std::endl;
+    oss << "- topk=" << _attributes->topk << std::endl;
+
+    oss << "- values: " << _attributes->values->info() << std::endl;
+    oss << "- indices: " << _attributes->indices->info() << std::endl;
+    oss << "- x: " << _attributes->x->info() << std::endl;
+    oss << "- correction_bias: " << _attributes->correction_bias->info() << std::endl;
+
+    oss << "- lable_values: " << _attributes->lable_values->info() << std::endl;
+    oss << "- lable_indices: " << _attributes->lable_indices->info() << std::endl;
+
+    oss << std::scientific << std::setprecision(2);
+    oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl;
+    return oss.str();
+}
+
+Test::~Test() {
+    delete _attributes;
+}
+
+} // namespace infiniop_test::topkrouter