Skip to content

Commit 84ee85f

Browse files
ZiminliZiminli
authored
Merge pull request InfiniTensor#1340 from InfiniTensor/issue/1320/infinirt-ops644
issue/1320 - support InfiniOps standalone InfiniRT across platforms
2 parents 81afd0b + 7098b3c commit 84ee85f

27 files changed

Lines changed: 177 additions & 82 deletions

‎src/infinicore/ops/add/add_infiniops.cc‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ struct PlannedMeta {
1616
} // namespace
1717

1818
void *plan(Tensor c, const Tensor &a, const Tensor &b) {
19-
INFINICORE_ASSERT(c->device().getType() == Device::Type::NVIDIA);
19+
INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(c->device().getType()));
2020
INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b);
2121

2222
return new PlannedMeta{
@@ -49,9 +49,9 @@ void cleanup(void **planned_meta_ptr) {
4949
}
5050

5151
static bool registered = []() {
52-
Add::plan_dispatcher().registerDevice(Device::Type::NVIDIA, &plan);
53-
Add::run_dispatcher().registerDevice(Device::Type::NVIDIA, &run);
54-
Add::cleanup_dispatcher().registerDevice(Device::Type::NVIDIA, &cleanup);
52+
::infinicore::op::infiniops::registerSupportedDevices(Add::plan_dispatcher(), &plan);
53+
::infinicore::op::infiniops::registerSupportedDevices(Add::run_dispatcher(), &run);
54+
::infinicore::op::infiniops::registerSupportedDevices(Add::cleanup_dispatcher(), &cleanup);
5555
return true;
5656
}();
5757

‎src/infinicore/ops/add_rms_norm/add_rms_norm_infiniops.cc‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ struct PlannedMeta {
2121
} // namespace
2222

2323
void *plan(Tensor out, Tensor residual, const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon) {
24-
INFINICORE_ASSERT(out->device().getType() == Device::Type::NVIDIA);
24+
INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(out->device().getType()));
2525
INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, residual, a, b, weight);
2626

2727
return new PlannedMeta{
@@ -62,9 +62,9 @@ void cleanup(void **planned_meta_ptr) {
6262
}
6363

6464
static bool registered = []() {
65-
AddRMSNorm::plan_dispatcher().registerDevice(Device::Type::NVIDIA, &plan);
66-
AddRMSNorm::run_dispatcher().registerDevice(Device::Type::NVIDIA, &run);
67-
AddRMSNorm::cleanup_dispatcher().registerDevice(Device::Type::NVIDIA, &cleanup);
65+
::infinicore::op::infiniops::registerSupportedDevices(AddRMSNorm::plan_dispatcher(), &plan);
66+
::infinicore::op::infiniops::registerSupportedDevices(AddRMSNorm::run_dispatcher(), &run);
67+
::infinicore::op::infiniops::registerSupportedDevices(AddRMSNorm::cleanup_dispatcher(), &cleanup);
6868
return true;
6969
}();
7070

‎src/infinicore/ops/causal_softmax/causal_softmax_infiniops.cc‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ struct PlannedMeta {
1818
} // namespace
1919

2020
void *plan(Tensor output, const Tensor &input) {
21-
INFINICORE_ASSERT(output->device().getType() == Device::Type::NVIDIA);
21+
INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType()));
2222
INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input);
2323

2424
return new PlannedMeta{
@@ -48,9 +48,9 @@ void cleanup(void **planned_meta_ptr) {
4848
}
4949

5050
static bool registered = []() {
51-
CausalSoftmax::plan_dispatcher().registerDevice(Device::Type::NVIDIA, &plan);
52-
CausalSoftmax::run_dispatcher().registerDevice(Device::Type::NVIDIA, &run);
53-
CausalSoftmax::cleanup_dispatcher().registerDevice(Device::Type::NVIDIA, &cleanup);
51+
::infinicore::op::infiniops::registerSupportedDevices(CausalSoftmax::plan_dispatcher(), &plan);
52+
::infinicore::op::infiniops::registerSupportedDevices(CausalSoftmax::run_dispatcher(), &run);
53+
::infinicore::op::infiniops::registerSupportedDevices(CausalSoftmax::cleanup_dispatcher(), &cleanup);
5454
return true;
5555
}();
5656

‎src/infinicore/ops/conv2d/conv2d_infiniops.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ void calculate(Tensor output,
3030
const size_t *strides,
3131
const size_t *dilations,
3232
size_t n) {
33-
INFINICORE_ASSERT(output->device().getType() == Device::Type::NVIDIA);
33+
INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType()));
3434
INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, weight);
3535
if (bias) {
3636
INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, bias);
@@ -64,7 +64,7 @@ void calculate(Tensor output,
6464
} // namespace
6565

