Skip to content

Commit ed530e1

Browse files
author
pengcheng888
committed
issue/427 - the sigmoid, topksoftmax, and topkrouter ops
1 parent 3959c94 commit ed530e1

38 files changed

+2523
-141
lines changed

‎include/infiniop.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
#include "infiniop/ops/sub.h"
2020
#include "infiniop/ops/swiglu.h"
2121
#include "infiniop/ops/topkrouter.h"
22+
#include "infiniop/ops/topksoftmax.h"
23+
#include "infiniop/ops/sigmoid.h"
2224
#include "infiniop/tensor_descriptor.h"
2325

2426
#endif // __INFINIOP_API_H__

‎include/infiniop/ops/sigmoid.h‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
#ifndef __INFINIOP_SIGMOID_API_H__
2+
#define __INFINIOP_SIGMOID_API_H__
3+
4+
#include "../operator_descriptor.h"
5+
6+
typedef struct InfiniopDescriptor *infiniopSigmoidDescriptor_t;
7+
8+
__C __export infiniStatus_t infiniopCreateSigmoidDescriptor(infiniopHandle_t handle,
9+
infiniopSigmoidDescriptor_t *desc_ptr,
10+
infiniopTensorDescriptor_t y,
11+
infiniopTensorDescriptor_t x);
12+
13+
__C __export infiniStatus_t infiniopGetSigmoidWorkspaceSize(infiniopSigmoidDescriptor_t desc, size_t *size);
14+
15+
__C __export infiniStatus_t infiniopSigmoid(infiniopSigmoidDescriptor_t desc,
16+
void *workspace,
17+
size_t workspace_size,
18+
void *y,
19+
const void *x,
20+
void *stream);
21+
22+
__C __export infiniStatus_t infiniopDestroySigmoidDescriptor(infiniopSigmoidDescriptor_t desc);
23+
24+
#endif

‎include/infiniop/ops/topkrouter.h‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,16 +5,23 @@
55

66
typedef struct InfiniopDescriptor *infiniopTopkrouterDescriptor_t;
77

8-
__C __export infiniStatus_t infiniopCreateTopkrouterDescriptor(
9-
infiniopHandle_t handle,
10-
infiniopTopkrouterDescriptor_t *desc_ptr,
11-
infiniopTensorDescriptor_t x_desc,
12-
infiniopTensorDescriptor_t correction_bias_desc);
8+
__C __export infiniStatus_t infiniopCreateTopkrouterDescriptor(infiniopHandle_t handle,
9+
infiniopTopkrouterDescriptor_t *desc_ptr,
10+
infiniopTensorDescriptor_t x_desc,
11+
infiniopTensorDescriptor_t correction_bias_desc);
1312

1413
__C __export infiniStatus_t infiniopGetTopkrouterWorkspaceSize(infiniopTopkrouterDescriptor_t desc, size_t *size);
1514

16-
__C __export infiniStatus_t infiniopTopkrouter(infiniopTopkrouterDescriptor_t desc, void *workspace, size_t workspace_size,
17-
void *values, void *indices, void *x, void *correction_bias, float routed_scaling_factor, size_t topk, void *stream);
15+
__C __export infiniStatus_t infiniopTopkrouter(infiniopTopkrouterDescriptor_t desc,
16+
void *workspace,
17+
size_t workspace_size,
18+
void *values,
19+
void *indices,
20+
const void *x,
21+
const void *correction_bias,
22+
const float routed_scaling_factor,
23+
const size_t topk,
24+
void *stream);
1825

1926
__C __export infiniStatus_t infiniopDestroyTopkrouterDescriptor(infiniopTopkrouterDescriptor_t desc);
2027

