Skip to content

Commit a4d938f

Browse files
Shl1015CS史昊林
Shl1015CS
authored and
史昊林
committed
homework
1 parent f3466ec commit a4d938f

22 files changed

Lines changed: 1125 additions & 16 deletions

‎src/ops/argmax/cpu/argmax_cpu.cpp‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
#include "argmax_cpu.hpp"
2+
3+
#include "../../../utils.hpp"
4+
5+
#include <limits>
6+
7+
template <typename T>
8+
void argmax_(int64_t *max_idx, T *max_val, const T *vals, size_t numel) {
9+
if (numel == 0) return;
10+
11+
size_t best_idx = 0;
12+
T best_val;
13+
14+
if constexpr (std::is_same_v<T, llaisys::bf16_t> || std::is_same_v<T, llaisys::fp16_t>) {
15+
best_val = vals[0];
16+
float best_float = llaisys::utils::cast<float>(best_val);
17+
18+
for (size_t i = 1; i < numel; i++) {
19+
float current_float = llaisys::utils::cast<float>(vals[i]);
20+
if (current_float > best_float) {
21+
best_float = current_float;
22+
best_val = vals[i];
23+
best_idx = i;
24+
}
25+
}
26+
} else {
27+
best_val = vals[0];
28+
for (size_t i = 1; i < numel; i++) {
29+
if (vals[i] > best_val) {
30+
best_val = vals[i];
31+
best_idx = i;
32+
}
33+
}
34+
}
35+
36+
*max_idx = static_cast<int64_t>(best_idx);
37+
*max_val = best_val;
38+
}
39+
40+
namespace llaisys::ops::cpu {
41+
void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, llaisysDataType_t type, size_t numel) {
42+
switch (type) {
43+
case LLAISYS_DTYPE_F32:
44+
return argmax_(
45+
reinterpret_cast<int64_t *>(max_idx),
46+
reinterpret_cast<float *>(max_val),
47+
reinterpret_cast<const float *>(vals),
48+
numel
49+
);
50+
case LLAISYS_DTYPE_BF16:
51+
return argmax_(
52+
reinterpret_cast<int64_t *>(max_idx),
53+
reinterpret_cast<llaisys::bf16_t *>(max_val),
54+
reinterpret_cast<const llaisys::bf16_t *>(vals),
55+
numel
56+
);
57+
case LLAISYS_DTYPE_F16:
58+
return argmax_(
59+
reinterpret_cast<int64_t *>(max_idx),
60+
reinterpret_cast<llaisys::fp16_t *>(max_val),
61+
reinterpret_cast<const llaisys::fp16_t *>(vals),
62+
numel
63+
);
64+
default:
65+
EXCEPTION_UNSUPPORTED_DATATYPE(type);
66+
}
67+
}
68+
} // namespace llaisys::ops::cpu

‎src/ops/argmax/cpu/argmax_cpu.hpp‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#pragma once
2+
#include "llaisys.h"
3+
4+
#include <cstddef>
5+
6+
namespace llaisys::ops::cpu {
7+
void argmax(std::byte *max_idx, std::byte *max_val, const std::byte *vals, llaisysDataType_t type, size_t numel);
8+
}

‎src/ops/argmax/op.cpp‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,45 @@
11
#include "op.hpp"
22

3+
#include "../../core/llaisys_core.hpp"
4+
#include "../../utils.hpp"
5+
6+
#include "cpu/argmax_cpu.hpp"
7+
38
namespace llaisys::ops {
49
void argmax(tensor_t max_idx, tensor_t max_val, tensor_t vals) {
5-
TO_BE_IMPLEMENTED();
10+
// Check that tensors are on same device
11+
CHECK_SAME_DEVICE(max_idx, max_val, vals);
12+
13+
// Check data types
14+
ASSERT(max_idx->dtype() == LLAISYS_DTYPE_I64, "Argmax: max_idx must be int64");
15+
CHECK_SAME_DTYPE(max_val->dtype(), vals->dtype());
16+
17+
// Check contiguity
18+
ASSERT(max_idx->isContiguous() && max_val->isContiguous() && vals->isContiguous(),
19+
"Argmax: all tensors must be contiguous");
20+
21+
// For now, assume vals is 1D and results are scalar (single element)
22+
ASSERT(vals->ndim() == 1, "Argmax: vals must be 1D tensor for now");
23+
ASSERT(max_idx->numel() == 1, "Argmax: max_idx must be scalar");
24+
ASSERT(max_val->numel() == 1, "Argmax: max_val must be scalar");
25+
26+
// always support cpu calculation
27+
if (vals->deviceType() == LLAISYS_DEVICE_CPU) {
28+
return cpu::argmax(max_idx->data(), max_val->data(), vals->data(), vals->dtype(), vals->numel());
29+
}
30+
31+
llaisys::core::context().setDevice(vals->deviceType(), vals->deviceId());
32+
33+
switch (vals->deviceType()) {
34+
case LLAISYS_DEVICE_CPU:
35+
return cpu::argmax(max_idx->data(), max_val->data(), vals->data(), vals->dtype(), vals->numel());
36+
#ifdef ENABLE_NVIDIA_API
37+
case LLAISYS_DEVICE_NVIDIA:
38+
TO_BE_IMPLEMENTED();
39+
return;
40+
#endif
41+
default:
42+
EXCEPTION_UNSUPPORTED_DEVICE;
43+
}
644
}
745
} // namespace llaisys::ops