6666
static bool registered = []() {
67-
Conv2d::dispatcher().registerDevice(Device::Type::NVIDIA, &calculate);
67+
::infinicore::op::infiniops::registerSupportedDevices(Conv2d::dispatcher(), &calculate);
6868
return true;
6969
}();
7070

‎src/infinicore/ops/embedding/embedding_infiniops.cc‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ struct PlannedMeta {
1818
} // namespace
1919

2020
void *plan(Tensor out, const Tensor &input, const Tensor &weight) {
21-
INFINICORE_ASSERT(out->device().getType() == Device::Type::NVIDIA);
21+
INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(out->device().getType()));
2222
INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight);
2323

2424
return new PlannedMeta{
@@ -51,9 +51,9 @@ void cleanup(void **planned_meta_ptr) {
5151
}
5252

5353
static bool registered = []() {
54-
Embedding::plan_dispatcher().registerDevice(Device::Type::NVIDIA, &plan);
55-
Embedding::run_dispatcher().registerDevice(Device::Type::NVIDIA, &run);
56-
Embedding::cleanup_dispatcher().registerDevice(Device::Type::NVIDIA, &cleanup);
54+
::infinicore::op::infiniops::registerSupportedDevices(Embedding::plan_dispatcher(), &plan);
55+
::infinicore::op::infiniops::registerSupportedDevices(Embedding::run_dispatcher(), &run);
56+
::infinicore::op::infiniops::registerSupportedDevices(Embedding::cleanup_dispatcher(), &cleanup);
5757
return true;
5858
}();
5959

‎src/infinicore/ops/gelu/gelu_infiniops.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ namespace {
1313
using TensorMeta = ::infinicore::op::infiniops::TensorMeta;
1414

1515
void calculate(Tensor output, Tensor input) {
16-
INFINICORE_ASSERT(output->device().getType() == Device::Type::NVIDIA);
16+
INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType()));
1717
INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input);
1818

1919
infini::ops::Handle handle;
@@ -33,7 +33,7 @@ void calculate(Tensor output, Tensor input) {
3333
} // namespace
3434

3535
static bool registered = []() {
36-
Gelu::dispatcher().registerDevice(Device::Type::NVIDIA, &calculate);
36+
::infinicore::op::infiniops::registerSupportedDevices(Gelu::dispatcher(), &calculate);
3737
return true;
3838
}();
3939

‎src/infinicore/ops/gelutanh/gelutanh_infiniops.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ namespace {
1111
using TensorMeta = ::infinicore::op::infiniops::TensorMeta;
1212

1313
void calculate(Tensor output, Tensor input) {
14-
INFINICORE_ASSERT(output->device().getType() == Device::Type::NVIDIA);
14+
INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType()));
1515
INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input);
1616

1717
infini::ops::Handle handle;
@@ -30,7 +30,7 @@ void calculate(Tensor output, Tensor input) {
3030
} // namespace
3131

3232
static bool registered = []() {
33-
GeluTanh::dispatcher().registerDevice(Device::Type::NVIDIA, &calculate);
33+
::infinicore::op::infiniops::registerSupportedDevices(GeluTanh::dispatcher(), &calculate);
3434
return true;
3535
}();
3636

‎src/infinicore/ops/gemm/gemm_infiniops.cc‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ struct PlannedMeta {
1919
} // namespace
2020

2121
void *plan(Tensor c, const Tensor &a, const Tensor &b, float alpha, float beta) {
22-
INFINICORE_ASSERT(c->device().getType() == Device::Type::NVIDIA);
22+
INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(c->device().getType()));
2323
INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b);
2424