‎include/infiniop/ops/topksoftmax.h‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
#ifndef __INFINIOP_TOPKSOFTMAX_API_H__
2+
#define __INFINIOP_TOPKSOFTMAX_API_H__
3+
4+
#include "../operator_descriptor.h"
5+
6+
typedef struct InfiniopDescriptor *infiniopTopksoftmaxDescriptor_t;
7+
8+
__C __export infiniStatus_t infiniopCreateTopksoftmaxDescriptor(infiniopHandle_t handle,
9+
infiniopTopksoftmaxDescriptor_t *desc_ptr,
10+
infiniopTensorDescriptor_t x_desc);
11+
12+
__C __export infiniStatus_t infiniopGetTopksoftmaxWorkspaceSize(infiniopTopksoftmaxDescriptor_t desc, size_t *size);
13+
14+
__C __export infiniStatus_t infiniopTopksoftmax(infiniopTopksoftmaxDescriptor_t desc,
15+
void *workspace,
16+
size_t workspace_size,
17+
void *values,
18+
void *indices,
19+
const void *x,
20+
const size_t topk,
21+
const int norm,
22+
void *stream);
23+
24+
__C __export infiniStatus_t infiniopDestroyTopksoftmaxDescriptor(infiniopTopksoftmaxDescriptor_t desc);
25+
26+
#endif

