Skip to content

Commit e9567ce

Browse files
kilinchangekilinchange
committed
test: migrate test_transformer_architecture to ctest framework
1 parent dbdf569 commit e9567ce

11 files changed

Lines changed: 222 additions & 690 deletions

File tree

infini_train/include/autograd/indexing.h renamed to infini_train/include/autograd/gather.h

Lines changed: 0 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -27,21 +27,4 @@ class Gather : public Function {
2727
std::vector<int64_t> input_dims_;
2828
};
2929

30-
class Slice : public Function {
31-
public:
32-
static constexpr char kType[] = "SliceFunction";
33-
34-
Slice(const std::vector<int64_t> &starts, const std::vector<int64_t> &ends, const std::vector<int64_t> &steps)
35-
: Function(kType), starts_(starts), ends_(ends), steps_(steps) {}
36-
std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
37-
void SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
38-
const std::vector<std::shared_ptr<Tensor>> &output_tensors) override;
39-
std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
40-
41-
private:
42-
const std::vector<int64_t> starts_;
43-
const std::vector<int64_t> ends_;
44-
const std::vector<int64_t> steps_;
45-
};
46-
4730
} // namespace infini_train::autograd

‎infini_train/include/autograd/transform.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,4 +127,21 @@ class Concat : public Function {
127127
std::vector<std::vector<int64_t>> input_dims_list_;
128128
};
129129

130+
class Slice : public Function {
131+
public:
132+
static constexpr char kType[] = "SliceFunction";
133+
134+
Slice(const std::vector<int64_t> &starts, const std::vector<int64_t> &ends, const std::vector<int64_t> &steps)
135+
: Function(kType), starts_(starts), ends_(ends), steps_(steps) {}
136+
std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
137+
void SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
138+
const std::vector<std::shared_ptr<Tensor>> &output_tensors) override;
139+
std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
140+
141+
private:
142+
const std::vector<int64_t> starts_;
143+
const std::vector<int64_t> ends_;
144+
const std::vector<int64_t> steps_;
145+
};
146+
130147
} // namespace infini_train::autograd

‎infini_train/include/datatype.h‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -100,18 +100,18 @@ enum class DataType : int8_t {
100100
};
101101

102102
inline const std::unordered_map<DataType, size_t> kDataTypeToSize = {
103-
{DataType::kBOOL, 1},
104-
{DataType::kUINT8, 1}, {DataType::kINT8, 1}, {DataType::kUINT16, 2}, {DataType::kINT16, 2},
105-
{DataType::kUINT32, 4}, {DataType::kINT32, 4}, {DataType::kUINT64, 8}, {DataType::kINT64, 8},
106-
{DataType::kBFLOAT16, 2}, {DataType::kFLOAT16, 2}, {DataType::kFLOAT32, 4}, {DataType::kFLOAT64, 8},
103+
{DataType::kBOOL, 1}, {DataType::kUINT8, 1}, {DataType::kINT8, 1}, {DataType::kUINT16, 2},
104+
{DataType::kINT16, 2}, {DataType::kUINT32, 4}, {DataType::kINT32, 4}, {DataType::kUINT64, 8},
105+
{DataType::kINT64, 8}, {DataType::kBFLOAT16, 2}, {DataType::kFLOAT16, 2}, {DataType::kFLOAT32, 4},
106+
{DataType::kFLOAT64, 8},
107107
};
108108

109109
inline const std::unordered_map<DataType, std::string> kDataTypeToDesc = {
110-
{DataType::kBOOL, "bool"},
111-
{DataType::kUINT8, "uint8"}, {DataType::kINT8, "int8"}, {DataType::kUINT16, "uint16"},
112-
{DataType::kINT16, "int16"}, {DataType::kUINT32, "uint32"}, {DataType::kINT32, "int32"},
113-
{DataType::kUINT64, "uint64"}, {DataType::kINT64, "int64"}, {DataType::kBFLOAT16, "bf16"},
114-
{DataType::kFLOAT16, "fp16"}, {DataType::kFLOAT32, "fp32"}, {DataType::kFLOAT64, "fp64"},
110+
{DataType::kBOOL, "bool"}, {DataType::kUINT8, "uint8"}, {DataType::kINT8, "int8"},
111+
{DataType::kUINT16, "uint16"}, {DataType::kINT16, "int16"}, {DataType::kUINT32, "uint32"},
112+
{DataType::kINT32, "int32"}, {DataType::kUINT64, "uint64"}, {DataType::kINT64, "int64"},
113+
{DataType::kBFLOAT16, "bf16"}, {DataType::kFLOAT16, "fp16"}, {DataType::kFLOAT32, "fp32"},
114+
{DataType::kFLOAT64, "fp64"},
115115
};
116116