‎src/ops/embedding/cpu/embedding_cpu.cpp‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
#include "embedding_cpu.hpp"
2+
3+
#include "../../../utils.hpp"
4+
5+
#include <cstring>
6+
7+
template <typename T>
8+
void embedding_(T *out, const int64_t *index, const T *weight, size_t idx_len, size_t embed_dim) {
9+
for (size_t i = 0; i < idx_len; i++) {
10+
int64_t idx = index[i];
11+
const T *src_row = weight + idx * embed_dim; // 源行指针
12+
T *dst_row = out + i * embed_dim; // 目标行指针
13+
14+
// 复制整行数据
15+
if constexpr (std::is_same_v<T, llaisys::bf16_t> || std::is_same_v<T, llaisys::fp16_t>) {
16+
// 对于半精度类型，逐元素复制
17+
for (size_t j = 0; j < embed_dim; j++) {
18+
dst_row[j] = src_row[j];
19+
}
20+
} else {
21+
// 对于标准类型，使用内存复制
22+
std::memcpy(dst_row, src_row, embed_dim * sizeof(T));
23+
}
24+
}
25+
}
26+
27+
namespace llaisys::ops::cpu {
28+
void embedding(std::byte *out, const std::byte *index, const std::byte *weight,
29+
llaisysDataType_t type, size_t idx_len, size_t embed_dim) {
30+
switch (type) {
31+
case LLAISYS_DTYPE_F32:
32+
return embedding_(
33+
reinterpret_cast<float *>(out),
34+
reinterpret_cast<const int64_t *>(index),
35+
reinterpret_cast<const float *>(weight),
36+
idx_len, embed_dim
37+
);
38+
case LLAISYS_DTYPE_BF16:
39+
return embedding_(
40+
reinterpret_cast<llaisys::bf16_t *>(out),
41+
reinterpret_cast<const int64_t *>(index),
42+
reinterpret_cast<const llaisys::bf16_t *>(weight),
43+
idx_len, embed_dim
44+
);
45+
case LLAISYS_DTYPE_F16:
46+
return embedding_(
47+
reinterpret_cast<llaisys::fp16_t *>(out),
48+
reinterpret_cast<const int64_t *>(index),
49+
reinterpret_cast<const llaisys::fp16_t *>(weight),
50+
idx_len, embed_dim
51+
);
52+
default:
53+
EXCEPTION_UNSUPPORTED_DATATYPE(type);
54+
}
55+
}
56+
} // namespace llaisys::ops::cpu

‎src/ops/embedding/cpu/embedding_cpu.hpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
#pragma once
2+
#include "llaisys.h"
3+
4+
#include <cstddef>
5+
6+
namespace llaisys::ops::cpu {
7+
void embedding(std::byte *out, const std::byte *index, const std::byte *weight,
8+
llaisysDataType_t type, size_t idx_len, size_t embed_dim);
9+
}

‎src/ops/embedding/op.cpp‎

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,53 @@
11
#include "op.hpp"
22

3+
#include "../../core/llaisys_core.hpp"
4+
#include "../../utils.hpp"
5+
6+
#include "cpu/embedding_cpu.hpp"
7+
38
namespace llaisys::ops {
49
void embedding(tensor_t out, tensor_t index, tensor_t weight) {
5-
TO_BE_IMPLEMENTED();
10+
// Check that tensors are on same device
11+
CHECK_SAME_DEVICE(out, index, weight);
12+
13+
// Check data types
14+
ASSERT(index->dtype() == LLAISYS_DTYPE_I64, "Embedding: index must be int64");
15+
CHECK_SAME_DTYPE(out->dtype(), weight->dtype());
16+
17+
// Check contiguity
18+
ASSERT(out->isContiguous() && index->isContiguous() && weight->isContiguous(),
19+
"Embedding: all tensors must be contiguous");
20+
21+
// Check dimensions
22+
ASSERT(index->ndim() == 1, "Embedding: index must be 1D tensor");
23+
ASSERT(weight->ndim() == 2, "Embedding: weight must be 2D tensor");
24+
ASSERT(out->ndim() == 2, "Embedding: out must be 2D tensor");
25+
26+
// Check shapes are compatible
27+
size_t idx_len = index->shape()[0];
28+
size_t embed_dim = weight->shape()[1];
29+
ASSERT(out->shape()[0] == idx_len, "Embedding: output length must match index length");
30+
ASSERT(out->shape()[1] == embed_dim, "Embedding: output embedding dimension must match weight");
31+
32+
// always support cpu calculation
33+
if (out->deviceType() == LLAISYS_DEVICE_CPU) {
34+
return cpu::embedding(out->data(), index->data(), weight->data(),
35+
out->dtype(), idx_len, embed_dim);
36+
}
37+
38+
llaisys::core::context().setDevice(out->deviceType(), out->deviceId());
39+
40+
switch (out->deviceType()) {
41+
case LLAISYS_DEVICE_CPU:
42+
return cpu::embedding(out->data(), index->data(), weight->data(),
43+
out->dtype(), idx_len, embed_dim);
44+
#ifdef ENABLE_NVIDIA_API
45+
case LLAISYS_DEVICE_NVIDIA:
46+
TO_BE_IMPLEMENTED();
47+
return;
48+
#endif
49+
default:
50+
EXCEPTION_UNSUPPORTED_DEVICE;
51+
}
652
}
753
} // namespace llaisys::ops