‎scripts/python_test.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,9 @@ def run_tests(args):
2525
"sub.py",
2626
"swiglu.py",
2727
"softplus.py",
28+
"sigmoid.py",
29+
"topkrouter.py",
30+
"topksoftmax.py",
2831
]:
2932
result = subprocess.run(
3033
f"python {test} {args} --debug", text=True, encoding="utf-8", shell=True

‎src/infiniop-test/include/ops.hpp‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@ DECLARE_INFINIOP_TEST(add)
1616
DECLARE_INFINIOP_TEST(causal_softmax)
1717
DECLARE_INFINIOP_TEST(rearrange)
1818
DECLARE_INFINIOP_TEST(sub)
19-
19+
DECLARE_INFINIOP_TEST(sigmoid)
20+
DECLARE_INFINIOP_TEST(topkrouter)
21+
DECLARE_INFINIOP_TEST(topksoftmax)
2022
#define REGISTER_INFINIOP_TEST(name) \
2123
{ \
2224
#name, \
@@ -43,6 +45,9 @@ DECLARE_INFINIOP_TEST(sub)
4345
REGISTER_INFINIOP_TEST(causal_softmax) \
4446
REGISTER_INFINIOP_TEST(rearrange) \
4547
REGISTER_INFINIOP_TEST(sub) \
48+
REGISTER_INFINIOP_TEST(sigmoid) \
49+
REGISTER_INFINIOP_TEST(topkrouter) \
50+
REGISTER_INFINIOP_TEST(topksoftmax) \
4651
}
4752

4853
namespace infiniop_test {

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

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
#include "ops.hpp"
2+
#include "utils.hpp"
3+
#include <infinirt.h>
4+
#include <iomanip>
5+
#include <iostream>
6+
7+
namespace infiniop_test::sigmoid {
8+
struct Test::Attributes {
9+
std::shared_ptr<Tensor> x;
10+
std::shared_ptr<Tensor> y;
11+
std::shared_ptr<Tensor> ans;
12+
};
13+
14+
std::shared_ptr<Test> Test::build(
15+
std::unordered_map<std::string, std::vector<uint8_t>> attributes,
16+
std::unordered_map<std::string, std::shared_ptr<Tensor>> tensors,
17+
double rtol, double atol) {
18+
auto test = std::shared_ptr<Test>(new Test(rtol, atol));
19+
test->_attributes = new Attributes();
20+
if (tensors.find("x") == tensors.end()
21+
|| tensors.find("y") == tensors.end()
22+
|| tensors.find("ans") == tensors.end()) {
23+
throw std::runtime_error("Invalid Test");
24+
}
25+
26+
test->_attributes->x = tensors["x"];
27+
test->_attributes->y = tensors["y"];
28+
test->_attributes->ans = tensors["ans"];
29+
30+
return test;
31+
}
32+
33+
std::shared_ptr<infiniop_test::Result> Test::run(
34+
infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) {
35+
infiniopSigmoidDescriptor_t op_desc;
36+
auto x = _attributes->x->to(device, device_id);
37+
auto y = _attributes->y->to(device, device_id);
38+
CHECK_OR(infiniopCreateSigmoidDescriptor(handle, &op_desc,
39+
y->desc(),
40+
x->desc()),
41+
return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor."));
42+
size_t workspace_size;
43+
CHECK_OR(infiniopGetSigmoidWorkspaceSize(op_desc, &workspace_size),
44+
return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size."));
45+
void *workspace;
46+
CHECK_OR(infinirtMalloc(&workspace, workspace_size),
47+
return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace."));
48+
CHECK_OR(infiniopSigmoid(op_desc, workspace, workspace_size,
49+
y->data(),
50+
x->data(),
51+
nullptr),
52+
return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution."));
53+
54+
try {
55+
allClose(y, _attributes->ans, _rtol, _atol);
56+
} catch (const std::exception &e) {
57+
return TEST_FAILED(RESULT_INCORRECT, e.what());
58+
}
59+
60+
double elapsed_time = 0.;
61+
62+
elapsed_time = benchmark(
63+
[=]() {
64+
infiniopSigmoid(
65+
op_desc, workspace, workspace_size,
66+
y->data(),
67+
x->data(),
68+
nullptr);
69+
},
70+
warm_ups, iterations);
71+
72+
infiniopDestroySigmoidDescriptor(op_desc);
73+
infinirtFree(workspace);
74+
return TEST_PASSED(elapsed_time);
75+
}
76+
77+
std::vector<std::string> Test::attribute_names() {
78+
return {};
79+
}
80+
81+
std::vector<std::string> Test::tensor_names() {
82+
return {"x", "y", "ans"};
83+
}
84+
85+
std::vector<std::string> Test::output_names() {
86+
return {"y"};
87+
}
88+
89+
std::string Test::toString() const {
90+
std::ostringstream oss;
91+
oss << op_name() << std::endl;
92+
oss << "- x: " << _attributes->x->info() << std::endl;
93+
oss << "- y: " << _attributes->y->info() << std::endl;
94+
oss << std::scientific << std::setprecision(2);
95+
oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl;
96+
return oss.str();
97+
}
98+
99+
Test::~Test() {
100+
delete _attributes;
101+
}
102+
103+
} // namespace infiniop_test::sigmoid

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

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
#include "ops.hpp"
2+
#include "utils.hpp"
3+
#include <infinirt.h>
4+
#include <iomanip>
5+
#include <iostream>
6+
7+
namespace infiniop_test::topkrouter {
8+
struct Test::Attributes {
9+
std::shared_ptr<Tensor> values;
10+
std::shared_ptr<Tensor> indices;
11+
std::shared_ptr<Tensor> x;
12+
std::shared_ptr<Tensor> correction_bias;
13+
float routed_scaling_factor;
14+
int topk;
15+
std::shared_ptr<Tensor> lable_values;
16+
std::shared_ptr<Tensor> lable_indices;
17+
};
18+
19+
std::shared_ptr<Test> Test::build(std::unordered_map<std::string, std::vector<uint8_t>> attributes,
20+
std::unordered_map<std::string, std::shared_ptr<Tensor>> tensors, double rtol,
21+
double atol) {
22+
auto test = std::shared_ptr<Test>(new Test(rtol, atol));
23+
test->_attributes = new Attributes();
24+
25+
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()) {
26+
throw std::runtime_error("Invalid Test: Missing attributes or tensors");
27+
}
28+
29+
test->_attributes->values = tensors["values"];
30+
test->_attributes->indices = tensors["indices"];
31+
test->_attributes->x = tensors["x"];
32+
test->_attributes->correction_bias = tensors["correction_bias"];
33+
34+
test->_attributes->routed_scaling_factor = *reinterpret_cast<float *>(attributes["routed_scaling_factor"].data());
35+
test->_attributes->topk = *reinterpret_cast<int *>(attributes["topk"].data());
36+
37+
test->_attributes->lable_values = tensors["lable_values"];
38+
test->_attributes->lable_indices = tensors["lable_indices"];
39+
40+
return test;
41+
}
42+
43+
std::shared_ptr<infiniop_test::Result> Test::run(infiniopHandle_t handle, infiniDevice_t device, int device_id,
44+
size_t warm_ups, size_t iterations) {
45+
infiniopTopkrouterDescriptor_t op_desc;
46+
CHECK_OR(infiniopCreateTopkrouterDescriptor(handle, &op_desc, _attributes->x->desc(),
47+
_attributes->correction_bias->desc()),
48+
return TEST_FAILED(OP_CREATION_FAILED, "Failed to create Topkrouter descriptor"));
49+
50+
//
51+
auto values = _attributes->values->to(device, device_id);
52+
auto indices = _attributes->indices->to(device, device_id);
53+
auto x = _attributes->x->to(device, device_id);
54+
auto correction_bias = _attributes->correction_bias->to(device, device_id);
55+
56+
float routed_scaling_factor = _attributes->routed_scaling_factor;
57+
int topk = _attributes->topk;
58+
59+
size_t workspace_size;
60+
CHECK_OR(infiniopGetTopkrouterWorkspaceSize(op_desc, &workspace_size),
61+
return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size"));
62+
void *workspace = nullptr;
63+
if (workspace_size > 0) {
64+
CHECK_OR(infinirtMalloc(&workspace, workspace_size),
65+
return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace"));
66+
}
67+
68+
CHECK_OR(infiniopTopkrouter(op_desc, workspace, workspace_size, values->data(), indices->data(), x->data(),
69+
correction_bias->data(), routed_scaling_factor, topk, nullptr),
70+
return TEST_FAILED(OP_EXECUTION_FAILED, "Topkrouter execution failed"));
71+
72+
try {
73+
allClose(values, _attributes->lable_values, _rtol, _atol);
74+
allClose(indices, _attributes->lable_indices, _rtol, _atol);
75+
} catch (const std::exception &e) {
76+
return TEST_FAILED(RESULT_INCORRECT, e.what());
77+
}
78+
79+
double elapsed_time = 0.;
80+
81+
elapsed_time = benchmark(
82+
[=]() {
83+
infiniopTopkrouter(op_desc, workspace, workspace_size, values->data(), indices->data(), x->data(),
84+
correction_bias->data(), routed_scaling_factor, topk, nullptr);
85+
},
86+
warm_ups, iterations);
87+
88+
if (workspace != nullptr) {
89+
infinirtFree(workspace);
90+
}
91+
92+
return TEST_PASSED(elapsed_time);
93+
}
94+
95+
std::vector<std::string> Test::attribute_names() {
96+
return {"routed_scaling_factor", "topk"};
97+
}
98+
99+
std::vector<std::string> Test::tensor_names() {
100+
return {"values", "indices", "x", "correction_bias", "lable_values", "lable_indices"};
101+
}
102+
103+
std::vector<std::string> Test::output_names() {
104+
return {"values", "indices"};
105+
}
106+
107+
std::string Test::toString() const {
108+
std::ostringstream oss;
109+
oss << op_name() << std::endl;
110+
oss << "- routed_scaling_factor=" << _attributes->routed_scaling_factor << std::endl;
111+
oss << "- topk=" << _attributes->topk << std::endl;
112+
113+
oss << "- values: " << _attributes->values->info() << std::endl;
114+
oss << "- indices: " << _attributes->indices->info() << std::endl;
115+
oss << "- x: " << _attributes->x->info() << std::endl;
116+
oss << "- correction_bias: " << _attributes->correction_bias->info() << std::endl;
117+
118+
oss << "- lable_values: " << _attributes->lable_values->info() << std::endl;
119+
oss << "- lable_indices: " << _attributes->lable_indices->info() << std::endl;
120+
121+
oss << std::scientific << std::setprecision(2);
122+
oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl;
123+
return oss.str();
124+
}
125+
126+
Test::~Test() {
127+
delete _attributes;
128+
}
129+
130+
} // namespace infiniop_test::topkrouter

0 commit comments

Comments
 (0)