117117
// =============================================================================

infini_train/src/autograd/indexing.cc renamed to infini_train/src/autograd/gather.cc

Lines changed: 1 addition & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
#include "infini_train/include/autograd/indexing.h"
1+
#include "infini_train/include/autograd/gather.h"
22

33
#include "glog/logging.h"
44

@@ -34,30 +34,4 @@ std::vector<std::shared_ptr<Tensor>> Gather::Backward(const std::vector<std::sha
3434
return {kernel.Call<std::shared_ptr<Tensor>>(grad_output, index, dim_, input_dims_), nullptr};
3535
}
3636

37-
std::vector<std::shared_ptr<Tensor>> Slice::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
38-
CHECK_EQ(input_tensors.size(), 1);
39-
const auto &input = input_tensors[0];
40-
41-
auto device = input->GetDevice().type();
42-
return {
43-
Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "SliceForward"}, input, starts_, ends_, steps_)};
44-
}
45-
46-
void Slice::SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
47-
const std::vector<std::shared_ptr<Tensor>> &) {
48-
// FIXME(dcj): only input's dim need to be saved
49-
const auto &input = input_tensors[0];
50-
saved_tensors_ = {input};
51-
}
52-
53-
std::vector<std::shared_ptr<Tensor>> Slice::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
54-
CHECK_EQ(saved_tensors_.size(), 1);
55-
const auto &input = saved_tensors_[0];
56-
const auto &grad_output = grad_outputs[0];
57-
58-
auto device = input->GetDevice().type();
59-
return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "SliceBackward"}, grad_output, input, starts_,
60-
ends_, steps_)};
61-
}
62-
6337
} // namespace infini_train::autograd

‎infini_train/src/autograd/transform.cc‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,4 +153,30 @@ std::vector<std::shared_ptr<Tensor>> Concat::Backward(const std::vector<std::sha
153153
return kernel.Call<std::vector<std::shared_ptr<Tensor>>>(grad_output, input_dims_list_, dim_);
154154
}
155155

156+
std::vector<std::shared_ptr<Tensor>> Slice::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
157+
CHECK_EQ(input_tensors.size(), 1);
158+
const auto &input = input_tensors[0];
159+
160+
auto device = input->GetDevice().type();
161+
return {
162+
Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "SliceForward"}, input, starts_, ends_, steps_)};
163+
}
164+
165+
void Slice::SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
166+
const std::vector<std::shared_ptr<Tensor>> &) {
167+
// FIXME(dcj): only input's dim need to be saved
168+
const auto &input = input_tensors[0];
169+
saved_tensors_ = {input};
170+
}
171+
172+
std::vector<std::shared_ptr<Tensor>> Slice::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
173+
CHECK_EQ(saved_tensors_.size(), 1);
174+
const auto &input = saved_tensors_[0];
175+
const auto &grad_output = grad_outputs[0];
176+
177+
auto device = input->GetDevice().type();
178+
return {Dispatcher::Instance().Call<std::shared_ptr<Tensor>>({device, "SliceBackward"}, grad_output, input, starts_,
179+
ends_, steps_)};
180+
}
181+
156182
} // namespace infini_train::autograd

‎infini_train/src/kernels/cpu/gather.cc‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99