2525
return new PlannedMeta{
@@ -58,9 +58,9 @@ void cleanup(void **planned_meta_ptr) {
5858
}
5959

6060
static bool registered = []() {
61-
Gemm::plan_dispatcher().registerDevice(Device::Type::NVIDIA, &plan);
62-
Gemm::run_dispatcher().registerDevice(Device::Type::NVIDIA, &run);
63-
Gemm::cleanup_dispatcher().registerDevice(Device::Type::NVIDIA, &cleanup);
61+
::infinicore::op::infiniops::registerSupportedDevices(Gemm::plan_dispatcher(), &plan);
62+
::infinicore::op::infiniops::registerSupportedDevices(Gemm::run_dispatcher(), &run);
63+
::infinicore::op::infiniops::registerSupportedDevices(Gemm::cleanup_dispatcher(), &cleanup);
6464
return true;
6565
}();
6666

‎src/infinicore/ops/infiniops_impl.hpp‎

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,38 @@ inline infini::ops::DataType toInfiniOpsDtype(DataType dtype) {
4646
}
4747

4848
inline infini::ops::Device toInfiniOpsDevice(const Device &device) {
49-
INFINICORE_ASSERT(device.getType() == Device::Type::NVIDIA);
50-
return infini::ops::Device{infini::ops::Device::Type::kNvidia, static_cast<int>(device.getIndex())};
49+
switch (device.getType()) {
50+
case Device::Type::NVIDIA:
51+
return infini::ops::Device{infini::ops::Device::Type::kNvidia, static_cast<int>(device.getIndex())};
52+
case Device::Type::METAX:
53+
return infini::ops::Device{infini::ops::Device::Type::kMetax, static_cast<int>(device.getIndex())};
54+
case Device::Type::MOORE:
55+
return infini::ops::Device{infini::ops::Device::Type::kMoore, static_cast<int>(device.getIndex())};
56+
case Device::Type::ILUVATAR:
57+
return infini::ops::Device{infini::ops::Device::Type::kIluvatar, static_cast<int>(device.getIndex())};
58+
default:
59+
throw std::runtime_error("InfiniOps backend does not support this device type.");
60+
}
61+
}
62+
63+
inline bool isSupportedDevice(Device::Type device_type) {
64+
switch (device_type) {
65+
case Device::Type::NVIDIA:
66+
case Device::Type::METAX:
67+
case Device::Type::MOORE:
68+
case Device::Type::ILUVATAR:
69+
return true;
70+
default:
71+
return false;
72+
}
73+
}
74+
75+
template <typename Dispatcher, typename Function>
76+
void registerSupportedDevices(Dispatcher &dispatcher, Function function) {
77+
dispatcher.registerDevice(Device::Type::NVIDIA, function);
78+
dispatcher.registerDevice(Device::Type::METAX, function);
79+
dispatcher.registerDevice(Device::Type::MOORE, function);
80+
dispatcher.registerDevice(Device::Type::ILUVATAR, function);
5181
}
5282

5383
struct TensorMeta {

‎src/infinicore/ops/kv_caching/kv_caching_infiniops.cc‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ struct PlannedMeta {
1515
} // namespace
1616

1717
void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &past_kv_lengths) {
18-
INFINICORE_ASSERT(k_cache->device().getType() == Device::Type::NVIDIA);
18+
INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(k_cache->device().getType()));
1919
INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, past_kv_lengths);
2020
return new PlannedMeta{TensorMeta(k_cache), TensorMeta(v_cache), TensorMeta(k), TensorMeta(v), TensorMeta(past_kv_lengths), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(past_kv_lengths)};
2121
}
@@ -41,9 +41,9 @@ void cleanup(void **planned_meta_ptr) {
4141
}
4242

4343
static bool registered = []() {
44-
KVCaching::plan_dispatcher().registerDevice(Device::Type::NVIDIA, &plan);
45-
KVCaching::run_dispatcher().registerDevice(Device::Type::NVIDIA, &run);
46-
KVCaching::cleanup_dispatcher().registerDevice(Device::Type::NVIDIA, &cleanup);
44+
::infinicore::op::infiniops::registerSupportedDevices(KVCaching::plan_dispatcher(), &plan);
45+
::infinicore::op::infiniops::registerSupportedDevices(KVCaching::run_dispatcher(), &run);
46+
::infinicore::op::infiniops::registerSupportedDevices(KVCaching::cleanup_dispatcher(), &cleanup);
4747
return true;
4848
}();
4949
} // namespace infinicore::op::kv_caching_impl::infiniops

0 commit comments

Comments
 (0)