‎src/ops/linear/cpu/linear_cpu.cpp‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
#include "linear_cpu.hpp"
2+
3+
#include "../../../utils.hpp"
4+
5+
// 矩阵乘法: Y = X * W^T + bias
6+
// X: (batch_size, in_features)
7+
// W: (out_features, in_features)
8+
// Y: (batch_size, out_features)
9+
// bias: (out_features) 可选
10+
template <typename T>
11+
void linear_(T *out, const T *in, const T *weight, const T *bias,
12+
size_t batch_size, size_t in_features, size_t out_features) {
13+
14+
// 对每个批次的每个输出特征计算
15+
for (size_t b = 0; b < batch_size; b++) {
16+
for (size_t o = 0; o < out_features; o++) {
17+
if constexpr (std::is_same_v<T, llaisys::bf16_t> || std::is_same_v<T, llaisys::fp16_t>) {
18+
// 对半精度类型，使用float进行累加以避免精度损失
19+
float sum_float = 0.0f;
20+
21+
// 计算点积: X[b,:] · W[o,:]
22+
for (size_t i = 0; i < in_features; i++) {
23+
float x_float = llaisys::utils::cast<float>(in[b * in_features + i]);
24+
float w_float = llaisys::utils::cast<float>(weight[o * in_features + i]);
25+
sum_float += x_float * w_float;
26+
}
27+
28+
// 添加偏置（如果有）
29+
if (bias != nullptr) {
30+
sum_float += llaisys::utils::cast<float>(bias[o]);
31+
}
32+
33+
out[b * out_features + o] = llaisys::utils::cast<T>(sum_float);
34+
} else {
35+
// 对全精度类型，直接计算
36+
T sum = T(0);
37+
38+
// 计算点积: X[b,:] · W[o,:]
39+
for (size_t i = 0; i < in_features; i++) {
40+
sum += in[b * in_features + i] * weight[o * in_features + i];
41+
}
42+
43+
// 添加偏置（如果有）
44+
if (bias != nullptr) {
45+
sum += bias[o];
46+
}
47+
48+
out[b * out_features + o] = sum;
49+
}
50+
}
51+
}
52+
}
53+
54+
namespace llaisys::ops::cpu {
55+
void linear(std::byte *out, const std::byte *in, const std::byte *weight, const std::byte *bias,
56+
llaisysDataType_t type, size_t batch_size, size_t in_features, size_t out_features) {
57+
switch (type) {
58+
case LLAISYS_DTYPE_F32:
59+
return linear_(
60+
reinterpret_cast<float *>(out),
61+
reinterpret_cast<const float *>(in),
62+
reinterpret_cast<const float *>(weight),
63+
bias ? reinterpret_cast<const float *>(bias) : nullptr,
64+
batch_size, in_features, out_features
65+
);
66+
case LLAISYS_DTYPE_BF16:
67+
return linear_(
68+
reinterpret_cast<llaisys::bf16_t *>(out),
69+
reinterpret_cast<const llaisys::bf16_t *>(in),
70+
reinterpret_cast<const llaisys::bf16_t *>(weight),
71+
bias ? reinterpret_cast<const llaisys::bf16_t *>(bias) : nullptr,
72+
batch_size, in_features, out_features
73+
);
74+
case LLAISYS_DTYPE_F16:
75+
return linear_(
76+
reinterpret_cast<llaisys::fp16_t *>(out),
77+
reinterpret_cast<const llaisys::fp16_t *>(in),
78+
reinterpret_cast<const llaisys::fp16_t *>(weight),
79+
bias ? reinterpret_cast<const llaisys::fp16_t *>(bias) : nullptr,
80+
batch_size, in_features, out_features
81+
);
82+
default:
83+
EXCEPTION_UNSUPPORTED_DATATYPE(type);
84+
}
85+
}
86+
} // namespace llaisys::ops::cpu

‎src/ops/linear/cpu/linear_cpu.hpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
#pragma once
2+
#include "llaisys.h"
3+
4+
#include <cstddef>
5+
6+
namespace llaisys::ops::cpu {
7+
void linear(std::byte *out, const std::byte *in, const std::byte *weight, const std::byte *bias,
8+
llaisysDataType_t type, size_t batch_size, size_t in_features, size_t out_features);
9+
}

0 commit comments

Comments
 (0)