1010
namespace infini_train::kernels::cpu {
1111
std::shared_ptr<Tensor> GatherForward(const std::shared_ptr<Tensor> &input, const std::shared_ptr<Tensor> &index,
12-
int64_t dim) {
12+
int64_t dim) {
1313
const auto &in_dims = input->Dims();
1414
const auto &idx_dims = index->Dims();
1515
CHECK_EQ(in_dims.size(), idx_dims.size());
@@ -100,9 +100,8 @@ std::shared_ptr<Tensor> GatherForward(const std::shared_ptr<Tensor> &input, cons
100100
return out;
101101
}
102102

103-
std::shared_ptr<Tensor> GatherBackward(const std::shared_ptr<Tensor> &grad_output,
104-
const std::shared_ptr<Tensor> &index, int64_t dim,
105-
const std::vector<int64_t> &input_dims) {
103+
std::shared_ptr<Tensor> GatherBackward(const std::shared_ptr<Tensor> &grad_output, const std::shared_ptr<Tensor> &index,
104+
int64_t dim, const std::vector<int64_t> &input_dims) {
106105
const auto &in_dims = input_dims;
107106
const auto &idx_dims = index->Dims();
108107
CHECK_EQ(in_dims.size(), idx_dims.size());

‎infini_train/src/kernels/cpu/scatter.cc‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,7 @@
1010

1111
namespace infini_train::kernels::cpu {
1212

13-
std::shared_ptr<Tensor> ScatterForward(const std::shared_ptr<Tensor> &values,
14-
const std::shared_ptr<Tensor> &indices,
13+
std::shared_ptr<Tensor> ScatterForward(const std::shared_ptr<Tensor> &values, const std::shared_ptr<Tensor> &indices,
1514
const std::vector<int64_t> &output_dims) {
1615
CHECK(indices->Dtype() == DataType::kINT64) << "CPU ScatterForward expects int64 indices";
1716
CHECK(values->Dims() == indices->Dims());
@@ -39,8 +38,8 @@ std::shared_ptr<Tensor> ScatterForward(const std::shared_ptr<Tensor> &values,
3938
const int64_t expert_idx = idx[row * topk + selected];
4039
CHECK_GE(expert_idx, 0);
4140
CHECK_LT(expert_idx, num_experts);
42-
std::memcpy(dst + (row * num_experts + expert_idx) * elem_size,
43-
src + (row * topk + selected) * elem_size, elem_size);
41+
std::memcpy(dst + (row * num_experts + expert_idx) * elem_size, src + (row * topk + selected) * elem_size,
42+
elem_size);
4443
}
4544
}
4645

@@ -68,8 +67,8 @@ std::shared_ptr<Tensor> ScatterBackward(const std::shared_ptr<Tensor> &grad_outp
6867
const int64_t expert_idx = idx[row * topk + selected];
6968
CHECK_GE(expert_idx, 0);
7069
CHECK_LT(expert_idx, num_experts);
71-
std::memcpy(dst + (row * topk + selected) * elem_size,
72-
src + (row * num_experts + expert_idx) * elem_size, elem_size);
70+
std::memcpy(dst + (row * topk + selected) * elem_size, src + (row * num_experts + expert_idx) * elem_size,
71+
elem_size);
7372
}
7473
}
7574

@@ -78,7 +77,7 @@ std::shared_ptr<Tensor> ScatterBackward(const std::shared_ptr<Tensor> &grad_outp
7877

7978
} // namespace infini_train::kernels::cpu
8079

81-
#define REGISTER_CPU_SCATTER_KERNEL(kernel_name) \
80+
#define REGISTER_CPU_SCATTER_KERNEL(kernel_name) \
8281
REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name)
8382

8483
REGISTER_CPU_SCATTER_KERNEL(ScatterForward)

‎infini_train/src/nn/functional.cc‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66

77
#include "infini_train/include/autograd/activations.h"
88
#include "infini_train/include/autograd/elementwise.h"
9-
#include "infini_train/include/autograd/misc.h"
109
#include "infini_train/include/autograd/reduction.h"
1110
#include "infini_train/include/autograd/softmax.h"
1211
#include "infini_train/include/autograd/transform.h"

‎infini_train/src/tensor.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@
1313
#include "infini_train/include/autograd/elementwise.h"
1414
#include "infini_train/include/autograd/function.h"
1515
#include "infini_train/include/autograd/function_hook.h"
16+
#include "infini_train/include/autograd/gather.h"
1617
#include "infini_train/include/autograd/matmul.h"
17-
#include "infini_train/include/autograd/indexing.h"
1818
#include "infini_train/include/autograd/no_op.h"
1919
#include "infini_train/include/autograd/outer.h"
2020
#include "infini_train/include/autograd/reduction.h"

0 commit comments

Comments
 (0)