diff --git a/README.md b/README.md index fc57b11..efd0f9b 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,687 @@ -# TinyInfiniTrain - -简化版的训练框架，全部由C++实现，包含了autograd机制、kernel层关键算子实现，测例中包含单点验证以及端到端训练GPT2模型流程，支持CPU和CUDA平台。 - -[环境部署文档](docs/项目部署.md) - -[训练营作业介绍文档](docs/训练营作业介绍.md) - -[训练营作业报告模版](docs/TinyInfiniTrain%20作业报告.md) \ No newline at end of file +# TinyInfiniTrain 作业报告 + +## 一、测试通过截图 + +按测试顺序排列，全部 8 个测试通过： + +### 1. test_elementwise（5分）✅ + +验证 autograd 机制调用 Neg kernel 的实现，依赖作业一、作业五。 + +``` +[==========] Running 2 tests from 1 test suite. +[----------] 2 tests from ElementwiseTest +[ RUN ] ElementwiseTest.NegForward +[ OK ] ElementwiseTest.NegForward (0 ms) +[ RUN ] ElementwiseTest.NegBackward +[ OK ] ElementwiseTest.NegBackward (0 ms) +[----------] 2 tests from ElementwiseTest (0 ms total) +[ PASSED ] 2 tests. +``` + +### 2. test_matmul（5分）✅ + +验证 Matmul kernel 的 CPU 实现，依赖作业二。 + +``` +[==========] Running 3 tests from 1 test suite. +[----------] 3 tests from MatmulTest +[ RUN ] MatmulTest.BasicMatrixMultiply +[ OK ] MatmulTest.BasicMatrixMultiply (0 ms) +[ RUN ] MatmulTest.BatchedMatrixMultiply +[ OK ] MatmulTest.BatchedMatrixMultiply (0 ms) +[ RUN ] MatmulTest.BackwardPass +[ OK ] MatmulTest.BackwardPass (0 ms) +[----------] 3 tests from MatmulTest (0 ms total) +[ PASSED ] 3 tests. +``` + +### 3. test_matmul_cuda（10分）✅ + +验证 Matmul kernel 的 CUDA 实现，依赖作业二。 + +``` +[==========] Running 3 tests from 1 test suite. +[----------] 3 tests from MatmulTest +[ RUN ] MatmulTest.BasicMatrixMultiplyCuda +[ OK ] MatmulTest.BasicMatrixMultiplyCuda (277 ms) +[ RUN ] MatmulTest.BatchedMatrixMultiplyCuda +[ OK ] MatmulTest.BatchedMatrixMultiplyCuda (2 ms) +[ RUN ] MatmulTest.BackwardPassCuda +[ OK ] MatmulTest.BackwardPassCuda (2 ms) +[----------] 3 tests from MatmulTest (282 ms total) +[ PASSED ] 3 tests. +``` + +### 4. test_adam（5分）✅ + +验证 Adam 优化器的 CPU 实现，依赖作业三。 + +``` +[==========] Running 2 tests from 1 test suite. +[----------] 2 tests from AdamOptimizerTest +[ RUN ] AdamOptimizerTest.BasicParameterUpdate +[ OK ] AdamOptimizerTest.BasicParameterUpdate (0 ms) +[ RUN ] AdamOptimizerTest.MomentumAccumulation +[ OK ] AdamOptimizerTest.MomentumAccumulation (0 ms) +[----------] 2 tests from AdamOptimizerTest (0 ms total) +[ PASSED ] 2 tests. +``` + +### 5. test_adam_cuda（10分）✅ + +验证 Adam 优化器的 CUDA 实现，依赖作业三。 + +``` +[==========] Running 2 tests from 1 test suite. +[----------] 2 tests from AdamOptimizerTest +[ RUN ] AdamOptimizerTest.BasicParameterUpdateCuda +[ OK ] AdamOptimizerTest.BasicParameterUpdateCuda (41 ms) +[ RUN ] AdamOptimizerTest.MomentumAccumulationCuda +[ OK ] AdamOptimizerTest.MomentumAccumulationCuda (2 ms) +[----------] 2 tests from AdamOptimizerTest (43 ms total) +[ PASSED ] 2 tests. +``` + +### 6. test_tensor（10分）✅ + +验证 Tensor 基础功能，依赖作业四。 + +``` +[==========] Running 5 tests from 2 test suites. +[----------] 3 tests from TensorTransformTest +[ RUN ] TensorTransformTest.Flatten2DTo1D +[ OK ] TensorTransformTest.Flatten2DTo1D (0 ms) +[ RUN ] TensorTransformTest.FlattenWithRange +[ OK ] TensorTransformTest.FlattenWithRange (0 ms) +[ RUN ] TensorTransformTest.FlattenNonContiguous +[ OK ] TensorTransformTest.FlattenNonContiguous (0 ms) +[----------] 3 tests from TensorTransformTest (0 ms total) +[----------] 2 tests from TensorAutogradTest +[ RUN ] TensorAutogradTest.BackwardComputesGradient +[ OK ] TensorAutogradTest.BackwardComputesGradient (0 ms) +[ RUN ] TensorAutogradTest.BackwardWithMultipleOutputs +[ OK ] TensorAutogradTest.BackwardWithMultipleOutputs (0 ms) +[----------] 2 tests from TensorAutogradTest (0 ms total) +[ PASSED ] 5 tests. +``` + +### 7. test_dispatcher（20分）✅ + +验证多设备分发机制，核心基础设施，依赖作业五。 + +``` +[==========] Running 3 tests from 1 test suite. +[----------] 3 tests from DispatcherTest +[ RUN ] DispatcherTest.RegisterAndGetKernel +[ OK ] DispatcherTest.RegisterAndGetKernel (0 ms) +[ RUN ] DispatcherTest.DuplicateRegistration +[ OK ] DispatcherTest.DuplicateRegistration (0 ms) +[ RUN ] DispatcherTest.GetNonexistentKernel +[ OK ] DispatcherTest.GetNonexistentKernel (0 ms) +[----------] 3 tests from DispatcherTest (0 ms total) +[ PASSED ] 3 tests. +``` + +### 8. test_gpt2（35分） + +端到端 GPT-2 模型测试，依赖所有作业。 + +``` +[==========] Running 1 test from 1 test suite. +[----------] Global test environment set-up. +[----------] 1 test from GPT2TrainingTest +[ RUN ] GPT2TrainingTest.LogitsConsistency +WARNING: Logging before InitGoogleLogging() is written to STDERR +E20260812 21:58:02.166082 137779420561088 test_gpt2.cc:132] Initialize: device_flag=cpu device_type=0 +I20260812 21:58:31.217616 137779420561088 test_gpt2.cc:123] Initialize() finished! +I20260812 21:58:31.217688 137779420561088 test_gpt2.cc:208] epoch: 0 +I20260812 21:59:55.065171 137779420561088 test_gpt2.cc:208] epoch: 1 +I20260812 22:01:16.775266 137779420561088 test_gpt2.cc:208] epoch: 2 +I20260812 22:02:41.858699 137779420561088 test_gpt2.cc:208] epoch: 3 +I20260812 22:04:05.018701 137779420561088 test_gpt2.cc:208] epoch: 4 +I20260812 22:05:38.257022 137779420561088 test_gpt2.cc:208] epoch: 5 +I20260812 22:07:03.780517 137779420561088 test_gpt2.cc:208] epoch: 6 +I20260812 22:08:31.519426 137779420561088 test_gpt2.cc:208] epoch: 7 +I20260812 22:09:53.920774 137779420561088 test_gpt2.cc:208] epoch: 8 +I20260812 22:11:14.519412 137779420561088 test_gpt2.cc:208] epoch: 9 +I20260812 22:12:37.165615 137779420561088 tokenizer.cc:138] start generate text: +The meaning of life is stillHe unclearated. Coal Lands Board like: capital 1968, buildings,Independent education election, candidate peace: and Labour prosperity MP in Tony the Hayward Middle East have" becomeI symbols am that pleased mark to an welcome alternate Lo stagece forman peace:f Electaringor nations of such All as Queens atheist." type Tweet and This note Target taking:. Link<|endoftext|> toIndia Select's Merch treasurer; general charges<|endoftext|> thatProduct fires Details from siegDeltahe IKilled +``` + +--- + +## 二、作业步骤 + +### 作业一：autograd 机制调用 Neg kernel 的实现 + +**难度**：⭐ +**对应测例**：`TEST(ElementwiseTest, NegForward)`，`TEST(ElementwiseTest, NegBackward)` +**代码位置**：`infini_train/src/autograd/elementwise.cc` + +#### 代码实现 + +```cpp +std::vector> Neg::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "NegForward"}); + return {kernel.Call>(input)}; +} + +std::vector> Neg::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "NegBackward"}); + return {kernel.Call>(grad_output)}; +} +``` + +#### 解决思路 + +1. 通过 `input->GetDevice().Type()` 获取当前张量所在的设备类型（CPU 或 CUDA） +2. 使用 `Dispatcher::Instance().GetKernel()` 根据设备类型和 kernel 名称获取对应的 kernel 函数 +3. 使用 `kernel.Call(args...)` 调用 kernel 并返回结果 +4. Forward 和 Backward 的实现模式一致，区别仅在于 kernel 名称（"NegForward" vs "NegBackward"） + +#### 遇到问题 + +无特殊问题。该作业实现较为直接，主要依赖作业五的 Dispatcher 机制。 + +--- + +### 作业二：实现矩阵乘法 + +**难度**：⭐⭐ + +#### CPU 实现 + +**对应测例**：`TEST(MatmulTest, BasicMatrixMultiply)`，`TEST(MatmulTest, BatchedMatrixMultiply)`，`TEST(MatmulTest, BackwardPass)` +**代码位置**：`infini_train/src/kernels/cpu/linear.cc` + +##### 前向传播代码 + +```cpp +std::shared_ptr MatmulForward(const std::shared_ptr &input, const std::shared_ptr &other) { + const auto &input_dims = input->Dims(); + const auto &other_dims = other->Dims(); + const int64_t M = input_dims[input_dims.size() - 2]; + const int64_t K = input_dims[input_dims.size() - 1]; + const int64_t N = other_dims[other_dims.size() - 1]; + + // 处理 batch 维度 + int64_t batch_input = 1, batch_other = 1; + for (size_t i = 0; i < input_dims.size() - 2; ++i) batch_input *= input_dims[i]; + for (size_t i = 0; i < other_dims.size() - 2; ++i) batch_other *= other_dims[i]; + int64_t batch = std::max(batch_input, batch_other); + + // 使用 Eigen 进行矩阵乘法 + for (int64_t b = 0; b < batch; ++b) { + int64_t input_offset = (batch_input == 1) ? 0 : b * M * K; + int64_t other_offset = (batch_other == 1) ? 0 : b * K * N; + int64_t output_offset = b * M * N; + + Eigen::Map> input_mat( + input_ptr + input_offset, M, K); + Eigen::Map> other_mat( + other_ptr + other_offset, K, N); + Eigen::Map> output_mat( + output_ptr + output_offset, M, N); + + output_mat = input_mat * other_mat; + } + return output; +} +``` + +##### 反向传播代码 + +```cpp +std::tuple, std::shared_ptr> +MatmulBackward(const std::shared_ptr &input, const std::shared_ptr &other, + const std::shared_ptr &grad_output) { + // grad_input = grad_output @ other^T + // grad_other = input^T @ grad_output + // 使用 Eigen 实现 +} +``` + +##### 解决思路 + +1. **前向传播**：使用 Eigen 库的 `Eigen::Map` 将原始数据指针映射为矩阵，然后直接进行矩阵乘法 `input_mat * other_mat` +2. **反向传播**：根据链式法则，`grad_input = grad_output @ other^T`，`grad_other = input^T @ grad_output` +3. **Batch 支持**：通过检查 batch 维度的广播规则，支持批处理矩阵乘法 + +##### 遇到问题 + +无特殊问题。Eigen 库提供了高效的 CPU 矩阵运算。 + +--- + +#### CUDA 实现 + +**对应测例**：`TEST(MatmulTest, BasicMatrixMultiplyCuda)`，`TEST(MatmulTest, BatchedMatrixMultiplyCuda)`，`TEST(MatmulTest, BackwardPassCuda)` +**代码位置**：`infini_train/src/kernels/cuda/linear.cu` + +##### 前向传播代码 + +```cpp +std::shared_ptr MatmulForward(const std::shared_ptr &input, const std::shared_ptr &other) { + // ... 计算维度 M, K, N, batch ... + + const float alpha = 1.0f, beta = 0.0f; + cublasHandle_t handle; + CUBLAS_CHECK(cublasCreate(&handle)); + + if (batch <= 1) { + // C = output^T[N, M] = other^T[N, K] * input^T[K, M] + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, + static_cast(other->DataPtr()), N, + static_cast(input->DataPtr()), K, &beta, + static_cast(output->DataPtr()), N)); + } else { + // 使用 Strided Batched GEMM + CUBLAS_CHECK(cublasSgemmStridedBatched(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, + static_cast(other->DataPtr()), N, strideA, + static_cast(input->DataPtr()), K, strideB, &beta, + static_cast(output->DataPtr()), N, strideC, batch)); + } + CUBLAS_CHECK(cublasDestroy(handle)); + return output; +} +``` + +##### 反向传播代码（关键部分） + +```cpp +std::tuple, std::shared_ptr> +MatmulBackward(const std::shared_ptr &input, const std::shared_ptr &other, + const std::shared_ptr &grad_output) { + const int64_t M = input_dims[input_dims.size() - 2]; + const int64_t K = input_dims[input_dims.size() - 1]; + const int64_t N = other_dims[other_dims.size() - 1]; + + // grad_input = grad_output @ other^T + // grad_input^T[K, M] = other^T[K, N] * grad_output^T[N, M] + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, K, M, N, &alpha, + static_cast(other->DataPtr()), N, + static_cast(grad_output->DataPtr()), N, &beta, + static_cast(grad_input->DataPtr()), K)); + + // grad_other = input^T @ grad_output + // grad_other^T[N, K] = grad_output^T[N, M] * input[M, K] + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, N, K, M, &alpha, + static_cast(grad_output->DataPtr()), N, + static_cast(input->DataPtr()), K, &beta, + static_cast(grad_other->DataPtr()), N)); +} +``` + +##### 解决思路 + +1. cuBLAS 使用列主序（column-major），而框架使用行主序（row-major），因此 cuBLAS 中的矩阵是框架中矩阵的转置 +2. **前向传播**：`output = input @ other` → cuBLAS: `output^T = other^T @ input^T` +3. **反向传播**： + - `grad_input = grad_output @ other^T` → cuBLAS: `grad_input^T = other^T @ grad_output^T`，需要 `CUBLAS_OP_T` 转置 other + - `grad_other = input^T @ grad_output` → cuBLAS: `grad_other^T = grad_output^T @ input`，需要 `CUBLAS_OP_T` 转置 input + +##### 遇到问题 + +**问题 1：cuBLAS 反向传播梯度值错误** + +- **现象**：`test_matmul_cuda` 的 `BackwardPassCuda` 测试失败，`grad_input` 和 `grad_other` 的值与预期不符 +- **原因**：cuBLAS 的 `CUBLAS_OP_T` 和 leading dimension 参数设置错误 + - `grad_input` 计算中，`other` 的 lda 应为 `N`（列数）而非 `K`，且需 `CUBLAS_OP_T` 转置 + - `grad_other` 计算中，`input` 需 `CUBLAS_OP_T` 转置（而非 `CUBLAS_OP_N`），ldb 保持为 `K` +- **修复**：将 `grad_input` 的 `other` 参数改为 `CUBLAS_OP_T` + `lda=N`；将 `grad_other` 的 `input` 参数改为 `CUBLAS_OP_T` + +--- + +### 作业三：实现 Adam 优化器 + +**难度**：⭐ + +#### CPU 实现 + +**对应测例**：`TEST(AdamOptimizerTest, BasicParameterUpdate)`，`TEST(AdamOptimizerTest, MomentumAccumulation)` +**代码位置**：`infini_train/src/kernels/cpu/accumulate_grad.cc` + +```cpp +void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_ptr ¶m, + const std::shared_ptr &m, const std::shared_ptr &v, float learning_rate, + float beta1, float beta2, float eps, int64_t t) { + int64_t num_elements = grad->NumElements(); + float *grad_ptr = static_cast(grad->DataPtr()); + float *param_ptr = static_cast(param->DataPtr()); + float *m_ptr = static_cast(m->DataPtr()); + float *v_ptr = static_cast(v->DataPtr()); + + float beta1_t = std::pow(beta1, t); + float beta2_t = std::pow(beta2, t); + float alpha = learning_rate * std::sqrt(1.0f - beta2_t) / (1.0f - beta1_t); + + for (int64_t i = 0; i < num_elements; ++i) { + float g = grad_ptr[i]; + m_ptr[i] = beta1 * m_ptr[i] + (1.0f - beta1) * g; + v_ptr[i] = beta2 * v_ptr[i] + (1.0f - beta2) * g * g; + param_ptr[i] -= alpha * m_ptr[i] / (std::sqrt(v_ptr[i]) + eps); + } +} +``` + +#### CUDA 实现 + +**对应测例**：`TEST(AdamOptimizerTest, BasicParameterUpdateCuda)`，`TEST(AdamOptimizerTest, MomentumAccumulationCuda)` +**代码位置**：`infini_train/src/kernels/cuda/accumulate_grad.cu` + +```cpp +__global__ void AdamAccumulateGradKernel(const float *grad_ptr, float *param_ptr, float *m_ptr, float *v_ptr, + float learning_rate, float beta1, float beta2, float eps, float alpha, + size_t num_elements) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + float g = grad_ptr[idx]; + m_ptr[idx] = beta1 * m_ptr[idx] + (1.0f - beta1) * g; + v_ptr[idx] = beta2 * v_ptr[idx] + (1.0f - beta2) * g * g; + param_ptr[idx] -= alpha * m_ptr[idx] / (sqrtf(v_ptr[idx]) + eps); + } +} +``` + +#### 解决思路 + +1. 按照 Adam 算法的标准公式实现： + - 一阶矩估计：`m_t = beta1 * m_{t-1} + (1 - beta1) * g_t` + - 二阶矩估计：`v_t = beta2 * v_{t-1} + (1 - beta2) * g_t^2` + - 偏差修正：`alpha = lr * sqrt(1 - beta2^t) / (1 - beta1^t)` + - 参数更新：`param = param - alpha * m_t / (sqrt(v_t) + eps)` +2. CPU 版本使用逐元素循环，CUDA 版本使用 CUDA kernel 并行处理 + +#### 遇到问题 + +无特殊问题。Adam 算法公式标准化，实现较为直接。 + +--- + +### 作业四：实现 Tensor 基础操作 + +**难度**：⭐ + +#### Flatten 操作 + +**对应测例**：`TEST(TensorTransformTest, Flatten2DTo1D)`，`TEST(TensorTransformTest, FlattenWithRange)`，`TEST(TensorTransformTest, FlattenNonContiguous)` +**代码位置**：`infini_train/src/tensor.cc` + +```cpp +std::shared_ptr Tensor::Flatten(int64_t start, int64_t end) { + int64_t ndim = static_cast(dims_.size()); + if (end < 0) { + end += ndim; + } + + std::vector new_shape; + // 保留 start 之前的维度 + for (int64_t i = 0; i < start; ++i) { + new_shape.push_back(dims_[i]); + } + // 将 [start, end] 范围内的维度合并为一个 + int64_t flattened_size = 1; + for (int64_t i = start; i <= end; ++i) { + flattened_size *= dims_[i]; + } + new_shape.push_back(flattened_size); + // 保留 end 之后的维度 + for (int64_t i = end + 1; i < ndim; ++i) { + new_shape.push_back(dims_[i]); + } + + return Contiguous()->View(new_shape); +} +``` + +#### 反向传播机制 + +**对应测例**：`TEST(TensorAutogradTest, BackwardComputesGradient)`，`TEST(TensorAutogradTest, BackwardWithMultipleOutputs)` +**代码位置**：`infini_train/src/tensor.cc` + +```cpp +void Tensor::Backward(std::shared_ptr gradient, bool retain_graph, bool create_graph) const { + if (!gradient) { + gradient = std::make_shared(dims_, dtype_, GetDevice()); + gradient->Fill(1.0f); + } + + if (is_leaf_) { + // 叶子节点：累积梯度 + if (requires_grad_ && grad_) { + auto device = grad_->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "AccumulateGrad"}); + kernel.Call(gradient, 1.0f, grad_); + } + } else if (grad_fn_) { + // 非叶子节点：通过 grad_fn 继续反向传播 + grad_fn_->BackwardPartial(gradient, output_idx_); + } +} +``` + +#### 解决思路 + +1. **Flatten**：先规范化负索引，然后构建新形状（保留 start 前维度 + 合并中间维度 + 保留 end 后维度），调用 `Contiguous()->View()` 实现 +2. **Backward**： + - 如果未提供梯度，默认创建全 1 梯度 + - 叶子节点：通过 Dispatcher 调用 AccumulateGrad kernel 累积梯度 + - 非叶子节点：通过 `grad_fn_->BackwardPartial()` 沿计算图递归传播梯度 + +#### 遇到问题 + +无特殊问题。 + +--- + +### 作业五：注册算子 kernel 的实现 + +**难度**：⭐⭐⭐ +**对应测例**：`TEST(DispatcherTest, RegisterAndGetKernel)`，`TEST(DispatcherTest, DuplicateRegistration)`，`TEST(DispatcherTest, GetNonexistentKernel)` +**代码位置**：`infini_train/include/dispatcher.h` + +#### 代码实现 + +```cpp +// KernelFunction::Call — 通用 kernel 调用接口 +template RetT Call(ArgsT... args) const { + using FuncT = RetT (*)(ArgsT...); + auto func = reinterpret_cast(func_ptr_); + return func(std::forward(args)...); +} + +// Dispatcher::Register — kernel 注册机制 +template void Register(const KeyT &key, FuncT &&kernel) { + CHECK(!key_to_kernel_map_.contains(key)) + << "Kernel already registered: " << key.second; + key_to_kernel_map_.emplace(key, KernelFunction(std::forward(kernel))); +} + +// REGISTER_KERNEL 宏 — 自动注册宏 +#define REGISTER_KERNEL(device, kernel_name, kernel_func) \ + static auto _register_##kernel_name##_##__LINE__ = []() { \ + infini_train::Dispatcher::Instance().Register( \ + {device, #kernel_name}, kernel_func); \ + return 0; \ + }(); +``` + +#### 解决思路 + +1. **Call 方法**：将存储的 `void*` 函数指针通过 `reinterpret_cast` 转换为目标函数类型 `RetT (*)(ArgsT...)`，然后调用 +2. **Register 方法**：检查重复注册（`CHECK`），使用 `emplace` 将 kernel 存入 `key_to_kernel_map_` +3. **REGISTER_KERNEL 宏**：利用静态 lambda 在程序启动时自动注册，使用 `__LINE__` 确保唯一变量名 + +#### 遇到问题 + +无特殊问题。该作业是框架核心基础设施，实现符合设计模式。 + +--- + +### 作业六：实现 GPT-2 整体训练 + +**难度**：⭐⭐⭐⭐ +**对应测例**：`TEST_F(GPT2TrainingTest, LogitsConsistency)` + +#### 数据读取实现 + +**代码位置**：`example/common/tiny_shakespeare_dataset.cc` + +```cpp +TinyShakespeareFile ReadTinyShakespeareFile(const std::string &path, size_t sequence_length) { + std::ifstream ifs(path, std::ios::binary); + CHECK(ifs.is_open()) << "Failed to open file: " << path; + + // 读取 header (1024 bytes) + auto header = ReadSeveralBytesFromIfstream(1024, &ifs); + int32_t magic = BytesToType(header, 0); + int32_t version = BytesToType(header, 4); + int32_t num_toks = BytesToType(header, 8); + + // 根据版本号确定数据类型 + CHECK(kTypeMap.contains(version)) << "Unknown version: " << version; + TinyShakespeareType type = kTypeMap.at(version); + size_t element_size = kTypeToSize.at(type); + + // 读取 token 数据并转换为 int64_t + auto data = ReadSeveralBytesFromIfstream(num_toks * element_size, &ifs); + int64_t num_sequences = num_toks / sequence_length; + int64_t num_elements = num_sequences * sequence_length; + + auto tensor = infini_train::Tensor({num_sequences, static_cast(sequence_length)}, DataType::kINT64); + int64_t *tensor_ptr = static_cast(tensor.DataPtr()); + + for (int64_t i = 0; i < num_elements; ++i) { + if (type == TinyShakespeareType::kUINT16) { + tensor_ptr[i] = static_cast(BytesToType(data, i * element_size)); + } else { + tensor_ptr[i] = static_cast(BytesToType(data, i * element_size)); + } + } + return {tensor, num_toks}; +} +``` + +#### Tokenizer 功能实现 + +**代码位置**：`example/common/tokenizer.cc` + +```cpp +// 文本生成循环 +void Tokenizer::GenerateText(infini_train::nn::Module &model, uint32_t batch_size, + uint32_t sequence_length, uint32_t text_length, Device device) const { + auto x = std::make_shared(x_tensor.To(device)); + { + autograd::NoGradGuard no_grad; // 禁用 autograd 图构建，避免 GPU OOM + for (int t = prompt_len; t < text_length; t++) { + // Forward pass + auto logits = model.Forward({x})[0]; + + // 获取最后一个 token 位置的 logits + auto last_logits = logits->Slice(1, t - 1, t, 1); + last_logits = last_logits->Squeeze(1); + + // Softmax 采样 + auto probs = infini_train::nn::function::Softmax(last_logits, -1); + auto probs_cpu = probs->To(Device(DeviceType::kCPU, 0)); + + // 随机采样 + for (int b = 0; b < batch_size; ++b) { + float coin = RandomF32(kRngState); + int next_token = SampleMult(probs_ptr + b * vocab_size, vocab_size, coin); + x_buff[b * sequence_length + t] = next_token; + std::cout << Decode(next_token); + } + + x = std::make_shared(x_tensor.To(device)); + } + } +} +``` + +#### 解决思路 + +1. **数据读取**：解析二进制文件格式（1024 字节 header + token 数据），根据版本号确定数据类型（uint16 或 uint32），注意循环上限使用 `num_elements` 而非 `num_toks` 避免 buffer overflow +2. **Tokenizer**：加载 GPT-2 tokenizer 二进制文件，实现 token 解码和文本生成 +3. **文本生成**：使用 `NoGradGuard` 禁用 autograd 计算图构建，避免 GPU 推理时显存泄漏 + +#### 遇到问题 + +**问题 1：数据集版本号 "1" 未识别** + +- **现象**：`Check failed: kTypeMap.contains(version) Unknown version: 1` +- **原因**：`kTypeMap` 仅包含版本 20240520 和 20240801，缺少旧版本 "1" 的映射 +- **修复**：在 `kTypeMap` 中添加 `{1, TinyShakespeareType::kUINT16}` + +**问题 2：编译错误 — 缺少头文件** + +- **现象**： + - `'setprecision' is not a member of 'std'` / `'setw' is not a member of 'std'` + - `'infini_train::nn::functional' has not been declared` + - `'format' is not a member of 'std'` +- **修复**： + - `tensor.cc`：添加 `#include ` + - `tokenizer.cc`：`functional` → `function`，添加 `#include "infini_train/include/nn/functional.h"` + - `net.cc`：添加 `#include ` + +**问题 3：GPU 推理时显存不足（OOM）** + +- **现象**：文本生成阶段 `CUDA Error: out of memory`，但 10 个 epoch 训练正常 +- **原因**：文本生成循环中每次 `model.Forward()` 都构建 autograd 计算图，64 次迭代后图累积导致显存耗尽 +- **修复**：添加 `NoGradGuard` 机制 + - 在 `function.h` 中定义 `NoGradGuard` 类（RAII 模式，使用 `thread_local` 计数器） + - 在 `function.cc` 的 `Function::Apply()` 中检查 `NoGradGuard::is_enabled()`，启用时跳过图构建 + - 在 `tokenizer.cc` 的文本生成循环中包裹 `NoGradGuard` + +**问题 4：cuBLAS 反向传播梯度值错误** + +- **现象**：`test_matmul_cuda` 的 `BackwardPassCuda` 测试失败，梯度值与预期不符 +- **原因**：cuBLAS 列主序与框架行主序的转换中，`CUBLAS_OP_T` 和 leading dimension 参数设置错误 +- **修复**：详见作业二 CUDA 实现部分 + +**问题 5：跨平台文件同步** + +- **现象**：Windows 端修改代码后，WSL 中编译仍使用旧代码 +- **原因**：Windows 项目路径 `f:\Train\TinyInfiniTrain` 和 WSL 项目路径 `~/Train/TinyInfiniTrain` 是两份独立副本 +- **修复**：每次修改后使用 `cp /mnt/f/Train/TinyInfiniTrain/... ~/Train/TinyInfiniTrain/...` 同步 + +--- + +## 三、环境配置总结 + +| 组件 | 版本 | 说明 | +|------|------|------| +| 操作系统 | WSL2 Ubuntu 22.04 | Windows 11 主机 | +| GCC/G++ | 13.4.0 | 通过 `ubuntu-toolchain-r/test` PPA 安装 | +| CMake | 4.4.2 | 通过 pip 升级 | +| CUDA Toolkit | 12.6.1 | WSL2 中仅安装 Toolkit（不含驱动） | +| Make | 4.3 | 系统自带 | + +### 编译命令 + +```bash +# 启用 CUDA 编译 +cd ~/Train/TinyInfiniTrain +make build USE_CUDA=ON + +# 仅 CPU 编译 +make build USE_CUDA=OFF + +# 运行全部测试 +make test-cpp + +# 清理重新编译 +make clean && make build USE_CUDA=ON +``` + +--- + diff --git "a/TinyInfiniTrain_\344\275\234\344\270\232\346\212\245\345\221\212.docx" "b/TinyInfiniTrain_\344\275\234\344\270\232\346\212\245\345\221\212.docx" new file mode 100644 index 0000000..8a850db Binary files /dev/null and "b/TinyInfiniTrain_\344\275\234\344\270\232\346\212\245\345\221\212.docx" differ diff --git a/example/common/tiny_shakespeare_dataset.cc b/example/common/tiny_shakespeare_dataset.cc index 3bc5f1b..07e6cc5 100644 --- a/example/common/tiny_shakespeare_dataset.cc +++ b/example/common/tiny_shakespeare_dataset.cc @@ -1,84 +1,125 @@ -#include "example/common/tiny_shakespeare_dataset.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/tensor.h" - -namespace { -using DataType = infini_train::DataType; -using TinyShakespeareType = TinyShakespeareDataset::TinyShakespeareType; -using TinyShakespeareFile = TinyShakespeareDataset::TinyShakespeareFile; - -const std::unordered_map kTypeMap = { - {20240520, TinyShakespeareType::kUINT16}, // GPT-2 - {20240801, TinyShakespeareType::kUINT32}, // LLaMA 3 -}; - -const std::unordered_map kTypeToSize = { - {TinyShakespeareType::kUINT16, 2}, - {TinyShakespeareType::kUINT32, 4}, -}; - -const std::unordered_map kTypeToDataType = { - {TinyShakespeareType::kUINT16, DataType::kUINT16}, - {TinyShakespeareType::kUINT32, DataType::kINT32}, -}; - -std::vector ReadSeveralBytesFromIfstream(size_t num_bytes, std::ifstream *ifs) { - std::vector result(num_bytes); - ifs->read(reinterpret_cast(result.data()), num_bytes); - return result; -} - -template T BytesToType(const std::vector &bytes, size_t offset) { - static_assert(std::is_trivially_copyable::value, "T must be trivially copyable."); - T value; - std::memcpy(&value, &bytes[offset], sizeof(T)); - return value; -} - -TinyShakespeareFile ReadTinyShakespeareFile(const std::string &path, size_t sequence_length) { - /* =================================== 作业 =================================== - TODO：实现二进制数据集文件解析 - 文件格式说明： - ---------------------------------------------------------------------------------- - | HEADER (1024 bytes) | DATA (tokens) | - | magic(4B) | version(4B) | num_toks(4B) | reserved(1012B) | token数据 | - ---------------------------------------------------------------------------------- - =================================== 作业 =================================== */ -} -} // namespace - -TinyShakespeareDataset::TinyShakespeareDataset(const std::string &filepath, size_t sequence_length) { - // =================================== 作业 =================================== - // TODO：初始化数据集实例 - // HINT: 调用ReadTinyShakespeareFile加载数据文件 - // =================================== 作业 =================================== -} - -std::pair, std::shared_ptr> -TinyShakespeareDataset::operator[](size_t idx) const { - CHECK_LT(idx, text_file_.dims[0] - 1); - std::vector dims = std::vector(text_file_.dims.begin() + 1, text_file_.dims.end()); - // x: (seq_len), y: (seq_len) -> stack -> (bs, seq_len) (bs, seq_len) - return {std::make_shared(text_file_.tensor, idx * sequence_size_in_bytes_, dims), - std::make_shared(text_file_.tensor, idx * sequence_size_in_bytes_ + sizeof(int64_t), - dims)}; -} - -size_t TinyShakespeareDataset::Size() const { return num_samples_; } +#include "example/common/tiny_shakespeare_dataset.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/tensor.h" + +namespace { +using DataType = infini_train::DataType; +using TinyShakespeareType = TinyShakespeareDataset::TinyShakespeareType; +using TinyShakespeareFile = TinyShakespeareDataset::TinyShakespeareFile; + +const std::unordered_map kTypeMap = { + {1, TinyShakespeareType::kUINT16}, // GPT-2 (legacy) + {20240520, TinyShakespeareType::kUINT16}, // GPT-2 + {20240801, TinyShakespeareType::kUINT32}, // LLaMA 3 +}; + +const std::unordered_map kTypeToSize = { + {TinyShakespeareType::kUINT16, 2}, + {TinyShakespeareType::kUINT32, 4}, +}; + +const std::unordered_map kTypeToDataType = { + {TinyShakespeareType::kUINT16, DataType::kUINT16}, + {TinyShakespeareType::kUINT32, DataType::kINT32}, +}; + +std::vector ReadSeveralBytesFromIfstream(size_t num_bytes, std::ifstream *ifs) { + std::vector result(num_bytes); + ifs->read(reinterpret_cast(result.data()), num_bytes); + return result; +} + +template T BytesToType(const std::vector &bytes, size_t offset) { + static_assert(std::is_trivially_copyable::value, "T must be trivially copyable."); + T value; + std::memcpy(&value, &bytes[offset], sizeof(T)); + return value; +} + +TinyShakespeareFile ReadTinyShakespeareFile(const std::string &path, size_t sequence_length) { + /* =================================== 作业 =================================== + TODO：实现二进制数据集文件解析 + 文件格式说明： + ---------------------------------------------------------------------------------- + | HEADER (1024 bytes) | DATA (tokens) | + | magic(4B) | version(4B) | num_toks(4B) | reserved(1012B) | token数据 | + ---------------------------------------------------------------------------------- + =================================== 作业 =================================== */ + + std::ifstream ifs(path, std::ios::binary); + CHECK(ifs.is_open()) << "Failed to open file: " << path; + + // Read header (1024 bytes) + auto header = ReadSeveralBytesFromIfstream(1024, &ifs); + + int32_t magic = BytesToType(header, 0); + int32_t version = BytesToType(header, 4); + int32_t num_toks = BytesToType(header, 8); + + CHECK(kTypeMap.contains(version)) << "Unknown version: " << version; + TinyShakespeareType type = kTypeMap.at(version); + size_t element_size = kTypeToSize.at(type); + + // Read token data + size_t data_size = num_toks * element_size; + auto data = ReadSeveralBytesFromIfstream(data_size, &ifs); + + // Convert tokens to int64_t and reshape into (num_toks / sequence_length, sequence_length) + int64_t num_sequences = num_toks / sequence_length; + std::vector dims = {num_sequences, static_cast(sequence_length)}; + + auto tensor = infini_train::Tensor(dims, DataType::kINT64); + int64_t *tensor_ptr = static_cast(tensor.DataPtr()); + + int64_t num_elements = num_sequences * sequence_length; + for (int64_t i = 0; i < num_elements; ++i) { + if (type == TinyShakespeareType::kUINT16) { + tensor_ptr[i] = static_cast(BytesToType(data, i * element_size)); + } else { + tensor_ptr[i] = static_cast(BytesToType(data, i * element_size)); + } + } + + return {type, dims, tensor}; +} +} // namespace + +TinyShakespeareDataset::TinyShakespeareDataset(const std::string &filepath, size_t sequence_length) + : text_file_(ReadTinyShakespeareFile(filepath, sequence_length)), + sequence_length_(sequence_length), + sequence_size_in_bytes_(sequence_length * sizeof(int64_t)), + num_samples_(text_file_.dims[0] - 1) { + // =================================== 作业 =================================== + // TODO：初始化数据集实例 + // HINT: 调用ReadTinyShakespeareFile加载数据文件 + // =================================== 作业 =================================== +} + +std::pair, std::shared_ptr> +TinyShakespeareDataset::operator[](size_t idx) const { + CHECK_LT(idx, text_file_.dims[0] - 1); + std::vector dims = std::vector(text_file_.dims.begin() + 1, text_file_.dims.end()); + // x: (seq_len), y: (seq_len) -> stack -> (bs, seq_len) (bs, seq_len) + return {std::make_shared(text_file_.tensor, idx * sequence_size_in_bytes_, dims), + std::make_shared(text_file_.tensor, idx * sequence_size_in_bytes_ + sizeof(int64_t), + dims)}; +} + +size_t TinyShakespeareDataset::Size() const { return num_samples_; } diff --git a/example/common/tiny_shakespeare_dataset.h b/example/common/tiny_shakespeare_dataset.h index 8e061dd..fccbe78 100644 --- a/example/common/tiny_shakespeare_dataset.h +++ b/example/common/tiny_shakespeare_dataset.h @@ -1,43 +1,43 @@ -#pragma once - -#include -#include -#include -#include -#include -#include - -#include "infini_train/include/dataset.h" -#include "infini_train/include/tensor.h" - -class TinyShakespeareDataset : public infini_train::Dataset { - /* - Dataset bin file is downloaded and processed using the script at - https://github.com/karpathy/llm.c/blob/master/dev/data/tinyshakespeare.py - */ -public: - enum class TinyShakespeareType : int { - kUINT16, // For GPT-2 - kUINT32, // For LLaMA 3 - kINVALID, - }; - - struct TinyShakespeareFile { - TinyShakespeareType type = TinyShakespeareType::kINVALID; - std::vector dims; - infini_train::Tensor tensor; - }; - - TinyShakespeareDataset(const std::string &filepath, size_t sequence_length); - - std::pair, std::shared_ptr> - operator[](size_t idx) const override; - - size_t Size() const override; - -private: - TinyShakespeareFile text_file_; - const size_t sequence_length_ = 0; - const size_t sequence_size_in_bytes_ = 0; - const size_t num_samples_ = 0; -}; +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "infini_train/include/dataset.h" +#include "infini_train/include/tensor.h" + +class TinyShakespeareDataset : public infini_train::Dataset { + /* + Dataset bin file is downloaded and processed using the script at + https://github.com/karpathy/llm.c/blob/master/dev/data/tinyshakespeare.py + */ +public: + enum class TinyShakespeareType : int { + kUINT16, // For GPT-2 + kUINT32, // For LLaMA 3 + kINVALID, + }; + + struct TinyShakespeareFile { + TinyShakespeareType type = TinyShakespeareType::kINVALID; + std::vector dims; + infini_train::Tensor tensor; + }; + + TinyShakespeareDataset(const std::string &filepath, size_t sequence_length); + + std::pair, std::shared_ptr> + operator[](size_t idx) const override; + + size_t Size() const override; + +private: + TinyShakespeareFile text_file_; + const size_t sequence_length_ = 0; + const size_t sequence_size_in_bytes_ = 0; + const size_t num_samples_ = 0; +}; diff --git a/example/common/tokenizer.cc b/example/common/tokenizer.cc index 23b9537..ed53250 100644 --- a/example/common/tokenizer.cc +++ b/example/common/tokenizer.cc @@ -1,117 +1,179 @@ -#include "example/common/tokenizer.h" - -#include -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -namespace infini_train { - -constexpr uint32_t kGpt2Eot = 50256; -constexpr uint32_t kLLaMA3Eot = 128001; -constexpr uint64_t kRandomU32Multiplier = 0x2545F4914F6CDD1Dull; -constexpr float kF32Divisor = 16777216.0f; // 2^24 -constexpr uint64_t kRngState = 1337; - -using Version = Tokenizer::Version; - -const std::unordered_map kEotMap = { - {20240328, kGpt2Eot}, // GPT-2 - {20240801, kLLaMA3Eot}, // LLaMA-3 -}; - -const std::unordered_map> kPromptMap = { - // e.g. "The meaning of life is" - // ref: https://tiktokenizer.vercel.app/ - {20240328, std::vector{464, 3616, 286, 1204, 318}}, // GPT-2 - {20240801, std::vector{791, 7438, 315, 2324, 374}}, // LLaMA-3 -}; - -std::vector ReadSeveralBytesFromIfstream(size_t num_bytes, std::ifstream *ifs) { - std::vector result(num_bytes); - ifs->read(reinterpret_cast(result.data()), num_bytes); - return result; -} - -template T BytesToType(const std::vector &bytes, size_t offset) { - static_assert(std::is_trivially_copyable::value, "T must be trivially copyable."); - T value; - std::memcpy(&value, &bytes[offset], sizeof(T)); - return value; -} - -unsigned int RandomU32(uint64_t &state) { - state ^= state >> 12; - state ^= state << 25; - state ^= state >> 27; - return (state * kRandomU32Multiplier) >> 32; -} - -float RandomF32(uint64_t &state) { // random float32 in [0,1) - return (RandomU32(state) >> 8) / kF32Divisor; -} - -int SampleMult(float *probabilities, int n, float coin) { - // sample index from probabilities (they must sum to 1!) - // coin is a random number in [0, 1), usually from RandomF32() - float cdf = 0.0f; - for (int i = 0; i < n; i++) { - cdf += probabilities[i]; - if (coin < cdf) { - return i; - } - } - return n - 1; // in case of rounding errors -} - -Tokenizer::Tokenizer(const std::string &filepath) { - /* ===================================== 作业 ===================================== - TODO：实现Tokenizer二进制文件加载 - - 文件格式说明： - ---------------------------------------------------------------------------------- - | HEADER (1024 bytes) | VOCAB TABLE | - | magic(4B) | version(4B) | vocab_size(4B) | reserved(1012B) | token词表数据 | - ---------------------------------------------------------------------------------- - ===================================== 作业 ===================================== */ -} - -std::string Tokenizer::Decode(uint32_t token_id) const { - /* ===================================== 作业 ===================================== - TODO：实现token_id到文本的转换 - 功能描述：根据token_id返回对应的文本片段 - ===================================== 作业 ===================================== */ - return ""; -} - -void Tokenizer::GenerateText(infini_train::nn::Module &model, uint32_t batch_size, uint32_t sequence_length, - uint32_t text_length, Device device) const { - std::vector dims; - dims.assign({batch_size, sequence_length}); - // x_tensor (FLAGS_batch_size, FLAGS_sequence_length) eq:(4, 64) - infini_train::Tensor x_tensor = infini_train::Tensor(dims, DataType::kINT64); - int64_t *x_buff = static_cast(x_tensor.DataPtr()); - for (int i = 0; i < batch_size * sequence_length; ++i) { x_buff[i] = eot_token_; } - - // Give some contexts: "The meaning of life is " - auto prompt = kPromptMap.at(magic_number_); - auto prompt_len = prompt.size(); - for (int i = 0; i < prompt_len; ++i) { x_buff[i] = prompt[i]; } - std::cout << "The meaning of life is"; - - auto x = std::make_shared(x_tensor.To(device)); - uint64_t kRngState = kRngState; - LOG(INFO) << "start generate text:"; - for (int t = prompt_len; t < text_length; t++) { - /* ===================================== 作业 ===================================== - TODO：实现单步文本生成逻辑 - HINT：调用model.Forward推理获取logits，根据推理结果进行随机采样，调用Decode获取文本结果 - ===================================== 作业 ===================================== */ - } - std::cout << std::endl; -} -} // namespace infini_train +#include "example/common/tokenizer.h" + +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/nn/functional.h" + +namespace infini_train { + +constexpr uint32_t kGpt2Eot = 50256; +constexpr uint32_t kLLaMA3Eot = 128001; +constexpr uint64_t kRandomU32Multiplier = 0x2545F4914F6CDD1Dull; +constexpr float kF32Divisor = 16777216.0f; // 2^24 +constexpr uint64_t kRngState = 1337; + +using Version = Tokenizer::Version; + +const std::unordered_map kEotMap = { + {20240328, kGpt2Eot}, // GPT-2 + {20240801, kLLaMA3Eot}, // LLaMA-3 +}; + +const std::unordered_map> kPromptMap = { + // e.g. "The meaning of life is" + // ref: https://tiktokenizer.vercel.app/ + {20240328, std::vector{464, 3616, 286, 1204, 318}}, // GPT-2 + {20240801, std::vector{791, 7438, 315, 2324, 374}}, // LLaMA-3 +}; + +std::vector ReadSeveralBytesFromIfstream(size_t num_bytes, std::ifstream *ifs) { + std::vector result(num_bytes); + ifs->read(reinterpret_cast(result.data()), num_bytes); + return result; +} + +template T BytesToType(const std::vector &bytes, size_t offset) { + static_assert(std::is_trivially_copyable::value, "T must be trivially copyable."); + T value; + std::memcpy(&value, &bytes[offset], sizeof(T)); + return value; +} + +unsigned int RandomU32(uint64_t &state) { + state ^= state >> 12; + state ^= state << 25; + state ^= state >> 27; + return (state * kRandomU32Multiplier) >> 32; +} + +float RandomF32(uint64_t &state) { // random float32 in [0,1) + return (RandomU32(state) >> 8) / kF32Divisor; +} + +int SampleMult(float *probabilities, int n, float coin) { + // sample index from probabilities (they must sum to 1!) + // coin is a random number in [0, 1), usually from RandomF32() + float cdf = 0.0f; + for (int i = 0; i < n; i++) { + cdf += probabilities[i]; + if (coin < cdf) { + return i; + } + } + return n - 1; // in case of rounding errors +} + +Tokenizer::Tokenizer(const std::string &filepath) { + /* ===================================== 作业 ===================================== + TODO：实现Tokenizer二进制文件加载 + + 文件格式说明： + ---------------------------------------------------------------------------------- + | HEADER (1024 bytes) | VOCAB TABLE | + | magic(4B) | version(4B) | vocab_size(4B) | reserved(1012B) | token词表数据 | + ---------------------------------------------------------------------------------- + ===================================== 作业 ===================================== */ + + std::ifstream ifs(filepath, std::ios::binary); + CHECK(ifs.is_open()) << "Failed to open tokenizer file: " << filepath; + + // Read header (1024 bytes) + auto header = ReadSeveralBytesFromIfstream(1024, &ifs); + + magic_number_ = BytesToType(header, 0); + uint32_t version = BytesToType(header, 4); + vocab_size_ = BytesToType(header, 8); + + CHECK(kEotMap.contains(magic_number_)) << "Unknown magic number: " << magic_number_; + eot_token_ = kEotMap.at(magic_number_); + + // Read vocab table + token_table_.resize(vocab_size_); + constexpr size_t kMaxTokenLength = 1024; + for (uint32_t i = 0; i < vocab_size_; ++i) { + auto len_bytes = ReadSeveralBytesFromIfstream(sizeof(uint8_t), &ifs); + uint32_t token_len = BytesToType(len_bytes, 0); + CHECK_LE(token_len, kMaxTokenLength) << "Token too long: " << token_len; + auto token_bytes = ReadSeveralBytesFromIfstream(token_len, &ifs); + token_table_[i] = std::string(reinterpret_cast(token_bytes.data()), token_len); + // Strip trailing whitespace + while (!token_table_[i].empty() && std::isspace(token_table_[i].back())) { + token_table_[i].pop_back(); + } + } +} + +std::string Tokenizer::Decode(uint32_t token_id) const { + /* ===================================== 作业 ===================================== + TODO：实现token_id到文本的转换 + 功能描述：根据token_id返回对应的文本片段 + ===================================== 作业 ===================================== */ + CHECK_LT(token_id, token_table_.size()) << "Token ID out of range: " << token_id; + return token_table_[token_id]; +} + +void Tokenizer::GenerateText(infini_train::nn::Module &model, uint32_t batch_size, uint32_t sequence_length, + uint32_t text_length, Device device) const { + std::vector dims; + dims.assign({batch_size, sequence_length}); + // x_tensor (FLAGS_batch_size, FLAGS_sequence_length) eq:(4, 64) + infini_train::Tensor x_tensor = infini_train::Tensor(dims, DataType::kINT64); + int64_t *x_buff = static_cast(x_tensor.DataPtr()); + for (int i = 0; i < batch_size * sequence_length; ++i) { x_buff[i] = eot_token_; } + + // Give some contexts: "The meaning of life is " + auto prompt = kPromptMap.at(magic_number_); + auto prompt_len = prompt.size(); + for (int i = 0; i < prompt_len; ++i) { x_buff[i] = prompt[i]; } + std::cout << "The meaning of life is"; + + auto x = std::make_shared(x_tensor.To(device)); + uint64_t kRngState = kRngState; + LOG(INFO) << "start generate text:"; + { + autograd::NoGradGuard no_grad; + for (int t = prompt_len; t < text_length; t++) { + /* ===================================== 作业 ===================================== + TODO：实现单步文本生成逻辑 + HINT：调用model.Forward推理获取logits，根据推理结果进行随机采样，调用Decode获取文本结果 + ===================================== 作业 ===================================== */ + + // Forward pass + auto logits = model.Forward({x})[0]; + + // Get logits for the last token position (position t-1) + auto last_logits = logits->Slice(1, t - 1, t, 1); + // Squeeze the sequence dimension: (batch_size, 1, vocab_size) -> (batch_size, vocab_size) + last_logits = last_logits->Squeeze(1); + + // Apply softmax to get probabilities + auto probs = infini_train::nn::function::Softmax(last_logits, -1); + + // Sample from the probability distribution for each batch + auto probs_cpu = probs->To(Device(DeviceType::kCPU, 0)); + float *probs_ptr = static_cast(probs_cpu.DataPtr()); + int64_t vocab_size = probs_cpu.Dims().back(); + + for (int b = 0; b < batch_size; ++b) { + float coin = RandomF32(kRngState); + int next_token = SampleMult(probs_ptr + b * vocab_size, vocab_size, coin); + x_buff[b * sequence_length + t] = next_token; + std::cout << Decode(next_token); + } + std::cout.flush(); + + // Update the device tensor with the new token + x = std::make_shared(x_tensor.To(device)); + } + } + std::cout << std::endl; +} +} // namespace infini_train diff --git a/example/common/tokenizer.h b/example/common/tokenizer.h index 1da539f..9f39c64 100644 --- a/example/common/tokenizer.h +++ b/example/common/tokenizer.h @@ -1,35 +1,35 @@ -#include -#include -#include -#include - -#include "infini_train/include/device.h" -#include "infini_train/include/nn/functional.h" -#include "infini_train/include/nn/modules/module.h" -#include "infini_train/include/tensor.h" - -namespace infini_train { - -class Tokenizer { -public: - enum class Version : uint32_t { - kV1 = 1, - kV2 = 2, - }; - - Tokenizer(const std::string &filepath); - - std::string Decode(uint32_t token_id) const; - - void GenerateText(infini_train::nn::Module &model, uint32_t batch_size, uint32_t sequence_length, - uint32_t text_length, Device device) const; - - uint32_t GetEndToken() const { return eot_token_; }; - -private: - uint32_t magic_number_ = 0; - uint32_t vocab_size_ = 0; - std::vector token_table_; - uint32_t eot_token_ = 0; -}; -} // namespace infini_train +#include +#include +#include +#include + +#include "infini_train/include/device.h" +#include "infini_train/include/nn/functional.h" +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/tensor.h" + +namespace infini_train { + +class Tokenizer { +public: + enum class Version : uint32_t { + kV1 = 1, + kV2 = 2, + }; + + Tokenizer(const std::string &filepath); + + std::string Decode(uint32_t token_id) const; + + void GenerateText(infini_train::nn::Module &model, uint32_t batch_size, uint32_t sequence_length, + uint32_t text_length, Device device) const; + + uint32_t GetEndToken() const { return eot_token_; }; + +private: + uint32_t magic_number_ = 0; + uint32_t vocab_size_ = 0; + std::vector token_table_; + uint32_t eot_token_ = 0; +}; +} // namespace infini_train diff --git a/example/gpt2/net.cc b/example/gpt2/net.cc index 441b121..8e614ea 100644 --- a/example/gpt2/net.cc +++ b/example/gpt2/net.cc @@ -1,377 +1,386 @@ -#include "example/gpt2/net.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/device.h" -#include "infini_train/include/nn/functional.h" -#include "infini_train/include/nn/init.h" -#include "infini_train/include/nn/modules/container.h" -#include "infini_train/include/nn/modules/linear.h" -#include "infini_train/include/nn/modules/module.h" -#include "infini_train/include/nn/modules/normalization.h" -#include "infini_train/include/nn/modules/sparse.h" -#include "infini_train/include/tensor.h" - -namespace nn = infini_train::nn; - -namespace { -constexpr int kRandomSeed = 42; - -class GPT2Linear : public nn::Linear { -public: - GPT2Linear(int64_t in_features, int64_t out_features, bool residual_scale = false, bool skip_init = false) - : Linear(in_features, out_features), residual_scale_(residual_scale), skip_init_(skip_init) {} - - bool residual_scale() const { return residual_scale_; } - bool skip_init() const { return skip_init_; } - -private: - const bool residual_scale_ = false; - const bool skip_init_ = false; -}; - -static std::mt19937 gen{kRandomSeed}; -} // namespace - -std::vector> -NewGELU::Forward(const std::vector> &x) { - auto &input = x[0]; - return {0.5 * input - * (1.0 + nn::function::Tanh(std::sqrt(2.0 / M_PI) * (input + 0.044715 * nn::function::Pow(input, 3.0))))}; -} - -CausalSelfAttention::CausalSelfAttention(const GPT2Config &config) - : config_(config), n_head_(config.n_head), n_embd_(config.n_embd) { - CHECK_EQ(config.n_embd % config.n_head, 0); - modules_[kCAttnLayerName] = std::make_unique(config.n_embd, config.n_embd * 3); - modules_[kCProjLayerName] = std::make_unique(config.n_embd, config.n_embd, true, false); - // (1, 1, block_size, block_size) - bias_ = nn::function::Tril(nn::function::Ones({config_.block_size, config_.block_size})) - ->View({1, 1, config_.block_size, config_.block_size}); -} - -void CausalSelfAttention::To(infini_train::Device device) { - nn::Module::To(device); - bias_ = std::make_shared(bias_->To(device)); -} - -std::vector> -CausalSelfAttention::Forward(const std::vector> &x) { - const auto B = x[0]->Dims()[0]; // bs - const auto T = x[0]->Dims()[1]; // seq_len - const auto C = x[0]->Dims()[2]; // n_embd - - // calculate query, key, values for all heads in batch and move head forward to be the batch dim - // (bs, seq_len, n_embd) -> Linear(n_embd, 3 * n_embd) -> (bs, seq_len, 3 * n_embd) - // -> Split -> (3, bs, seq_len, n_embd) - auto qkv = modules_[kCAttnLayerName]->Forward(x)[0]->Split(n_embd_, 2); - // (bs, seq_len, n_embd) - auto q = qkv[0]; - // (bs, seq_len, n_embd) - auto k = qkv[1]; - // (bs, seq_len, n_embd) - auto v = qkv[2]; - // (bs, seq_len, n_embd) -> (bs, seq_len, n_head, n_embd / n_head) -> (bs, n_head, seq_len, n_embd / n_head) - k = k->View({B, T, n_head_, C / n_head_})->Transpose(1, 2); - // (bs, seq_len, n_embd) -> (bs, seq_len, n_head, n_embd / n_head) -> (bs, n_head, seq_len, n_embd / n_head) - q = q->View({B, T, n_head_, C / n_head_})->Transpose(1, 2); - // (bs, seq_len, n_embd) -> (bs, seq_len, n_head, n_embd / n_head) -> (bs, n_head, seq_len, n_embd / n_head) - v = v->View({B, T, n_head_, C / n_head_})->Transpose(1, 2); - - // manual implementation of attention - // this materializes the large (T,T) matrix for all the queries and keys - - // q: (bs, n_head, seq_len, n_embd / n_head) - // k: (bs, n_head, seq_len, n_embd / n_head) -> (bs, n_head, n_embd / n_head, seq_len) - // q matmul k: (bs, n_head, seq_len, seq_len) -> mul 1.0 / sqrt(n_embd / n_head) -> (bs, n_head, seq_len, seq_len) - auto att = q->Matmul(k->Transpose(-2, -1)) * (1.0 / std::sqrt(*k->Dims().rbegin())); - // (1, 1, seq_len, seq_len) - auto mask = bias_->Slice({0, 0, 0, 0}, {1, 1, T, T}, {1, 1, 1, 1}); - // (1, 1, seq_len, seq_len) -> eq 0 -> (1, 1, seq_len, seq_len) -> masked_fill -> (bs, n_head, seq_len, seq_len) - att = att->MaskedFill(mask == 0, -std::numeric_limits::infinity()); - // (bs, n_head, seq_len, seq_len) - att = nn::function::Softmax(att, -1); - // (bs, n_head, seq_len, n_embd / n_head) - auto y = att->Matmul(v); - // (bs, n_head, seq_len, n_embd / n_head) -> Transpose(1, 2) -> (bs, seq_len, n_head, n_embd / n_head) - // -> (bs, seq_len, n_embd) - y = y->Transpose(1, 2)->Contiguous()->View({B, T, C}); - // output projection - // (bs, seq_len, n_embd) -> Linear(n_embd, n_embd) -> (bs, seq_len, n_embd) - y = modules_[kCProjLayerName]->Forward({y})[0]; - // (bs, seq_len, n_embd) - return {y}; -} - -MLP::MLP(const GPT2Config &config) { - modules_[kCFclayerName] = std::make_unique(config.n_embd, config.n_embd * 4); - modules_[kGeluLayerName] = std::make_unique(); - modules_[kCProjLayerName] = std::make_unique(config.n_embd * 4, config.n_embd, true, false); -} - -std::vector> -MLP::Forward(const std::vector> &x) { - // (bs, seq_len, n_embd) -> Linear(n_embd, 4 * n_embd) -> (bs, seq_len, 4 * n_embd) - auto x1 = modules_[kCFclayerName]->Forward(x); - // (bs, seq_len, 4 * n_embd) -> GELU -> (bs, seq_len, 4 * n_embd) - auto x2 = modules_[kGeluLayerName]->Forward(x1); - // (bs, seq_len, 4 * n_embd) -> Linear(4 * n_embd, n_embd) -> (bs, seq_len, n_embd) - auto x3 = modules_[kCProjLayerName]->Forward(x2); - // (bs, seq_len, n_embd) - return x3; -} - -Block::Block(const GPT2Config &config) { - modules_[kLn1LayerName] = std::make_unique(std::vector{config.n_embd}); - modules_[kAttnLayerName] = std::make_unique(config); - modules_[kLn2LayerName] = std::make_unique(std::vector{config.n_embd}); - modules_[kMlpLayerName] = std::make_unique(config); -} - -std::vector> -Block::Forward(const std::vector> &x) { - // (bs, seq_len, n_embd) -> Layernorm -> (bs, seq_len, n_embd) -> attention -> (bs, seq_len, n_embd) - // -> Add -> (bs, seq_len, n_embd) - auto x1 = x[0] + modules_[kAttnLayerName]->Forward(modules_[kLn1LayerName]->Forward(x))[0]; - // (bs, seq_len, n_embd) -> Layernorm -> (bs, seq_len, n_embd) -> MLP -> (bs, seq_len, n_embd) - // -> Add -> (bs, seq_len, n_embd) - auto x2 = x1 + modules_[kMlpLayerName]->Forward(modules_[kLn2LayerName]->Forward({x1}))[0]; - // (bs, seq_len, n_embd) - return {x2}; -} - -GPT2::GPT2(const GPT2Config &config) : config_(config) { - { - std::unordered_map> transformer; - transformer[kWTELayerName] = std::make_unique(config.vocab_size, config.n_embd); - transformer[kWPELayerName] = std::make_unique(config.block_size, config.n_embd); - { - std::vector> h; - for (int64_t i = 0; i < config.n_layer; i++) { h.push_back(std::make_unique(config)); } - transformer[kHLayerName] = std::make_unique(std::move(h)); - } - transformer[kLnFLayerName] = std::make_unique(std::vector{config.n_embd}); - modules_[kTransformerLayerName] = std::make_unique(std::move(transformer)); - } - // don't init this one, we will tie weights - modules_[kLMHeadLayerName] = std::make_unique(config.n_embd, config.vocab_size, false, true); - // https://paperswithcode.com/method/weight-tying - *mutable_module(kTransformerLayerName) - ->mutable_module(kWTELayerName) - ->mutable_parameter(GPT2Linear::kParamWeightName) - = module(kLMHeadLayerName).parameter(GPT2Linear::kParamWeightName); - - // init all weights - Apply([&](Module *module) { - if (module->type() == nn::Linear::kType) { - auto *linear = static_cast(module); - const float std = linear->residual_scale() ? 0.02 / std::sqrt(2 * config_.n_layer) : 0.02; - if (!linear->skip_init()) { - nn::init::Normal(*linear->mutable_parameter(GPT2Linear::kParamWeightName), 0.0f, std, gen); - } - if (linear->has_parameter(nn::Linear::kParamWeightName)) { - nn::init::Zeros(*linear->mutable_parameter(nn::Linear::kParamBiasName)); - } - } else if (module->type() == nn::Embedding::kType) { - nn::init::Normal(*module->mutable_parameter(nn::Embedding::kParamWeightName), 0.0f, 0.02, gen); - } - }); -} - -std::vector> -GPT2::Forward(const std::vector> &x) { - // (bs, seq_len) - auto &idx = x[0]; - const auto device = idx->GetDevice(); - const auto t = idx->Dims()[1]; // seq_len - CHECK_LE(t, config_.block_size) << "Cannot forward sequence of length " << t << ", block size is only " - << config_.block_size; - // (seq_len) - auto pos = nn::init::Arange(0, t, infini_train::DataType::kINT64, device); - - // forward the GPT2 model itself - auto &transformer = modules_[kTransformerLayerName]; - // (bs, seq_len) -> Embedding(vocab_size, n_embd) -> (bs, seq_len, n_embd) - auto tok_emb = transformer->mutable_module(kWTELayerName)->Forward({idx})[0]; - // (seq_len) -> Embedding(block_size, n_embd) -> (seq_len, n_embd) - auto pos_emb = transformer->mutable_module(kWPELayerName)->Forward({pos})[0]; - // (bs, seq_len, n_embd) - auto x1 = tok_emb + pos_emb; - - // (bs, seq_len, n_embd) -> transformer -> (bs, seq_len, n_embd) - auto x2 = transformer->mutable_module(kHLayerName)->Forward({x1}); - // (bs, seq_len, n_embd) -> Layernorm -> (bs, seq_len, n_embd) - auto x3 = transformer->mutable_module(kLnFLayerName)->Forward(x2); - - // (bs, seq_len, n_embd) -> Linear(n_embd, vocab_size) -> (bs, seq_len, vocab_size) - auto logits = modules_[kLMHeadLayerName]->Forward(x3); - - // (bs, seq_len, vocab_size) - return logits; -} - -std::unique_ptr GPT2::FromPretrained(ModelType model_type) { - LOG(FATAL) << "Not implemented yet"; - return nullptr; -} - -namespace { -std::vector ReadSeveralBytesFromIfstream(size_t num_bytes, std::ifstream *ifs) { - std::vector result(num_bytes); - ifs->read(reinterpret_cast(result.data()), num_bytes); - return result; -} - -template T BytesToType(const std::vector &bytes, size_t offset) { - static_assert(std::is_trivially_copyable::value, "T must be trivially copyable."); - T value; - std::memcpy(&value, &bytes[offset], sizeof(T)); - return value; -} - -constexpr int32_t kHeaderMagic = 20240326; -constexpr int32_t kHeaderFP32Version = 3; -} // namespace - -std::unique_ptr GPT2::FromLLMC(const std::string &filepath) { - if (!std::filesystem::exists(filepath)) { - LOG(FATAL) << "File not found: " << filepath; - } - - std::ifstream ifs(filepath, std::ios::binary); - const auto header = ReadSeveralBytesFromIfstream(256 * sizeof(int32_t), &ifs); - - const auto magic = BytesToType(header, 0); - CHECK_EQ(magic, kHeaderMagic); - const auto version = BytesToType(header, 4); - CHECK_EQ(version, kHeaderFP32Version); - - const auto block_size = BytesToType(header, 8); - const auto vocab_size = BytesToType(header, 12); - const auto n_layer = BytesToType(header, 16); - const auto n_head = BytesToType(header, 20); - const auto n_embd = BytesToType(header, 24); - auto gpt2 = std::make_unique(GPT2Config{ - .block_size = block_size, .vocab_size = vocab_size, .n_layer = n_layer, .n_head = n_head, .n_embd = n_embd}); - - const auto padded_vocab_size = BytesToType(header, 28); - LOG(ERROR) << "magic: " << magic << " version: " << version << " block_size: " << block_size - << " vocab_size: " << vocab_size << " n_layer: " << n_layer << " n_head: " << n_head - << " n_embd: " << n_embd << " padded_vocab_size: " << padded_vocab_size; - - auto state_dict = gpt2->StateDict(); - // transformer.wte.weight - // (padded_vocab_size, n_embd) -> un_pad -> (vocab_size, n_embd) - auto &transformer_wte_weight = state_dict[std::format("{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kWTELayerName, - nn::Embedding::kParamWeightName)]; - ifs.read(reinterpret_cast(transformer_wte_weight->DataPtr()), transformer_wte_weight->SizeInBytes()); - ifs.ignore((padded_vocab_size - vocab_size) * n_embd * sizeof(float)); - // transformer.wpe.weight - auto &transformer_wpe_weight = state_dict[std::format("{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kWPELayerName, - nn::Embedding::kParamWeightName)]; - ifs.read(reinterpret_cast(transformer_wpe_weight->DataPtr()), transformer_wpe_weight->SizeInBytes()); - // transformer.h.{i}.ln_1.weight - for (int idx = 0; idx < n_layer; idx++) { - auto &tensor - = state_dict[std::format("{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, - std::to_string(idx), Block::kLn1LayerName, nn::LayerNorm::kParamWeightName)]; - ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); - } - // transformer.h.{i}.ln_1.bias - for (int idx = 0; idx < n_layer; idx++) { - auto &tensor - = state_dict[std::format("{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, - std::to_string(idx), Block::kLn1LayerName, nn::LayerNorm::kParamBiasName)]; - ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); - } - // transformer.h.{i}.attn.c_attn.weight - for (int idx = 0; idx < n_layer; idx++) { - auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, - std::to_string(idx), Block::kAttnLayerName, - CausalSelfAttention::kCAttnLayerName, GPT2Linear::kParamWeightName)]; - ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); - } - // transformer.h.{i}.attn.c_attn.bias - for (int idx = 0; idx < n_layer; idx++) { - auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, - std::to_string(idx), Block::kAttnLayerName, - CausalSelfAttention::kCAttnLayerName, GPT2Linear::kParamBiasName)]; - ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); - } - // transformer.h.{i}.attn.c_proj.weight - for (int idx = 0; idx < n_layer; idx++) { - auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, - std::to_string(idx), Block::kAttnLayerName, - CausalSelfAttention::kCProjLayerName, GPT2Linear::kParamWeightName)]; - ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); - } - // transformer.h.{i}.attn.c_proj.bias - for (int idx = 0; idx < n_layer; idx++) { - auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, - std::to_string(idx), Block::kAttnLayerName, - CausalSelfAttention::kCProjLayerName, GPT2Linear::kParamBiasName)]; - ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); - } - // transformer.h.{i}.ln_2.weight - for (int idx = 0; idx < n_layer; idx++) { - auto &tensor - = state_dict[std::format("{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, - std::to_string(idx), Block::kLn2LayerName, nn::LayerNorm::kParamWeightName)]; - ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); - } - // transformer.h.{i}.ln_2.bias - for (int idx = 0; idx < n_layer; idx++) { - auto &tensor - = state_dict[std::format("{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, - std::to_string(idx), Block::kLn2LayerName, nn::LayerNorm::kParamBiasName)]; - ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); - } - // transformer.h.{i}.mlp.c_fc.weight - for (int idx = 0; idx < n_layer; idx++) { - auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, - std::to_string(idx), Block::kMlpLayerName, MLP::kCFclayerName, - GPT2Linear::kParamWeightName)]; - ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); - } - // transformer.h.{i}.mlp.c_fc.bias - for (int idx = 0; idx < n_layer; idx++) { - auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, - std::to_string(idx), Block::kMlpLayerName, MLP::kCFclayerName, - GPT2Linear::kParamBiasName)]; - ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); - } - // transformer.h.{i}.mlp.c_proj.weight - for (int idx = 0; idx < n_layer; idx++) { - auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, - std::to_string(idx), Block::kMlpLayerName, MLP::kCProjLayerName, - GPT2Linear::kParamWeightName)]; - ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); - } - // transformer.h.{i}.mlp.c_proj.bias - for (int idx = 0; idx < n_layer; idx++) { - auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, - std::to_string(idx), Block::kMlpLayerName, MLP::kCProjLayerName, - GPT2Linear::kParamBiasName)]; - ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); - } - // transformer.ln_f.weight - auto &transformer_ln_f_weight = state_dict[std::format("{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kLnFLayerName, - nn::LayerNorm::kParamWeightName)]; - ifs.read(reinterpret_cast(transformer_ln_f_weight->DataPtr()), transformer_ln_f_weight->SizeInBytes()); - // transformer.ln_f.bias - auto &transformer_ln_f_bias = state_dict[std::format("{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kLnFLayerName, - nn::LayerNorm::kParamBiasName)]; - ifs.read(reinterpret_cast(transformer_ln_f_bias->DataPtr()), transformer_ln_f_bias->SizeInBytes()); - - return gpt2; -} +#include "example/gpt2/net.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/device.h" +#include "infini_train/include/nn/functional.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/nn/modules/container.h" +#include "infini_train/include/nn/modules/linear.h" +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/nn/modules/normalization.h" +#include "infini_train/include/nn/modules/sparse.h" +#include "infini_train/include/tensor.h" + +namespace nn = infini_train::nn; + +namespace { +constexpr int kRandomSeed = 42; + +class GPT2Linear : public nn::Linear { +public: + GPT2Linear(int64_t in_features, int64_t out_features, bool residual_scale = false, bool skip_init = false) + : Linear(in_features, out_features), residual_scale_(residual_scale), skip_init_(skip_init) {} + + bool residual_scale() const { return residual_scale_; } + bool skip_init() const { return skip_init_; } + +private: + const bool residual_scale_ = false; + const bool skip_init_ = false; +}; + +static std::mt19937 gen{kRandomSeed}; +} // namespace + +std::vector> +NewGELU::Forward(const std::vector> &x) { + auto &input = x[0]; + return {0.5 * input + * (1.0 + nn::function::Tanh(std::sqrt(2.0 / M_PI) * (input + 0.044715 * nn::function::Pow(input, 3.0))))}; +} + +CausalSelfAttention::CausalSelfAttention(const GPT2Config &config) + : config_(config), n_head_(config.n_head), n_embd_(config.n_embd) { + CHECK_EQ(config.n_embd % config.n_head, 0); + modules_[kCAttnLayerName] = std::make_unique(config.n_embd, config.n_embd * 3); + modules_[kCProjLayerName] = std::make_unique(config.n_embd, config.n_embd, true, false); + // (1, 1, block_size, block_size) + bias_ = nn::function::Tril(nn::function::Ones({config_.block_size, config_.block_size})) + ->View({1, 1, config_.block_size, config_.block_size}); +} + +void CausalSelfAttention::To(infini_train::Device device) { + nn::Module::To(device); + bias_ = std::make_shared(bias_->To(device)); +} + +std::vector> +CausalSelfAttention::Forward(const std::vector> &x) { + const auto B = x[0]->Dims()[0]; // bs + const auto T = x[0]->Dims()[1]; // seq_len + const auto C = x[0]->Dims()[2]; // n_embd + + // calculate query, key, values for all heads in batch and move head forward to be the batch dim + // (bs, seq_len, n_embd) -> Linear(n_embd, 3 * n_embd) -> (bs, seq_len, 3 * n_embd) + // -> Split -> (3, bs, seq_len, n_embd) + auto qkv = modules_[kCAttnLayerName]->Forward(x)[0]->Split(n_embd_, 2); + // (bs, seq_len, n_embd) + auto q = qkv[0]; + // (bs, seq_len, n_embd) + auto k = qkv[1]; + // (bs, seq_len, n_embd) + auto v = qkv[2]; + // (bs, seq_len, n_embd) -> (bs, seq_len, n_head, n_embd / n_head) -> (bs, n_head, seq_len, n_embd / n_head) + k = k->View({B, T, n_head_, C / n_head_})->Transpose(1, 2); + // (bs, seq_len, n_embd) -> (bs, seq_len, n_head, n_embd / n_head) -> (bs, n_head, seq_len, n_embd / n_head) + q = q->View({B, T, n_head_, C / n_head_})->Transpose(1, 2); + // (bs, seq_len, n_embd) -> (bs, seq_len, n_head, n_embd / n_head) -> (bs, n_head, seq_len, n_embd / n_head) + v = v->View({B, T, n_head_, C / n_head_})->Transpose(1, 2); + + // manual implementation of attention + // this materializes the large (T,T) matrix for all the queries and keys + + // q: (bs, n_head, seq_len, n_embd / n_head) + // k: (bs, n_head, seq_len, n_embd / n_head) -> (bs, n_head, n_embd / n_head, seq_len) + // q matmul k: (bs, n_head, seq_len, seq_len) -> mul 1.0 / sqrt(n_embd / n_head) -> (bs, n_head, seq_len, seq_len) + auto att = q->Matmul(k->Transpose(-2, -1)) * (1.0 / std::sqrt(*k->Dims().rbegin())); + // (1, 1, seq_len, seq_len) + auto mask = bias_->Slice({0, 0, 0, 0}, {1, 1, T, T}, {1, 1, 1, 1}); + // (1, 1, seq_len, seq_len) -> eq 0 -> (1, 1, seq_len, seq_len) -> masked_fill -> (bs, n_head, seq_len, seq_len) + att = att->MaskedFill(mask == 0, -std::numeric_limits::infinity()); + // (bs, n_head, seq_len, seq_len) + att = nn::function::Softmax(att, -1); + // (bs, n_head, seq_len, n_embd / n_head) + auto y = att->Matmul(v); + // (bs, n_head, seq_len, n_embd / n_head) -> Transpose(1, 2) -> (bs, seq_len, n_head, n_embd / n_head) + // -> (bs, seq_len, n_embd) + y = y->Transpose(1, 2)->Contiguous()->View({B, T, C}); + // output projection + // (bs, seq_len, n_embd) -> Linear(n_embd, n_embd) -> (bs, seq_len, n_embd) + y = modules_[kCProjLayerName]->Forward({y})[0]; + // (bs, seq_len, n_embd) + return {y}; +} + +MLP::MLP(const GPT2Config &config) { + modules_[kCFclayerName] = std::make_unique(config.n_embd, config.n_embd * 4); + modules_[kGeluLayerName] = std::make_unique(); + modules_[kCProjLayerName] = std::make_unique(config.n_embd * 4, config.n_embd, true, false); +} + +std::vector> +MLP::Forward(const std::vector> &x) { + // (bs, seq_len, n_embd) -> Linear(n_embd, 4 * n_embd) -> (bs, seq_len, 4 * n_embd) + auto x1 = modules_[kCFclayerName]->Forward(x); + // (bs, seq_len, 4 * n_embd) -> GELU -> (bs, seq_len, 4 * n_embd) + auto x2 = modules_[kGeluLayerName]->Forward(x1); + // (bs, seq_len, 4 * n_embd) -> Linear(4 * n_embd, n_embd) -> (bs, seq_len, n_embd) + auto x3 = modules_[kCProjLayerName]->Forward(x2); + // (bs, seq_len, n_embd) + return x3; +} + +Block::Block(const GPT2Config &config) { + modules_[kLn1LayerName] = std::make_unique(std::vector{config.n_embd}); + modules_[kAttnLayerName] = std::make_unique(config); + modules_[kLn2LayerName] = std::make_unique(std::vector{config.n_embd}); + modules_[kMlpLayerName] = std::make_unique(config); +} + +std::vector> +Block::Forward(const std::vector> &x) { + // (bs, seq_len, n_embd) -> Layernorm -> (bs, seq_len, n_embd) -> attention -> (bs, seq_len, n_embd) + // -> Add -> (bs, seq_len, n_embd) + auto x1 = x[0] + modules_[kAttnLayerName]->Forward(modules_[kLn1LayerName]->Forward(x))[0]; + // (bs, seq_len, n_embd) -> Layernorm -> (bs, seq_len, n_embd) -> MLP -> (bs, seq_len, n_embd) + // -> Add -> (bs, seq_len, n_embd) + auto x2 = x1 + modules_[kMlpLayerName]->Forward(modules_[kLn2LayerName]->Forward({x1}))[0]; + // (bs, seq_len, n_embd) + return {x2}; +} + +GPT2::GPT2(const GPT2Config &config) : config_(config) { + { + std::unordered_map> transformer; + transformer[kWTELayerName] = std::make_unique(config.vocab_size, config.n_embd); + transformer[kWPELayerName] = std::make_unique(config.block_size, config.n_embd); + { + std::vector> h; + for (int64_t i = 0; i < config.n_layer; i++) { h.push_back(std::make_unique(config)); } + transformer[kHLayerName] = std::make_unique(std::move(h)); + } + transformer[kLnFLayerName] = std::make_unique(std::vector{config.n_embd}); + modules_[kTransformerLayerName] = std::make_unique(std::move(transformer)); + } + // don't init this one, we will tie weights + modules_[kLMHeadLayerName] = std::make_unique(config.n_embd, config.vocab_size, false, true); + // https://paperswithcode.com/method/weight-tying + *mutable_module(kTransformerLayerName) + ->mutable_module(kWTELayerName) + ->mutable_parameter(GPT2Linear::kParamWeightName) + = module(kLMHeadLayerName).parameter(GPT2Linear::kParamWeightName); + + // init all weights + Apply([&](Module *module) { + if (module->type() == nn::Linear::kType) { + auto *linear = static_cast(module); + const float std = linear->residual_scale() ? 0.02 / std::sqrt(2 * config_.n_layer) : 0.02; + if (!linear->skip_init()) { + nn::init::Normal(*linear->mutable_parameter(GPT2Linear::kParamWeightName), 0.0f, std, gen); + } + if (linear->has_parameter(nn::Linear::kParamWeightName)) { + nn::init::Zeros(*linear->mutable_parameter(nn::Linear::kParamBiasName)); + } + } else if (module->type() == nn::Embedding::kType) { + nn::init::Normal(*module->mutable_parameter(nn::Embedding::kParamWeightName), 0.0f, 0.02, gen); + } + }); +} + +std::vector> +GPT2::Forward(const std::vector> &x) { + // (bs, seq_len) + auto &idx = x[0]; + const auto device = idx->GetDevice(); + const auto t = idx->Dims()[1]; // seq_len + CHECK_LE(t, config_.block_size) << "Cannot forward sequence of length " << t << ", block size is only " + << config_.block_size; + // (seq_len) + auto pos = nn::init::Arange(0, t, infini_train::DataType::kINT64, device); + + // forward the GPT2 model itself + auto &transformer = modules_[kTransformerLayerName]; + // (bs, seq_len) -> Embedding(vocab_size, n_embd) -> (bs, seq_len, n_embd) + auto tok_emb = transformer->mutable_module(kWTELayerName)->Forward({idx})[0]; + // (seq_len) -> Embedding(block_size, n_embd) -> (seq_len, n_embd) + auto pos_emb = transformer->mutable_module(kWPELayerName)->Forward({pos})[0]; + // (bs, seq_len, n_embd) + auto x1 = tok_emb + pos_emb; + + // (bs, seq_len, n_embd) -> transformer -> (bs, seq_len, n_embd) + auto x2 = transformer->mutable_module(kHLayerName)->Forward({x1}); + // (bs, seq_len, n_embd) -> Layernorm -> (bs, seq_len, n_embd) + auto x3 = transformer->mutable_module(kLnFLayerName)->Forward(x2); + + // (bs, seq_len, n_embd) -> Linear(n_embd, vocab_size) -> (bs, seq_len, vocab_size) + auto logits = modules_[kLMHeadLayerName]->Forward(x3); + + // (bs, seq_len, vocab_size) + return logits; +} + +std::unique_ptr GPT2::FromPretrained(ModelType model_type) { + LOG(FATAL) << "Not implemented yet"; + return nullptr; +} + +namespace { +std::vector ReadSeveralBytesFromIfstream(size_t num_bytes, std::ifstream *ifs) { + std::vector result(num_bytes); + ifs->read(reinterpret_cast(result.data()), num_bytes); + return result; +} + +template T BytesToType(const std::vector &bytes, size_t offset) { + static_assert(std::is_trivially_copyable::value, "T must be trivially copyable."); + T value; + std::memcpy(&value, &bytes[offset], sizeof(T)); + return value; +} + +constexpr int32_t kHeaderMagic = 20240326; +constexpr int32_t kHeaderFP32Version = 3; +} // namespace + +std::unique_ptr GPT2::FromLLMC(const std::string &filepath) { + if (!std::filesystem::exists(filepath)) { + LOG(FATAL) << "File not found: " << filepath; + } + + std::ifstream ifs(filepath, std::ios::binary); + const auto header = ReadSeveralBytesFromIfstream(256 * sizeof(int32_t), &ifs); + + const auto magic = BytesToType(header, 0); + CHECK_EQ(magic, kHeaderMagic); + const auto version = BytesToType(header, 4); + CHECK_EQ(version, kHeaderFP32Version); + + const auto block_size = BytesToType(header, 8); + const auto vocab_size = BytesToType(header, 12); + const auto n_layer = BytesToType(header, 16); + const auto n_head = BytesToType(header, 20); + const auto n_embd = BytesToType(header, 24); + auto gpt2 = std::make_unique(GPT2Config{ + .block_size = block_size, .vocab_size = vocab_size, .n_layer = n_layer, .n_head = n_head, .n_embd = n_embd}); + + const auto padded_vocab_size = BytesToType(header, 28); + + auto state_dict = gpt2->StateDict(); + // transformer.wte.weight + // (padded_vocab_size, n_embd) -> un_pad -> (vocab_size, n_embd) + auto key = std::format("{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kWTELayerName, + nn::Embedding::kParamWeightName); + auto &transformer_wte_weight = state_dict.at(key); + ifs.read(reinterpret_cast(transformer_wte_weight->DataPtr()), transformer_wte_weight->SizeInBytes()); + // Use a small stack buffer to skip the padded part instead of ifs.ignore() + { + const std::streamsize skip_size = (padded_vocab_size - vocab_size) * n_embd * sizeof(float); + char skip_buf[4096]; + std::streamsize remaining = skip_size; + while (remaining > 0) { + std::streamsize to_read = std::min(remaining, static_cast(sizeof(skip_buf))); + ifs.read(skip_buf, to_read); + remaining -= to_read; + } + } + // transformer.wpe.weight + auto &transformer_wpe_weight = state_dict[std::format("{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kWPELayerName, + nn::Embedding::kParamWeightName)]; + ifs.read(reinterpret_cast(transformer_wpe_weight->DataPtr()), transformer_wpe_weight->SizeInBytes()); + // transformer.h.{i}.ln_1.weight + for (int idx = 0; idx < n_layer; idx++) { + auto &tensor + = state_dict[std::format("{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, + std::to_string(idx), Block::kLn1LayerName, nn::LayerNorm::kParamWeightName)]; + ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); + } + // transformer.h.{i}.ln_1.bias + for (int idx = 0; idx < n_layer; idx++) { + auto &tensor + = state_dict[std::format("{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, + std::to_string(idx), Block::kLn1LayerName, nn::LayerNorm::kParamBiasName)]; + ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); + } + // transformer.h.{i}.attn.c_attn.weight + for (int idx = 0; idx < n_layer; idx++) { + auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, + std::to_string(idx), Block::kAttnLayerName, + CausalSelfAttention::kCAttnLayerName, GPT2Linear::kParamWeightName)]; + ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); + } + // transformer.h.{i}.attn.c_attn.bias + for (int idx = 0; idx < n_layer; idx++) { + auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, + std::to_string(idx), Block::kAttnLayerName, + CausalSelfAttention::kCAttnLayerName, GPT2Linear::kParamBiasName)]; + ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); + } + // transformer.h.{i}.attn.c_proj.weight + for (int idx = 0; idx < n_layer; idx++) { + auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, + std::to_string(idx), Block::kAttnLayerName, + CausalSelfAttention::kCProjLayerName, GPT2Linear::kParamWeightName)]; + ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); + } + // transformer.h.{i}.attn.c_proj.bias + for (int idx = 0; idx < n_layer; idx++) { + auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, + std::to_string(idx), Block::kAttnLayerName, + CausalSelfAttention::kCProjLayerName, GPT2Linear::kParamBiasName)]; + ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); + } + // transformer.h.{i}.ln_2.weight + for (int idx = 0; idx < n_layer; idx++) { + auto &tensor + = state_dict[std::format("{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, + std::to_string(idx), Block::kLn2LayerName, nn::LayerNorm::kParamWeightName)]; + ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); + } + // transformer.h.{i}.ln_2.bias + for (int idx = 0; idx < n_layer; idx++) { + auto &tensor + = state_dict[std::format("{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, + std::to_string(idx), Block::kLn2LayerName, nn::LayerNorm::kParamBiasName)]; + ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); + } + // transformer.h.{i}.mlp.c_fc.weight + for (int idx = 0; idx < n_layer; idx++) { + auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, + std::to_string(idx), Block::kMlpLayerName, MLP::kCFclayerName, + GPT2Linear::kParamWeightName)]; + ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); + } + // transformer.h.{i}.mlp.c_fc.bias + for (int idx = 0; idx < n_layer; idx++) { + auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, + std::to_string(idx), Block::kMlpLayerName, MLP::kCFclayerName, + GPT2Linear::kParamBiasName)]; + ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); + } + // transformer.h.{i}.mlp.c_proj.weight + for (int idx = 0; idx < n_layer; idx++) { + auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, + std::to_string(idx), Block::kMlpLayerName, MLP::kCProjLayerName, + GPT2Linear::kParamWeightName)]; + ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); + } + // transformer.h.{i}.mlp.c_proj.bias + for (int idx = 0; idx < n_layer; idx++) { + auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kHLayerName, + std::to_string(idx), Block::kMlpLayerName, MLP::kCProjLayerName, + GPT2Linear::kParamBiasName)]; + ifs.read(reinterpret_cast(tensor->DataPtr()), tensor->SizeInBytes()); + } + // transformer.ln_f.weight + auto &transformer_ln_f_weight = state_dict[std::format("{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kLnFLayerName, + nn::LayerNorm::kParamWeightName)]; + ifs.read(reinterpret_cast(transformer_ln_f_weight->DataPtr()), transformer_ln_f_weight->SizeInBytes()); + // transformer.ln_f.bias + auto &transformer_ln_f_bias = state_dict[std::format("{}.{}.{}", GPT2::kTransformerLayerName, GPT2::kLnFLayerName, + nn::LayerNorm::kParamBiasName)]; + ifs.read(reinterpret_cast(transformer_ln_f_bias->DataPtr()), transformer_ln_f_bias->SizeInBytes()); + + return gpt2; +} diff --git a/example/gpt2/net.h b/example/gpt2/net.h index 619cbdc..13cee23 100644 --- a/example/gpt2/net.h +++ b/example/gpt2/net.h @@ -1,97 +1,97 @@ -#pragma once - -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/nn/modules/module.h" -#include "infini_train/include/tensor.h" - -struct GPT2Config { - int64_t block_size = 1024; - int64_t vocab_size = 50257; - int64_t n_layer = 12; - int64_t n_head = 12; - int64_t n_embd = 768; -}; - -class NewGELU : public infini_train::nn::Module { -public: - std::vector> - Forward(const std::vector> &x) override; -}; - -class CausalSelfAttention : public infini_train::nn::Module { -public: - static constexpr char kCAttnLayerName[] = "c_attn"; - static constexpr char kCProjLayerName[] = "c_proj"; - - explicit CausalSelfAttention(const GPT2Config &config); - - std::vector> - Forward(const std::vector> &x) override; - - void To(infini_train::Device device) override; - -private: - GPT2Config config_; - int64_t n_head_ = 0; - int64_t n_embd_ = 0; - std::shared_ptr bias_ = nullptr; -}; - -class MLP : public infini_train::nn::Module { -public: - static constexpr char kCFclayerName[] = "c_fc"; - static constexpr char kGeluLayerName[] = "gelu"; - static constexpr char kCProjLayerName[] = "c_proj"; - - explicit MLP(const GPT2Config &config); - - std::vector> - Forward(const std::vector> &x) override; -}; - -class Block : public infini_train::nn::Module { -public: - static constexpr char kLn1LayerName[] = "ln_1"; - static constexpr char kAttnLayerName[] = "attn"; - static constexpr char kLn2LayerName[] = "ln_2"; - static constexpr char kMlpLayerName[] = "mlp"; - - explicit Block(const GPT2Config &config); - - std::vector> - Forward(const std::vector> &x) override; -}; - -class GPT2 : public infini_train::nn::Module { -public: - static constexpr char kWTELayerName[] = "wte"; - static constexpr char kWPELayerName[] = "wpe"; - static constexpr char kHLayerName[] = "h"; - static constexpr char kLnFLayerName[] = "ln_f"; - static constexpr char kTransformerLayerName[] = "transformer"; - static constexpr char kLMHeadLayerName[] = "lm_head"; - - enum class ModelType : int8_t { - kGPT2, - kGPT2Medium, - kGPT2Large, - kGPT2XL, - }; - - explicit GPT2(const GPT2Config &config); - - std::vector> - Forward(const std::vector> &x) override; - - static std::unique_ptr FromPretrained(ModelType model_type); - static std::unique_ptr FromLLMC(const std::string &filepath); - -private: - GPT2Config config_; -}; +#pragma once + +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/tensor.h" + +struct GPT2Config { + int64_t block_size = 1024; + int64_t vocab_size = 50257; + int64_t n_layer = 12; + int64_t n_head = 12; + int64_t n_embd = 768; +}; + +class NewGELU : public infini_train::nn::Module { +public: + std::vector> + Forward(const std::vector> &x) override; +}; + +class CausalSelfAttention : public infini_train::nn::Module { +public: + static constexpr char kCAttnLayerName[] = "c_attn"; + static constexpr char kCProjLayerName[] = "c_proj"; + + explicit CausalSelfAttention(const GPT2Config &config); + + std::vector> + Forward(const std::vector> &x) override; + + void To(infini_train::Device device) override; + +private: + GPT2Config config_; + int64_t n_head_ = 0; + int64_t n_embd_ = 0; + std::shared_ptr bias_ = nullptr; +}; + +class MLP : public infini_train::nn::Module { +public: + static constexpr char kCFclayerName[] = "c_fc"; + static constexpr char kGeluLayerName[] = "gelu"; + static constexpr char kCProjLayerName[] = "c_proj"; + + explicit MLP(const GPT2Config &config); + + std::vector> + Forward(const std::vector> &x) override; +}; + +class Block : public infini_train::nn::Module { +public: + static constexpr char kLn1LayerName[] = "ln_1"; + static constexpr char kAttnLayerName[] = "attn"; + static constexpr char kLn2LayerName[] = "ln_2"; + static constexpr char kMlpLayerName[] = "mlp"; + + explicit Block(const GPT2Config &config); + + std::vector> + Forward(const std::vector> &x) override; +}; + +class GPT2 : public infini_train::nn::Module { +public: + static constexpr char kWTELayerName[] = "wte"; + static constexpr char kWPELayerName[] = "wpe"; + static constexpr char kHLayerName[] = "h"; + static constexpr char kLnFLayerName[] = "ln_f"; + static constexpr char kTransformerLayerName[] = "transformer"; + static constexpr char kLMHeadLayerName[] = "lm_head"; + + enum class ModelType : int8_t { + kGPT2, + kGPT2Medium, + kGPT2Large, + kGPT2XL, + }; + + explicit GPT2(const GPT2Config &config); + + std::vector> + Forward(const std::vector> &x) override; + + static std::unique_ptr FromPretrained(ModelType model_type); + static std::unique_ptr FromLLMC(const std::string &filepath); + +private: + GPT2Config config_; +}; diff --git a/infini_train/include/autograd/activations.h b/infini_train/include/autograd/activations.h index 952acc4..b7e3aae 100644 --- a/infini_train/include/autograd/activations.h +++ b/infini_train/include/autograd/activations.h @@ -1,21 +1,21 @@ -#pragma once - -#include -#include - -#include "infini_train/include/autograd/function.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class Sigmoid : public Function { -public: - static constexpr char kType[] = "SigmoidFunction"; - - Sigmoid() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; -}; -} // namespace infini_train::autograd +#pragma once + +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +class Sigmoid : public Function { +public: + static constexpr char kType[] = "SigmoidFunction"; + + Sigmoid() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/elementwise.h b/infini_train/include/autograd/elementwise.h index 9f5a3e6..8fa73e0 100644 --- a/infini_train/include/autograd/elementwise.h +++ b/infini_train/include/autograd/elementwise.h @@ -1,192 +1,192 @@ -#pragma once - -#include -#include - -#include "infini_train/include/autograd/function.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class Neg : public Function { -public: - static constexpr char kType[] = "NegFunction"; - - Neg() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; -}; - -class Reciprocal : public Function { -public: - static constexpr char kType[] = "ReciprocalFunction"; - - Reciprocal() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; -}; - -class Sin : public Function { -public: - static constexpr char kType[] = "SinFunction"; - - Sin() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; -}; - -class Cos : public Function { -public: - static constexpr char kType[] = "CosFunction"; - - Cos() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; -}; - -class Tanh : public Function { -public: - static constexpr char kType[] = "TanhFunction"; - - Tanh() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; -}; - -class Pow : public Function { -public: - static constexpr char kType[] = "PowFunction"; - - explicit Pow(float exponent, bool scalar_is_base = false) - : Function(kType), exponent_(exponent), scalar_is_base_(scalar_is_base) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - const float exponent_ = 0.0f; - const bool scalar_is_base_ = false; -}; - -class Rsqrt : public Function { -public: - static constexpr char kType[] = "RsqrtFunction"; - - explicit Rsqrt() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; -}; - -class EqualsScalar : public Function { -public: - static constexpr char kType[] = "EqualsScalarFunction"; - - explicit EqualsScalar(float scalar) : Function(kType), scalar_(scalar) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - const float scalar_ = 0.0f; -}; - -class Add : public Function { -public: - static constexpr char kType[] = "AddFunction"; - - explicit Add() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - std::vector a_dims_; - std::vector b_dims_; -}; - -class AddScalar : public Function { -public: - static constexpr char kType[] = "AddScalarFunction"; - - explicit AddScalar(float scalar) : Function(kType), scalar_(scalar) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - const float scalar_ = 0.0f; -}; - -class Sub : public Function { -public: - static constexpr char kType[] = "SubFunction"; - - explicit Sub() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - std::vector a_dims_; - std::vector b_dims_; -}; - -class Mul : public Function { -public: - static constexpr char kType[] = "MulFunction"; - - explicit Mul() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; -}; - -class MulScalar : public Function { -public: - static constexpr char kType[] = "MulScalarFunction"; - - explicit MulScalar(float scalar) : Function(kType), scalar_(scalar) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - const float scalar_ = 0.0f; -}; - -class Div : public Function { -public: - static constexpr char kType[] = "DivFunction"; - - explicit Div() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; -}; - -} // namespace infini_train::autograd +#pragma once + +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +class Neg : public Function { +public: + static constexpr char kType[] = "NegFunction"; + + Neg() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; + +class Reciprocal : public Function { +public: + static constexpr char kType[] = "ReciprocalFunction"; + + Reciprocal() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; + +class Sin : public Function { +public: + static constexpr char kType[] = "SinFunction"; + + Sin() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; + +class Cos : public Function { +public: + static constexpr char kType[] = "CosFunction"; + + Cos() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; + +class Tanh : public Function { +public: + static constexpr char kType[] = "TanhFunction"; + + Tanh() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; + +class Pow : public Function { +public: + static constexpr char kType[] = "PowFunction"; + + explicit Pow(float exponent, bool scalar_is_base = false) + : Function(kType), exponent_(exponent), scalar_is_base_(scalar_is_base) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + const float exponent_ = 0.0f; + const bool scalar_is_base_ = false; +}; + +class Rsqrt : public Function { +public: + static constexpr char kType[] = "RsqrtFunction"; + + explicit Rsqrt() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; + +class EqualsScalar : public Function { +public: + static constexpr char kType[] = "EqualsScalarFunction"; + + explicit EqualsScalar(float scalar) : Function(kType), scalar_(scalar) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + const float scalar_ = 0.0f; +}; + +class Add : public Function { +public: + static constexpr char kType[] = "AddFunction"; + + explicit Add() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + std::vector a_dims_; + std::vector b_dims_; +}; + +class AddScalar : public Function { +public: + static constexpr char kType[] = "AddScalarFunction"; + + explicit AddScalar(float scalar) : Function(kType), scalar_(scalar) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + const float scalar_ = 0.0f; +}; + +class Sub : public Function { +public: + static constexpr char kType[] = "SubFunction"; + + explicit Sub() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + std::vector a_dims_; + std::vector b_dims_; +}; + +class Mul : public Function { +public: + static constexpr char kType[] = "MulFunction"; + + explicit Mul() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; + +class MulScalar : public Function { +public: + static constexpr char kType[] = "MulScalarFunction"; + + explicit MulScalar(float scalar) : Function(kType), scalar_(scalar) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + const float scalar_ = 0.0f; +}; + +class Div : public Function { +public: + static constexpr char kType[] = "DivFunction"; + + explicit Div() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; + +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/function.h b/infini_train/include/autograd/function.h index b3bc71a..cd8645c 100644 --- a/infini_train/include/autograd/function.h +++ b/infini_train/include/autograd/function.h @@ -1,40 +1,51 @@ -#pragma once - -#include -#include -#include - -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class Function : public std::enable_shared_from_this { -public: - static constexpr char kUndefinedType[] = "Undefined"; - - Function() : type_(kUndefinedType) {} - explicit Function(const std::string &type) : type_(type) {} - - virtual ~Function() = default; - - virtual std::vector> Forward(const std::vector> &input_tensors) = 0; - virtual void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) {} - virtual std::vector> Backward(const std::vector> &grad_outputs) = 0; - - std::vector> Apply(const std::vector> &input_tensors); - virtual void BackwardPartial(const std::shared_ptr &grad_output, int idx); - - void IncreaseDependenciesNumber(); - -protected: - std::vector> saved_tensors_; - -private: - std::vector, int>> next_functions_; - int dependencies_number_ = 0; - int dependencies_reached_ = 0; - int grad_outputs_reached_ = 0; - std::vector> grad_outputs_; - const std::string type_ = kUndefinedType; -}; -} // namespace infini_train::autograd +#pragma once + +#include +#include +#include + +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { + +class NoGradGuard { +public: + NoGradGuard() { ++depth_; } + ~NoGradGuard() { --depth_; } + static bool is_enabled() { return depth_ > 0; } + +private: + static thread_local int depth_; +}; + +class Function : public std::enable_shared_from_this { +public: + static constexpr char kUndefinedType[] = "Undefined"; + + Function() : type_(kUndefinedType) {} + explicit Function(const std::string &type) : type_(type) {} + + virtual ~Function() = default; + + virtual std::vector> Forward(const std::vector> &input_tensors) = 0; + virtual void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) {} + virtual std::vector> Backward(const std::vector> &grad_outputs) = 0; + + std::vector> Apply(const std::vector> &input_tensors); + virtual void BackwardPartial(const std::shared_ptr &grad_output, int idx); + + void IncreaseDependenciesNumber(); + +protected: + std::vector> saved_tensors_; + +private: + std::vector, int>> next_functions_; + int dependencies_number_ = 0; + int dependencies_reached_ = 0; + int grad_outputs_reached_ = 0; + std::vector> grad_outputs_; + const std::string type_ = kUndefinedType; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/linear.h b/infini_train/include/autograd/linear.h index ee7dc3c..1f3431b 100644 --- a/infini_train/include/autograd/linear.h +++ b/infini_train/include/autograd/linear.h @@ -1,25 +1,25 @@ -#pragma once - -#include -#include - -#include "infini_train/include/autograd/function.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class Linear : public Function { -public: - static constexpr char kType[] = "LinearFunction"; - - Linear() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - int64_t out_features_ = 0; - bool bias_ = true; -}; -} // namespace infini_train::autograd +#pragma once + +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +class Linear : public Function { +public: + static constexpr char kType[] = "LinearFunction"; + + Linear() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + int64_t out_features_ = 0; + bool bias_ = true; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/loss.h b/infini_train/include/autograd/loss.h index 93487f3..c98751c 100644 --- a/infini_train/include/autograd/loss.h +++ b/infini_train/include/autograd/loss.h @@ -1,21 +1,21 @@ -#pragma once - -#include -#include - -#include "infini_train/include/autograd/function.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class CrossEntropy : public Function { -public: - static constexpr char kType[] = "CrossEntropyFunction"; - - CrossEntropy() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; -}; -} // namespace infini_train::autograd +#pragma once + +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +class CrossEntropy : public Function { +public: + static constexpr char kType[] = "CrossEntropyFunction"; + + CrossEntropy() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/matmul.h b/infini_train/include/autograd/matmul.h index af4ffda..d768471 100644 --- a/infini_train/include/autograd/matmul.h +++ b/infini_train/include/autograd/matmul.h @@ -1,24 +1,24 @@ -#pragma once - -#include -#include - -#include "infini_train/include/autograd/function.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class Matmul : public Function { -public: - static constexpr char kType[] = "MatmulFunction"; - - Matmul() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - int64_t out_features_ = 0; -}; -} // namespace infini_train::autograd +#pragma once + +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +class Matmul : public Function { +public: + static constexpr char kType[] = "MatmulFunction"; + + Matmul() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + int64_t out_features_ = 0; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/misc.h b/infini_train/include/autograd/misc.h index a6516d4..7857c1d 100644 --- a/infini_train/include/autograd/misc.h +++ b/infini_train/include/autograd/misc.h @@ -1,75 +1,75 @@ -#pragma once - -#include -#include - -#include "infini_train/include/autograd/function.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class Split : public Function { -public: - static constexpr char kType[] = "SplitFunction"; - - Split(int64_t split_size, int dim = 0) : Function(kType), split_size_(split_size), dim_(dim) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - const int64_t split_size_ = 0; - const int dim_ = 0; - std::vector input_dims_; -}; - -class NoOp : public Function { -public: - static constexpr char kType[] = "NoOpFunction"; - - explicit NoOp(const std::vector &output_dims) : Function(kType), output_dims_(output_dims) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - const std::vector output_dims_; - std::vector input_dims_; -}; - -class Slice : public Function { -public: - static constexpr char kType[] = "SliceFunction"; - - Slice(const std::vector &starts, const std::vector &ends, const std::vector &steps) - : Function(kType), starts_(starts), ends_(ends), steps_(steps) {} - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - const std::vector starts_; - const std::vector ends_; - const std::vector steps_; -}; - -class Stack : public Function { -public: - static constexpr char kType[] = "StackFunction"; - - Stack(int64_t dim) : Function(kType), dim_(dim) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - int64_t dim_ = 0; - std::vector input_dims_; -}; -} // namespace infini_train::autograd +#pragma once + +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +class Split : public Function { +public: + static constexpr char kType[] = "SplitFunction"; + + Split(int64_t split_size, int dim = 0) : Function(kType), split_size_(split_size), dim_(dim) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + const int64_t split_size_ = 0; + const int dim_ = 0; + std::vector input_dims_; +}; + +class NoOp : public Function { +public: + static constexpr char kType[] = "NoOpFunction"; + + explicit NoOp(const std::vector &output_dims) : Function(kType), output_dims_(output_dims) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + const std::vector output_dims_; + std::vector input_dims_; +}; + +class Slice : public Function { +public: + static constexpr char kType[] = "SliceFunction"; + + Slice(const std::vector &starts, const std::vector &ends, const std::vector &steps) + : Function(kType), starts_(starts), ends_(ends), steps_(steps) {} + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + const std::vector starts_; + const std::vector ends_; + const std::vector steps_; +}; + +class Stack : public Function { +public: + static constexpr char kType[] = "StackFunction"; + + Stack(int64_t dim) : Function(kType), dim_(dim) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + int64_t dim_ = 0; + std::vector input_dims_; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/normalization.h b/infini_train/include/autograd/normalization.h index c8432b9..6c858bb 100644 --- a/infini_train/include/autograd/normalization.h +++ b/infini_train/include/autograd/normalization.h @@ -1,24 +1,24 @@ -#pragma once - -#include -#include - -#include "infini_train/include/autograd/function.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class LayerNorm : public Function { -public: - static constexpr char kType[] = "LayerNormFunction"; - - explicit LayerNorm(float eps) : Function(kType), eps_(eps) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - const float eps_ = 1e-5f; -}; -} // namespace infini_train::autograd +#pragma once + +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +class LayerNorm : public Function { +public: + static constexpr char kType[] = "LayerNormFunction"; + + explicit LayerNorm(float eps) : Function(kType), eps_(eps) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + const float eps_ = 1e-5f; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/outer.h b/infini_train/include/autograd/outer.h index 4661f74..65899e9 100644 --- a/infini_train/include/autograd/outer.h +++ b/infini_train/include/autograd/outer.h @@ -1,21 +1,21 @@ -#pragma once - -#include -#include - -#include "infini_train/include/autograd/function.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class Outer : public Function { -public: - static constexpr char kType[] = "OuterFunction"; - - Outer() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; -}; -} // namespace infini_train::autograd +#pragma once + +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +class Outer : public Function { +public: + static constexpr char kType[] = "OuterFunction"; + + Outer() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/reduction.h b/infini_train/include/autograd/reduction.h index a7eea8e..7da6923 100644 --- a/infini_train/include/autograd/reduction.h +++ b/infini_train/include/autograd/reduction.h @@ -1,75 +1,75 @@ -#pragma once - -#include -#include - -#include "infini_train/include/autograd/function.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class Mean : public Function { -public: - static constexpr char kType[] = "MeanFunction"; - - explicit Mean(int64_t dim, bool keep_dim = false) : Function(kType), dim_(dim), keep_dim_(keep_dim) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - std::vector input_dims_; - int64_t dim_ = 0; - bool keep_dim_ = false; -}; - -class Sum : public Function { -public: - static constexpr char kType[] = "SumFunction"; - - explicit Sum(int64_t dim, bool keep_dim = false) : Function(kType), dim_(dim), keep_dim_(keep_dim) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - std::vector input_dims_; - int64_t dim_ = 0; - bool keep_dim_ = false; -}; - -class Max : public Function { -public: - static constexpr char kType[] = "MaxFunction"; - - explicit Max(int64_t dim, bool keep_dim = false) : Function(kType), dim_(dim), keep_dim_(keep_dim) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - int64_t dim_ = 0; - bool keep_dim_ = false; -}; - -class Min : public Function { -public: - static constexpr char kType[] = "MinFunction"; - - explicit Min(int64_t dim, bool keep_dim = false) : Function(kType), dim_(dim), keep_dim_(keep_dim) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - int64_t dim_ = 0; - bool keep_dim_ = false; -}; -} // namespace infini_train::autograd +#pragma once + +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +class Mean : public Function { +public: + static constexpr char kType[] = "MeanFunction"; + + explicit Mean(int64_t dim, bool keep_dim = false) : Function(kType), dim_(dim), keep_dim_(keep_dim) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + std::vector input_dims_; + int64_t dim_ = 0; + bool keep_dim_ = false; +}; + +class Sum : public Function { +public: + static constexpr char kType[] = "SumFunction"; + + explicit Sum(int64_t dim, bool keep_dim = false) : Function(kType), dim_(dim), keep_dim_(keep_dim) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + std::vector input_dims_; + int64_t dim_ = 0; + bool keep_dim_ = false; +}; + +class Max : public Function { +public: + static constexpr char kType[] = "MaxFunction"; + + explicit Max(int64_t dim, bool keep_dim = false) : Function(kType), dim_(dim), keep_dim_(keep_dim) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + int64_t dim_ = 0; + bool keep_dim_ = false; +}; + +class Min : public Function { +public: + static constexpr char kType[] = "MinFunction"; + + explicit Min(int64_t dim, bool keep_dim = false) : Function(kType), dim_(dim), keep_dim_(keep_dim) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + int64_t dim_ = 0; + bool keep_dim_ = false; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/softmax.h b/infini_train/include/autograd/softmax.h index 025a4e3..2bfe5ec 100644 --- a/infini_train/include/autograd/softmax.h +++ b/infini_train/include/autograd/softmax.h @@ -1,25 +1,25 @@ -#pragma once - -#include -#include -#include - -#include "infini_train/include/autograd/function.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class Softmax : public Function { -public: - static constexpr char kType[] = "SoftmaxFunction"; - - explicit Softmax(int64_t dim = -1) : Function(kType), dim_(dim){}; - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - const int64_t dim_ = -1; -}; -} // namespace infini_train::autograd +#pragma once + +#include +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +class Softmax : public Function { +public: + static constexpr char kType[] = "SoftmaxFunction"; + + explicit Softmax(int64_t dim = -1) : Function(kType), dim_(dim){}; + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + const int64_t dim_ = -1; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/sparse.h b/infini_train/include/autograd/sparse.h index 76c2a6b..de676db 100644 --- a/infini_train/include/autograd/sparse.h +++ b/infini_train/include/autograd/sparse.h @@ -1,24 +1,24 @@ -#pragma once - -#include -#include - -#include "infini_train/include/autograd/function.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class Embedding : public Function { -public: - static constexpr char kType[] = "EmbeddingFunction"; - - explicit Embedding() : Function(kType) {} - - std::vector> Forward(const std::vector> &input_tensors) override; - void SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - std::vector weight_dims_; -}; -} // namespace infini_train::autograd +#pragma once + +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +class Embedding : public Function { +public: + static constexpr char kType[] = "EmbeddingFunction"; + + explicit Embedding() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + std::vector weight_dims_; +}; +} // namespace infini_train::autograd diff --git a/infini_train/include/autograd/transform.h b/infini_train/include/autograd/transform.h index 3ef4908..5e5dbcc 100644 --- a/infini_train/include/autograd/transform.h +++ b/infini_train/include/autograd/transform.h @@ -1,78 +1,78 @@ -#pragma once - -#include -#include -#include - -#include "infini_train/include/autograd/function.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -class Tril : public Function { -public: - static constexpr char kType[] = "TrilFunction"; - - Tril(int64_t diagonal) : Function(kType), diagonal_(diagonal) {} - std::vector> Forward(const std::vector> &input_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - int64_t diagonal_ = 0; -}; - -class Triu : public Function { -public: - static constexpr char kType[] = "TriuFunction"; - - Triu(int64_t diagonal) : Function(kType), diagonal_(diagonal) {} - std::vector> Forward(const std::vector> &input_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - int64_t diagonal_ = 0; -}; - -class Transpose : public Function { -public: - static constexpr char kType[] = "TransposeFunction"; - - Transpose(int64_t dim0, int64_t dim1) : Function(kType), dim0_(dim0), dim1_(dim1) {} - std::vector> Forward(const std::vector> &input_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - int64_t dim0_ = 0; - int64_t dim1_ = 0; -}; - -class Mask : public Function { -public: - static constexpr char kType[] = "MaskFunction"; - - Mask(std::shared_ptr mask, float value) : Function(kType), mask_(mask), value_(value) {} - std::vector> Forward(const std::vector> &input_tensors) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - std::shared_ptr mask_; - float value_ = 0.f; -}; - -class RepeatInterleave : public Function { -public: - static constexpr char kType[] = "RepeatInterleaveFunction"; - - RepeatInterleave(int64_t repeat, int64_t dim) : Function(kType), repeat_(repeat), dim_(dim) {} - - std::vector> Forward(const std::vector> &inputs) override; - void SetupContext(const std::vector> &inputs, - const std::vector> &outputs) override; - std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - int64_t repeat_ = 0; - int64_t dim_ = 0; - std::vector input_dims_; -}; - -} // namespace infini_train::autograd +#pragma once + +#include +#include +#include + +#include "infini_train/include/autograd/function.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +class Tril : public Function { +public: + static constexpr char kType[] = "TrilFunction"; + + Tril(int64_t diagonal) : Function(kType), diagonal_(diagonal) {} + std::vector> Forward(const std::vector> &input_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + int64_t diagonal_ = 0; +}; + +class Triu : public Function { +public: + static constexpr char kType[] = "TriuFunction"; + + Triu(int64_t diagonal) : Function(kType), diagonal_(diagonal) {} + std::vector> Forward(const std::vector> &input_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + int64_t diagonal_ = 0; +}; + +class Transpose : public Function { +public: + static constexpr char kType[] = "TransposeFunction"; + + Transpose(int64_t dim0, int64_t dim1) : Function(kType), dim0_(dim0), dim1_(dim1) {} + std::vector> Forward(const std::vector> &input_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + int64_t dim0_ = 0; + int64_t dim1_ = 0; +}; + +class Mask : public Function { +public: + static constexpr char kType[] = "MaskFunction"; + + Mask(std::shared_ptr mask, float value) : Function(kType), mask_(mask), value_(value) {} + std::vector> Forward(const std::vector> &input_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + std::shared_ptr mask_; + float value_ = 0.f; +}; + +class RepeatInterleave : public Function { +public: + static constexpr char kType[] = "RepeatInterleaveFunction"; + + RepeatInterleave(int64_t repeat, int64_t dim) : Function(kType), repeat_(repeat), dim_(dim) {} + + std::vector> Forward(const std::vector> &inputs) override; + void SetupContext(const std::vector> &inputs, + const std::vector> &outputs) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + +private: + int64_t repeat_ = 0; + int64_t dim_ = 0; + std::vector input_dims_; +}; + +} // namespace infini_train::autograd diff --git a/infini_train/include/dataloader.h b/infini_train/include/dataloader.h index 5e9cd43..74fd98f 100644 --- a/infini_train/include/dataloader.h +++ b/infini_train/include/dataloader.h @@ -1,43 +1,43 @@ -#pragma once - -#include -#include -#include - -#include "infini_train/include/dataset.h" -#include "infini_train/include/tensor.h" - -namespace infini_train { -class DataLoaderIterator { -public: - DataLoaderIterator(const Dataset &dataset, size_t batch_size, size_t batch_idx, size_t max_batch_idx); - - std::pair, std::shared_ptr> operator*() const; - - DataLoaderIterator &operator++(); - DataLoaderIterator operator++(int); - - friend bool operator<(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs); - friend bool operator!=(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs); - friend bool operator==(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs); - -private: - const Dataset *dataset_ = nullptr; // not owned - size_t batch_size_ = 0; - size_t batch_idx_ = 0; - size_t max_batch_idx_ = 0; -}; - -class DataLoader { -public: - DataLoader(const std::shared_ptr &dataset, size_t batch_size); - - DataLoaderIterator begin() const; - DataLoaderIterator end() const; - -private: - std::shared_ptr dataset_; - size_t batch_size_ = 0; - size_t max_batch_idx_ = 0; -}; -} // namespace infini_train +#pragma once + +#include +#include +#include + +#include "infini_train/include/dataset.h" +#include "infini_train/include/tensor.h" + +namespace infini_train { +class DataLoaderIterator { +public: + DataLoaderIterator(const Dataset &dataset, size_t batch_size, size_t batch_idx, size_t max_batch_idx); + + std::pair, std::shared_ptr> operator*() const; + + DataLoaderIterator &operator++(); + DataLoaderIterator operator++(int); + + friend bool operator<(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs); + friend bool operator!=(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs); + friend bool operator==(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs); + +private: + const Dataset *dataset_ = nullptr; // not owned + size_t batch_size_ = 0; + size_t batch_idx_ = 0; + size_t max_batch_idx_ = 0; +}; + +class DataLoader { +public: + DataLoader(const std::shared_ptr &dataset, size_t batch_size); + + DataLoaderIterator begin() const; + DataLoaderIterator end() const; + +private: + std::shared_ptr dataset_; + size_t batch_size_ = 0; + size_t max_batch_idx_ = 0; +}; +} // namespace infini_train diff --git a/infini_train/include/dataset.h b/infini_train/include/dataset.h index 8769f45..a1ab248 100644 --- a/infini_train/include/dataset.h +++ b/infini_train/include/dataset.h @@ -1,17 +1,17 @@ -#pragma once - -#include -#include -#include - -#include "infini_train/include/tensor.h" - -namespace infini_train { -class Dataset { -public: - virtual ~Dataset() = default; - - virtual std::pair, std::shared_ptr> operator[](size_t idx) const = 0; - virtual size_t Size() const = 0; -}; -} // namespace infini_train +#pragma once + +#include +#include +#include + +#include "infini_train/include/tensor.h" + +namespace infini_train { +class Dataset { +public: + virtual ~Dataset() = default; + + virtual std::pair, std::shared_ptr> operator[](size_t idx) const = 0; + virtual size_t Size() const = 0; +}; +} // namespace infini_train diff --git a/infini_train/include/device.h b/infini_train/include/device.h index ebc49e3..ed464ff 100644 --- a/infini_train/include/device.h +++ b/infini_train/include/device.h @@ -1,38 +1,38 @@ -#pragma once - -#include - -#include "glog/logging.h" - -namespace infini_train { -enum class DeviceType : int8_t { - kCPU = 0, - kCUDA = 1, -}; - -class Device { -public: - Device(); - - Device(DeviceType type, int8_t index); - - bool operator==(const Device &other) const; - - bool operator!=(const Device &other) const; - - DeviceType Type() const; - int8_t Index() const; - - bool IsCPU() const; - bool IsCUDA() const; - - std::string ToString() const; - - friend std::ostream &operator<<(std::ostream &os, const Device &device); - -private: - DeviceType type_; - int8_t index_; -}; - -} // namespace infini_train +#pragma once + +#include + +#include "glog/logging.h" + +namespace infini_train { +enum class DeviceType : int8_t { + kCPU = 0, + kCUDA = 1, +}; + +class Device { +public: + Device(); + + Device(DeviceType type, int8_t index); + + bool operator==(const Device &other) const; + + bool operator!=(const Device &other) const; + + DeviceType Type() const; + int8_t Index() const; + + bool IsCPU() const; + bool IsCUDA() const; + + std::string ToString() const; + + friend std::ostream &operator<<(std::ostream &os, const Device &device); + +private: + DeviceType type_; + int8_t index_; +}; + +} // namespace infini_train diff --git a/infini_train/include/dispatcher.h b/infini_train/include/dispatcher.h index 5b91d85..ac85eeb 100644 --- a/infini_train/include/dispatcher.h +++ b/infini_train/include/dispatcher.h @@ -1,62 +1,66 @@ -#pragma once - -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/device.h" - -namespace infini_train { -class KernelFunction { -public: - template explicit KernelFunction(FuncT &&func) : func_ptr_(reinterpret_cast(func)) {} - - template RetT Call(ArgsT... args) const { - // =================================== 作业 =================================== - // TODO：实现通用kernel调用接口 - // 功能描述：将存储的函数指针转换为指定类型并调用 - // =================================== 作业 =================================== - - using FuncT = RetT (*)(ArgsT...); - // TODO: 实现函数调用逻辑 - } - -private: - void *func_ptr_ = nullptr; -}; - -class Dispatcher { -public: - using KeyT = std::pair; - - static Dispatcher &Instance() { - static Dispatcher instance; - return instance; - } - - const KernelFunction &GetKernel(KeyT key) const { - CHECK(key_to_kernel_map_.contains(key)) - << "Kernel not found: " << key.second << " on device: " << static_cast(key.first); - return key_to_kernel_map_.at(key); - } - - template void Register(const KeyT &key, FuncT &&kernel) { - // =================================== 作业 =================================== - // TODO：实现kernel注册机制 - // 功能描述：将kernel函数与设备类型、名称绑定 - // =================================== 作业 =================================== - } - -private: - std::map key_to_kernel_map_; -}; -} // namespace infini_train - -#define REGISTER_KERNEL(device, kernel_name, kernel_func) \ - // =================================== 作业 =================================== - // TODO：实现自动注册宏 - // 功能描述：在全局静态区注册kernel，避免显式初始化代码 - // =================================== 作业 =================================== +#pragma once + +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/device.h" + +namespace infini_train { +class KernelFunction { +public: + template explicit KernelFunction(FuncT &&func) : func_ptr_(reinterpret_cast(func)) {} + + template RetT Call(ArgsT... args) const { + // =================================== 作业 =================================== + // TODO：实现通用kernel调用接口 + // 功能描述：将存储的函数指针转换为指定类型并调用 + // =================================== 作业 =================================== + + using FuncT = RetT (*)(ArgsT...); + auto func = reinterpret_cast(func_ptr_); + return func(std::forward(args)...); + } + +private: + void *func_ptr_ = nullptr; +}; + +class Dispatcher { +public: + using KeyT = std::pair; + + static Dispatcher &Instance() { + static Dispatcher instance; + return instance; + } + + const KernelFunction &GetKernel(KeyT key) const { + CHECK(key_to_kernel_map_.contains(key)) + << "Kernel not found: " << key.second << " on device: " << static_cast(key.first); + return key_to_kernel_map_.at(key); + } + + template void Register(const KeyT &key, FuncT &&kernel) { + // =================================== 作业 =================================== + // TODO：实现kernel注册机制 + // 功能描述：将kernel函数与设备类型、名称绑定 + // =================================== 作业 =================================== + CHECK(!key_to_kernel_map_.contains(key)) << "Kernel already registered: " << key.second + << " on device: " << static_cast(key.first); + key_to_kernel_map_.emplace(key, KernelFunction(std::forward(kernel))); + } + +private: + std::map key_to_kernel_map_; +}; +} // namespace infini_train + +#define REGISTER_KERNEL(device, kernel_name, kernel_func) \ + static auto _register_##kernel_name##_##__LINE__ = []() { \ + infini_train::Dispatcher::Instance().Register({device, #kernel_name}, kernel_func); \ + return 0; \ + }(); diff --git a/infini_train/include/nn/functional.h b/infini_train/include/nn/functional.h index 7713a40..f4ebf83 100644 --- a/infini_train/include/nn/functional.h +++ b/infini_train/include/nn/functional.h @@ -1,170 +1,170 @@ -#pragma once - -#include -#include -#include - -#include "infini_train/include/tensor.h" - -namespace infini_train::nn::function { - -// Returns the lower triangular part of a 2D tensor or a batch of matrices. -// -// The lower triangular part includes elements on and below the specified -// diagonal. Elements above the diagonal are set to zero. -// -// Args: -// input: The input tensor. -// diagonal: Diagonal offset (default 0). Positive means above the main diagonal, -// negative means below. -// -// Returns: -// A tensor with the same shape as input, with upper-triangular values zeroed. -std::shared_ptr Tril(const std::shared_ptr &input, int64_t diagonal = 0); - -// Returns the upper triangular part of a 2D tensor or a batch of matrices. -// -// The upper triangular part includes elements on and above the specified -// diagonal. Elements below the diagonal are set to zero. -// -// Args: -// input: The input tensor. -// diagonal: Diagonal offset (default 0). Positive means above the main diagonal, -// negative means below. -// -// Returns: -// A tensor with the same shape as input, with lower-triangular values zeroed. -std::shared_ptr Triu(const std::shared_ptr &input, int64_t diagonal = 0); - -// Returns a tensor filled with ones of the specified shape. -// -// Args: -// size: A vector specifying the shape of the output tensor. -// -// Returns: -// A tensor of the given shape filled with the scalar value 1. -std::shared_ptr Ones(const std::vector size); - -// Returns a new tensor with the reciprocal of the elements of input. -// -// Args: -// input: The input tensor. -// -// Returns: -// A tensor containing reciprocal applied element-wise to the input. -std::shared_ptr Reciprocal(const std::shared_ptr &input); - -// Returns a new tensor with the sine of each element in the input. -// -// Args: -// input: The input tensor. -// -// Returns: -// A tensor containing sin applied element-wise to the input. -std::shared_ptr Sin(const std::shared_ptr &input); - -// Returns a new tensor with the cosine of each element in the input. -// -// Args: -// input: The input tensor. -// -// Returns: -// A tensor containing cos applied element-wise to the input. -std::shared_ptr Cos(const std::shared_ptr &input); - -// Returns a new tensor with the hyperbolic tangent of each element in the input. -// -// Args: -// input: The input tensor. -// -// Returns: -// A tensor containing tanh applied element-wise to the input. -std::shared_ptr Tanh(const std::shared_ptr &input); - -// Raises each element of the input tensor to the specified power. -// -// Args: -// input: The input tensor. -// exponent: The exponent to apply to each element. -// -// Returns: -// A tensor with each element raised to the given exponent. -std::shared_ptr Pow(const std::shared_ptr &input, float exponent); - -// Raises the specified base to the power of each element in the input tensor. -// -// Args: -// base: The scalar base value. -// input: The input tensor providing the exponents. -// -// Returns: -// A tensor where each element is computed as base raised to the power of the corresponding input element. -std::shared_ptr Pow(float base, const std::shared_ptr &input); - -// Returns a new tensor with reciprocal of the square-root of each element in the input. -// -// Args: -// input: The input tensor. -// -// Returns: -// A tensor containing reciprocal square-root applied element-wise to the input. -std::shared_ptr Rsqrt(const std::shared_ptr &input); - -// Returns the mean of all elements in the input tensor. -// -// Args: -// input: The input tensor. -// dim: The dimension to reduce. -// keep_dim: Whether the output tensor has dim retained or not (default false). -// -// Returns: -// A new tensor with the mean values computed along the specified dimension. -std::shared_ptr Mean(const std::shared_ptr &input, int64_t dim, bool keep_dim = false); - -// Returns a new tensor with the sigmoid of each element in the input. -// -// The sigmoid function is defined as 1 / (1 + exp(-x)). -// -// Args: -// input: The input tensor. -// -// Returns: -// A tensor containing sigmoid applied element-wise to the input. -std::shared_ptr Sigmoid(const std::shared_ptr &input); - -// Applies the softmax function along the specified dimension. -// -// The softmax function maps input values to the range [0, 1] and ensures they sum to 1. -// -// Args: -// input: The input tensor. -// dim: The dimension along which softmax is computed (default -1). -// -// Returns: -// A tensor with softmax applied along the specified dimension. -std::shared_ptr Softmax(const std::shared_ptr &input, int64_t dim = -1); - -// Returns a slice of the input tensor defined by start, end, and step per dimension. -// -// Args: -// input: The input tensor. -// starts: Start indices for each dimension. -// ends: End indices for each dimension (exclusive). -// steps: Step sizes for each dimension. -// -// Returns: -// A sliced view of the input tensor. -std::shared_ptr Slice(const std::shared_ptr &input, const std::vector &starts, - const std::vector &ends, const std::vector &steps); - -// Concatenates a sequence of tensors along a new dimension. -// -// Args: -// inputs: The sequence of tensors to concatenate. All tensors need to be of the same size. -// dim: dimension to insert (defualt 0). -// -// Returns: -// Concatenation of the input tensors. -std::shared_ptr Stack(const std::vector> &inputs, int64_t dim = 0); - -} // namespace infini_train::nn::function +#pragma once + +#include +#include +#include + +#include "infini_train/include/tensor.h" + +namespace infini_train::nn::function { + +// Returns the lower triangular part of a 2D tensor or a batch of matrices. +// +// The lower triangular part includes elements on and below the specified +// diagonal. Elements above the diagonal are set to zero. +// +// Args: +// input: The input tensor. +// diagonal: Diagonal offset (default 0). Positive means above the main diagonal, +// negative means below. +// +// Returns: +// A tensor with the same shape as input, with upper-triangular values zeroed. +std::shared_ptr Tril(const std::shared_ptr &input, int64_t diagonal = 0); + +// Returns the upper triangular part of a 2D tensor or a batch of matrices. +// +// The upper triangular part includes elements on and above the specified +// diagonal. Elements below the diagonal are set to zero. +// +// Args: +// input: The input tensor. +// diagonal: Diagonal offset (default 0). Positive means above the main diagonal, +// negative means below. +// +// Returns: +// A tensor with the same shape as input, with lower-triangular values zeroed. +std::shared_ptr Triu(const std::shared_ptr &input, int64_t diagonal = 0); + +// Returns a tensor filled with ones of the specified shape. +// +// Args: +// size: A vector specifying the shape of the output tensor. +// +// Returns: +// A tensor of the given shape filled with the scalar value 1. +std::shared_ptr Ones(const std::vector size); + +// Returns a new tensor with the reciprocal of the elements of input. +// +// Args: +// input: The input tensor. +// +// Returns: +// A tensor containing reciprocal applied element-wise to the input. +std::shared_ptr Reciprocal(const std::shared_ptr &input); + +// Returns a new tensor with the sine of each element in the input. +// +// Args: +// input: The input tensor. +// +// Returns: +// A tensor containing sin applied element-wise to the input. +std::shared_ptr Sin(const std::shared_ptr &input); + +// Returns a new tensor with the cosine of each element in the input. +// +// Args: +// input: The input tensor. +// +// Returns: +// A tensor containing cos applied element-wise to the input. +std::shared_ptr Cos(const std::shared_ptr &input); + +// Returns a new tensor with the hyperbolic tangent of each element in the input. +// +// Args: +// input: The input tensor. +// +// Returns: +// A tensor containing tanh applied element-wise to the input. +std::shared_ptr Tanh(const std::shared_ptr &input); + +// Raises each element of the input tensor to the specified power. +// +// Args: +// input: The input tensor. +// exponent: The exponent to apply to each element. +// +// Returns: +// A tensor with each element raised to the given exponent. +std::shared_ptr Pow(const std::shared_ptr &input, float exponent); + +// Raises the specified base to the power of each element in the input tensor. +// +// Args: +// base: The scalar base value. +// input: The input tensor providing the exponents. +// +// Returns: +// A tensor where each element is computed as base raised to the power of the corresponding input element. +std::shared_ptr Pow(float base, const std::shared_ptr &input); + +// Returns a new tensor with reciprocal of the square-root of each element in the input. +// +// Args: +// input: The input tensor. +// +// Returns: +// A tensor containing reciprocal square-root applied element-wise to the input. +std::shared_ptr Rsqrt(const std::shared_ptr &input); + +// Returns the mean of all elements in the input tensor. +// +// Args: +// input: The input tensor. +// dim: The dimension to reduce. +// keep_dim: Whether the output tensor has dim retained or not (default false). +// +// Returns: +// A new tensor with the mean values computed along the specified dimension. +std::shared_ptr Mean(const std::shared_ptr &input, int64_t dim, bool keep_dim = false); + +// Returns a new tensor with the sigmoid of each element in the input. +// +// The sigmoid function is defined as 1 / (1 + exp(-x)). +// +// Args: +// input: The input tensor. +// +// Returns: +// A tensor containing sigmoid applied element-wise to the input. +std::shared_ptr Sigmoid(const std::shared_ptr &input); + +// Applies the softmax function along the specified dimension. +// +// The softmax function maps input values to the range [0, 1] and ensures they sum to 1. +// +// Args: +// input: The input tensor. +// dim: The dimension along which softmax is computed (default -1). +// +// Returns: +// A tensor with softmax applied along the specified dimension. +std::shared_ptr Softmax(const std::shared_ptr &input, int64_t dim = -1); + +// Returns a slice of the input tensor defined by start, end, and step per dimension. +// +// Args: +// input: The input tensor. +// starts: Start indices for each dimension. +// ends: End indices for each dimension (exclusive). +// steps: Step sizes for each dimension. +// +// Returns: +// A sliced view of the input tensor. +std::shared_ptr Slice(const std::shared_ptr &input, const std::vector &starts, + const std::vector &ends, const std::vector &steps); + +// Concatenates a sequence of tensors along a new dimension. +// +// Args: +// inputs: The sequence of tensors to concatenate. All tensors need to be of the same size. +// dim: dimension to insert (defualt 0). +// +// Returns: +// Concatenation of the input tensors. +std::shared_ptr Stack(const std::vector> &inputs, int64_t dim = 0); + +} // namespace infini_train::nn::function diff --git a/infini_train/include/nn/init.h b/infini_train/include/nn/init.h index 169aba5..948b820 100644 --- a/infini_train/include/nn/init.h +++ b/infini_train/include/nn/init.h @@ -1,50 +1,50 @@ -#pragma once - -#include -#include -#include -#include -#include - -#include "infini_train/include/tensor.h" - -namespace infini_train::nn::init { -std::shared_ptr Normal(const std::shared_ptr &tensor, float mean = 0.0, float std = 1.0, - std::optional generator = std::nullopt); - -std::pair CalculateFanInAndFanOut(const std::shared_ptr &tensor); - -enum class KaimingMode : int8_t { - kFanIn, - kFanOut, -}; - -enum class NonLinearityType : int8_t { - kLinear, - kConv1D, - kConv2D, - kConv3D, - kConvTransposed1d, - kConvTransposed2d, - kConvTransposed3d, - kSigmoid, - kTanh, - kReLU, - kLeakyReLU, - kSELU, -}; - -std::shared_ptr KaimingUniform(const std::shared_ptr &tensor, float a = 0.0f, - KaimingMode mode = KaimingMode::kFanIn, - NonLinearityType non_linearity = NonLinearityType::kLeakyReLU, - std::optional generator = std::nullopt); - -std::shared_ptr Uniform(const std::shared_ptr &tensor, float a = 0.0f, float b = 1.0f, - std::optional generator = std::nullopt); - -std::shared_ptr Ones(const std::shared_ptr &tensor); - -std::shared_ptr Zeros(const std::shared_ptr &tensor); - -std::shared_ptr Arange(int64_t start, int64_t end, DataType dtype, Device device = Device()); -} // namespace infini_train::nn::init +#pragma once + +#include +#include +#include +#include +#include + +#include "infini_train/include/tensor.h" + +namespace infini_train::nn::init { +std::shared_ptr Normal(const std::shared_ptr &tensor, float mean = 0.0, float std = 1.0, + std::optional generator = std::nullopt); + +std::pair CalculateFanInAndFanOut(const std::shared_ptr &tensor); + +enum class KaimingMode : int8_t { + kFanIn, + kFanOut, +}; + +enum class NonLinearityType : int8_t { + kLinear, + kConv1D, + kConv2D, + kConv3D, + kConvTransposed1d, + kConvTransposed2d, + kConvTransposed3d, + kSigmoid, + kTanh, + kReLU, + kLeakyReLU, + kSELU, +}; + +std::shared_ptr KaimingUniform(const std::shared_ptr &tensor, float a = 0.0f, + KaimingMode mode = KaimingMode::kFanIn, + NonLinearityType non_linearity = NonLinearityType::kLeakyReLU, + std::optional generator = std::nullopt); + +std::shared_ptr Uniform(const std::shared_ptr &tensor, float a = 0.0f, float b = 1.0f, + std::optional generator = std::nullopt); + +std::shared_ptr Ones(const std::shared_ptr &tensor); + +std::shared_ptr Zeros(const std::shared_ptr &tensor); + +std::shared_ptr Arange(int64_t start, int64_t end, DataType dtype, Device device = Device()); +} // namespace infini_train::nn::init diff --git a/infini_train/include/nn/modules/activations.h b/infini_train/include/nn/modules/activations.h index e6f65ae..f887ceb 100644 --- a/infini_train/include/nn/modules/activations.h +++ b/infini_train/include/nn/modules/activations.h @@ -1,15 +1,15 @@ -#pragma once - -#include -#include - -#include "infini_train/include/nn/modules/module.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -class Sigmoid : public Module { -public: - Sigmoid() = default; - std::vector> Forward(const std::vector> &input_tensors) override; -}; -} // namespace infini_train::nn +#pragma once + +#include +#include + +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +class Sigmoid : public Module { +public: + Sigmoid() = default; + std::vector> Forward(const std::vector> &input_tensors) override; +}; +} // namespace infini_train::nn diff --git a/infini_train/include/nn/modules/container.h b/infini_train/include/nn/modules/container.h index fe8697c..7b774f2 100644 --- a/infini_train/include/nn/modules/container.h +++ b/infini_train/include/nn/modules/container.h @@ -1,24 +1,24 @@ -#pragma once - -#include -#include -#include - -#include "infini_train/include/nn/modules/module.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -class Sequential : public Module { -public: - explicit Sequential(std::vector> &&layers); - - std::vector> Forward(const std::vector> &input_tensors) override; -}; - -class ModuleDict : public Module { -public: - explicit ModuleDict(std::unordered_map> &&modules); - - std::vector> Forward(const std::vector> &input_tensors) override; -}; -} // namespace infini_train::nn +#pragma once + +#include +#include +#include + +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +class Sequential : public Module { +public: + explicit Sequential(std::vector> &&layers); + + std::vector> Forward(const std::vector> &input_tensors) override; +}; + +class ModuleDict : public Module { +public: + explicit ModuleDict(std::unordered_map> &&modules); + + std::vector> Forward(const std::vector> &input_tensors) override; +}; +} // namespace infini_train::nn diff --git a/infini_train/include/nn/modules/linear.h b/infini_train/include/nn/modules/linear.h index 500ae71..b86d3bc 100644 --- a/infini_train/include/nn/modules/linear.h +++ b/infini_train/include/nn/modules/linear.h @@ -1,25 +1,25 @@ -#pragma once - -#include -#include - -#include "infini_train/include/device.h" -#include "infini_train/include/nn/modules/module.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -class Linear : public Module { -public: - static constexpr char kType[] = "Linear"; - - static constexpr char kParamWeightName[] = "weight"; - static constexpr char kParamBiasName[] = "bias"; - - Linear(int64_t in_features, int64_t out_features, bool bias = true, Device device = Device()); - std::vector> Forward(const std::vector> &input_tensors) override; - -private: - void ResetParameters(); - bool bias_ = true; -}; -} // namespace infini_train::nn +#pragma once + +#include +#include + +#include "infini_train/include/device.h" +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +class Linear : public Module { +public: + static constexpr char kType[] = "Linear"; + + static constexpr char kParamWeightName[] = "weight"; + static constexpr char kParamBiasName[] = "bias"; + + Linear(int64_t in_features, int64_t out_features, bool bias = true, Device device = Device()); + std::vector> Forward(const std::vector> &input_tensors) override; + +private: + void ResetParameters(); + bool bias_ = true; +}; +} // namespace infini_train::nn diff --git a/infini_train/include/nn/modules/loss.h b/infini_train/include/nn/modules/loss.h index 0b4a831..e4ab912 100644 --- a/infini_train/include/nn/modules/loss.h +++ b/infini_train/include/nn/modules/loss.h @@ -1,16 +1,16 @@ -#pragma once - -#include -#include - -#include "infini_train/include/nn/modules/module.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -class CrossEntropyLoss : public Module { -public: - CrossEntropyLoss() = default; - - std::vector> Forward(const std::vector> &input_tensors) override; -}; -} // namespace infini_train::nn +#pragma once + +#include +#include + +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +class CrossEntropyLoss : public Module { +public: + CrossEntropyLoss() = default; + + std::vector> Forward(const std::vector> &input_tensors) override; +}; +} // namespace infini_train::nn diff --git a/infini_train/include/nn/modules/module.h b/infini_train/include/nn/modules/module.h index 3722bd7..e3b2c59 100644 --- a/infini_train/include/nn/modules/module.h +++ b/infini_train/include/nn/modules/module.h @@ -1,46 +1,46 @@ -#pragma once - -#include -#include -#include -#include - -#include "infini_train/include/device.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -class Module { -public: - static constexpr char kUndefinedType[] = "Undefined"; - - Module() : type_(kUndefinedType) {} - explicit Module(const std::string &type) : type_(type) {} - - virtual ~Module(){}; - - const std::string &type() const; - - std::vector> Parameters() const; - bool has_parameter(const std::string &name) const; - std::shared_ptr *mutable_parameter(const std::string &name); - const std::shared_ptr ¶meter(const std::string &name) const; - - std::vector modules() const; - Module *mutable_module(const std::string &name); - const Module &module(const std::string &name) const; - - std::unordered_map> StateDict() const; - - virtual std::vector> Forward(const std::vector> &input_tensors) = 0; - - virtual void To(Device device); - - void Apply(std::function fn); - -protected: - Device device_; // CPU by default - const std::string type_ = kUndefinedType; - std::unordered_map> modules_; - std::unordered_map> parameters_; -}; -} // namespace infini_train::nn +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/device.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +class Module { +public: + static constexpr char kUndefinedType[] = "Undefined"; + + Module() : type_(kUndefinedType) {} + explicit Module(const std::string &type) : type_(type) {} + + virtual ~Module(){}; + + const std::string &type() const; + + std::vector> Parameters() const; + bool has_parameter(const std::string &name) const; + std::shared_ptr *mutable_parameter(const std::string &name); + const std::shared_ptr ¶meter(const std::string &name) const; + + std::vector modules() const; + Module *mutable_module(const std::string &name); + const Module &module(const std::string &name) const; + + std::unordered_map> StateDict() const; + + virtual std::vector> Forward(const std::vector> &input_tensors) = 0; + + virtual void To(Device device); + + void Apply(std::function fn); + +protected: + Device device_; // CPU by default + const std::string type_ = kUndefinedType; + std::unordered_map> modules_; + std::unordered_map> parameters_; +}; +} // namespace infini_train::nn diff --git a/infini_train/include/nn/modules/normalization.h b/infini_train/include/nn/modules/normalization.h index cf536c6..3279d8c 100644 --- a/infini_train/include/nn/modules/normalization.h +++ b/infini_train/include/nn/modules/normalization.h @@ -1,24 +1,24 @@ -#pragma once - -#include -#include - -#include "infini_train/include/device.h" -#include "infini_train/include/nn/modules/module.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -class LayerNorm : public Module { -public: - static constexpr char kParamWeightName[] = "weight"; - static constexpr char kParamBiasName[] = "bias"; - - LayerNorm(const std::vector &normalized_shape, float eps = 1e-5f, Device device = Device()); - std::vector> Forward(const std::vector> &input_tensors) override; - -private: - void ResetParameters(); - - const float eps_ = 1e-5f; -}; -} // namespace infini_train::nn +#pragma once + +#include +#include + +#include "infini_train/include/device.h" +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +class LayerNorm : public Module { +public: + static constexpr char kParamWeightName[] = "weight"; + static constexpr char kParamBiasName[] = "bias"; + + LayerNorm(const std::vector &normalized_shape, float eps = 1e-5f, Device device = Device()); + std::vector> Forward(const std::vector> &input_tensors) override; + +private: + void ResetParameters(); + + const float eps_ = 1e-5f; +}; +} // namespace infini_train::nn diff --git a/infini_train/include/nn/modules/sparse.h b/infini_train/include/nn/modules/sparse.h index ea9b654..98ffd5c 100644 --- a/infini_train/include/nn/modules/sparse.h +++ b/infini_train/include/nn/modules/sparse.h @@ -1,22 +1,22 @@ -#pragma once - -#include -#include - -#include "infini_train/include/nn/modules/module.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -class Embedding : public Module { -public: - static constexpr char kType[] = "Embedding"; - - static constexpr char kParamWeightName[] = "weight"; - - Embedding(int num_embeddings, int embedding_dim, Device device = Device()); - std::vector> Forward(const std::vector> &input_tensors) override; - -private: - void ResetParameters(); -}; -} // namespace infini_train::nn +#pragma once + +#include +#include + +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +class Embedding : public Module { +public: + static constexpr char kType[] = "Embedding"; + + static constexpr char kParamWeightName[] = "weight"; + + Embedding(int num_embeddings, int embedding_dim, Device device = Device()); + std::vector> Forward(const std::vector> &input_tensors) override; + +private: + void ResetParameters(); +}; +} // namespace infini_train::nn diff --git a/infini_train/include/optimizer.h b/infini_train/include/optimizer.h index 241aeb9..2712bcc 100644 --- a/infini_train/include/optimizer.h +++ b/infini_train/include/optimizer.h @@ -1,49 +1,49 @@ -#pragma once - -#include -#include - -#include "infini_train/include/tensor.h" - -namespace infini_train { -class Optimizer { -public: - explicit Optimizer(const std::vector> ¶ms); - - void ZeroGrad(); - - virtual void Step() = 0; - -protected: - std::vector> params_; -}; - -namespace optimizers { -class SGD : public Optimizer { -public: - SGD(const std::vector> ¶ms, float learning_rate); - - void Step() override; - -private: - const float learning_rate_ = 0.0; -}; - -class Adam : public Optimizer { -public: - Adam(const std::vector> ¶ms, float learning_rate = 1e-3, float beta1 = 0.9, - float beta2 = 0.999, float eps = 1e-8); - - void Step() override; - -private: - int64_t t_; - const float learning_rate_; - const float beta1_; - const float beta2_; - const float eps_; - std::vector> m_; - std::vector> v_; -}; -} // namespace optimizers -} // namespace infini_train +#pragma once + +#include +#include + +#include "infini_train/include/tensor.h" + +namespace infini_train { +class Optimizer { +public: + explicit Optimizer(const std::vector> ¶ms); + + void ZeroGrad(); + + virtual void Step() = 0; + +protected: + std::vector> params_; +}; + +namespace optimizers { +class SGD : public Optimizer { +public: + SGD(const std::vector> ¶ms, float learning_rate); + + void Step() override; + +private: + const float learning_rate_ = 0.0; +}; + +class Adam : public Optimizer { +public: + Adam(const std::vector> ¶ms, float learning_rate = 1e-3, float beta1 = 0.9, + float beta2 = 0.999, float eps = 1e-8); + + void Step() override; + +private: + int64_t t_; + const float learning_rate_; + const float beta1_; + const float beta2_; + const float eps_; + std::vector> m_; + std::vector> v_; +}; +} // namespace optimizers +} // namespace infini_train diff --git a/infini_train/include/tensor.h b/infini_train/include/tensor.h index a6479de..9398de4 100644 --- a/infini_train/include/tensor.h +++ b/infini_train/include/tensor.h @@ -1,198 +1,198 @@ -#pragma once - -#include -#include -#include -#include -#include -#include - -#include "Eigen/Dense" -#include "glog/logging.h" - -#include "infini_train/include/device.h" - -namespace infini_train { -namespace autograd { -class Function; -} - -namespace { -struct PrintOptions { - // ref: https://pytorch.org/docs/stable/generated/torch.set_printoptions.html#torch.set_printoptions - int64_t precision = 4; - int64_t threshold = 1000; - int64_t edge_items = 3; - int64_t linewidth = 80; - std::optional sci_mode; - - static PrintOptions &Get() { - static PrintOptions instance; - return instance; - } -}; -} // namespace - -enum class DataType : int8_t { - kUINT8, - kINT8, - kUINT16, - kINT16, - kUINT32, - kINT32, - kUINT64, - kINT64, - kBFLOAT16, - kFLOAT16, - kFLOAT32, - kFLOAT64, -}; - -class TensorBuffer { -public: - TensorBuffer(Device device, size_t size); - ~TensorBuffer(); - - void *DataPtr(); - const void *DataPtr() const; - - Device GetDevice() const; - size_t Size() const; - -private: - Device device_; - size_t size_ = 0; - void *data_ = nullptr; -}; - -class Tensor : public std::enable_shared_from_this { -public: - Tensor() = default; - - Tensor(const std::vector &dims, DataType dtype, Device device); - Tensor(const std::vector &dims, DataType dtype) : Tensor(dims, dtype, Device(DeviceType::kCPU, 0)) {} - Tensor(const Tensor &tensor, size_t offset, const std::vector &dims); - - Device GetDevice() const; - - void *DataPtr(); - const void *DataPtr() const; - - size_t SizeInBytes() const; - - const std::vector &Dims() const; - size_t NumElements() const; - DataType Dtype() const; - - template void Fill(T value); - - Eigen::Map> EigenMatrix(); - Eigen::Map> EigenVector(); - - Tensor To(Device device); - - // operator overloading - std::shared_ptr Equals(float scalar); - std::shared_ptr Add(const std::shared_ptr &other); - std::shared_ptr Add(float scalar); - std::shared_ptr Sub(const std::shared_ptr &other); - std::shared_ptr Mul(const std::shared_ptr &other); - std::shared_ptr Mul(float scalar); - std::shared_ptr Div(const std::shared_ptr &other); - std::shared_ptr Neg(); - std::shared_ptr Reciprocal(); - std::shared_ptr Sin(); - std::shared_ptr Cos(); - std::shared_ptr Tanh(); - std::shared_ptr Pow(float exponent); - std::shared_ptr Rsqrt(); - - std::vector> Split(int split_size, int dim = 0); - std::shared_ptr Transpose(int dim0, int dim1); - std::shared_ptr Slice(const std::vector &starts, const std::vector &ends, - const std::vector &steps); - std::shared_ptr Slice(int64_t dim, int64_t start, int64_t end, int64_t step = 1); - - std::shared_ptr RepeatInterleave(int64_t repeat, int64_t dim); - - std::shared_ptr View(const std::vector &dims); - std::shared_ptr Contiguous(); - std::shared_ptr Flatten(int64_t start = 0, int64_t end = -1); - std::shared_ptr Squeeze(int64_t dim); - - // distribution - std::shared_ptr Uniform(float from = 0.0f, float to = 1.0f, - std::optional generator = std::nullopt); - - std::shared_ptr Matmul(const std::shared_ptr &other); - std::shared_ptr Outer(const std::shared_ptr &other); - std::shared_ptr MaskedFill(const std::shared_ptr &mask, float value); - - friend std::shared_ptr operator==(const std::shared_ptr &t, float scalar); - friend std::shared_ptr operator+(const std::shared_ptr &t1, const std::shared_ptr &t2); - friend std::shared_ptr operator+(float scalar, const std::shared_ptr &t); - friend std::shared_ptr operator+(const std::shared_ptr &t, float scalar); - friend std::shared_ptr operator-(const std::shared_ptr &t1, const std::shared_ptr &t2); - friend std::shared_ptr operator-(float scalar, const std::shared_ptr &t); - friend std::shared_ptr operator-(const std::shared_ptr &t, float scalar); - friend std::shared_ptr operator-(const std::shared_ptr &t); - friend std::shared_ptr operator*(const std::shared_ptr &t1, const std::shared_ptr &t2); - friend std::shared_ptr operator*(float scalar, const std::shared_ptr &t); - friend std::shared_ptr operator*(const std::shared_ptr &t, float scalar); - friend std::shared_ptr operator/(const std::shared_ptr &t1, const std::shared_ptr &t2); - friend std::shared_ptr operator/(float scalar, const std::shared_ptr &t); - friend std::shared_ptr operator/(const std::shared_ptr &t, float scalar); - - friend std::ostream &operator<<(std::ostream &os, const Tensor &tensor); - - void SaveAsNpy(const std::string &path) const; - void Print(std::ostream &os = std::cout) const; - static void - SetPrintOptions(std::optional precision = std::nullopt, std::optional threshold = std::nullopt, - std::optional edge_items = std::nullopt, std::optional linewidth = std::nullopt, - std::optional profile = std::nullopt, std::optional sci_mode = std::nullopt); - -private: - std::shared_ptr buffer_; - size_t offset_ = 0; - std::vector dims_; - size_t num_elements_ = 0; - DataType dtype_; - - // autograd related -public: - std::shared_ptr RequiresGrad(); - - std::shared_ptr grad() const { return grad_; }; - bool requires_grad() const { return requires_grad_; } - void set_requires_grad(bool requires_grad) { requires_grad_ = requires_grad; } - - bool is_leaf() const { return is_leaf_; } - void set_is_leaf(bool is_leaf) { is_leaf_ = is_leaf; } - - std::shared_ptr grad_fn() const { return grad_fn_; } - void set_grad_fn(std::shared_ptr grad_fn) { grad_fn_ = grad_fn; } - - int output_idx() const { return output_idx_; } - void set_output_idx(int output_idx) { output_idx_ = output_idx; } - - void ZeroGrad(); - - void Backward(std::shared_ptr gradient = nullptr, bool retain_graph = false, - bool create_graph = false) const; - -private: - std::shared_ptr grad_ = nullptr; - bool requires_grad_ = false; - bool is_leaf_ = true; - std::shared_ptr grad_fn_ = nullptr; - int output_idx_ = -1; -}; - -std::shared_ptr operator==(const std::shared_ptr &t, float scalar); -std::shared_ptr operator+(const std::shared_ptr &t1, const std::shared_ptr &t2); -std::shared_ptr operator+(float scalar, const std::shared_ptr &t); -std::shared_ptr operator*(const std::shared_ptr &t1, const std::shared_ptr &t2); -std::shared_ptr operator*(float scalar, const std::shared_ptr &t); -std::shared_ptr operator*(const std::shared_ptr &t, float scalar); -} // namespace infini_train +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "Eigen/Dense" +#include "glog/logging.h" + +#include "infini_train/include/device.h" + +namespace infini_train { +namespace autograd { +class Function; +} + +namespace { +struct PrintOptions { + // ref: https://pytorch.org/docs/stable/generated/torch.set_printoptions.html#torch.set_printoptions + int64_t precision = 4; + int64_t threshold = 1000; + int64_t edge_items = 3; + int64_t linewidth = 80; + std::optional sci_mode; + + static PrintOptions &Get() { + static PrintOptions instance; + return instance; + } +}; +} // namespace + +enum class DataType : int8_t { + kUINT8, + kINT8, + kUINT16, + kINT16, + kUINT32, + kINT32, + kUINT64, + kINT64, + kBFLOAT16, + kFLOAT16, + kFLOAT32, + kFLOAT64, +}; + +class TensorBuffer { +public: + TensorBuffer(Device device, size_t size); + ~TensorBuffer(); + + void *DataPtr(); + const void *DataPtr() const; + + Device GetDevice() const; + size_t Size() const; + +private: + Device device_; + size_t size_ = 0; + void *data_ = nullptr; +}; + +class Tensor : public std::enable_shared_from_this { +public: + Tensor() = default; + + Tensor(const std::vector &dims, DataType dtype, Device device); + Tensor(const std::vector &dims, DataType dtype) : Tensor(dims, dtype, Device(DeviceType::kCPU, 0)) {} + Tensor(const Tensor &tensor, size_t offset, const std::vector &dims); + + Device GetDevice() const; + + void *DataPtr(); + const void *DataPtr() const; + + size_t SizeInBytes() const; + + const std::vector &Dims() const; + size_t NumElements() const; + DataType Dtype() const; + + template void Fill(T value); + + Eigen::Map> EigenMatrix(); + Eigen::Map> EigenVector(); + + Tensor To(Device device); + + // operator overloading + std::shared_ptr Equals(float scalar); + std::shared_ptr Add(const std::shared_ptr &other); + std::shared_ptr Add(float scalar); + std::shared_ptr Sub(const std::shared_ptr &other); + std::shared_ptr Mul(const std::shared_ptr &other); + std::shared_ptr Mul(float scalar); + std::shared_ptr Div(const std::shared_ptr &other); + std::shared_ptr Neg(); + std::shared_ptr Reciprocal(); + std::shared_ptr Sin(); + std::shared_ptr Cos(); + std::shared_ptr Tanh(); + std::shared_ptr Pow(float exponent); + std::shared_ptr Rsqrt(); + + std::vector> Split(int split_size, int dim = 0); + std::shared_ptr Transpose(int dim0, int dim1); + std::shared_ptr Slice(const std::vector &starts, const std::vector &ends, + const std::vector &steps); + std::shared_ptr Slice(int64_t dim, int64_t start, int64_t end, int64_t step = 1); + + std::shared_ptr RepeatInterleave(int64_t repeat, int64_t dim); + + std::shared_ptr View(const std::vector &dims); + std::shared_ptr Contiguous(); + std::shared_ptr Flatten(int64_t start = 0, int64_t end = -1); + std::shared_ptr Squeeze(int64_t dim); + + // distribution + std::shared_ptr Uniform(float from = 0.0f, float to = 1.0f, + std::optional generator = std::nullopt); + + std::shared_ptr Matmul(const std::shared_ptr &other); + std::shared_ptr Outer(const std::shared_ptr &other); + std::shared_ptr MaskedFill(const std::shared_ptr &mask, float value); + + friend std::shared_ptr operator==(const std::shared_ptr &t, float scalar); + friend std::shared_ptr operator+(const std::shared_ptr &t1, const std::shared_ptr &t2); + friend std::shared_ptr operator+(float scalar, const std::shared_ptr &t); + friend std::shared_ptr operator+(const std::shared_ptr &t, float scalar); + friend std::shared_ptr operator-(const std::shared_ptr &t1, const std::shared_ptr &t2); + friend std::shared_ptr operator-(float scalar, const std::shared_ptr &t); + friend std::shared_ptr operator-(const std::shared_ptr &t, float scalar); + friend std::shared_ptr operator-(const std::shared_ptr &t); + friend std::shared_ptr operator*(const std::shared_ptr &t1, const std::shared_ptr &t2); + friend std::shared_ptr operator*(float scalar, const std::shared_ptr &t); + friend std::shared_ptr operator*(const std::shared_ptr &t, float scalar); + friend std::shared_ptr operator/(const std::shared_ptr &t1, const std::shared_ptr &t2); + friend std::shared_ptr operator/(float scalar, const std::shared_ptr &t); + friend std::shared_ptr operator/(const std::shared_ptr &t, float scalar); + + friend std::ostream &operator<<(std::ostream &os, const Tensor &tensor); + + void SaveAsNpy(const std::string &path) const; + void Print(std::ostream &os = std::cout) const; + static void + SetPrintOptions(std::optional precision = std::nullopt, std::optional threshold = std::nullopt, + std::optional edge_items = std::nullopt, std::optional linewidth = std::nullopt, + std::optional profile = std::nullopt, std::optional sci_mode = std::nullopt); + +private: + std::shared_ptr buffer_; + size_t offset_ = 0; + std::vector dims_; + size_t num_elements_ = 0; + DataType dtype_; + + // autograd related +public: + std::shared_ptr RequiresGrad(); + + std::shared_ptr grad() const { return grad_; }; + bool requires_grad() const { return requires_grad_; } + void set_requires_grad(bool requires_grad) { requires_grad_ = requires_grad; } + + bool is_leaf() const { return is_leaf_; } + void set_is_leaf(bool is_leaf) { is_leaf_ = is_leaf; } + + std::shared_ptr grad_fn() const { return grad_fn_; } + void set_grad_fn(std::shared_ptr grad_fn) { grad_fn_ = grad_fn; } + + int output_idx() const { return output_idx_; } + void set_output_idx(int output_idx) { output_idx_ = output_idx; } + + void ZeroGrad(); + + void Backward(std::shared_ptr gradient = nullptr, bool retain_graph = false, + bool create_graph = false) const; + +private: + std::shared_ptr grad_ = nullptr; + bool requires_grad_ = false; + bool is_leaf_ = true; + std::shared_ptr grad_fn_ = nullptr; + int output_idx_ = -1; +}; + +std::shared_ptr operator==(const std::shared_ptr &t, float scalar); +std::shared_ptr operator+(const std::shared_ptr &t1, const std::shared_ptr &t2); +std::shared_ptr operator+(float scalar, const std::shared_ptr &t); +std::shared_ptr operator*(const std::shared_ptr &t1, const std::shared_ptr &t2); +std::shared_ptr operator*(float scalar, const std::shared_ptr &t); +std::shared_ptr operator*(const std::shared_ptr &t, float scalar); +} // namespace infini_train diff --git a/infini_train/src/autograd/activations.cc b/infini_train/src/autograd/activations.cc index 7fdb4d2..c4283dc 100644 --- a/infini_train/src/autograd/activations.cc +++ b/infini_train/src/autograd/activations.cc @@ -1,34 +1,34 @@ -#include "infini_train/include/autograd/activations.h" - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -std::vector> Sigmoid::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SigmoidForward"}); - return {kernel.Call>(input)}; -} - -void Sigmoid::SetupContext(const std::vector> &, - const std::vector> &output_tensors) { - const auto &output = output_tensors[0]; - saved_tensors_ = {output}; -} - -std::vector> Sigmoid::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 1); - const auto &output = saved_tensors_[0]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SigmoidBackward"}); - return {kernel.Call>(output, grad_output)}; -} -} // namespace infini_train::autograd +#include "infini_train/include/autograd/activations.h" + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +std::vector> Sigmoid::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SigmoidForward"}); + return {kernel.Call>(input)}; +} + +void Sigmoid::SetupContext(const std::vector> &, + const std::vector> &output_tensors) { + const auto &output = output_tensors[0]; + saved_tensors_ = {output}; +} + +std::vector> Sigmoid::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 1); + const auto &output = saved_tensors_[0]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SigmoidBackward"}); + return {kernel.Call>(output, grad_output)}; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/elementwise.cc b/infini_train/src/autograd/elementwise.cc index 5a790a5..e82e3d0 100644 --- a/infini_train/src/autograd/elementwise.cc +++ b/infini_train/src/autograd/elementwise.cc @@ -1,345 +1,355 @@ -#include "infini_train/include/autograd/elementwise.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" -#include - -namespace infini_train::autograd { -std::vector> Neg::Forward(const std::vector> &input_tensors) { - // =================================== 作业 =================================== - // TODO：通过Dispatcher获取设备专属kernel，对输入张量进行取反操作 - // NOTES: 依赖test_dispatcher，Neg kernel实现已给出 - // =================================== 作业 =================================== - - return std::vector>(); -} - -std::vector> Neg::Backward(const std::vector> &grad_outputs) { - // =================================== 作业 =================================== - // TODO：通过Dispatcher获取设备专属的反向传播kernel，计算梯度 - // NOTES: 依赖test_dispatcher，Neg的kernel实现已给出 - // =================================== 作业 =================================== - - return std::vector>(); -} - -std::vector> Reciprocal::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "ReciprocalForward"}); - return {kernel.Call>(input)}; -} - -void Reciprocal::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &input = input_tensors[0]; - saved_tensors_ = {input}; -} - -std::vector> Reciprocal::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 1); - const auto &input = saved_tensors_[0]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "ReciprocalBackward"}); - return {kernel.Call>(grad_output, input)}; -} - -std::vector> Sin::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SinForward"}); - return {kernel.Call>(input)}; -} - -void Sin::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &input = input_tensors[0]; - saved_tensors_ = {input}; -} - -std::vector> Sin::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 1); - const auto &input = saved_tensors_[0]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SinBackward"}); - return {kernel.Call>(grad_output, input)}; -} - -std::vector> Cos::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "CosForward"}); - return {kernel.Call>(input)}; -} - -void Cos::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &input = input_tensors[0]; - saved_tensors_ = {input}; -} - -std::vector> Cos::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 1); - const auto &input = saved_tensors_[0]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "CosBackward"}); - return {kernel.Call>(grad_output, input)}; -} - -std::vector> Tanh::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "TanhForward"}); - return {kernel.Call>(input)}; -} - -void Tanh::SetupContext(const std::vector> &, - const std::vector> &output_tensors) { - const auto &output = output_tensors[0]; - saved_tensors_ = {output}; -} - -std::vector> Tanh::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 1); - const auto &output = saved_tensors_[0]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "TanhBackward"}); - return {kernel.Call>(grad_output, output)}; -} - -std::vector> Pow::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "PowForward"}); - return {kernel.Call>(input, exponent_, scalar_is_base_)}; -} - -void Pow::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &input = input_tensors[0]; - saved_tensors_ = {input}; -} - -std::vector> Pow::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 1); - const auto &input = saved_tensors_[0]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "PowBackward"}); - return {kernel.Call>(grad_output, input, exponent_, scalar_is_base_)}; -} - -std::vector> Rsqrt::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "RsqrtForward"}); - return {kernel.Call>(input)}; -} - -void Rsqrt::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &input = input_tensors[0]; - saved_tensors_ = {input}; -} - -std::vector> Rsqrt::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 1); - const auto &input = saved_tensors_[0]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "RsqrtBackward"}); - return {kernel.Call>(grad_output, input)}; -} - -std::vector> EqualsScalar::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "EqualsScalarForward"}); - return {kernel.Call>(input, scalar_)}; -} - -std::vector> EqualsScalar::Backward(const std::vector> &grad_outputs) { - LOG(FATAL) << "EqualsScalar::Backward shall not be called anytime"; - return {}; -} - -std::vector> Add::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 2); - const auto &a = input_tensors[0]; - const auto &b = input_tensors[1]; - - auto device = a->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "AddForward"}); - return {kernel.Call>(a, b)}; -} - -void Add::SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) { - a_dims_ = input_tensors[0]->Dims(); - b_dims_ = input_tensors[1]->Dims(); -} - -std::vector> Add::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "AddBackward"}); - auto [grad_a, grad_b] - = kernel.Call, std::shared_ptr>>(grad_output, a_dims_, b_dims_); - return {grad_a, grad_b}; -} - -std::vector> AddScalar::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "AddScalarForward"}); - return {kernel.Call>(input, scalar_)}; -} - -std::vector> AddScalar::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "AddScalarBackward"}); - return {kernel.Call>(grad_output)}; -} - -std::vector> Sub::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 2); - const auto &a = input_tensors[0]; - const auto &b = input_tensors[1]; - - auto device = a->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SubForward"}); - return {kernel.Call>(a, b)}; -} - -void Sub::SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) { - a_dims_ = input_tensors[0]->Dims(); - b_dims_ = input_tensors[1]->Dims(); -} - -std::vector> Sub::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SubBackward"}); - auto [grad_a, grad_b] - = kernel.Call, std::shared_ptr>>(grad_output, a_dims_, b_dims_); - return {grad_a, grad_b}; -} - -std::vector> Mul::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 2); - const auto &a = input_tensors[0]; - const auto &b = input_tensors[1]; - - auto device = a->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MulForward"}); - return {kernel.Call>(a, b)}; -} - -void Mul::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &a = input_tensors[0]; - const auto &b = input_tensors[1]; - saved_tensors_ = {a, b}; -} - -std::vector> Mul::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 2); - const auto &a = saved_tensors_[0]; - const auto &b = saved_tensors_[1]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MulBackward"}); - auto [grad_a, grad_b] = kernel.Call, std::shared_ptr>>(grad_output, a, b); - return {grad_a, grad_b}; -} - -std::vector> MulScalar::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MulScalarForward"}); - return {kernel.Call>(input, scalar_)}; -} - -std::vector> MulScalar::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MulScalarBackward"}); - return {kernel.Call>(grad_output, scalar_)}; -} - -std::vector> Div::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 2); - const auto &a = input_tensors[0]; - const auto &b = input_tensors[1]; - - auto device = a->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "DivForward"}); - return {kernel.Call>(a, b)}; -} - -void Div::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &a = input_tensors[0]; - const auto &b = input_tensors[1]; - saved_tensors_ = {a, b}; -} - -std::vector> Div::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 2); - const auto &a = saved_tensors_[0]; - const auto &b = saved_tensors_[1]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "DivBackward"}); - auto [grad_a, grad_b] = kernel.Call, std::shared_ptr>>(grad_output, a, b); - return {grad_a, grad_b}; -} -} // namespace infini_train::autograd +#include "infini_train/include/autograd/elementwise.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" +#include + +namespace infini_train::autograd { +std::vector> Neg::Forward(const std::vector> &input_tensors) { + // =================================== 作业 =================================== + // TODO：通过Dispatcher获取设备专属kernel，对输入张量进行取反操作 + // NOTES: 依赖test_dispatcher，Neg kernel实现已给出 + // =================================== 作业 =================================== + + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "NegForward"}); + return {kernel.Call>(input)}; +} + +std::vector> Neg::Backward(const std::vector> &grad_outputs) { + // =================================== 作业 =================================== + // TODO：通过Dispatcher获取设备专属的反向传播kernel，计算梯度 + // NOTES: 依赖test_dispatcher，Neg的kernel实现已给出 + // =================================== 作业 =================================== + + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "NegBackward"}); + return {kernel.Call>(grad_output)}; +} + +std::vector> Reciprocal::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "ReciprocalForward"}); + return {kernel.Call>(input)}; +} + +void Reciprocal::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + saved_tensors_ = {input}; +} + +std::vector> Reciprocal::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 1); + const auto &input = saved_tensors_[0]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "ReciprocalBackward"}); + return {kernel.Call>(grad_output, input)}; +} + +std::vector> Sin::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SinForward"}); + return {kernel.Call>(input)}; +} + +void Sin::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + saved_tensors_ = {input}; +} + +std::vector> Sin::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 1); + const auto &input = saved_tensors_[0]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SinBackward"}); + return {kernel.Call>(grad_output, input)}; +} + +std::vector> Cos::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "CosForward"}); + return {kernel.Call>(input)}; +} + +void Cos::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + saved_tensors_ = {input}; +} + +std::vector> Cos::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 1); + const auto &input = saved_tensors_[0]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "CosBackward"}); + return {kernel.Call>(grad_output, input)}; +} + +std::vector> Tanh::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "TanhForward"}); + return {kernel.Call>(input)}; +} + +void Tanh::SetupContext(const std::vector> &, + const std::vector> &output_tensors) { + const auto &output = output_tensors[0]; + saved_tensors_ = {output}; +} + +std::vector> Tanh::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 1); + const auto &output = saved_tensors_[0]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "TanhBackward"}); + return {kernel.Call>(grad_output, output)}; +} + +std::vector> Pow::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "PowForward"}); + return {kernel.Call>(input, exponent_, scalar_is_base_)}; +} + +void Pow::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + saved_tensors_ = {input}; +} + +std::vector> Pow::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 1); + const auto &input = saved_tensors_[0]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "PowBackward"}); + return {kernel.Call>(grad_output, input, exponent_, scalar_is_base_)}; +} + +std::vector> Rsqrt::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "RsqrtForward"}); + return {kernel.Call>(input)}; +} + +void Rsqrt::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + saved_tensors_ = {input}; +} + +std::vector> Rsqrt::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 1); + const auto &input = saved_tensors_[0]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "RsqrtBackward"}); + return {kernel.Call>(grad_output, input)}; +} + +std::vector> EqualsScalar::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "EqualsScalarForward"}); + return {kernel.Call>(input, scalar_)}; +} + +std::vector> EqualsScalar::Backward(const std::vector> &grad_outputs) { + LOG(FATAL) << "EqualsScalar::Backward shall not be called anytime"; + return {}; +} + +std::vector> Add::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 2); + const auto &a = input_tensors[0]; + const auto &b = input_tensors[1]; + + auto device = a->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "AddForward"}); + return {kernel.Call>(a, b)}; +} + +void Add::SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) { + a_dims_ = input_tensors[0]->Dims(); + b_dims_ = input_tensors[1]->Dims(); +} + +std::vector> Add::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "AddBackward"}); + auto [grad_a, grad_b] + = kernel.Call, std::shared_ptr>>(grad_output, a_dims_, b_dims_); + return {grad_a, grad_b}; +} + +std::vector> AddScalar::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "AddScalarForward"}); + return {kernel.Call>(input, scalar_)}; +} + +std::vector> AddScalar::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "AddScalarBackward"}); + return {kernel.Call>(grad_output)}; +} + +std::vector> Sub::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 2); + const auto &a = input_tensors[0]; + const auto &b = input_tensors[1]; + + auto device = a->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SubForward"}); + return {kernel.Call>(a, b)}; +} + +void Sub::SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) { + a_dims_ = input_tensors[0]->Dims(); + b_dims_ = input_tensors[1]->Dims(); +} + +std::vector> Sub::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SubBackward"}); + auto [grad_a, grad_b] + = kernel.Call, std::shared_ptr>>(grad_output, a_dims_, b_dims_); + return {grad_a, grad_b}; +} + +std::vector> Mul::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 2); + const auto &a = input_tensors[0]; + const auto &b = input_tensors[1]; + + auto device = a->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MulForward"}); + return {kernel.Call>(a, b)}; +} + +void Mul::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &a = input_tensors[0]; + const auto &b = input_tensors[1]; + saved_tensors_ = {a, b}; +} + +std::vector> Mul::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 2); + const auto &a = saved_tensors_[0]; + const auto &b = saved_tensors_[1]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MulBackward"}); + auto [grad_a, grad_b] = kernel.Call, std::shared_ptr>>(grad_output, a, b); + return {grad_a, grad_b}; +} + +std::vector> MulScalar::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MulScalarForward"}); + return {kernel.Call>(input, scalar_)}; +} + +std::vector> MulScalar::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MulScalarBackward"}); + return {kernel.Call>(grad_output, scalar_)}; +} + +std::vector> Div::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 2); + const auto &a = input_tensors[0]; + const auto &b = input_tensors[1]; + + auto device = a->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "DivForward"}); + return {kernel.Call>(a, b)}; +} + +void Div::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &a = input_tensors[0]; + const auto &b = input_tensors[1]; + saved_tensors_ = {a, b}; +} + +std::vector> Div::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 2); + const auto &a = saved_tensors_[0]; + const auto &b = saved_tensors_[1]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "DivBackward"}); + auto [grad_a, grad_b] = kernel.Call, std::shared_ptr>>(grad_output, a, b); + return {grad_a, grad_b}; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/function.cc b/infini_train/src/autograd/function.cc index 6edf723..6cc9f9e 100644 --- a/infini_train/src/autograd/function.cc +++ b/infini_train/src/autograd/function.cc @@ -1,93 +1,101 @@ -#include "infini_train/include/autograd/function.h" - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -namespace { -class AccumulateGrad final : public Function { -public: - explicit AccumulateGrad(std::shared_ptr grad) : grad_(grad) {} - - std::vector> Forward(const std::vector> &) override { - LOG(FATAL) << "AccumulateGrad::Forward shall not be called directly!"; - return {}; - } - - std::vector> Backward(const std::vector> &) override { - LOG(FATAL) << "AccumulateGrad::Backward shall not be called directly!"; - return {}; - } - - void BackwardPartial(const std::shared_ptr &grad_output, int) override { - if (grad_output) { - auto device = grad_->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "AccumulateGrad"}); - kernel.Call(grad_output, 1.0f, grad_); - } - } - -private: - std::shared_ptr grad_ = nullptr; -}; -} // namespace - -std::vector> Function::Apply(const std::vector> &input_tensors) { - auto output_tensors = Forward(input_tensors); - SetupContext(input_tensors, output_tensors); - - bool output_requires_grad = false; - for (int idx = 0; idx < input_tensors.size(); ++idx) { - const auto &input_tensor = input_tensors[idx]; - if (input_tensor->requires_grad() && input_tensor->is_leaf()) { - next_functions_.emplace_back(std::make_shared(input_tensor->grad()), 0); - } else { - next_functions_.emplace_back(input_tensor->grad_fn(), input_tensor->output_idx()); - if (input_tensor->grad_fn()) { - input_tensor->grad_fn()->IncreaseDependenciesNumber(); - } - } - output_requires_grad |= input_tensor->requires_grad(); - } - - grad_outputs_reached_ = 0; - grad_outputs_.resize(output_tensors.size(), nullptr); - for (int output_idx = 0; output_idx < output_tensors.size(); ++output_idx) { - auto &output_tensor = output_tensors[output_idx]; - output_tensor->set_requires_grad(output_requires_grad); - output_tensor->set_is_leaf(false); - output_tensor->set_grad_fn(output_requires_grad ? shared_from_this() : nullptr); - output_tensor->set_output_idx(output_idx); - } - - return output_tensors; -} - -void Function::BackwardPartial(const std::shared_ptr &grad_output, int grad_output_idx) { - if (!grad_outputs_[grad_output_idx]) { - grad_outputs_[grad_output_idx] = grad_output; - ++grad_outputs_reached_; - } else { - auto accumulate_function = std::make_shared(grad_outputs_[grad_output_idx]); - accumulate_function->BackwardPartial(grad_output, 0); - } - ++dependencies_reached_; - if (grad_outputs_reached_ == grad_outputs_.size() - && (dependencies_reached_ == dependencies_number_ || dependencies_number_ == 0)) { - auto grad_inputs = Backward(grad_outputs_); - saved_tensors_.clear(); - grad_outputs_.clear(); - CHECK_EQ(grad_inputs.size(), next_functions_.size()); - for (int idx = 0; idx < grad_inputs.size(); ++idx) { - auto &grad_input = grad_inputs[idx]; - auto &[next_function, output_idx] = next_functions_[idx]; - if (grad_input && next_function) { - next_function->BackwardPartial(grad_input, output_idx); - } - } - } -} -void Function::IncreaseDependenciesNumber() { ++dependencies_number_; } -} // namespace infini_train::autograd +#include "infini_train/include/autograd/function.h" + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { + +thread_local int NoGradGuard::depth_ = 0; + +namespace { + +class AccumulateGrad final : public Function { +public: + explicit AccumulateGrad(std::shared_ptr grad) : grad_(grad) {} + + std::vector> Forward(const std::vector> &) override { + LOG(FATAL) << "AccumulateGrad::Forward shall not be called directly!"; + return {}; + } + + std::vector> Backward(const std::vector> &) override { + LOG(FATAL) << "AccumulateGrad::Backward shall not be called directly!"; + return {}; + } + + void BackwardPartial(const std::shared_ptr &grad_output, int) override { + if (grad_output) { + auto device = grad_->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "AccumulateGrad"}); + kernel.Call(grad_output, 1.0f, grad_); + } + } + +private: + std::shared_ptr grad_ = nullptr; +}; +} // namespace + +std::vector> Function::Apply(const std::vector> &input_tensors) { + if (NoGradGuard::is_enabled()) { + return Forward(input_tensors); + } + + auto output_tensors = Forward(input_tensors); + SetupContext(input_tensors, output_tensors); + + bool output_requires_grad = false; + for (int idx = 0; idx < input_tensors.size(); ++idx) { + const auto &input_tensor = input_tensors[idx]; + if (input_tensor->requires_grad() && input_tensor->is_leaf()) { + next_functions_.emplace_back(std::make_shared(input_tensor->grad()), 0); + } else { + next_functions_.emplace_back(input_tensor->grad_fn(), input_tensor->output_idx()); + if (input_tensor->grad_fn()) { + input_tensor->grad_fn()->IncreaseDependenciesNumber(); + } + } + output_requires_grad |= input_tensor->requires_grad(); + } + + grad_outputs_reached_ = 0; + grad_outputs_.resize(output_tensors.size(), nullptr); + for (int output_idx = 0; output_idx < output_tensors.size(); ++output_idx) { + auto &output_tensor = output_tensors[output_idx]; + output_tensor->set_requires_grad(output_requires_grad); + output_tensor->set_is_leaf(false); + output_tensor->set_grad_fn(output_requires_grad ? shared_from_this() : nullptr); + output_tensor->set_output_idx(output_idx); + } + + return output_tensors; +} + +void Function::BackwardPartial(const std::shared_ptr &grad_output, int grad_output_idx) { + if (!grad_outputs_[grad_output_idx]) { + grad_outputs_[grad_output_idx] = grad_output; + ++grad_outputs_reached_; + } else { + auto accumulate_function = std::make_shared(grad_outputs_[grad_output_idx]); + accumulate_function->BackwardPartial(grad_output, 0); + } + ++dependencies_reached_; + if (grad_outputs_reached_ == grad_outputs_.size() + && (dependencies_reached_ == dependencies_number_ || dependencies_number_ == 0)) { + auto grad_inputs = Backward(grad_outputs_); + saved_tensors_.clear(); + grad_outputs_.clear(); + CHECK_EQ(grad_inputs.size(), next_functions_.size()); + for (int idx = 0; idx < grad_inputs.size(); ++idx) { + auto &grad_input = grad_inputs[idx]; + auto &[next_function, output_idx] = next_functions_[idx]; + if (grad_input && next_function) { + next_function->BackwardPartial(grad_input, output_idx); + } + } + } +} +void Function::IncreaseDependenciesNumber() { ++dependencies_number_; } +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/linear.cc b/infini_train/src/autograd/linear.cc index 7748e59..bac8620 100644 --- a/infini_train/src/autograd/linear.cc +++ b/infini_train/src/autograd/linear.cc @@ -1,45 +1,45 @@ -#include "infini_train/include/autograd/linear.h" - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -std::vector> Linear::Forward(const std::vector> &input_tensors) { - CHECK_GE(input_tensors.size(), 2); - const auto &input = input_tensors[0]; - const auto &weight = input_tensors[1]; - const auto &bias = input_tensors.size() == 3 ? input_tensors[2] : nullptr; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "LinearForward"}); - return {kernel.Call>(input, weight, true, bias)}; -} - -void Linear::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &input = input_tensors[0]; - const auto &weight = input_tensors[1]; - saved_tensors_ = {input, weight}; - bias_ = input_tensors.size() == 3; - out_features_ = weight->Dims()[0]; -} - -std::vector> Linear::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 2); - const auto &input = saved_tensors_[0]; - const auto &weight = saved_tensors_[1]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "LinearBackward"}); - auto [grad_input, grad_weight, grad_bias] - = kernel.Call, std::shared_ptr, std::shared_ptr>>( - input, weight, true, out_features_, grad_output, bias_); - return bias_ ? std::vector>{grad_input, grad_weight, grad_bias} - : std::vector>{grad_input, grad_weight}; - ; -} -} // namespace infini_train::autograd +#include "infini_train/include/autograd/linear.h" + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +std::vector> Linear::Forward(const std::vector> &input_tensors) { + CHECK_GE(input_tensors.size(), 2); + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + const auto &bias = input_tensors.size() == 3 ? input_tensors[2] : nullptr; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "LinearForward"}); + return {kernel.Call>(input, weight, true, bias)}; +} + +void Linear::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + saved_tensors_ = {input, weight}; + bias_ = input_tensors.size() == 3; + out_features_ = weight->Dims()[0]; +} + +std::vector> Linear::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 2); + const auto &input = saved_tensors_[0]; + const auto &weight = saved_tensors_[1]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "LinearBackward"}); + auto [grad_input, grad_weight, grad_bias] + = kernel.Call, std::shared_ptr, std::shared_ptr>>( + input, weight, true, out_features_, grad_output, bias_); + return bias_ ? std::vector>{grad_input, grad_weight, grad_bias} + : std::vector>{grad_input, grad_weight}; + ; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/loss.cc b/infini_train/src/autograd/loss.cc index b3824bd..a9145ab 100644 --- a/infini_train/src/autograd/loss.cc +++ b/infini_train/src/autograd/loss.cc @@ -1,38 +1,38 @@ -#include "infini_train/include/autograd/loss.h" - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -std::vector> CrossEntropy::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 2); - const auto &input = input_tensors[0]; - const auto &target = input_tensors[1]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "CrossEntropyForward"}); - return {kernel.Call>(input, target)}; -} - -void CrossEntropy::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &input = input_tensors[0]; - const auto &target = input_tensors[1]; - saved_tensors_ = {input, target}; -} - -std::vector> CrossEntropy::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 2); - const auto &input = saved_tensors_[0]; - const auto &target = saved_tensors_[1]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "CrossEntropyBackward"}); - auto grad_input = kernel.Call>(input, target, grad_output); - return {grad_input, nullptr}; -} -} // namespace infini_train::autograd +#include "infini_train/include/autograd/loss.h" + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +std::vector> CrossEntropy::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 2); + const auto &input = input_tensors[0]; + const auto &target = input_tensors[1]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "CrossEntropyForward"}); + return {kernel.Call>(input, target)}; +} + +void CrossEntropy::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + const auto &target = input_tensors[1]; + saved_tensors_ = {input, target}; +} + +std::vector> CrossEntropy::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 2); + const auto &input = saved_tensors_[0]; + const auto &target = saved_tensors_[1]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "CrossEntropyBackward"}); + auto grad_input = kernel.Call>(input, target, grad_output); + return {grad_input, nullptr}; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/matmul.cc b/infini_train/src/autograd/matmul.cc index 9bca403..198f904 100644 --- a/infini_train/src/autograd/matmul.cc +++ b/infini_train/src/autograd/matmul.cc @@ -1,41 +1,41 @@ -#include "infini_train/include/autograd/matmul.h" - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -std::vector> Matmul::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 2); - const auto &input1 = input_tensors[0]; - const auto &input2 = input_tensors[1]; - - auto device = input1->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MatmulForward"}); - return {kernel.Call>(input1, input2)}; -} - -void Matmul::SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) { - const auto &input1 = input_tensors[0]; - const auto &input2 = input_tensors[1]; - const auto &output = output_tensors[0]; - saved_tensors_ = {input1, input2}; - out_features_ = output->Dims()[0]; -} - -std::vector> Matmul::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 2); - const auto &input1 = saved_tensors_[0]; - const auto &input2 = saved_tensors_[1]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = input1->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MatmulBackward"}); - auto [grad_input1, grad_input2] - = kernel.Call, std::shared_ptr>>(input1, input2, grad_output); - return {grad_input1, grad_input2}; -} -} // namespace infini_train::autograd +#include "infini_train/include/autograd/matmul.h" + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +std::vector> Matmul::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 2); + const auto &input1 = input_tensors[0]; + const auto &input2 = input_tensors[1]; + + auto device = input1->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MatmulForward"}); + return {kernel.Call>(input1, input2)}; +} + +void Matmul::SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) { + const auto &input1 = input_tensors[0]; + const auto &input2 = input_tensors[1]; + const auto &output = output_tensors[0]; + saved_tensors_ = {input1, input2}; + out_features_ = output->Dims()[0]; +} + +std::vector> Matmul::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 2); + const auto &input1 = saved_tensors_[0]; + const auto &input2 = saved_tensors_[1]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input1->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MatmulBackward"}); + auto [grad_input1, grad_input2] + = kernel.Call, std::shared_ptr>>(input1, input2, grad_output); + return {grad_input1, grad_input2}; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/misc.cc b/infini_train/src/autograd/misc.cc index 59d395e..759d0c6 100644 --- a/infini_train/src/autograd/misc.cc +++ b/infini_train/src/autograd/misc.cc @@ -1,101 +1,101 @@ -#include "infini_train/include/autograd/misc.h" - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -std::vector> Split::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SplitForward"}); - return {kernel.Call>>(input, split_size_, dim_)}; -} - -void Split::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &input = input_tensors[0]; - input_dims_ = input->Dims(); -} - -std::vector> Split::Backward(const std::vector> &grad_outputs) { - auto device = grad_outputs[0]->GetDevice(); - auto kernel = Dispatcher::Instance().GetKernel({device.Type(), "SplitBackward"}); - return {kernel.Call>(input_dims_, split_size_, dim_, grad_outputs)}; -} - -std::vector> NoOp::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "NoOpForward"}); - return {kernel.Call>(input, output_dims_)}; -} - -void NoOp::SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) { - const auto &input = input_tensors[0]; - input_dims_ = input->Dims(); -} - -std::vector> NoOp::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "NoOpBackward"}); - return {kernel.Call>(input_dims_, grad_output)}; -} - -std::vector> Slice::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SliceForward"}); - return {kernel.Call>(input, starts_, ends_, steps_)}; -} - -void Slice::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - // FIXME(dcj): only input's dim need to be saved - const auto &input = input_tensors[0]; - saved_tensors_ = {input}; -} - -std::vector> Slice::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 1); - const auto &input = saved_tensors_[0]; - const auto &grad_output = grad_outputs[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SliceBackward"}); - return {kernel.Call>(grad_output, input, starts_, ends_, steps_)}; -} - -std::vector> Stack::Forward(const std::vector> &input_tensors) { - CHECK_GE(input_tensors.size(), 2); - const auto device = input_tensors[0]->GetDevice().Type(); - - auto kernel = Dispatcher::Instance().GetKernel({device, "StackForward"}); - return {kernel.Call>(input_tensors, dim_)}; -} - -void Stack::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &input = input_tensors[0]; - input_dims_ = input->Dims(); -} - -std::vector> Stack::Backward(const std::vector> &grad_outputs) { - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "StackBackward"}); - return kernel.Call>>(input_dims_, dim_, grad_output); -} -} // namespace infini_train::autograd +#include "infini_train/include/autograd/misc.h" + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +std::vector> Split::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SplitForward"}); + return {kernel.Call>>(input, split_size_, dim_)}; +} + +void Split::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + input_dims_ = input->Dims(); +} + +std::vector> Split::Backward(const std::vector> &grad_outputs) { + auto device = grad_outputs[0]->GetDevice(); + auto kernel = Dispatcher::Instance().GetKernel({device.Type(), "SplitBackward"}); + return {kernel.Call>(input_dims_, split_size_, dim_, grad_outputs)}; +} + +std::vector> NoOp::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "NoOpForward"}); + return {kernel.Call>(input, output_dims_)}; +} + +void NoOp::SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) { + const auto &input = input_tensors[0]; + input_dims_ = input->Dims(); +} + +std::vector> NoOp::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "NoOpBackward"}); + return {kernel.Call>(input_dims_, grad_output)}; +} + +std::vector> Slice::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SliceForward"}); + return {kernel.Call>(input, starts_, ends_, steps_)}; +} + +void Slice::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + // FIXME(dcj): only input's dim need to be saved + const auto &input = input_tensors[0]; + saved_tensors_ = {input}; +} + +std::vector> Slice::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 1); + const auto &input = saved_tensors_[0]; + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SliceBackward"}); + return {kernel.Call>(grad_output, input, starts_, ends_, steps_)}; +} + +std::vector> Stack::Forward(const std::vector> &input_tensors) { + CHECK_GE(input_tensors.size(), 2); + const auto device = input_tensors[0]->GetDevice().Type(); + + auto kernel = Dispatcher::Instance().GetKernel({device, "StackForward"}); + return {kernel.Call>(input_tensors, dim_)}; +} + +void Stack::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + input_dims_ = input->Dims(); +} + +std::vector> Stack::Backward(const std::vector> &grad_outputs) { + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "StackBackward"}); + return kernel.Call>>(input_dims_, dim_, grad_output); +} +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/normalization.cc b/infini_train/src/autograd/normalization.cc index 95b64f9..a7adce3 100644 --- a/infini_train/src/autograd/normalization.cc +++ b/infini_train/src/autograd/normalization.cc @@ -1,50 +1,50 @@ -#include "infini_train/include/autograd/normalization.h" - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { - -std::vector> LayerNorm::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 3); - const auto &input = input_tensors[0]; - const auto &weight = input_tensors[1]; - const auto &bias = input_tensors[2]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "LayerNormForward"}); - auto [output, mean, rstd] - = kernel.Call, std::shared_ptr, std::shared_ptr>>( - input, weight, bias, eps_); - saved_tensors_ = {mean, rstd}; - return {output}; -} - -void LayerNorm::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &input = input_tensors[0]; - const auto &weight = input_tensors[1]; - const auto &bias = input_tensors[2]; - saved_tensors_.insert(saved_tensors_.begin(), {input, weight, bias}); -} - -std::vector> LayerNorm::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 5); - const auto &input = saved_tensors_[0]; - const auto &weight = saved_tensors_[1]; - const auto &bias = saved_tensors_[2]; - const auto &mean = saved_tensors_[3]; - const auto &rstd = saved_tensors_[4]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "LayerNormBackward"}); - auto [grad_input, grad_weight, grad_bias] - = kernel.Call, std::shared_ptr, std::shared_ptr>>( - input, weight, bias, mean, rstd, grad_output); - return {grad_input, grad_weight, grad_bias}; -} -} // namespace infini_train::autograd +#include "infini_train/include/autograd/normalization.h" + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { + +std::vector> LayerNorm::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 3); + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + const auto &bias = input_tensors[2]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "LayerNormForward"}); + auto [output, mean, rstd] + = kernel.Call, std::shared_ptr, std::shared_ptr>>( + input, weight, bias, eps_); + saved_tensors_ = {mean, rstd}; + return {output}; +} + +void LayerNorm::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + const auto &bias = input_tensors[2]; + saved_tensors_.insert(saved_tensors_.begin(), {input, weight, bias}); +} + +std::vector> LayerNorm::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 5); + const auto &input = saved_tensors_[0]; + const auto &weight = saved_tensors_[1]; + const auto &bias = saved_tensors_[2]; + const auto &mean = saved_tensors_[3]; + const auto &rstd = saved_tensors_[4]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "LayerNormBackward"}); + auto [grad_input, grad_weight, grad_bias] + = kernel.Call, std::shared_ptr, std::shared_ptr>>( + input, weight, bias, mean, rstd, grad_output); + return {grad_input, grad_weight, grad_bias}; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/outer.cc b/infini_train/src/autograd/outer.cc index 7d9339f..4b414f2 100644 --- a/infini_train/src/autograd/outer.cc +++ b/infini_train/src/autograd/outer.cc @@ -1,43 +1,43 @@ -#include "infini_train/include/autograd/outer.h" - -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -std::vector> Outer::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 2); - const auto &input1 = input_tensors[0]; - const auto &input2 = input_tensors[1]; - - auto device = input1->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "OuterForward"}); - return {kernel.Call>(input1, input2)}; -} - -void Outer::SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) { - const auto &input1 = input_tensors[0]; - const auto &input2 = input_tensors[1]; - saved_tensors_ = {input1, input2}; -} - -std::vector> Outer::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 2); - const auto &input1 = saved_tensors_[0]; - const auto &input2 = saved_tensors_[1]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = input1->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "OuterBackward"}); - auto [grad_input1, grad_input2] - = kernel.Call, std::shared_ptr>>(input1, input2, grad_output); - return {grad_input1, grad_input2}; -} -} // namespace infini_train::autograd +#include "infini_train/include/autograd/outer.h" + +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +std::vector> Outer::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 2); + const auto &input1 = input_tensors[0]; + const auto &input2 = input_tensors[1]; + + auto device = input1->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "OuterForward"}); + return {kernel.Call>(input1, input2)}; +} + +void Outer::SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) { + const auto &input1 = input_tensors[0]; + const auto &input2 = input_tensors[1]; + saved_tensors_ = {input1, input2}; +} + +std::vector> Outer::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 2); + const auto &input1 = saved_tensors_[0]; + const auto &input2 = saved_tensors_[1]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input1->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "OuterBackward"}); + auto [grad_input1, grad_input2] + = kernel.Call, std::shared_ptr>>(input1, input2, grad_output); + return {grad_input1, grad_input2}; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/reduction.cc b/infini_train/src/autograd/reduction.cc index cbe3b21..fde90d5 100644 --- a/infini_train/src/autograd/reduction.cc +++ b/infini_train/src/autograd/reduction.cc @@ -1,115 +1,115 @@ -#include "infini_train/include/autograd/reduction.h" - -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -std::vector> Mean::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MeanForward"}); - return {kernel.Call>(input, dim_, keep_dim_)}; -} - -void Mean::SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) { - const auto &input = input_tensors[0]; - input_dims_ = input->Dims(); -} - -std::vector> Mean::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MeanBackward"}); - return {kernel.Call>(grad_output, input_dims_, dim_, keep_dim_)}; -} - -std::vector> Sum::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SumForward"}); - return {kernel.Call>(input, dim_, keep_dim_)}; -} - -void Sum::SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) { - const auto &input = input_tensors[0]; - input_dims_ = input->Dims(); -} - -std::vector> Sum::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SumBackward"}); - return {kernel.Call>(grad_output, input_dims_, dim_, keep_dim_)}; -} - -std::vector> Max::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MaxForward"}); - return {kernel.Call>(input, dim_, keep_dim_)}; -} - -void Max::SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) { - const auto &input = input_tensors[0]; - const auto &output = output_tensors[0]; - saved_tensors_ = {input, output}; -} - -std::vector> Max::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(grad_outputs.size(), 1); - CHECK_EQ(saved_tensors_.size(), 2); - const auto &grad_output = grad_outputs[0]; - const auto &input = saved_tensors_[0]; - const auto &reduced = saved_tensors_[1]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MaxBackward"}); - return {kernel.Call>(grad_output, input, reduced, dim_, keep_dim_)}; -} - -std::vector> Min::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MinForward"}); - return {kernel.Call>(input, dim_, keep_dim_)}; -} - -void Min::SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) { - const auto &input = input_tensors[0]; - const auto &output = output_tensors[0]; - saved_tensors_ = {input, output}; -} - -std::vector> Min::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(grad_outputs.size(), 1); - CHECK_EQ(saved_tensors_.size(), 2); - const auto &grad_output = grad_outputs[0]; - const auto &input = saved_tensors_[0]; - const auto &reduced = saved_tensors_[1]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MinBackward"}); - return {kernel.Call>(grad_output, input, reduced, dim_, keep_dim_)}; -} -} // namespace infini_train::autograd +#include "infini_train/include/autograd/reduction.h" + +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +std::vector> Mean::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MeanForward"}); + return {kernel.Call>(input, dim_, keep_dim_)}; +} + +void Mean::SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) { + const auto &input = input_tensors[0]; + input_dims_ = input->Dims(); +} + +std::vector> Mean::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MeanBackward"}); + return {kernel.Call>(grad_output, input_dims_, dim_, keep_dim_)}; +} + +std::vector> Sum::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SumForward"}); + return {kernel.Call>(input, dim_, keep_dim_)}; +} + +void Sum::SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) { + const auto &input = input_tensors[0]; + input_dims_ = input->Dims(); +} + +std::vector> Sum::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SumBackward"}); + return {kernel.Call>(grad_output, input_dims_, dim_, keep_dim_)}; +} + +std::vector> Max::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MaxForward"}); + return {kernel.Call>(input, dim_, keep_dim_)}; +} + +void Max::SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) { + const auto &input = input_tensors[0]; + const auto &output = output_tensors[0]; + saved_tensors_ = {input, output}; +} + +std::vector> Max::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + CHECK_EQ(saved_tensors_.size(), 2); + const auto &grad_output = grad_outputs[0]; + const auto &input = saved_tensors_[0]; + const auto &reduced = saved_tensors_[1]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MaxBackward"}); + return {kernel.Call>(grad_output, input, reduced, dim_, keep_dim_)}; +} + +std::vector> Min::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MinForward"}); + return {kernel.Call>(input, dim_, keep_dim_)}; +} + +void Min::SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) { + const auto &input = input_tensors[0]; + const auto &output = output_tensors[0]; + saved_tensors_ = {input, output}; +} + +std::vector> Min::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + CHECK_EQ(saved_tensors_.size(), 2); + const auto &grad_output = grad_outputs[0]; + const auto &input = saved_tensors_[0]; + const auto &reduced = saved_tensors_[1]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MinBackward"}); + return {kernel.Call>(grad_output, input, reduced, dim_, keep_dim_)}; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/softmax.cc b/infini_train/src/autograd/softmax.cc index cd0e42a..45f770f 100644 --- a/infini_train/src/autograd/softmax.cc +++ b/infini_train/src/autograd/softmax.cc @@ -1,34 +1,34 @@ -#include "infini_train/include/autograd/softmax.h" - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -std::vector> Softmax::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SoftmaxForward"}); - return {kernel.Call>(input, dim_)}; -} - -void Softmax::SetupContext(const std::vector> &, - const std::vector> &output_tensors) { - const auto &output = output_tensors[0]; - saved_tensors_ = {output}; -} - -std::vector> Softmax::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(saved_tensors_.size(), 1); - const auto &output = saved_tensors_[0]; - CHECK_EQ(grad_outputs.size(), 1); - const auto &grad_output = grad_outputs[0]; - - auto device = output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "SoftmaxBackward"}); - return {kernel.Call>(grad_output, output, dim_)}; -} -} // namespace infini_train::autograd +#include "infini_train/include/autograd/softmax.h" + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +std::vector> Softmax::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SoftmaxForward"}); + return {kernel.Call>(input, dim_)}; +} + +void Softmax::SetupContext(const std::vector> &, + const std::vector> &output_tensors) { + const auto &output = output_tensors[0]; + saved_tensors_ = {output}; +} + +std::vector> Softmax::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(saved_tensors_.size(), 1); + const auto &output = saved_tensors_[0]; + CHECK_EQ(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "SoftmaxBackward"}); + return {kernel.Call>(grad_output, output, dim_)}; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/sparse.cc b/infini_train/src/autograd/sparse.cc index 39a0e4f..a439bc4 100644 --- a/infini_train/src/autograd/sparse.cc +++ b/infini_train/src/autograd/sparse.cc @@ -1,37 +1,37 @@ -#include "infini_train/include/autograd/sparse.h" - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::autograd { -std::vector> Embedding::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 2); - const auto &input = input_tensors[0]; - const auto &weight = input_tensors[1]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "EmbeddingForward"}); - return {kernel.Call>(input, weight)}; -} - -void Embedding::SetupContext(const std::vector> &input_tensors, - const std::vector> &output_tensors) { - const auto &input = input_tensors[0]; - const auto &weight = input_tensors[1]; - weight_dims_ = weight->Dims(); - saved_tensors_ = {input}; -} - -std::vector> Embedding::Backward(const std::vector> &grad_outputs) { - CHECK_EQ(grad_outputs.size(), 1); - const auto &input = saved_tensors_[0]; - const auto &grad_output = grad_outputs[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "EmbeddingBackward"}); - auto grad_weight = kernel.Call>(input, weight_dims_, grad_output); - return {nullptr, grad_weight}; -} -} // namespace infini_train::autograd +#include "infini_train/include/autograd/sparse.h" + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::autograd { +std::vector> Embedding::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 2); + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "EmbeddingForward"}); + return {kernel.Call>(input, weight)}; +} + +void Embedding::SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) { + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + weight_dims_ = weight->Dims(); + saved_tensors_ = {input}; +} + +std::vector> Embedding::Backward(const std::vector> &grad_outputs) { + CHECK_EQ(grad_outputs.size(), 1); + const auto &input = saved_tensors_[0]; + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "EmbeddingBackward"}); + auto grad_weight = kernel.Call>(input, weight_dims_, grad_output); + return {nullptr, grad_weight}; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/autograd/transform.cc b/infini_train/src/autograd/transform.cc index aad7edc..cdcfb8d 100644 --- a/infini_train/src/autograd/transform.cc +++ b/infini_train/src/autograd/transform.cc @@ -1,98 +1,98 @@ -#include "infini_train/include/autograd/transform.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" -namespace infini_train::autograd { -std::vector> Tril::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "TrilForward"}); - return {kernel.Call>(input, diagonal_)}; -} - -std::vector> Tril::Backward(const std::vector> &grad_outputs) { - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "TrilBackward"}); - return {kernel.Call>(grad_output, diagonal_)}; -} - -std::vector> Triu::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "TriuForward"}); - return {kernel.Call>(input, diagonal_)}; -} - -std::vector> Triu::Backward(const std::vector> &grad_outputs) { - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "TriuBackward"}); - return {kernel.Call>(grad_output, diagonal_)}; -} - -std::vector> Transpose::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "TransposeForward"}); - return {kernel.Call>(input, dim0_, dim1_)}; -} - -std::vector> Transpose::Backward(const std::vector> &grad_outputs) { - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "TransposeBackward"}); - return {kernel.Call>(grad_output, dim0_, dim1_)}; -} - -std::vector> Mask::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MaskForward"}); - return {kernel.Call>(input, mask_, value_)}; -} - -std::vector> Mask::Backward(const std::vector> &grad_outputs) { - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "MaskBackward"}); - return {kernel.Call>(grad_output, mask_)}; -} - -std::vector> -RepeatInterleave::Forward(const std::vector> &input_tensors) { - CHECK_EQ(input_tensors.size(), 1); - const auto &input = input_tensors[0]; - - auto device = input->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "RepeatInterleaveForward"}); - return {kernel.Call>(input, repeat_, dim_)}; -} - -void RepeatInterleave::SetupContext(const std::vector> &input_tensors, - const std::vector> &) { - const auto &input = input_tensors[0]; - input_dims_ = input->Dims(); -} - -std::vector> -RepeatInterleave::Backward(const std::vector> &grad_outputs) { - const auto &grad_output = grad_outputs[0]; - - auto device = grad_output->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "RepeatInterleaveBackward"}); - return {kernel.Call>(grad_output, input_dims_, dim_)}; -} -} // namespace infini_train::autograd +#include "infini_train/include/autograd/transform.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" +namespace infini_train::autograd { +std::vector> Tril::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "TrilForward"}); + return {kernel.Call>(input, diagonal_)}; +} + +std::vector> Tril::Backward(const std::vector> &grad_outputs) { + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "TrilBackward"}); + return {kernel.Call>(grad_output, diagonal_)}; +} + +std::vector> Triu::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "TriuForward"}); + return {kernel.Call>(input, diagonal_)}; +} + +std::vector> Triu::Backward(const std::vector> &grad_outputs) { + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "TriuBackward"}); + return {kernel.Call>(grad_output, diagonal_)}; +} + +std::vector> Transpose::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "TransposeForward"}); + return {kernel.Call>(input, dim0_, dim1_)}; +} + +std::vector> Transpose::Backward(const std::vector> &grad_outputs) { + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "TransposeBackward"}); + return {kernel.Call>(grad_output, dim0_, dim1_)}; +} + +std::vector> Mask::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MaskForward"}); + return {kernel.Call>(input, mask_, value_)}; +} + +std::vector> Mask::Backward(const std::vector> &grad_outputs) { + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "MaskBackward"}); + return {kernel.Call>(grad_output, mask_)}; +} + +std::vector> +RepeatInterleave::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + + auto device = input->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "RepeatInterleaveForward"}); + return {kernel.Call>(input, repeat_, dim_)}; +} + +void RepeatInterleave::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + const auto &input = input_tensors[0]; + input_dims_ = input->Dims(); +} + +std::vector> +RepeatInterleave::Backward(const std::vector> &grad_outputs) { + const auto &grad_output = grad_outputs[0]; + + auto device = grad_output->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "RepeatInterleaveBackward"}); + return {kernel.Call>(grad_output, input_dims_, dim_)}; +} +} // namespace infini_train::autograd diff --git a/infini_train/src/dataloader.cc b/infini_train/src/dataloader.cc index 9b022f6..e827bc6 100644 --- a/infini_train/src/dataloader.cc +++ b/infini_train/src/dataloader.cc @@ -1,87 +1,87 @@ -#include "infini_train/include/dataloader.h" - -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dataset.h" -#include "infini_train/include/tensor.h" - -namespace infini_train { -namespace { -std::shared_ptr Stack(const std::vector> &tensors) { - const int batch_size = tensors.size(); - const auto &dims = tensors[0]->Dims(); - const int stacked_dim = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); - auto stacked_tensor = std::make_shared(std::vector{batch_size, stacked_dim}, tensors[0]->Dtype()); - for (const auto &tensor : tensors) { - CHECK_EQ(static_cast(tensors[0]->Dtype()), static_cast(tensor->Dtype())); - const auto &dims = tensor->Dims(); - CHECK_EQ(stacked_dim, std::accumulate(dims.begin(), dims.end(), 1, std::multiplies())); - } - - size_t offset = 0; - for (const auto &tensor : tensors) { - memcpy(reinterpret_cast(stacked_tensor->DataPtr()) + offset, tensor->DataPtr(), - tensor->SizeInBytes()); - offset += tensor->SizeInBytes(); - } - return stacked_tensor; -} -} // namespace - -DataLoaderIterator::DataLoaderIterator(const Dataset &dataset, size_t batch_size, size_t batch_idx, - size_t max_batch_idx) - : dataset_(&dataset), batch_size_(batch_size), batch_idx_(batch_idx), max_batch_idx_(max_batch_idx){}; - -std::pair, std::shared_ptr> DataLoaderIterator::operator*() const { - /* - 0, 1, ..., x, ... - [0, bs-1], [bs, 2*bs-1], ..., [x*bs, (x+1)*bs-1], ... - ^ - batch_idx - */ - std::vector> data_vec; - std::vector> label_vec; - for (int idx = batch_idx_ * batch_size_; idx < (batch_idx_ + 1) * batch_size_ && idx < dataset_->Size(); ++idx) { - auto &&[data, label] = dataset_->operator[](idx); - data_vec.push_back(std::move(data)); - label_vec.push_back(std::move(label)); - } - return {Stack(std::move(data_vec)), Stack(std::move(label_vec))}; -} - -DataLoaderIterator &DataLoaderIterator::operator++() { - batch_idx_ = std::min(batch_idx_ + 1, max_batch_idx_); - return *this; -} - -DataLoaderIterator DataLoaderIterator::operator++(int) { - DataLoaderIterator tmp(*this); - ++(*this); - return tmp; -} - -bool operator<(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs) { return lhs.batch_idx_ < rhs.batch_idx_; } - -bool operator!=(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs) { - return lhs.batch_idx_ != rhs.batch_idx_; -} - -bool operator==(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs) { - return lhs.batch_idx_ == rhs.batch_idx_; -} - -DataLoader::DataLoader(const std::shared_ptr &dataset, size_t batch_size) - : dataset_(dataset), batch_size_(batch_size), max_batch_idx_((dataset_->Size() + batch_size_ - 1) / batch_size_) {} - -DataLoaderIterator DataLoader::begin() const { return DataLoaderIterator(*dataset_, batch_size_, 0, max_batch_idx_); } - -DataLoaderIterator DataLoader::end() const { - return DataLoaderIterator(*dataset_, batch_size_, max_batch_idx_, max_batch_idx_); -} -} // namespace infini_train +#include "infini_train/include/dataloader.h" + +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dataset.h" +#include "infini_train/include/tensor.h" + +namespace infini_train { +namespace { +std::shared_ptr Stack(const std::vector> &tensors) { + const int batch_size = tensors.size(); + const auto &dims = tensors[0]->Dims(); + const int stacked_dim = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); + auto stacked_tensor = std::make_shared(std::vector{batch_size, stacked_dim}, tensors[0]->Dtype()); + for (const auto &tensor : tensors) { + CHECK_EQ(static_cast(tensors[0]->Dtype()), static_cast(tensor->Dtype())); + const auto &dims = tensor->Dims(); + CHECK_EQ(stacked_dim, std::accumulate(dims.begin(), dims.end(), 1, std::multiplies())); + } + + size_t offset = 0; + for (const auto &tensor : tensors) { + memcpy(reinterpret_cast(stacked_tensor->DataPtr()) + offset, tensor->DataPtr(), + tensor->SizeInBytes()); + offset += tensor->SizeInBytes(); + } + return stacked_tensor; +} +} // namespace + +DataLoaderIterator::DataLoaderIterator(const Dataset &dataset, size_t batch_size, size_t batch_idx, + size_t max_batch_idx) + : dataset_(&dataset), batch_size_(batch_size), batch_idx_(batch_idx), max_batch_idx_(max_batch_idx){}; + +std::pair, std::shared_ptr> DataLoaderIterator::operator*() const { + /* + 0, 1, ..., x, ... + [0, bs-1], [bs, 2*bs-1], ..., [x*bs, (x+1)*bs-1], ... + ^ + batch_idx + */ + std::vector> data_vec; + std::vector> label_vec; + for (int idx = batch_idx_ * batch_size_; idx < (batch_idx_ + 1) * batch_size_ && idx < dataset_->Size(); ++idx) { + auto &&[data, label] = dataset_->operator[](idx); + data_vec.push_back(std::move(data)); + label_vec.push_back(std::move(label)); + } + return {Stack(std::move(data_vec)), Stack(std::move(label_vec))}; +} + +DataLoaderIterator &DataLoaderIterator::operator++() { + batch_idx_ = std::min(batch_idx_ + 1, max_batch_idx_); + return *this; +} + +DataLoaderIterator DataLoaderIterator::operator++(int) { + DataLoaderIterator tmp(*this); + ++(*this); + return tmp; +} + +bool operator<(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs) { return lhs.batch_idx_ < rhs.batch_idx_; } + +bool operator!=(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs) { + return lhs.batch_idx_ != rhs.batch_idx_; +} + +bool operator==(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs) { + return lhs.batch_idx_ == rhs.batch_idx_; +} + +DataLoader::DataLoader(const std::shared_ptr &dataset, size_t batch_size) + : dataset_(dataset), batch_size_(batch_size), max_batch_idx_((dataset_->Size() + batch_size_ - 1) / batch_size_) {} + +DataLoaderIterator DataLoader::begin() const { return DataLoaderIterator(*dataset_, batch_size_, 0, max_batch_idx_); } + +DataLoaderIterator DataLoader::end() const { + return DataLoaderIterator(*dataset_, batch_size_, max_batch_idx_, max_batch_idx_); +} +} // namespace infini_train diff --git a/infini_train/src/device.cc b/infini_train/src/device.cc index cc6842f..00bfe2d 100644 --- a/infini_train/src/device.cc +++ b/infini_train/src/device.cc @@ -1,41 +1,41 @@ -#include "infini_train/include/device.h" - -#include - -#include "glog/logging.h" - -namespace infini_train { -Device::Device() : type_(DeviceType::kCPU), index_(0) {} - -Device::Device(DeviceType type, int8_t index) : type_(type), index_(index) { - if (type_ == DeviceType::kCPU && index_ != 0) { - LOG(FATAL) << "CPU device index should be 0"; - } - - if (type_ == DeviceType::kCUDA && index_ != 0) { - LOG(FATAL) << "CUDA device index should be 0"; - } -} - -bool Device::operator==(const Device &other) const { return type_ == other.type_ && index_ == other.index_; } - -bool Device::operator!=(const Device &other) const { return !(*this == other); } - -DeviceType Device::Type() const { return type_; } -int8_t Device::Index() const { return index_; } - -bool Device::IsCPU() const { return type_ == DeviceType::kCPU; } -bool Device::IsCUDA() const { return type_ == DeviceType::kCUDA; } - -std::string Device::ToString() const { - std::ostringstream oss; - oss << "Device(" << (type_ == DeviceType::kCPU ? "CPU" : "CUDA") << ", " << static_cast(index_) << ")"; - return oss.str(); -} - -std::ostream &operator<<(std::ostream &os, const Device &device) { - os << device.ToString(); - return os; -} - -} // namespace infini_train +#include "infini_train/include/device.h" + +#include + +#include "glog/logging.h" + +namespace infini_train { +Device::Device() : type_(DeviceType::kCPU), index_(0) {} + +Device::Device(DeviceType type, int8_t index) : type_(type), index_(index) { + if (type_ == DeviceType::kCPU && index_ != 0) { + LOG(FATAL) << "CPU device index should be 0"; + } + + if (type_ == DeviceType::kCUDA && index_ != 0) { + LOG(FATAL) << "CUDA device index should be 0"; + } +} + +bool Device::operator==(const Device &other) const { return type_ == other.type_ && index_ == other.index_; } + +bool Device::operator!=(const Device &other) const { return !(*this == other); } + +DeviceType Device::Type() const { return type_; } +int8_t Device::Index() const { return index_; } + +bool Device::IsCPU() const { return type_ == DeviceType::kCPU; } +bool Device::IsCUDA() const { return type_ == DeviceType::kCUDA; } + +std::string Device::ToString() const { + std::ostringstream oss; + oss << "Device(" << (type_ == DeviceType::kCPU ? "CPU" : "CUDA") << ", " << static_cast(index_) << ")"; + return oss.str(); +} + +std::ostream &operator<<(std::ostream &os, const Device &device) { + os << device.ToString(); + return os; +} + +} // namespace infini_train diff --git a/infini_train/src/dispatcher.cc b/infini_train/src/dispatcher.cc index fe33b5d..b252ad9 100644 --- a/infini_train/src/dispatcher.cc +++ b/infini_train/src/dispatcher.cc @@ -1,23 +1,23 @@ -#include "infini_train/include/dispatcher.h" - -namespace infini_train { - -// Dispatcher &Dispatcher::Instance() { -// static Dispatcher instance; -// return instance; -// } - -// bool Dispatcher::RegisterImpl(const DispatchKey &key, const std::function &kernel) { -// CHECK(dispatch_table_.find(key) == dispatch_table_.end()) -// << "Kernel already registered for key " << std::get<1>(key); -// dispatch_table_.emplace(key, kernel); -// return true; -// } - -// std::function Dispatcher::GetImpl(const DispatchKey &key) const { -// auto it = dispatch_table_.find(key); -// CHECK(it != dispatch_table_.end()) << "No kernel registered for key " << std::get<1>(key); -// return it->second; -// } - -} // namespace infini_train +#include "infini_train/include/dispatcher.h" + +namespace infini_train { + +// Dispatcher &Dispatcher::Instance() { +// static Dispatcher instance; +// return instance; +// } + +// bool Dispatcher::RegisterImpl(const DispatchKey &key, const std::function &kernel) { +// CHECK(dispatch_table_.find(key) == dispatch_table_.end()) +// << "Kernel already registered for key " << std::get<1>(key); +// dispatch_table_.emplace(key, kernel); +// return true; +// } + +// std::function Dispatcher::GetImpl(const DispatchKey &key) const { +// auto it = dispatch_table_.find(key); +// CHECK(it != dispatch_table_.end()) << "No kernel registered for key " << std::get<1>(key); +// return it->second; +// } + +} // namespace infini_train diff --git a/infini_train/src/kernels/cpu/accumulate_grad.cc b/infini_train/src/kernels/cpu/accumulate_grad.cc index 55637cd..0232959 100644 --- a/infini_train/src/kernels/cpu/accumulate_grad.cc +++ b/infini_train/src/kernels/cpu/accumulate_grad.cc @@ -1,31 +1,48 @@ -#include -#include - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -void AccumulateGrad(const std::shared_ptr &gradient, float rate, const std::shared_ptr &tensor) { - for (int64_t idx = 0; idx < gradient->NumElements(); ++idx) { - static_cast(tensor->DataPtr())[idx] += rate * static_cast(gradient->DataPtr())[idx]; - } -} - -void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_ptr ¶m, - const std::shared_ptr &m, const std::shared_ptr &v, float learning_rate, - float beta1, float beta2, float eps, int64_t t) { - // =================================== 作业 =================================== - // TODO：实现Adam优化器的梯度累积和参数更新 - // REF: - // =================================== 作业 =================================== -} - -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(AccumulateGrad) -REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(AdamAccumulateGrad) - -#undef REGISTER_CPU_ACCUMULATE_GRAD_KERNEL +#include +#include + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +void AccumulateGrad(const std::shared_ptr &gradient, float rate, const std::shared_ptr &tensor) { + for (int64_t idx = 0; idx < gradient->NumElements(); ++idx) { + static_cast(tensor->DataPtr())[idx] += rate * static_cast(gradient->DataPtr())[idx]; + } +} + +void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_ptr ¶m, + const std::shared_ptr &m, const std::shared_ptr &v, float learning_rate, + float beta1, float beta2, float eps, int64_t t) { + // =================================== 作业 =================================== + // TODO：实现Adam优化器的梯度累积和参数更新 + // REF: + // =================================== 作业 =================================== + + int64_t num_elements = grad->NumElements(); + float *grad_ptr = static_cast(grad->DataPtr()); + float *param_ptr = static_cast(param->DataPtr()); + float *m_ptr = static_cast(m->DataPtr()); + float *v_ptr = static_cast(v->DataPtr()); + + float beta1_t = std::pow(beta1, t); + float beta2_t = std::pow(beta2, t); + float alpha = learning_rate * std::sqrt(1.0f - beta2_t) / (1.0f - beta1_t); + + for (int64_t i = 0; i < num_elements; ++i) { + float g = grad_ptr[i]; + m_ptr[i] = beta1 * m_ptr[i] + (1.0f - beta1) * g; + v_ptr[i] = beta2 * v_ptr[i] + (1.0f - beta2) * g * g; + param_ptr[i] -= alpha * m_ptr[i] / (std::sqrt(v_ptr[i]) + eps); + } +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(AccumulateGrad) +REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(AdamAccumulateGrad) + +#undef REGISTER_CPU_ACCUMULATE_GRAD_KERNEL diff --git a/infini_train/src/kernels/cpu/cross_entropy.cc b/infini_train/src/kernels/cpu/cross_entropy.cc index 2f6b4cc..427c57a 100644 --- a/infini_train/src/kernels/cpu/cross_entropy.cc +++ b/infini_train/src/kernels/cpu/cross_entropy.cc @@ -1,105 +1,105 @@ -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -namespace { -constexpr float kNegativeInfinity = -std::numeric_limits::infinity(); -} - -std::shared_ptr CrossEntropyForward(const std::shared_ptr &input, - const std::shared_ptr &target) { - const auto &input_dims = input->Dims(); - CHECK_GE(input_dims.size(), 2); - const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); - const int num_classes = *input_dims.rbegin(); - - auto output = std::make_shared(std::vector{}, DataType::kFLOAT32); - static_cast(output->DataPtr())[0] = 0.0f; - for (int64_t i = 0; i < bs; ++i) { - float max_logit = kNegativeInfinity; - for (int64_t j = 0; j < num_classes; ++j) { - max_logit = std::max(max_logit, static_cast(input->DataPtr())[i * num_classes + j]); - } - float sum_exp = 0.0f; - for (int64_t j = 0; j < num_classes; ++j) { - sum_exp += exp(static_cast(input->DataPtr())[i * num_classes + j] - max_logit); - } - if (target->Dtype() == DataType::kUINT8) { - static_cast(output->DataPtr())[0] - -= log(exp(static_cast( - input->DataPtr())[i * num_classes + static_cast(target->DataPtr())[i]] - - max_logit) - / sum_exp); - } else if (target->Dtype() == DataType::kINT64) { - static_cast(output->DataPtr())[0] - -= log(exp(static_cast( - input->DataPtr())[i * num_classes + static_cast(target->DataPtr())[i]] - - max_logit) - / sum_exp); - } else { - LOG(FATAL) << "Unsupported target data type: " << static_cast(target->Dtype()); - } - } - static_cast(output->DataPtr())[0] /= bs; - return {output}; -} - -std::shared_ptr CrossEntropyBackward(const std::shared_ptr &input, - const std::shared_ptr &target, - const std::shared_ptr &grad_output) { - const auto &input_dims = input->Dims(); - CHECK_GE(input_dims.size(), 2); - const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); - const int num_classes = *input_dims.rbegin(); - - CHECK_EQ(grad_output->Dims().size(), 0); - auto grad_input = std::make_shared(input->Dims(), DataType::kFLOAT32); - std::vector softmax(bs * num_classes, 0.0f); - for (int64_t i = 0; i < bs; ++i) { - float max_logit = kNegativeInfinity; - for (int64_t j = 0; j < num_classes; ++j) { - max_logit = std::max(max_logit, static_cast(input->DataPtr())[i * num_classes + j]); - } - float sum_exp = 0.0f; - for (int64_t j = 0; j < num_classes; ++j) { - sum_exp += exp(static_cast(input->DataPtr())[i * num_classes + j] - max_logit); - } - for (int64_t j = 0; j < num_classes; ++j) { - const auto idx = i * num_classes + j; - softmax[idx] = exp(static_cast(input->DataPtr())[idx] - max_logit) / sum_exp; - } - } - for (int64_t i = 0; i < bs; ++i) { - auto target_idx = 0; - if (target->Dtype() == DataType::kUINT8) { - target_idx = static_cast(target->DataPtr())[i]; - } else if (target->Dtype() == DataType::kINT64) { - target_idx = static_cast(target->DataPtr())[i]; - } else { - LOG(FATAL) << "Unsupported target data type: " << static_cast(target->Dtype()); - } - for (int64_t j = 0; j < num_classes; ++j) { - const auto idx = i * num_classes + j; - static_cast(grad_input->DataPtr())[idx] = static_cast(grad_output->DataPtr())[0] - * (softmax[idx] - (j == target_idx ? 1.0f : 0.0f)) / bs; - } - } - return {grad_input}; -} -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_CROSS_ENTROPY_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_CROSS_ENTROPY_KERNEL(CrossEntropyForward) -REGISTER_CPU_CROSS_ENTROPY_KERNEL(CrossEntropyBackward) - -#undef REGISTER_CPU_CROSS_ENTROPY_KERNEL +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +namespace { +constexpr float kNegativeInfinity = -std::numeric_limits::infinity(); +} + +std::shared_ptr CrossEntropyForward(const std::shared_ptr &input, + const std::shared_ptr &target) { + const auto &input_dims = input->Dims(); + CHECK_GE(input_dims.size(), 2); + const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int num_classes = *input_dims.rbegin(); + + auto output = std::make_shared(std::vector{}, DataType::kFLOAT32); + static_cast(output->DataPtr())[0] = 0.0f; + for (int64_t i = 0; i < bs; ++i) { + float max_logit = kNegativeInfinity; + for (int64_t j = 0; j < num_classes; ++j) { + max_logit = std::max(max_logit, static_cast(input->DataPtr())[i * num_classes + j]); + } + float sum_exp = 0.0f; + for (int64_t j = 0; j < num_classes; ++j) { + sum_exp += exp(static_cast(input->DataPtr())[i * num_classes + j] - max_logit); + } + if (target->Dtype() == DataType::kUINT8) { + static_cast(output->DataPtr())[0] + -= log(exp(static_cast( + input->DataPtr())[i * num_classes + static_cast(target->DataPtr())[i]] + - max_logit) + / sum_exp); + } else if (target->Dtype() == DataType::kINT64) { + static_cast(output->DataPtr())[0] + -= log(exp(static_cast( + input->DataPtr())[i * num_classes + static_cast(target->DataPtr())[i]] + - max_logit) + / sum_exp); + } else { + LOG(FATAL) << "Unsupported target data type: " << static_cast(target->Dtype()); + } + } + static_cast(output->DataPtr())[0] /= bs; + return {output}; +} + +std::shared_ptr CrossEntropyBackward(const std::shared_ptr &input, + const std::shared_ptr &target, + const std::shared_ptr &grad_output) { + const auto &input_dims = input->Dims(); + CHECK_GE(input_dims.size(), 2); + const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int num_classes = *input_dims.rbegin(); + + CHECK_EQ(grad_output->Dims().size(), 0); + auto grad_input = std::make_shared(input->Dims(), DataType::kFLOAT32); + std::vector softmax(bs * num_classes, 0.0f); + for (int64_t i = 0; i < bs; ++i) { + float max_logit = kNegativeInfinity; + for (int64_t j = 0; j < num_classes; ++j) { + max_logit = std::max(max_logit, static_cast(input->DataPtr())[i * num_classes + j]); + } + float sum_exp = 0.0f; + for (int64_t j = 0; j < num_classes; ++j) { + sum_exp += exp(static_cast(input->DataPtr())[i * num_classes + j] - max_logit); + } + for (int64_t j = 0; j < num_classes; ++j) { + const auto idx = i * num_classes + j; + softmax[idx] = exp(static_cast(input->DataPtr())[idx] - max_logit) / sum_exp; + } + } + for (int64_t i = 0; i < bs; ++i) { + auto target_idx = 0; + if (target->Dtype() == DataType::kUINT8) { + target_idx = static_cast(target->DataPtr())[i]; + } else if (target->Dtype() == DataType::kINT64) { + target_idx = static_cast(target->DataPtr())[i]; + } else { + LOG(FATAL) << "Unsupported target data type: " << static_cast(target->Dtype()); + } + for (int64_t j = 0; j < num_classes; ++j) { + const auto idx = i * num_classes + j; + static_cast(grad_input->DataPtr())[idx] = static_cast(grad_output->DataPtr())[0] + * (softmax[idx] - (j == target_idx ? 1.0f : 0.0f)) / bs; + } + } + return {grad_input}; +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_CROSS_ENTROPY_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_CROSS_ENTROPY_KERNEL(CrossEntropyForward) +REGISTER_CPU_CROSS_ENTROPY_KERNEL(CrossEntropyBackward) + +#undef REGISTER_CPU_CROSS_ENTROPY_KERNEL diff --git a/infini_train/src/kernels/cpu/elementwise.cc b/infini_train/src/kernels/cpu/elementwise.cc index b9e02d4..00a21de 100644 --- a/infini_train/src/kernels/cpu/elementwise.cc +++ b/infini_train/src/kernels/cpu/elementwise.cc @@ -1,295 +1,295 @@ -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/device.h" -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -namespace { -std::vector ComputeStrides(const std::vector &dims) { - std::vector strides(dims.size()); - int64_t stride = 1; - for (int i = dims.size() - 1; i >= 0; --i) { - strides[i] = stride; - stride *= dims[i]; - } - return strides; -} - -std::shared_ptr UnaryForward(const std::shared_ptr &input, std::function unary_fn) { - auto output = std::make_shared(input->Dims(), DataType::kFLOAT32); - for (int64_t idx = 0; idx < output->NumElements(); ++idx) { - static_cast(output->DataPtr())[idx] = unary_fn(static_cast(input->DataPtr())[idx]); - } - return output; -} - -std::shared_ptr UnaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr &a, - std::function unary_fn) { - auto grad_input = std::make_shared(grad_output->Dims(), DataType::kFLOAT32); - for (int idx = 0; idx < grad_input->NumElements(); ++idx) { - const float x = a ? static_cast(a->DataPtr())[idx] : 0.0f; - const float grad = static_cast(grad_output->DataPtr())[idx]; - static_cast(grad_input->DataPtr())[idx] = grad * unary_fn(x); - } - return grad_input; -} - -std::shared_ptr BinaryForward(const std::shared_ptr &a, const std::shared_ptr &b, - std::function binary_fn) { - // Currently a and b should have the same data type and only one-way broadcasting from b to a is assumed by default - CHECK(a->Dtype() == b->Dtype() && a->NumElements() >= b->NumElements() && a->NumElements() % b->NumElements() == 0); - auto out_dims = a->Dims(); - auto output = std::make_shared(out_dims, DataType::kFLOAT32); - int64_t num_elements = output->NumElements(); - - auto out_strides = ComputeStrides(out_dims); - auto b_strides = ComputeStrides(b->Dims()); - - int ndim = out_dims.size(); - std::vector b_padded_dims = b->Dims(); - b_padded_dims.insert(b_padded_dims.begin(), ndim - b_padded_dims.size(), 1); - b_strides.insert(b_strides.begin(), ndim - b_strides.size(), 0); - - float *out_ptr = static_cast(output->DataPtr()); - const float *a_ptr = static_cast(a->DataPtr()); - const float *b_ptr = static_cast(b->DataPtr()); - - for (int64_t idx = 0; idx < num_elements; ++idx) { - int64_t tmp = idx, b_offset = 0; - for (int i = 0; i < ndim; ++i) { - int64_t index = tmp / out_strides[i]; - tmp %= out_strides[i]; - b_offset += (b_padded_dims[i] == 1 ? 0 : index) * b_strides[i]; - } - out_ptr[idx] = binary_fn(a_ptr[idx], b_ptr[b_offset]); - } - - return output; -} - -std::pair, std::shared_ptr> -BinaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr &a, - const std::shared_ptr &b, const std::vector &a_dims, const std::vector &b_dims, - std::function fn_a, std::function fn_b) { - // Currently a and b should have the same data type and only one-way broadcasting from b to a is assumed by default - CHECK(a_dims.size() >= b_dims.size()); - - const int64_t a_num_elements = std::accumulate(a_dims.begin(), a_dims.end(), 1LL, std::multiplies()); - const int64_t b_num_elements = std::accumulate(b_dims.begin(), b_dims.end(), 1LL, std::multiplies()); - - CHECK(a_num_elements == grad_output->NumElements()); - if (a) { - CHECK_EQ(a_num_elements, a->NumElements()); - } - if (b) { - CHECK_EQ(b_num_elements, b->NumElements()); - } - - auto grad_a = std::make_shared(a_dims, DataType::kFLOAT32); - auto grad_b = std::make_shared(b_dims, DataType::kFLOAT32); - grad_a->Fill(0.0f); - grad_b->Fill(0.0f); - - int ndim = a_dims.size(); - auto out_strides = ComputeStrides(a_dims); - auto b_strides = ComputeStrides(b_dims); - std::vector b_padded_dims = b_dims; - b_padded_dims.insert(b_padded_dims.begin(), ndim - b_dims.size(), 1); - b_strides.insert(b_strides.begin(), ndim - b_strides.size(), 0); - - for (int64_t idx = 0; idx < a_num_elements; ++idx) { - int64_t tmp = idx, b_offset = 0; - for (size_t i = 0; i < ndim; ++i) { - int64_t index = tmp / out_strides[i]; - tmp %= out_strides[i]; - b_offset += (b_padded_dims[i] == 1 ? 0 : index) * b_strides[i]; - } - - const float x = a ? static_cast(a->DataPtr())[idx] : 0.0f; - const float y = b ? static_cast(b->DataPtr())[b_offset] : 0.0f; - const float grad = static_cast(grad_output->DataPtr())[idx]; - - static_cast(grad_a->DataPtr())[idx] = grad * fn_a(x, y); - static_cast(grad_b->DataPtr())[b_offset] += grad * fn_b(x, y); - } - return {grad_a, grad_b}; -} -} // namespace - -std::shared_ptr NegForward(const std::shared_ptr &input) { - return UnaryForward(input, [](float x) { return -x; }); -} - -std::shared_ptr NegBackward(const std::shared_ptr &grad_output) { - return UnaryBackward(grad_output, nullptr, [](float) { return -1.0f; }); -} - -std::shared_ptr ReciprocalForward(const std::shared_ptr &input) { - return UnaryForward(input, [](float x) { return 1.0f / x; }); -} - -std::shared_ptr ReciprocalBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &input) { - return UnaryBackward(grad_output, input, [](float x) { return -1.0f / (x * x); }); -} - -std::shared_ptr SinForward(const std::shared_ptr &input) { - return UnaryForward(input, [](float x) { return sinf(x); }); -} - -std::shared_ptr SinBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { - return UnaryBackward(grad_output, input, [](float x) { return cosf(x); }); -} - -std::shared_ptr CosForward(const std::shared_ptr &input) { - return UnaryForward(input, [](float x) { return cosf(x); }); -} - -std::shared_ptr CosBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { - return UnaryBackward(grad_output, input, [](float x) { return -sinf(x); }); -} - -std::shared_ptr TanhForward(const std::shared_ptr &input) { - return UnaryForward(input, [](float x) { return tanhf(x); }); -} - -std::shared_ptr TanhBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &output) { - return UnaryBackward(grad_output, output, [](float x) { return 1.0 - x * x; }); -} - -std::shared_ptr PowForward(const std::shared_ptr &input, float scalar, bool scalar_is_base) { - if (scalar_is_base) { - return UnaryForward(input, [scalar](float x) { return powf(scalar, x); }); - } else { - return UnaryForward(input, [scalar](float x) { return powf(x, scalar); }); - } -} - -std::shared_ptr PowBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, - float scalar, bool scalar_is_base) { - if (scalar_is_base) { - return UnaryBackward(grad_output, input, [scalar](float x) { return std::log(scalar) * powf(scalar, x); }); - } else { - return UnaryBackward(grad_output, input, [scalar](float x) { return scalar * powf(x, scalar - 1.0f); }); - } -} - -std::shared_ptr RsqrtForward(const std::shared_ptr &input) { - return UnaryForward(input, [](float x) { return 1.0f / std::sqrt(x); }); -} - -std::shared_ptr RsqrtBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &input) { - return UnaryBackward(grad_output, input, [](float x) { return -0.5f / (x * std::sqrt(x)); }); -} - -std::shared_ptr EqualsScalarForward(const std::shared_ptr &a, float scalar) { - return UnaryForward(a, [scalar](float x) { return x == scalar ? 1.0f : 0.0f; }); -} - -std::shared_ptr AddForward(const std::shared_ptr &a, const std::shared_ptr &b) { - return BinaryForward(a, b, [](float x, float y) { return x + y; }); -} - -std::pair, std::shared_ptr> AddBackward(const std::shared_ptr &grad_output, - const std::vector &a_dims, - const std::vector &b_dims) { - - return BinaryBackward( - grad_output, nullptr, nullptr, a_dims, b_dims, [](float, float) { return 1.0f; }, - [](float, float) { return 1.0f; }); -} - -std::shared_ptr AddScalarForward(const std::shared_ptr &a, float scalar) { - return UnaryForward(a, [scalar](float x) { return x + scalar; }); -} - -std::shared_ptr AddScalarBackward(const std::shared_ptr &grad_output) { - return UnaryBackward(grad_output, nullptr, [](float) { return 1.0f; }); -} - -std::shared_ptr SubForward(const std::shared_ptr &a, const std::shared_ptr &b) { - return BinaryForward(a, b, [](float x, float y) { return x - y; }); -} - -std::pair, std::shared_ptr> SubBackward(const std::shared_ptr &grad_output, - const std::vector &a_dims, - const std::vector &b_dims) { - return BinaryBackward( - grad_output, nullptr, nullptr, a_dims, b_dims, [](float, float) { return 1.f; }, - [](float, float) { return -1.f; }); -} - -std::shared_ptr MulForward(const std::shared_ptr &a, const std::shared_ptr &b) { - return BinaryForward(a, b, [](float x, float y) { return x * y; }); -} - -std::pair, std::shared_ptr> MulBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &a, - const std::shared_ptr &b) { - return BinaryBackward( - grad_output, a, b, a->Dims(), b->Dims(), [](float, float y) { return y; }, [](float x, float) { return x; }); -} - -std::shared_ptr MulScalarForward(const std::shared_ptr &a, float scalar) { - return UnaryForward(a, [scalar](float x) { return x * scalar; }); -} - -std::shared_ptr MulScalarBackward(const std::shared_ptr &grad_output, float scalar) { - return UnaryBackward(grad_output, nullptr, [scalar](float) { return scalar; }); -} - -std::shared_ptr DivForward(const std::shared_ptr &a, const std::shared_ptr &b) { - return BinaryForward(a, b, [](float x, float y) { return x / y; }); -} - -std::pair, std::shared_ptr> DivBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &a, - const std::shared_ptr &b) { - return BinaryBackward( - grad_output, a, b, a->Dims(), b->Dims(), [](float, float y) { return 1 / y; }, - [](float x, float y) { return -x / (y * y); }); -} - -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_ELEMENTWISE_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_ELEMENTWISE_KERNEL(NegForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(NegBackward) -REGISTER_CPU_ELEMENTWISE_KERNEL(ReciprocalForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(ReciprocalBackward) -REGISTER_CPU_ELEMENTWISE_KERNEL(SinForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(SinBackward) -REGISTER_CPU_ELEMENTWISE_KERNEL(CosForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(CosBackward) -REGISTER_CPU_ELEMENTWISE_KERNEL(TanhForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(TanhBackward) -REGISTER_CPU_ELEMENTWISE_KERNEL(PowForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(PowBackward) -REGISTER_CPU_ELEMENTWISE_KERNEL(RsqrtForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(RsqrtBackward) -REGISTER_CPU_ELEMENTWISE_KERNEL(EqualsScalarForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(AddForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(AddBackward) -REGISTER_CPU_ELEMENTWISE_KERNEL(AddScalarForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(AddScalarBackward) -REGISTER_CPU_ELEMENTWISE_KERNEL(SubForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(SubBackward) -REGISTER_CPU_ELEMENTWISE_KERNEL(MulForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(MulBackward) -REGISTER_CPU_ELEMENTWISE_KERNEL(MulScalarForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(MulScalarBackward) -REGISTER_CPU_ELEMENTWISE_KERNEL(DivForward) -REGISTER_CPU_ELEMENTWISE_KERNEL(DivBackward) - -#undef REGISTER_CPU_ELEMENTWISE_KERNEL +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +namespace { +std::vector ComputeStrides(const std::vector &dims) { + std::vector strides(dims.size()); + int64_t stride = 1; + for (int i = dims.size() - 1; i >= 0; --i) { + strides[i] = stride; + stride *= dims[i]; + } + return strides; +} + +std::shared_ptr UnaryForward(const std::shared_ptr &input, std::function unary_fn) { + auto output = std::make_shared(input->Dims(), DataType::kFLOAT32); + for (int64_t idx = 0; idx < output->NumElements(); ++idx) { + static_cast(output->DataPtr())[idx] = unary_fn(static_cast(input->DataPtr())[idx]); + } + return output; +} + +std::shared_ptr UnaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr &a, + std::function unary_fn) { + auto grad_input = std::make_shared(grad_output->Dims(), DataType::kFLOAT32); + for (int idx = 0; idx < grad_input->NumElements(); ++idx) { + const float x = a ? static_cast(a->DataPtr())[idx] : 0.0f; + const float grad = static_cast(grad_output->DataPtr())[idx]; + static_cast(grad_input->DataPtr())[idx] = grad * unary_fn(x); + } + return grad_input; +} + +std::shared_ptr BinaryForward(const std::shared_ptr &a, const std::shared_ptr &b, + std::function binary_fn) { + // Currently a and b should have the same data type and only one-way broadcasting from b to a is assumed by default + CHECK(a->Dtype() == b->Dtype() && a->NumElements() >= b->NumElements() && a->NumElements() % b->NumElements() == 0); + auto out_dims = a->Dims(); + auto output = std::make_shared(out_dims, DataType::kFLOAT32); + int64_t num_elements = output->NumElements(); + + auto out_strides = ComputeStrides(out_dims); + auto b_strides = ComputeStrides(b->Dims()); + + int ndim = out_dims.size(); + std::vector b_padded_dims = b->Dims(); + b_padded_dims.insert(b_padded_dims.begin(), ndim - b_padded_dims.size(), 1); + b_strides.insert(b_strides.begin(), ndim - b_strides.size(), 0); + + float *out_ptr = static_cast(output->DataPtr()); + const float *a_ptr = static_cast(a->DataPtr()); + const float *b_ptr = static_cast(b->DataPtr()); + + for (int64_t idx = 0; idx < num_elements; ++idx) { + int64_t tmp = idx, b_offset = 0; + for (int i = 0; i < ndim; ++i) { + int64_t index = tmp / out_strides[i]; + tmp %= out_strides[i]; + b_offset += (b_padded_dims[i] == 1 ? 0 : index) * b_strides[i]; + } + out_ptr[idx] = binary_fn(a_ptr[idx], b_ptr[b_offset]); + } + + return output; +} + +std::pair, std::shared_ptr> +BinaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr &a, + const std::shared_ptr &b, const std::vector &a_dims, const std::vector &b_dims, + std::function fn_a, std::function fn_b) { + // Currently a and b should have the same data type and only one-way broadcasting from b to a is assumed by default + CHECK(a_dims.size() >= b_dims.size()); + + const int64_t a_num_elements = std::accumulate(a_dims.begin(), a_dims.end(), 1LL, std::multiplies()); + const int64_t b_num_elements = std::accumulate(b_dims.begin(), b_dims.end(), 1LL, std::multiplies()); + + CHECK(a_num_elements == grad_output->NumElements()); + if (a) { + CHECK_EQ(a_num_elements, a->NumElements()); + } + if (b) { + CHECK_EQ(b_num_elements, b->NumElements()); + } + + auto grad_a = std::make_shared(a_dims, DataType::kFLOAT32); + auto grad_b = std::make_shared(b_dims, DataType::kFLOAT32); + grad_a->Fill(0.0f); + grad_b->Fill(0.0f); + + int ndim = a_dims.size(); + auto out_strides = ComputeStrides(a_dims); + auto b_strides = ComputeStrides(b_dims); + std::vector b_padded_dims = b_dims; + b_padded_dims.insert(b_padded_dims.begin(), ndim - b_dims.size(), 1); + b_strides.insert(b_strides.begin(), ndim - b_strides.size(), 0); + + for (int64_t idx = 0; idx < a_num_elements; ++idx) { + int64_t tmp = idx, b_offset = 0; + for (size_t i = 0; i < ndim; ++i) { + int64_t index = tmp / out_strides[i]; + tmp %= out_strides[i]; + b_offset += (b_padded_dims[i] == 1 ? 0 : index) * b_strides[i]; + } + + const float x = a ? static_cast(a->DataPtr())[idx] : 0.0f; + const float y = b ? static_cast(b->DataPtr())[b_offset] : 0.0f; + const float grad = static_cast(grad_output->DataPtr())[idx]; + + static_cast(grad_a->DataPtr())[idx] = grad * fn_a(x, y); + static_cast(grad_b->DataPtr())[b_offset] += grad * fn_b(x, y); + } + return {grad_a, grad_b}; +} +} // namespace + +std::shared_ptr NegForward(const std::shared_ptr &input) { + return UnaryForward(input, [](float x) { return -x; }); +} + +std::shared_ptr NegBackward(const std::shared_ptr &grad_output) { + return UnaryBackward(grad_output, nullptr, [](float) { return -1.0f; }); +} + +std::shared_ptr ReciprocalForward(const std::shared_ptr &input) { + return UnaryForward(input, [](float x) { return 1.0f / x; }); +} + +std::shared_ptr ReciprocalBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &input) { + return UnaryBackward(grad_output, input, [](float x) { return -1.0f / (x * x); }); +} + +std::shared_ptr SinForward(const std::shared_ptr &input) { + return UnaryForward(input, [](float x) { return sinf(x); }); +} + +std::shared_ptr SinBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { + return UnaryBackward(grad_output, input, [](float x) { return cosf(x); }); +} + +std::shared_ptr CosForward(const std::shared_ptr &input) { + return UnaryForward(input, [](float x) { return cosf(x); }); +} + +std::shared_ptr CosBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { + return UnaryBackward(grad_output, input, [](float x) { return -sinf(x); }); +} + +std::shared_ptr TanhForward(const std::shared_ptr &input) { + return UnaryForward(input, [](float x) { return tanhf(x); }); +} + +std::shared_ptr TanhBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &output) { + return UnaryBackward(grad_output, output, [](float x) { return 1.0 - x * x; }); +} + +std::shared_ptr PowForward(const std::shared_ptr &input, float scalar, bool scalar_is_base) { + if (scalar_is_base) { + return UnaryForward(input, [scalar](float x) { return powf(scalar, x); }); + } else { + return UnaryForward(input, [scalar](float x) { return powf(x, scalar); }); + } +} + +std::shared_ptr PowBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + float scalar, bool scalar_is_base) { + if (scalar_is_base) { + return UnaryBackward(grad_output, input, [scalar](float x) { return std::log(scalar) * powf(scalar, x); }); + } else { + return UnaryBackward(grad_output, input, [scalar](float x) { return scalar * powf(x, scalar - 1.0f); }); + } +} + +std::shared_ptr RsqrtForward(const std::shared_ptr &input) { + return UnaryForward(input, [](float x) { return 1.0f / std::sqrt(x); }); +} + +std::shared_ptr RsqrtBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &input) { + return UnaryBackward(grad_output, input, [](float x) { return -0.5f / (x * std::sqrt(x)); }); +} + +std::shared_ptr EqualsScalarForward(const std::shared_ptr &a, float scalar) { + return UnaryForward(a, [scalar](float x) { return x == scalar ? 1.0f : 0.0f; }); +} + +std::shared_ptr AddForward(const std::shared_ptr &a, const std::shared_ptr &b) { + return BinaryForward(a, b, [](float x, float y) { return x + y; }); +} + +std::pair, std::shared_ptr> AddBackward(const std::shared_ptr &grad_output, + const std::vector &a_dims, + const std::vector &b_dims) { + + return BinaryBackward( + grad_output, nullptr, nullptr, a_dims, b_dims, [](float, float) { return 1.0f; }, + [](float, float) { return 1.0f; }); +} + +std::shared_ptr AddScalarForward(const std::shared_ptr &a, float scalar) { + return UnaryForward(a, [scalar](float x) { return x + scalar; }); +} + +std::shared_ptr AddScalarBackward(const std::shared_ptr &grad_output) { + return UnaryBackward(grad_output, nullptr, [](float) { return 1.0f; }); +} + +std::shared_ptr SubForward(const std::shared_ptr &a, const std::shared_ptr &b) { + return BinaryForward(a, b, [](float x, float y) { return x - y; }); +} + +std::pair, std::shared_ptr> SubBackward(const std::shared_ptr &grad_output, + const std::vector &a_dims, + const std::vector &b_dims) { + return BinaryBackward( + grad_output, nullptr, nullptr, a_dims, b_dims, [](float, float) { return 1.f; }, + [](float, float) { return -1.f; }); +} + +std::shared_ptr MulForward(const std::shared_ptr &a, const std::shared_ptr &b) { + return BinaryForward(a, b, [](float x, float y) { return x * y; }); +} + +std::pair, std::shared_ptr> MulBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &a, + const std::shared_ptr &b) { + return BinaryBackward( + grad_output, a, b, a->Dims(), b->Dims(), [](float, float y) { return y; }, [](float x, float) { return x; }); +} + +std::shared_ptr MulScalarForward(const std::shared_ptr &a, float scalar) { + return UnaryForward(a, [scalar](float x) { return x * scalar; }); +} + +std::shared_ptr MulScalarBackward(const std::shared_ptr &grad_output, float scalar) { + return UnaryBackward(grad_output, nullptr, [scalar](float) { return scalar; }); +} + +std::shared_ptr DivForward(const std::shared_ptr &a, const std::shared_ptr &b) { + return BinaryForward(a, b, [](float x, float y) { return x / y; }); +} + +std::pair, std::shared_ptr> DivBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &a, + const std::shared_ptr &b) { + return BinaryBackward( + grad_output, a, b, a->Dims(), b->Dims(), [](float, float y) { return 1 / y; }, + [](float x, float y) { return -x / (y * y); }); +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_ELEMENTWISE_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_ELEMENTWISE_KERNEL(NegForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(NegBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(ReciprocalForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(ReciprocalBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(SinForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(SinBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(CosForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(CosBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(TanhForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(TanhBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(PowForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(PowBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(RsqrtForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(RsqrtBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(EqualsScalarForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(AddForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(AddBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(AddScalarForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(AddScalarBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(SubForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(SubBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(MulForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(MulBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(MulScalarForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(MulScalarBackward) +REGISTER_CPU_ELEMENTWISE_KERNEL(DivForward) +REGISTER_CPU_ELEMENTWISE_KERNEL(DivBackward) + +#undef REGISTER_CPU_ELEMENTWISE_KERNEL diff --git a/infini_train/src/kernels/cpu/embedding.cc b/infini_train/src/kernels/cpu/embedding.cc index f0462be..e742810 100644 --- a/infini_train/src/kernels/cpu/embedding.cc +++ b/infini_train/src/kernels/cpu/embedding.cc @@ -1,64 +1,64 @@ -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -std::shared_ptr EmbeddingForward(const std::shared_ptr &input, const std::shared_ptr &weight) { - /* - x: [*] - -> Embedding (weight: [num_embeddings, embedding_dim]) - -> o: [*, embedding_dim] - */ - CHECK(input->Dtype() == DataType::kINT64); - const auto &input_dims = input->Dims(); - CHECK_EQ(weight->Dims().size(), 2); - const int embedding_dim = weight->Dims()[1]; - auto output_dims = input_dims; - output_dims.push_back(embedding_dim); - auto output = std::make_shared(output_dims, DataType::kFLOAT32); - - for (int i = 0; i < input->NumElements(); i++) { - int idx = static_cast(static_cast(input->DataPtr())[i]); - for (int j = 0; j < embedding_dim; j++) { - static_cast(output->DataPtr())[i * embedding_dim + j] - = static_cast(weight->DataPtr())[idx * embedding_dim + j]; - } - } - - return output; -} - -std::shared_ptr EmbeddingBackward(const std::shared_ptr &input, const std::vector &weight_dims, - const std::shared_ptr &grad_output) { - CHECK(input->Dtype() == DataType::kINT64); - CHECK_EQ(weight_dims.size(), 2); - const int embedding_dim = weight_dims[1]; - CHECK_EQ(input->Dims().size() + 1, grad_output->Dims().size()); - for (int idx = 0; idx < input->Dims().size(); ++idx) { CHECK_EQ(input->Dims()[idx], grad_output->Dims()[idx]); } - CHECK_EQ(*grad_output->Dims().rbegin(), embedding_dim); - - auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32); - grad_weight->Fill(0.0f); - - for (int i = 0; i < input->NumElements(); i++) { - int idx = static_cast(static_cast(input->DataPtr())[i]); - for (int j = 0; j < embedding_dim; j++) { - static_cast(grad_weight->DataPtr())[idx * embedding_dim + j] // <-- 修复这里 - += static_cast(grad_output->DataPtr())[i * embedding_dim + j]; - } - } - return grad_weight; -} - -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_EMBEDDING_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_EMBEDDING_KERNEL(EmbeddingForward) -REGISTER_CPU_EMBEDDING_KERNEL(EmbeddingBackward) - -#undef REGISTER_CPU_EMBEDDING_KERNEL +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr EmbeddingForward(const std::shared_ptr &input, const std::shared_ptr &weight) { + /* + x: [*] + -> Embedding (weight: [num_embeddings, embedding_dim]) + -> o: [*, embedding_dim] + */ + CHECK(input->Dtype() == DataType::kINT64); + const auto &input_dims = input->Dims(); + CHECK_EQ(weight->Dims().size(), 2); + const int embedding_dim = weight->Dims()[1]; + auto output_dims = input_dims; + output_dims.push_back(embedding_dim); + auto output = std::make_shared(output_dims, DataType::kFLOAT32); + + for (int i = 0; i < input->NumElements(); i++) { + int idx = static_cast(static_cast(input->DataPtr())[i]); + for (int j = 0; j < embedding_dim; j++) { + static_cast(output->DataPtr())[i * embedding_dim + j] + = static_cast(weight->DataPtr())[idx * embedding_dim + j]; + } + } + + return output; +} + +std::shared_ptr EmbeddingBackward(const std::shared_ptr &input, const std::vector &weight_dims, + const std::shared_ptr &grad_output) { + CHECK(input->Dtype() == DataType::kINT64); + CHECK_EQ(weight_dims.size(), 2); + const int embedding_dim = weight_dims[1]; + CHECK_EQ(input->Dims().size() + 1, grad_output->Dims().size()); + for (int idx = 0; idx < input->Dims().size(); ++idx) { CHECK_EQ(input->Dims()[idx], grad_output->Dims()[idx]); } + CHECK_EQ(*grad_output->Dims().rbegin(), embedding_dim); + + auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32); + grad_weight->Fill(0.0f); + + for (int i = 0; i < input->NumElements(); i++) { + int idx = static_cast(static_cast(input->DataPtr())[i]); + for (int j = 0; j < embedding_dim; j++) { + static_cast(grad_weight->DataPtr())[idx * embedding_dim + j] // <-- 修复这里 + += static_cast(grad_output->DataPtr())[i * embedding_dim + j]; + } + } + return grad_weight; +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_EMBEDDING_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_EMBEDDING_KERNEL(EmbeddingForward) +REGISTER_CPU_EMBEDDING_KERNEL(EmbeddingBackward) + +#undef REGISTER_CPU_EMBEDDING_KERNEL diff --git a/infini_train/src/kernels/cpu/fill.cc b/infini_train/src/kernels/cpu/fill.cc index 2e8fdbc..e2d4946 100644 --- a/infini_train/src/kernels/cpu/fill.cc +++ b/infini_train/src/kernels/cpu/fill.cc @@ -1,19 +1,19 @@ -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -void Fill(std::shared_ptr tensor, void *value_ptr) { - // FIXME(zbl): support other data types - auto data = reinterpret_cast(tensor->DataPtr()); - std::fill(data, data + tensor->NumElements(), *(static_cast(value_ptr))); -} -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_FILL_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_FILL_KERNEL(Fill) - -#undef REGISTER_CPU_FILL_KERNEL +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +void Fill(std::shared_ptr tensor, void *value_ptr) { + // FIXME(zbl): support other data types + auto data = reinterpret_cast(tensor->DataPtr()); + std::fill(data, data + tensor->NumElements(), *(static_cast(value_ptr))); +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_FILL_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_FILL_KERNEL(Fill) + +#undef REGISTER_CPU_FILL_KERNEL diff --git a/infini_train/src/kernels/cpu/layernorm.cc b/infini_train/src/kernels/cpu/layernorm.cc index f8fa458..192c262 100644 --- a/infini_train/src/kernels/cpu/layernorm.cc +++ b/infini_train/src/kernels/cpu/layernorm.cc @@ -1,146 +1,146 @@ -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -std::tuple, std::shared_ptr, std::shared_ptr> -LayerNormForward(const std::shared_ptr &input, const std::shared_ptr &weight, - const std::shared_ptr &bias, const float eps) { - /* - x: [bs, seq_len, embed_dim] - -> LayerNorm (w: [embed_dim], b: [embed_dim]) - -> o: [bs, seq_len, embed_dim] - */ - CHECK_EQ(input->Dims().size(), 3); - CHECK_LE(input->Dims()[2], weight->Dims()[0]); - CHECK_LE(input->Dims()[2], bias->Dims()[0]); - - const int batch_size = input->Dims()[0]; - const int max_seqlen = input->Dims()[1]; - const int embed_dim = input->Dims()[2]; - - auto output = std::make_shared(std::vector{batch_size, max_seqlen, embed_dim}, DataType::kFLOAT32); - auto mean = std::make_shared(std::vector{batch_size, max_seqlen}, DataType::kFLOAT32); - auto rstd = std::make_shared(std::vector{batch_size, max_seqlen}, DataType::kFLOAT32); - mean->Fill(0.0f); - rstd->Fill(0.0f); - - for (int b = 0; b < batch_size; b++) { - for (int t = 0; t < max_seqlen; t++) { - float m = 0.0f; - for (int i = 0; i < embed_dim; i++) { - m += static_cast(input->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i]; - } - m = m / embed_dim; - - float v = 0.0f; - for (int i = 0; i < embed_dim; i++) { - float xshift - = static_cast(input->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i] - m; - v += xshift * xshift; - } - v = v / embed_dim; - - float s = 1.0f / sqrtf(v + eps); - - for (int i = 0; i < embed_dim; i++) { - float n = (s - * (static_cast(input->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i] - - m)); // normalize - float o = n * static_cast(weight->DataPtr())[i] - + static_cast(bias->DataPtr())[i]; // scale and shift - static_cast(output->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i] = o; // write - } - // cache the mean and rstd for the backward pass later - static_cast(mean->DataPtr())[b * max_seqlen + t] = m; - static_cast(rstd->DataPtr())[b * max_seqlen + t] = s; - } - } - - return {output, mean, rstd}; -} - -std::tuple, std::shared_ptr, std::shared_ptr> -LayerNormBackward(const std::shared_ptr &input, const std::shared_ptr &weight, - const std::shared_ptr &bias, const std::shared_ptr &mean, - const std::shared_ptr &rstd, const std::shared_ptr &grad_output) { - CHECK_EQ(input->Dims().size(), 3); - CHECK_LE(input->Dims()[2], weight->Dims()[0]); - CHECK_LE(input->Dims()[2], bias->Dims()[0]); - CHECK_NE(mean, nullptr); - CHECK_NE(rstd, nullptr); - - const int batch_size = input->Dims()[0]; - const int max_seqlen = input->Dims()[1]; - const int embed_dim = input->Dims()[2]; - - auto grad_input = std::make_shared(input->Dims(), DataType::kFLOAT32); - auto grad_weight = std::make_shared(weight->Dims(), DataType::kFLOAT32); - auto grad_bias = std::make_shared(bias->Dims(), DataType::kFLOAT32); - - grad_input->Fill(0.0f); - grad_weight->Fill(0.0f); - grad_bias->Fill(0.0f); - - for (int b = 0; b < batch_size; b++) { - for (int t = 0; t < max_seqlen; t++) { - float mean_bt = static_cast(mean->DataPtr())[b * max_seqlen + t]; - float rstd_bt = static_cast(rstd->DataPtr())[b * max_seqlen + t]; - - // first: two reduce operations - float dnorm_mean = 0.0f; - float dnorm_norm_mean = 0.0f; - for (int i = 0; i < embed_dim; i++) { - float norm_bti - = (static_cast(input->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i] - mean_bt) - * rstd_bt; - float dnorm_i - = static_cast(weight->DataPtr())[i] - * static_cast(grad_output->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i]; - dnorm_mean += dnorm_i; - dnorm_norm_mean += dnorm_i * norm_bti; - } - dnorm_mean = dnorm_mean / embed_dim; - dnorm_norm_mean = dnorm_norm_mean / embed_dim; - - // now iterate again and accumulate all the gradients - for (int i = 0; i < embed_dim; i++) { - float norm_bti - = (static_cast(input->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i] - mean_bt) - * rstd_bt; - float dnorm_i - = static_cast(weight->DataPtr())[i] - * static_cast(grad_output->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i]; - // gradient contribution to bias - static_cast(grad_bias->DataPtr())[i] - += static_cast(grad_output->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i]; - // gradient contribution to weight - static_cast(grad_weight->DataPtr())[i] - += norm_bti - * static_cast(grad_output->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i]; - // gradient contribution to input - float dval = 0.0f; - dval += dnorm_i; // term 1 - dval -= dnorm_mean; // term 2 - dval -= norm_bti * dnorm_norm_mean; // term 3 - dval *= rstd_bt; // final scale - static_cast(grad_input->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i] += dval; - } - } - } - return {grad_input, grad_weight, grad_bias}; -} -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_LAYERNORM_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_LAYERNORM_KERNEL(LayerNormForward) -REGISTER_CPU_LAYERNORM_KERNEL(LayerNormBackward) - -#undef REGISTER_CPU_LAYERNORM_KERNEL +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::tuple, std::shared_ptr, std::shared_ptr> +LayerNormForward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias, const float eps) { + /* + x: [bs, seq_len, embed_dim] + -> LayerNorm (w: [embed_dim], b: [embed_dim]) + -> o: [bs, seq_len, embed_dim] + */ + CHECK_EQ(input->Dims().size(), 3); + CHECK_LE(input->Dims()[2], weight->Dims()[0]); + CHECK_LE(input->Dims()[2], bias->Dims()[0]); + + const int batch_size = input->Dims()[0]; + const int max_seqlen = input->Dims()[1]; + const int embed_dim = input->Dims()[2]; + + auto output = std::make_shared(std::vector{batch_size, max_seqlen, embed_dim}, DataType::kFLOAT32); + auto mean = std::make_shared(std::vector{batch_size, max_seqlen}, DataType::kFLOAT32); + auto rstd = std::make_shared(std::vector{batch_size, max_seqlen}, DataType::kFLOAT32); + mean->Fill(0.0f); + rstd->Fill(0.0f); + + for (int b = 0; b < batch_size; b++) { + for (int t = 0; t < max_seqlen; t++) { + float m = 0.0f; + for (int i = 0; i < embed_dim; i++) { + m += static_cast(input->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i]; + } + m = m / embed_dim; + + float v = 0.0f; + for (int i = 0; i < embed_dim; i++) { + float xshift + = static_cast(input->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i] - m; + v += xshift * xshift; + } + v = v / embed_dim; + + float s = 1.0f / sqrtf(v + eps); + + for (int i = 0; i < embed_dim; i++) { + float n = (s + * (static_cast(input->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i] + - m)); // normalize + float o = n * static_cast(weight->DataPtr())[i] + + static_cast(bias->DataPtr())[i]; // scale and shift + static_cast(output->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i] = o; // write + } + // cache the mean and rstd for the backward pass later + static_cast(mean->DataPtr())[b * max_seqlen + t] = m; + static_cast(rstd->DataPtr())[b * max_seqlen + t] = s; + } + } + + return {output, mean, rstd}; +} + +std::tuple, std::shared_ptr, std::shared_ptr> +LayerNormBackward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias, const std::shared_ptr &mean, + const std::shared_ptr &rstd, const std::shared_ptr &grad_output) { + CHECK_EQ(input->Dims().size(), 3); + CHECK_LE(input->Dims()[2], weight->Dims()[0]); + CHECK_LE(input->Dims()[2], bias->Dims()[0]); + CHECK_NE(mean, nullptr); + CHECK_NE(rstd, nullptr); + + const int batch_size = input->Dims()[0]; + const int max_seqlen = input->Dims()[1]; + const int embed_dim = input->Dims()[2]; + + auto grad_input = std::make_shared(input->Dims(), DataType::kFLOAT32); + auto grad_weight = std::make_shared(weight->Dims(), DataType::kFLOAT32); + auto grad_bias = std::make_shared(bias->Dims(), DataType::kFLOAT32); + + grad_input->Fill(0.0f); + grad_weight->Fill(0.0f); + grad_bias->Fill(0.0f); + + for (int b = 0; b < batch_size; b++) { + for (int t = 0; t < max_seqlen; t++) { + float mean_bt = static_cast(mean->DataPtr())[b * max_seqlen + t]; + float rstd_bt = static_cast(rstd->DataPtr())[b * max_seqlen + t]; + + // first: two reduce operations + float dnorm_mean = 0.0f; + float dnorm_norm_mean = 0.0f; + for (int i = 0; i < embed_dim; i++) { + float norm_bti + = (static_cast(input->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i] - mean_bt) + * rstd_bt; + float dnorm_i + = static_cast(weight->DataPtr())[i] + * static_cast(grad_output->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i]; + dnorm_mean += dnorm_i; + dnorm_norm_mean += dnorm_i * norm_bti; + } + dnorm_mean = dnorm_mean / embed_dim; + dnorm_norm_mean = dnorm_norm_mean / embed_dim; + + // now iterate again and accumulate all the gradients + for (int i = 0; i < embed_dim; i++) { + float norm_bti + = (static_cast(input->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i] - mean_bt) + * rstd_bt; + float dnorm_i + = static_cast(weight->DataPtr())[i] + * static_cast(grad_output->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i]; + // gradient contribution to bias + static_cast(grad_bias->DataPtr())[i] + += static_cast(grad_output->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i]; + // gradient contribution to weight + static_cast(grad_weight->DataPtr())[i] + += norm_bti + * static_cast(grad_output->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i]; + // gradient contribution to input + float dval = 0.0f; + dval += dnorm_i; // term 1 + dval -= dnorm_mean; // term 2 + dval -= norm_bti * dnorm_norm_mean; // term 3 + dval *= rstd_bt; // final scale + static_cast(grad_input->DataPtr())[b * max_seqlen * embed_dim + t * embed_dim + i] += dval; + } + } + } + return {grad_input, grad_weight, grad_bias}; +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_LAYERNORM_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_LAYERNORM_KERNEL(LayerNormForward) +REGISTER_CPU_LAYERNORM_KERNEL(LayerNormBackward) + +#undef REGISTER_CPU_LAYERNORM_KERNEL diff --git a/infini_train/src/kernels/cpu/linear.cc b/infini_train/src/kernels/cpu/linear.cc index 140e756..d72983e 100644 --- a/infini_train/src/kernels/cpu/linear.cc +++ b/infini_train/src/kernels/cpu/linear.cc @@ -1,134 +1,231 @@ -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -std::shared_ptr MatmulForward(const std::shared_ptr &input, const std::shared_ptr &other) { - // =================================== 作业 =================================== - // TODO：实现CPU上的矩阵乘法前向计算 - // REF: - // =================================== 作业 =================================== - - auto output = std::make_shared(); - return {output}; -} - -std::tuple, std::shared_ptr> -MatmulBackward(const std::shared_ptr &input, const std::shared_ptr &other, - const std::shared_ptr &grad_output) { - // =================================== 作业 =================================== - // TODO：实现CPU上的矩阵乘法反向传播 - // REF: - // =================================== 作业 =================================== - - auto grad_input = std::make_shared(); - auto grad_other = std::make_shared(); - return {grad_input, grad_other}; -} - -std::shared_ptr LinearForward(const std::shared_ptr &input, const std::shared_ptr &weight, - bool transpose, const std::shared_ptr &bias) { - /* - transpose: output = input * weight^T + bias - output[*, out_features] = input[*, in_features] * weight[out_features, in_features]^T + bias[out_features] - - !transpose: output = input * weight + bias - output[*, out_features] = input[*, in_features] * weight[in_features, out_features] + bias[out_features] - */ - - const auto &input_dims = input->Dims(); - CHECK_GE(input_dims.size(), 2); - const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); - const int64_t in_features = *input_dims.rbegin(); - - const auto &weight_dims = weight->Dims(); - CHECK_EQ(weight_dims.size(), 2); - CHECK_EQ(in_features, weight_dims[transpose ? 1 : 0]); - const int out_features = weight_dims[transpose ? 0 : 1]; - - if (bias) { - const auto &bias_dims = bias->Dims(); - CHECK_EQ(bias_dims.size(), 1); - CHECK_EQ(bias_dims[0], out_features); - } - - auto output_dims = input_dims; - *output_dims.rbegin() = out_features; - auto output = std::make_shared(output_dims, DataType::kFLOAT32); - - if (transpose) { - output->EigenMatrix() = input->EigenMatrix() * weight->EigenMatrix().transpose(); - } else { - output->EigenMatrix() = input->EigenMatrix() * weight->EigenMatrix(); - } - - if (bias) { - output->EigenMatrix().rowwise() += bias->EigenVector(); - } - - return output; -} - -std::tuple, std::shared_ptr, std::shared_ptr> -LinearBackward(const std::shared_ptr &input, const std::shared_ptr &weight, bool transpose, - int64_t out_features, const std::shared_ptr &grad_output, const bool bias) { - /* - transpose: grad_input = grad_output * weight - grad_input[*, in_features] = grad_output[*, out_features] * weight[out_features, in_features] - grad_weight[out_features, in_features] = grad_output[*, out_features]^T * input[*, in_features] - grad_bias[out_features] = grad_output[*, out_features].sum(axis=0) - - !transpose: grad_input = grad_output * weight^T - grad_input[*, in_features] = grad_output[_, out_features] * weight[in_features, out_features]^T - grad_weight[in_features, out_features] = input[*, in_features]^T * grad_output[*, out_features] - grad_bias[out_features] = grad_output[*, out_features].sum(axis=0) - */ - - const auto &input_dims = input->Dims(); - CHECK_GE(input_dims.size(), 2); - const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); - const int64_t in_features = *input_dims.rbegin(); - - const auto &weight_dims = weight->Dims(); - CHECK_EQ(weight_dims.size(), 2); - CHECK_EQ(in_features, weight_dims[transpose ? 1 : 0]); - CHECK_EQ(out_features, weight_dims[transpose ? 0 : 1]); - - auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32); - auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32); - std::shared_ptr grad_bias = nullptr; - if (bias) { - grad_bias = std::make_shared(std::vector{out_features}, DataType::kFLOAT32); - } - - if (transpose) { - grad_input->EigenMatrix() = grad_output->EigenMatrix() * weight->EigenMatrix(); - grad_weight->EigenMatrix() = grad_output->EigenMatrix().transpose() * input->EigenMatrix(); - } else { - grad_input->EigenMatrix() = grad_output->EigenMatrix() * weight->EigenMatrix().transpose(); - grad_weight->EigenMatrix() = input->EigenMatrix().transpose() * grad_output->EigenMatrix(); - } - if (bias) { - grad_bias->EigenVector() = grad_output->EigenMatrix().colwise().sum(); - } - - return {grad_input, grad_weight, grad_bias}; -} -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_LINEAR_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_LINEAR_KERNEL(MatmulForward) -REGISTER_CPU_LINEAR_KERNEL(MatmulBackward) -REGISTER_CPU_LINEAR_KERNEL(LinearForward) -REGISTER_CPU_LINEAR_KERNEL(LinearBackward) - -#undef REGISTER_CPU_LINEAR_KERNEL +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr MatmulForward(const std::shared_ptr &input, const std::shared_ptr &other) { + // =================================== 作业 =================================== + // TODO：实现CPU上的矩阵乘法前向计算 + // REF: + // =================================== 作业 =================================== + + const auto &input_dims = input->Dims(); + const auto &other_dims = other->Dims(); + CHECK_GE(input_dims.size(), 2); + CHECK_GE(other_dims.size(), 2); + const int64_t M = input_dims[input_dims.size() - 2]; + const int64_t K = input_dims[input_dims.size() - 1]; + const int64_t N = other_dims[other_dims.size() - 1]; + CHECK_EQ(K, other_dims[other_dims.size() - 2]); + + // Determine batch dimensions + int64_t batch_input = 1; + for (size_t i = 0; i < input_dims.size() - 2; ++i) { + batch_input *= input_dims[i]; + } + int64_t batch_other = 1; + for (size_t i = 0; i < other_dims.size() - 2; ++i) { + batch_other *= other_dims[i]; + } + int64_t batch = std::max(batch_input, batch_other); + + // Build output dims + std::vector output_dims; + if (batch > 1) { + output_dims = (input_dims.size() >= other_dims.size()) ? input_dims : other_dims; + output_dims[output_dims.size() - 2] = M; + output_dims[output_dims.size() - 1] = N; + } else { + output_dims = {M, N}; + } + + auto output = std::make_shared(output_dims, DataType::kFLOAT32); + + float *input_ptr = static_cast(input->DataPtr()); + float *other_ptr = static_cast(other->DataPtr()); + float *output_ptr = static_cast(output->DataPtr()); + + for (int64_t b = 0; b < batch; ++b) { + int64_t input_offset = (batch_input == 1) ? 0 : b * M * K; + int64_t other_offset = (batch_other == 1) ? 0 : b * K * N; + int64_t output_offset = b * M * N; + + Eigen::Map> input_mat( + input_ptr + input_offset, M, K); + Eigen::Map> other_mat( + other_ptr + other_offset, K, N); + Eigen::Map> output_mat( + output_ptr + output_offset, M, N); + + output_mat = input_mat * other_mat; + } + + return output; +} + +std::tuple, std::shared_ptr> +MatmulBackward(const std::shared_ptr &input, const std::shared_ptr &other, + const std::shared_ptr &grad_output) { + // =================================== 作业 =================================== + // TODO：实现CPU上的矩阵乘法反向传播 + // REF: + // =================================== 作业 =================================== + + const auto &input_dims = input->Dims(); + const auto &other_dims = other->Dims(); + const int64_t M = input_dims[input_dims.size() - 2]; + const int64_t K = input_dims[input_dims.size() - 1]; + const int64_t N = other_dims[other_dims.size() - 1]; + + int64_t batch_input = 1; + for (size_t i = 0; i < input_dims.size() - 2; ++i) { + batch_input *= input_dims[i]; + } + int64_t batch_other = 1; + for (size_t i = 0; i < other_dims.size() - 2; ++i) { + batch_other *= other_dims[i]; + } + int64_t batch = std::max(batch_input, batch_other); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32); + auto grad_other = std::make_shared(other_dims, DataType::kFLOAT32); + + float *input_ptr = static_cast(input->DataPtr()); + float *other_ptr = static_cast(other->DataPtr()); + float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + float *grad_other_ptr = static_cast(grad_other->DataPtr()); + + for (int64_t b = 0; b < batch; ++b) { + int64_t input_offset = (batch_input == 1) ? 0 : b * M * K; + int64_t other_offset = (batch_other == 1) ? 0 : b * K * N; + int64_t grad_output_offset = b * M * N; + int64_t grad_input_offset = (batch_input == 1) ? 0 : b * M * K; + int64_t grad_other_offset = (batch_other == 1) ? 0 : b * K * N; + + Eigen::Map> input_mat( + input_ptr + input_offset, M, K); + Eigen::Map> other_mat( + other_ptr + other_offset, K, N); + Eigen::Map> grad_output_mat( + grad_output_ptr + grad_output_offset, M, N); + Eigen::Map> grad_input_mat( + grad_input_ptr + grad_input_offset, M, K); + Eigen::Map> grad_other_mat( + grad_other_ptr + grad_other_offset, K, N); + + // dX = dY * W^T + grad_input_mat = grad_output_mat * other_mat.transpose(); + // dW = X^T * dY + grad_other_mat = input_mat.transpose() * grad_output_mat; + } + + return {grad_input, grad_other}; +} + +std::shared_ptr LinearForward(const std::shared_ptr &input, const std::shared_ptr &weight, + bool transpose, const std::shared_ptr &bias) { + /* + transpose: output = input * weight^T + bias + output[*, out_features] = input[*, in_features] * weight[out_features, in_features]^T + bias[out_features] + + !transpose: output = input * weight + bias + output[*, out_features] = input[*, in_features] * weight[in_features, out_features] + bias[out_features] + */ + + const auto &input_dims = input->Dims(); + CHECK_GE(input_dims.size(), 2); + const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int64_t in_features = *input_dims.rbegin(); + + const auto &weight_dims = weight->Dims(); + CHECK_EQ(weight_dims.size(), 2); + CHECK_EQ(in_features, weight_dims[transpose ? 1 : 0]); + const int out_features = weight_dims[transpose ? 0 : 1]; + + if (bias) { + const auto &bias_dims = bias->Dims(); + CHECK_EQ(bias_dims.size(), 1); + CHECK_EQ(bias_dims[0], out_features); + } + + auto output_dims = input_dims; + *output_dims.rbegin() = out_features; + auto output = std::make_shared(output_dims, DataType::kFLOAT32); + + if (transpose) { + output->EigenMatrix() = input->EigenMatrix() * weight->EigenMatrix().transpose(); + } else { + output->EigenMatrix() = input->EigenMatrix() * weight->EigenMatrix(); + } + + if (bias) { + output->EigenMatrix().rowwise() += bias->EigenVector(); + } + + return output; +} + +std::tuple, std::shared_ptr, std::shared_ptr> +LinearBackward(const std::shared_ptr &input, const std::shared_ptr &weight, bool transpose, + int64_t out_features, const std::shared_ptr &grad_output, const bool bias) { + /* + transpose: grad_input = grad_output * weight + grad_input[*, in_features] = grad_output[*, out_features] * weight[out_features, in_features] + grad_weight[out_features, in_features] = grad_output[*, out_features]^T * input[*, in_features] + grad_bias[out_features] = grad_output[*, out_features].sum(axis=0) + + !transpose: grad_input = grad_output * weight^T + grad_input[*, in_features] = grad_output[_, out_features] * weight[in_features, out_features]^T + grad_weight[in_features, out_features] = input[*, in_features]^T * grad_output[*, out_features] + grad_bias[out_features] = grad_output[*, out_features].sum(axis=0) + */ + + const auto &input_dims = input->Dims(); + CHECK_GE(input_dims.size(), 2); + const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int64_t in_features = *input_dims.rbegin(); + + const auto &weight_dims = weight->Dims(); + CHECK_EQ(weight_dims.size(), 2); + CHECK_EQ(in_features, weight_dims[transpose ? 1 : 0]); + CHECK_EQ(out_features, weight_dims[transpose ? 0 : 1]); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32); + auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32); + std::shared_ptr grad_bias = nullptr; + if (bias) { + grad_bias = std::make_shared(std::vector{out_features}, DataType::kFLOAT32); + } + + if (transpose) { + grad_input->EigenMatrix() = grad_output->EigenMatrix() * weight->EigenMatrix(); + grad_weight->EigenMatrix() = grad_output->EigenMatrix().transpose() * input->EigenMatrix(); + } else { + grad_input->EigenMatrix() = grad_output->EigenMatrix() * weight->EigenMatrix().transpose(); + grad_weight->EigenMatrix() = input->EigenMatrix().transpose() * grad_output->EigenMatrix(); + } + if (bias) { + grad_bias->EigenVector() = grad_output->EigenMatrix().colwise().sum(); + } + + return {grad_input, grad_weight, grad_bias}; +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_LINEAR_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_LINEAR_KERNEL(MatmulForward) +REGISTER_CPU_LINEAR_KERNEL(MatmulBackward) +REGISTER_CPU_LINEAR_KERNEL(LinearForward) +REGISTER_CPU_LINEAR_KERNEL(LinearBackward) + +#undef REGISTER_CPU_LINEAR_KERNEL diff --git a/infini_train/src/kernels/cpu/no_op.cc b/infini_train/src/kernels/cpu/no_op.cc index 343d7a2..a99d12b 100644 --- a/infini_train/src/kernels/cpu/no_op.cc +++ b/infini_train/src/kernels/cpu/no_op.cc @@ -1,33 +1,33 @@ -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -std::shared_ptr NoOpForward(const std::shared_ptr &input, const std::vector &dims) { - const int64_t num_elements = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); - CHECK_EQ(input->NumElements(), num_elements); - - auto output = std::make_shared(*input, 0, dims); - return output; -} - -std::shared_ptr NoOpBackward(const std::vector &dims, const std::shared_ptr &grad_output) { - auto num_elements = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); - CHECK_EQ(num_elements, grad_output->NumElements()); - - auto grad_input = std::make_shared(*grad_output, 0, dims); - return grad_input; -} -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_NO_OP_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_NO_OP_KERNEL(NoOpForward) -REGISTER_CPU_NO_OP_KERNEL(NoOpBackward) - -#undef REGISTER_CPU_NO_OP_KERNEL +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr NoOpForward(const std::shared_ptr &input, const std::vector &dims) { + const int64_t num_elements = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); + CHECK_EQ(input->NumElements(), num_elements); + + auto output = std::make_shared(*input, 0, dims); + return output; +} + +std::shared_ptr NoOpBackward(const std::vector &dims, const std::shared_ptr &grad_output) { + auto num_elements = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); + CHECK_EQ(num_elements, grad_output->NumElements()); + + auto grad_input = std::make_shared(*grad_output, 0, dims); + return grad_input; +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_NO_OP_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_NO_OP_KERNEL(NoOpForward) +REGISTER_CPU_NO_OP_KERNEL(NoOpBackward) + +#undef REGISTER_CPU_NO_OP_KERNEL diff --git a/infini_train/src/kernels/cpu/outer.cc b/infini_train/src/kernels/cpu/outer.cc index 1a8d392..2ac4c95 100644 --- a/infini_train/src/kernels/cpu/outer.cc +++ b/infini_train/src/kernels/cpu/outer.cc @@ -1,66 +1,66 @@ -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -std::shared_ptr OuterForward(const std::shared_ptr &input, const std::shared_ptr &other) { - /* - output[i, j] = input[i] * other[j] - output shape: [input.size(0), other.size(0)] - */ - - const auto &input_dims = input->Dims(); - const auto &other_dims = other->Dims(); - CHECK_EQ(input_dims.size(), 1); - CHECK_EQ(other_dims.size(), 1); - - std::vector out_shape = {input_dims[0], other_dims[0]}; - auto output = std::make_shared(out_shape, DataType::kFLOAT32); - - // [M, N] = [M, 1] * [1, N] - output->EigenMatrix() = input->EigenVector().transpose() * other->EigenVector(); - - return output; -} - -std::tuple, std::shared_ptr> OuterBackward(const std::shared_ptr &input, - const std::shared_ptr &other, - const std::shared_ptr &grad_output) { - /* - grad_input[i] = sum_j(grad_output[i, j] * other[j]) - grad_other[j] = sum_i(grad_output[i, j] * input[i]) - */ - - const int64_t m = input->Dims()[0]; - const int64_t n = other->Dims()[0]; - CHECK_EQ(grad_output->Dims().size(), 2); - CHECK_EQ(grad_output->Dims()[0], m); - CHECK_EQ(grad_output->Dims()[1], n); - - auto grad_input = std::make_shared(std::vector{m}, DataType::kFLOAT32); - auto grad_other = std::make_shared(std::vector{n}, DataType::kFLOAT32); - grad_input->Fill(0.0f); - grad_other->Fill(0.0f); - - grad_input->EigenVector() = grad_output->EigenMatrix() * other->EigenVector().transpose(); - grad_other->EigenVector() = grad_output->EigenMatrix().transpose() * input->EigenVector().transpose(); - - return {grad_input, grad_other}; -} - -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_OUTER_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_OUTER_KERNEL(OuterForward) -REGISTER_CPU_OUTER_KERNEL(OuterBackward) - -#undef REGISTER_CPU_OUTER_KERNEL +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr OuterForward(const std::shared_ptr &input, const std::shared_ptr &other) { + /* + output[i, j] = input[i] * other[j] + output shape: [input.size(0), other.size(0)] + */ + + const auto &input_dims = input->Dims(); + const auto &other_dims = other->Dims(); + CHECK_EQ(input_dims.size(), 1); + CHECK_EQ(other_dims.size(), 1); + + std::vector out_shape = {input_dims[0], other_dims[0]}; + auto output = std::make_shared(out_shape, DataType::kFLOAT32); + + // [M, N] = [M, 1] * [1, N] + output->EigenMatrix() = input->EigenVector().transpose() * other->EigenVector(); + + return output; +} + +std::tuple, std::shared_ptr> OuterBackward(const std::shared_ptr &input, + const std::shared_ptr &other, + const std::shared_ptr &grad_output) { + /* + grad_input[i] = sum_j(grad_output[i, j] * other[j]) + grad_other[j] = sum_i(grad_output[i, j] * input[i]) + */ + + const int64_t m = input->Dims()[0]; + const int64_t n = other->Dims()[0]; + CHECK_EQ(grad_output->Dims().size(), 2); + CHECK_EQ(grad_output->Dims()[0], m); + CHECK_EQ(grad_output->Dims()[1], n); + + auto grad_input = std::make_shared(std::vector{m}, DataType::kFLOAT32); + auto grad_other = std::make_shared(std::vector{n}, DataType::kFLOAT32); + grad_input->Fill(0.0f); + grad_other->Fill(0.0f); + + grad_input->EigenVector() = grad_output->EigenMatrix() * other->EigenVector().transpose(); + grad_other->EigenVector() = grad_output->EigenMatrix().transpose() * input->EigenVector().transpose(); + + return {grad_input, grad_other}; +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_OUTER_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_OUTER_KERNEL(OuterForward) +REGISTER_CPU_OUTER_KERNEL(OuterBackward) + +#undef REGISTER_CPU_OUTER_KERNEL diff --git a/infini_train/src/kernels/cpu/reduction.cc b/infini_train/src/kernels/cpu/reduction.cc index 87ed538..00be1ec 100644 --- a/infini_train/src/kernels/cpu/reduction.cc +++ b/infini_train/src/kernels/cpu/reduction.cc @@ -1,183 +1,183 @@ -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -std::shared_ptr ReduceOpForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim, - const std::function &reduce_fn) { - const auto &input_dims = input->Dims(); - int64_t actual_dim = dim < 0 ? dim + input_dims.size() : dim; - CHECK_GE(actual_dim, 0); - CHECK_LT(actual_dim, input_dims.size()); - - std::vector output_dims = input_dims; - const int64_t reduce_size = input_dims[dim]; - if (keep_dim) { - output_dims[actual_dim] = 1; - } else { - output_dims.erase(output_dims.begin() + actual_dim); - } - - auto output = std::make_shared(output_dims, DataType::kFLOAT32); - - int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + actual_dim, 1, std::multiplies()); - int64_t H = input_dims[actual_dim]; - int64_t W = std::accumulate(input_dims.begin() + actual_dim + 1, input_dims.end(), 1, std::multiplies()); - - const float *input_ptr = static_cast(input->DataPtr()); - float *output_ptr = static_cast(output->DataPtr()); - - for (int64_t n = 0; n < N; ++n) { - for (int64_t w = 0; w < W; ++w) { - const float *segment = &input_ptr[(n * H) * W + w]; - output_ptr[n * W + w] = reduce_fn(segment, H); - } - } - - return output; -} - -std::shared_ptr ReduceOpBackwardMask(const std::shared_ptr &input, - const std::shared_ptr &grad_output, - const std::shared_ptr &reduced, - const std::vector &input_dims, const int64_t dim, - const bool keep_dim, const std::function &mask_fn) { - std::vector grad_output_dims = grad_output->Dims(); - int64_t actual_dim = dim < 0 ? dim + input_dims.size() : dim; - if (!keep_dim) { - grad_output_dims.insert(grad_output_dims.begin() + actual_dim, 1); - } - - auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32); - - int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + actual_dim, 1, std::multiplies()); - int64_t H = input_dims[actual_dim]; - int64_t W = std::accumulate(input_dims.begin() + actual_dim + 1, input_dims.end(), 1, std::multiplies()); - - const float *input_ptr = static_cast(input->DataPtr()); - const float *grad_output_ptr = static_cast(grad_output->DataPtr()); - const float *reduced_ptr = static_cast(reduced->DataPtr()); - float *grad_input_ptr = static_cast(grad_input->DataPtr()); - - for (int64_t n = 0; n < N; ++n) { - for (int64_t w = 0; w < W; ++w) { - int64_t reduced_idx = n * W + w; - float target = reduced_ptr[reduced_idx]; - - for (int64_t h = 0; h < H; ++h) { - int64_t input_idx = (n * H + h) * W + w; - float value = input_ptr[input_idx]; - grad_input_ptr[input_idx] = (value == target) ? grad_output_ptr[reduced_idx] : 0.0f; - } - } - } - - return grad_input; -} - -std::shared_ptr ReduceOpBackward(const std::shared_ptr &grad_output, - const std::vector &input_dims, const int64_t dim, const bool keep_dim, - const std::function &scale_fn) { - int64_t actual_dim = dim < 0 ? dim + input_dims.size() : dim; - CHECK_GE(actual_dim, 0); - CHECK_LT(actual_dim, input_dims.size()); - - auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32); - - int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + actual_dim, 1, std::multiplies()); - int64_t H = input_dims[actual_dim]; - int64_t W = std::accumulate(input_dims.begin() + actual_dim + 1, input_dims.end(), 1, std::multiplies()); - - const float *grad_output_ptr = static_cast(grad_output->DataPtr()); - float *grad_input_ptr = static_cast(grad_input->DataPtr()); - - for (int64_t n = 0; n < N; ++n) { - for (int64_t h = 0; h < H; ++h) { - for (int64_t w = 0; w < W; ++w) { - int64_t output_idx = n * W + w; - int64_t input_idx = (n * H + h) * W + w; - grad_input_ptr[input_idx] = scale_fn(grad_output_ptr[output_idx], H); - } - } - } - - return grad_input; -} - -std::shared_ptr MeanForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { - return ReduceOpForward(input, dim, keep_dim, [](const float *data, int64_t len) { - float sum = 0.0f; - for (int64_t i = 0; i < len; ++i) { sum += data[i]; } - return sum / static_cast(len); - }); -} - -std::shared_ptr MeanBackward(const std::shared_ptr &grad_output, const std::vector &input_dims, - const int64_t dim, const bool keep_dim) { - return ReduceOpBackward(grad_output, input_dims, dim, keep_dim, - [](float grad, int64_t len) { return grad / static_cast(len); }); -} - -std::shared_ptr SumForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { - return ReduceOpForward(input, dim, keep_dim, [](const float *data, int64_t len) { - float sum = 0.0f; - for (int64_t i = 0; i < len; ++i) { sum += data[i]; } - return sum; - }); -} - -std::shared_ptr SumBackward(const std::shared_ptr &grad_output, const std::vector &input_dims, - const int64_t dim, const bool keep_dim) { - return ReduceOpBackward(grad_output, input_dims, dim, keep_dim, [](float grad, int64_t) { return grad; }); -} - -std::shared_ptr MaxForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { - return ReduceOpForward(input, dim, keep_dim, [](const float *data, int64_t len) { - float max_val = -std::numeric_limits::infinity(); - for (int64_t i = 0; i < len; ++i) { max_val = std::max(max_val, data[i]); } - return max_val; - }); -} - -std::shared_ptr MaxBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, - const std::shared_ptr &reduced, const int64_t dim, bool keep_dim) { - return ReduceOpBackwardMask(input, grad_output, reduced, input->Dims(), dim, keep_dim, - [](float val, float current) { return val > current; }); -} - -std::shared_ptr MinForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { - return ReduceOpForward(input, dim, keep_dim, [](const float *data, int64_t len) { - float min_val = std::numeric_limits::infinity(); - for (int64_t i = 0; i < len; ++i) { min_val = std::min(min_val, data[i]); } - return min_val; - }); -} - -std::shared_ptr MinBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, - const std::shared_ptr &reduced, const int64_t dim, bool keep_dim) { - return ReduceOpBackwardMask(input, grad_output, reduced, input->Dims(), dim, keep_dim, - [](float val, float current) { return val < current; }); -} - -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_REDUCTION_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_REDUCTION_KERNEL(MeanForward) -REGISTER_CPU_REDUCTION_KERNEL(MeanBackward) -REGISTER_CPU_REDUCTION_KERNEL(SumForward) -REGISTER_CPU_REDUCTION_KERNEL(SumBackward) -REGISTER_CPU_REDUCTION_KERNEL(MaxForward) -REGISTER_CPU_REDUCTION_KERNEL(MaxBackward) -REGISTER_CPU_REDUCTION_KERNEL(MinForward) -REGISTER_CPU_REDUCTION_KERNEL(MinBackward) - -#undef REGISTER_CPU_REDUCTION_KERNEL +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr ReduceOpForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim, + const std::function &reduce_fn) { + const auto &input_dims = input->Dims(); + int64_t actual_dim = dim < 0 ? dim + input_dims.size() : dim; + CHECK_GE(actual_dim, 0); + CHECK_LT(actual_dim, input_dims.size()); + + std::vector output_dims = input_dims; + const int64_t reduce_size = input_dims[dim]; + if (keep_dim) { + output_dims[actual_dim] = 1; + } else { + output_dims.erase(output_dims.begin() + actual_dim); + } + + auto output = std::make_shared(output_dims, DataType::kFLOAT32); + + int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + actual_dim, 1, std::multiplies()); + int64_t H = input_dims[actual_dim]; + int64_t W = std::accumulate(input_dims.begin() + actual_dim + 1, input_dims.end(), 1, std::multiplies()); + + const float *input_ptr = static_cast(input->DataPtr()); + float *output_ptr = static_cast(output->DataPtr()); + + for (int64_t n = 0; n < N; ++n) { + for (int64_t w = 0; w < W; ++w) { + const float *segment = &input_ptr[(n * H) * W + w]; + output_ptr[n * W + w] = reduce_fn(segment, H); + } + } + + return output; +} + +std::shared_ptr ReduceOpBackwardMask(const std::shared_ptr &input, + const std::shared_ptr &grad_output, + const std::shared_ptr &reduced, + const std::vector &input_dims, const int64_t dim, + const bool keep_dim, const std::function &mask_fn) { + std::vector grad_output_dims = grad_output->Dims(); + int64_t actual_dim = dim < 0 ? dim + input_dims.size() : dim; + if (!keep_dim) { + grad_output_dims.insert(grad_output_dims.begin() + actual_dim, 1); + } + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32); + + int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + actual_dim, 1, std::multiplies()); + int64_t H = input_dims[actual_dim]; + int64_t W = std::accumulate(input_dims.begin() + actual_dim + 1, input_dims.end(), 1, std::multiplies()); + + const float *input_ptr = static_cast(input->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + const float *reduced_ptr = static_cast(reduced->DataPtr()); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + + for (int64_t n = 0; n < N; ++n) { + for (int64_t w = 0; w < W; ++w) { + int64_t reduced_idx = n * W + w; + float target = reduced_ptr[reduced_idx]; + + for (int64_t h = 0; h < H; ++h) { + int64_t input_idx = (n * H + h) * W + w; + float value = input_ptr[input_idx]; + grad_input_ptr[input_idx] = (value == target) ? grad_output_ptr[reduced_idx] : 0.0f; + } + } + } + + return grad_input; +} + +std::shared_ptr ReduceOpBackward(const std::shared_ptr &grad_output, + const std::vector &input_dims, const int64_t dim, const bool keep_dim, + const std::function &scale_fn) { + int64_t actual_dim = dim < 0 ? dim + input_dims.size() : dim; + CHECK_GE(actual_dim, 0); + CHECK_LT(actual_dim, input_dims.size()); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32); + + int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + actual_dim, 1, std::multiplies()); + int64_t H = input_dims[actual_dim]; + int64_t W = std::accumulate(input_dims.begin() + actual_dim + 1, input_dims.end(), 1, std::multiplies()); + + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + + for (int64_t n = 0; n < N; ++n) { + for (int64_t h = 0; h < H; ++h) { + for (int64_t w = 0; w < W; ++w) { + int64_t output_idx = n * W + w; + int64_t input_idx = (n * H + h) * W + w; + grad_input_ptr[input_idx] = scale_fn(grad_output_ptr[output_idx], H); + } + } + } + + return grad_input; +} + +std::shared_ptr MeanForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + return ReduceOpForward(input, dim, keep_dim, [](const float *data, int64_t len) { + float sum = 0.0f; + for (int64_t i = 0; i < len; ++i) { sum += data[i]; } + return sum / static_cast(len); + }); +} + +std::shared_ptr MeanBackward(const std::shared_ptr &grad_output, const std::vector &input_dims, + const int64_t dim, const bool keep_dim) { + return ReduceOpBackward(grad_output, input_dims, dim, keep_dim, + [](float grad, int64_t len) { return grad / static_cast(len); }); +} + +std::shared_ptr SumForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + return ReduceOpForward(input, dim, keep_dim, [](const float *data, int64_t len) { + float sum = 0.0f; + for (int64_t i = 0; i < len; ++i) { sum += data[i]; } + return sum; + }); +} + +std::shared_ptr SumBackward(const std::shared_ptr &grad_output, const std::vector &input_dims, + const int64_t dim, const bool keep_dim) { + return ReduceOpBackward(grad_output, input_dims, dim, keep_dim, [](float grad, int64_t) { return grad; }); +} + +std::shared_ptr MaxForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + return ReduceOpForward(input, dim, keep_dim, [](const float *data, int64_t len) { + float max_val = -std::numeric_limits::infinity(); + for (int64_t i = 0; i < len; ++i) { max_val = std::max(max_val, data[i]); } + return max_val; + }); +} + +std::shared_ptr MaxBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + const std::shared_ptr &reduced, const int64_t dim, bool keep_dim) { + return ReduceOpBackwardMask(input, grad_output, reduced, input->Dims(), dim, keep_dim, + [](float val, float current) { return val > current; }); +} + +std::shared_ptr MinForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + return ReduceOpForward(input, dim, keep_dim, [](const float *data, int64_t len) { + float min_val = std::numeric_limits::infinity(); + for (int64_t i = 0; i < len; ++i) { min_val = std::min(min_val, data[i]); } + return min_val; + }); +} + +std::shared_ptr MinBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + const std::shared_ptr &reduced, const int64_t dim, bool keep_dim) { + return ReduceOpBackwardMask(input, grad_output, reduced, input->Dims(), dim, keep_dim, + [](float val, float current) { return val < current; }); +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_REDUCTION_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_REDUCTION_KERNEL(MeanForward) +REGISTER_CPU_REDUCTION_KERNEL(MeanBackward) +REGISTER_CPU_REDUCTION_KERNEL(SumForward) +REGISTER_CPU_REDUCTION_KERNEL(SumBackward) +REGISTER_CPU_REDUCTION_KERNEL(MaxForward) +REGISTER_CPU_REDUCTION_KERNEL(MaxBackward) +REGISTER_CPU_REDUCTION_KERNEL(MinForward) +REGISTER_CPU_REDUCTION_KERNEL(MinBackward) + +#undef REGISTER_CPU_REDUCTION_KERNEL diff --git a/infini_train/src/kernels/cpu/sigmoid.cc b/infini_train/src/kernels/cpu/sigmoid.cc index d4bc05d..7ef6e4b 100644 --- a/infini_train/src/kernels/cpu/sigmoid.cc +++ b/infini_train/src/kernels/cpu/sigmoid.cc @@ -1,43 +1,43 @@ -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -std::shared_ptr SigmoidForward(const std::shared_ptr &input) { - auto output = std::make_shared(input->Dims(), DataType::kFLOAT32); - const float *input_ptr = static_cast(input->DataPtr()); - float *output_ptr = static_cast(output->DataPtr()); - - const int64_t numel = input->NumElements(); - for (int64_t idx = 0; idx < numel; ++idx) { output_ptr[idx] = 1.0f / (1.0f + exp(-input_ptr[idx])); } - - return output; -} - -std::shared_ptr SigmoidBackward(const std::shared_ptr &output, - const std::shared_ptr &grad_output) { - auto grad_input = std::make_shared(output->Dims(), DataType::kFLOAT32); - const float *output_ptr = static_cast(output->DataPtr()); - const float *grad_output_ptr = static_cast(grad_output->DataPtr()); - float *grad_input_ptr = static_cast(grad_input->DataPtr()); - - const int64_t numel = output->NumElements(); - for (int64_t idx = 0; idx < numel; ++idx) { - const float y = output_ptr[idx]; - grad_input_ptr[idx] = grad_output_ptr[idx] * y * (1.0f - y); - } - return grad_input; -} -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_SIGMOID_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_SIGMOID_KERNEL(SigmoidForward) -REGISTER_CPU_SIGMOID_KERNEL(SigmoidBackward) - -#undef REGISTER_CPU_SIGMOID_KERNEL +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr SigmoidForward(const std::shared_ptr &input) { + auto output = std::make_shared(input->Dims(), DataType::kFLOAT32); + const float *input_ptr = static_cast(input->DataPtr()); + float *output_ptr = static_cast(output->DataPtr()); + + const int64_t numel = input->NumElements(); + for (int64_t idx = 0; idx < numel; ++idx) { output_ptr[idx] = 1.0f / (1.0f + exp(-input_ptr[idx])); } + + return output; +} + +std::shared_ptr SigmoidBackward(const std::shared_ptr &output, + const std::shared_ptr &grad_output) { + auto grad_input = std::make_shared(output->Dims(), DataType::kFLOAT32); + const float *output_ptr = static_cast(output->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + + const int64_t numel = output->NumElements(); + for (int64_t idx = 0; idx < numel; ++idx) { + const float y = output_ptr[idx]; + grad_input_ptr[idx] = grad_output_ptr[idx] * y * (1.0f - y); + } + return grad_input; +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_SIGMOID_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_SIGMOID_KERNEL(SigmoidForward) +REGISTER_CPU_SIGMOID_KERNEL(SigmoidBackward) + +#undef REGISTER_CPU_SIGMOID_KERNEL diff --git a/infini_train/src/kernels/cpu/slice.cc b/infini_train/src/kernels/cpu/slice.cc index d0f4946..3d2d0de 100644 --- a/infini_train/src/kernels/cpu/slice.cc +++ b/infini_train/src/kernels/cpu/slice.cc @@ -1,138 +1,138 @@ -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -std::shared_ptr SliceForward(const std::shared_ptr &input, const std::vector &starts, - const std::vector &ends, const std::vector &steps) { - CHECK_EQ(starts.size(), ends.size()); - CHECK_EQ(starts.size(), steps.size()); - auto &dims = input->Dims(); - CHECK_EQ(starts.size(), dims.size()); - - std::vector new_dims; - for (int i = 0; i < starts.size(); i++) { - CHECK_LE(starts[i], ends[i]); - CHECK_LE(0, steps[i]); - new_dims.push_back((ends[i] - starts[i] + steps[i] - 1) / steps[i]); - } - - auto new_tensor = std::make_shared(new_dims, input->Dtype(), input->GetDevice()); - - std::vector src_strides(dims.size()); - int64_t stride = 1; - for (int i = src_strides.size() - 1; i >= 0; --i) { - src_strides[i] = stride; - stride *= dims[i]; - } - - std::vector dst_strides(new_dims.size()); - stride = 1; - for (int i = dst_strides.size() - 1; i >= 0; --i) { - dst_strides[i] = stride; - stride *= new_dims[i]; - } - - std::vector dst_index(new_dims.size(), 0); - std::vector src_index(dims.size(), 0); - - std::function recurse = [&](int d) { - if (d == dims.size()) { - int64_t src_offset = 0; - int64_t dst_offset = 0; - for (int i = dims.size() - 1; i >= 0; --i) { - src_offset += src_index[i] * src_strides[i]; - dst_offset += dst_index[i] * dst_strides[i]; - } - static_cast(new_tensor->DataPtr())[dst_offset] - = static_cast(input->DataPtr())[src_offset]; - return; - } - - // fill in the src_index and dst_index - int64_t out_i = 0; - for (int64_t i = starts[d]; i < ends[d]; i += steps[d]) { - src_index[d] = i; - dst_index[d] = out_i++; - recurse(d + 1); - } - }; - - recurse(0); - return new_tensor; -} - -std::shared_ptr SliceBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, - const std::vector &starts, const std::vector &ends, - const std::vector &steps) { - CHECK_EQ(starts.size(), ends.size()); - CHECK_EQ(starts.size(), steps.size()); - auto &dims = input->Dims(); - CHECK_EQ(starts.size(), dims.size()); - - std::vector new_dims; - for (int i = 0; i < starts.size(); i++) { - CHECK_LE(starts[i], ends[i]); - CHECK_LE(0, steps[i]); - new_dims.push_back((ends[i] - starts[i] + steps[i] - 1) / steps[i]); - } - - auto new_tensor = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); - new_tensor->Fill(0.0); - - std::vector src_strides(dims.size()); - int64_t stride = 1; - for (int i = src_strides.size() - 1; i >= 0; --i) { - src_strides[i] = stride; - stride *= dims[i]; - } - - std::vector dst_strides(new_dims.size()); - stride = 1; - for (int i = dst_strides.size() - 1; i >= 0; --i) { - dst_strides[i] = stride; - stride *= new_dims[i]; - } - - std::vector dst_index(new_dims.size(), 0); - std::vector src_index(dims.size(), 0); - - std::function recurse = [&](int d) { - if (d == dims.size()) { - int64_t src_offset = 0; - int64_t dst_offset = 0; - for (int i = dims.size() - 1; i >= 0; --i) { - src_offset += src_index[i] * src_strides[i]; - dst_offset += dst_index[i] * dst_strides[i]; - } - static_cast(new_tensor->DataPtr())[src_offset] - = static_cast(grad_output->DataPtr())[dst_offset]; - return; - } - - // fill in the src_index and dst_index - int64_t out_i = 0; - for (int64_t i = starts[d]; i < ends[d]; i += steps[d]) { - src_index[d] = i; - dst_index[d] = out_i++; - recurse(d + 1); - } - }; - - recurse(0); - return new_tensor; -} -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_SLICE_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_SLICE_KERNEL(SliceForward) -REGISTER_CPU_SLICE_KERNEL(SliceBackward) - -#undef REGISTER_CPU_SLICE_KERNEL +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr SliceForward(const std::shared_ptr &input, const std::vector &starts, + const std::vector &ends, const std::vector &steps) { + CHECK_EQ(starts.size(), ends.size()); + CHECK_EQ(starts.size(), steps.size()); + auto &dims = input->Dims(); + CHECK_EQ(starts.size(), dims.size()); + + std::vector new_dims; + for (int i = 0; i < starts.size(); i++) { + CHECK_LE(starts[i], ends[i]); + CHECK_LE(0, steps[i]); + new_dims.push_back((ends[i] - starts[i] + steps[i] - 1) / steps[i]); + } + + auto new_tensor = std::make_shared(new_dims, input->Dtype(), input->GetDevice()); + + std::vector src_strides(dims.size()); + int64_t stride = 1; + for (int i = src_strides.size() - 1; i >= 0; --i) { + src_strides[i] = stride; + stride *= dims[i]; + } + + std::vector dst_strides(new_dims.size()); + stride = 1; + for (int i = dst_strides.size() - 1; i >= 0; --i) { + dst_strides[i] = stride; + stride *= new_dims[i]; + } + + std::vector dst_index(new_dims.size(), 0); + std::vector src_index(dims.size(), 0); + + std::function recurse = [&](int d) { + if (d == dims.size()) { + int64_t src_offset = 0; + int64_t dst_offset = 0; + for (int i = dims.size() - 1; i >= 0; --i) { + src_offset += src_index[i] * src_strides[i]; + dst_offset += dst_index[i] * dst_strides[i]; + } + static_cast(new_tensor->DataPtr())[dst_offset] + = static_cast(input->DataPtr())[src_offset]; + return; + } + + // fill in the src_index and dst_index + int64_t out_i = 0; + for (int64_t i = starts[d]; i < ends[d]; i += steps[d]) { + src_index[d] = i; + dst_index[d] = out_i++; + recurse(d + 1); + } + }; + + recurse(0); + return new_tensor; +} + +std::shared_ptr SliceBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + const std::vector &starts, const std::vector &ends, + const std::vector &steps) { + CHECK_EQ(starts.size(), ends.size()); + CHECK_EQ(starts.size(), steps.size()); + auto &dims = input->Dims(); + CHECK_EQ(starts.size(), dims.size()); + + std::vector new_dims; + for (int i = 0; i < starts.size(); i++) { + CHECK_LE(starts[i], ends[i]); + CHECK_LE(0, steps[i]); + new_dims.push_back((ends[i] - starts[i] + steps[i] - 1) / steps[i]); + } + + auto new_tensor = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + new_tensor->Fill(0.0); + + std::vector src_strides(dims.size()); + int64_t stride = 1; + for (int i = src_strides.size() - 1; i >= 0; --i) { + src_strides[i] = stride; + stride *= dims[i]; + } + + std::vector dst_strides(new_dims.size()); + stride = 1; + for (int i = dst_strides.size() - 1; i >= 0; --i) { + dst_strides[i] = stride; + stride *= new_dims[i]; + } + + std::vector dst_index(new_dims.size(), 0); + std::vector src_index(dims.size(), 0); + + std::function recurse = [&](int d) { + if (d == dims.size()) { + int64_t src_offset = 0; + int64_t dst_offset = 0; + for (int i = dims.size() - 1; i >= 0; --i) { + src_offset += src_index[i] * src_strides[i]; + dst_offset += dst_index[i] * dst_strides[i]; + } + static_cast(new_tensor->DataPtr())[src_offset] + = static_cast(grad_output->DataPtr())[dst_offset]; + return; + } + + // fill in the src_index and dst_index + int64_t out_i = 0; + for (int64_t i = starts[d]; i < ends[d]; i += steps[d]) { + src_index[d] = i; + dst_index[d] = out_i++; + recurse(d + 1); + } + }; + + recurse(0); + return new_tensor; +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_SLICE_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_SLICE_KERNEL(SliceForward) +REGISTER_CPU_SLICE_KERNEL(SliceBackward) + +#undef REGISTER_CPU_SLICE_KERNEL diff --git a/infini_train/src/kernels/cpu/softmax.cc b/infini_train/src/kernels/cpu/softmax.cc index 454bdc2..1ea24c7 100644 --- a/infini_train/src/kernels/cpu/softmax.cc +++ b/infini_train/src/kernels/cpu/softmax.cc @@ -1,89 +1,89 @@ -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -std::shared_ptr SoftmaxForward(const std::shared_ptr &input, int64_t dim) { - dim = dim < 0 ? input->Dims().size() + dim : dim; - auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); - - const float *input_data = static_cast(input->DataPtr()); - float *output_data = static_cast(output->DataPtr()); - - const auto &dims = input->Dims(); - int outer = 1; - int axis = dims[dim]; - int inner = 1; - - for (int i = 0; i < dim; ++i) { outer *= dims[i]; } - for (int i = dim + 1; i < dims.size(); ++i) { inner *= dims[i]; } - - for (int o = 0; o < outer; ++o) { - for (int i = 0; i < inner; ++i) { - int offset = o * axis * inner + i; - - // find max value in the axis dimension - float max_val = input_data[offset]; - for (int j = 1; j < axis; ++j) { max_val = std::max(max_val, input_data[offset + j * inner]); } - - // compute sum_exp - float sum_exp = 0.0f; - for (int j = 0; j < axis; ++j) { - output_data[offset + j * inner] = std::exp(input_data[offset + j * inner] - max_val); - sum_exp += output_data[offset + j * inner]; - } - - // normalize - for (int j = 0; j < axis; ++j) { output_data[offset + j * inner] /= sum_exp; } - } - } - return output; -} - -std::shared_ptr SoftmaxBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &output, int64_t dim) { - dim = dim < 0 ? output->Dims().size() + dim : dim; - auto grad_input = std::make_shared(output->Dims(), output->Dtype(), output->GetDevice()); - const float *grad_output_data = static_cast(grad_output->DataPtr()); - const float *output_data = static_cast(output->DataPtr()); - float *grad_input_data = static_cast(grad_input->DataPtr()); - - const auto &dims = output->Dims(); - int outer = 1; - int inner = 1; - int axis = dims[dim]; - - for (int i = 0; i < dim; ++i) { outer *= dims[i]; } - for (int i = dim + 1; i < dims.size(); ++i) { inner *= dims[i]; } - - for (int o = 0; o < outer; ++o) { - for (int i = 0; i < inner; ++i) { - int offset = o * axis * inner + i; - - // dot product: y * dy - float dot = 0.0f; - for (int j = 0; j < axis; ++j) { - dot += output_data[offset + j * inner] * grad_output_data[offset + j * inner]; - } - for (int j = 0; j < axis; ++j) { - grad_input_data[offset + j * inner] - = output_data[offset + j * inner] * (grad_output_data[offset + j * inner] - dot); - } - } - } - return grad_input; -} -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_SOFTMAX_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_SOFTMAX_KERNEL(SoftmaxForward) -REGISTER_CPU_SOFTMAX_KERNEL(SoftmaxBackward) - -#undef REGISTER_CPU_SOFTMAX_KERNEL +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr SoftmaxForward(const std::shared_ptr &input, int64_t dim) { + dim = dim < 0 ? input->Dims().size() + dim : dim; + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + + const float *input_data = static_cast(input->DataPtr()); + float *output_data = static_cast(output->DataPtr()); + + const auto &dims = input->Dims(); + int outer = 1; + int axis = dims[dim]; + int inner = 1; + + for (int i = 0; i < dim; ++i) { outer *= dims[i]; } + for (int i = dim + 1; i < dims.size(); ++i) { inner *= dims[i]; } + + for (int o = 0; o < outer; ++o) { + for (int i = 0; i < inner; ++i) { + int offset = o * axis * inner + i; + + // find max value in the axis dimension + float max_val = input_data[offset]; + for (int j = 1; j < axis; ++j) { max_val = std::max(max_val, input_data[offset + j * inner]); } + + // compute sum_exp + float sum_exp = 0.0f; + for (int j = 0; j < axis; ++j) { + output_data[offset + j * inner] = std::exp(input_data[offset + j * inner] - max_val); + sum_exp += output_data[offset + j * inner]; + } + + // normalize + for (int j = 0; j < axis; ++j) { output_data[offset + j * inner] /= sum_exp; } + } + } + return output; +} + +std::shared_ptr SoftmaxBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &output, int64_t dim) { + dim = dim < 0 ? output->Dims().size() + dim : dim; + auto grad_input = std::make_shared(output->Dims(), output->Dtype(), output->GetDevice()); + const float *grad_output_data = static_cast(grad_output->DataPtr()); + const float *output_data = static_cast(output->DataPtr()); + float *grad_input_data = static_cast(grad_input->DataPtr()); + + const auto &dims = output->Dims(); + int outer = 1; + int inner = 1; + int axis = dims[dim]; + + for (int i = 0; i < dim; ++i) { outer *= dims[i]; } + for (int i = dim + 1; i < dims.size(); ++i) { inner *= dims[i]; } + + for (int o = 0; o < outer; ++o) { + for (int i = 0; i < inner; ++i) { + int offset = o * axis * inner + i; + + // dot product: y * dy + float dot = 0.0f; + for (int j = 0; j < axis; ++j) { + dot += output_data[offset + j * inner] * grad_output_data[offset + j * inner]; + } + for (int j = 0; j < axis; ++j) { + grad_input_data[offset + j * inner] + = output_data[offset + j * inner] * (grad_output_data[offset + j * inner] - dot); + } + } + } + return grad_input; +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_SOFTMAX_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_SOFTMAX_KERNEL(SoftmaxForward) +REGISTER_CPU_SOFTMAX_KERNEL(SoftmaxBackward) + +#undef REGISTER_CPU_SOFTMAX_KERNEL diff --git a/infini_train/src/kernels/cpu/split.cc b/infini_train/src/kernels/cpu/split.cc index ddd335e..1aad785 100644 --- a/infini_train/src/kernels/cpu/split.cc +++ b/infini_train/src/kernels/cpu/split.cc @@ -1,80 +1,80 @@ -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -std::vector> SplitForward(const std::shared_ptr &input, int64_t split_size, int dim) { - CHECK_GT(split_size, 0); - CHECK_GE(dim, 0) << "Currently we do not support negative dimension"; - const auto &input_dims = input->Dims(); - CHECK_LT(dim, input_dims.size()); - - std::vector> outputs; - for (int64_t start = 0; start < input_dims[dim]; start += split_size) { - auto output_dims = input_dims; - output_dims[dim] = std::min(split_size, input_dims[dim] - start); - auto output = std::make_shared(output_dims, DataType::kFLOAT32); - const int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); - const int64_t W - = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); - for (int64_t n = 0; n < N; ++n) { - // out[0:N, :, 0:W] = in[0:N, start:end, 0:W] - const int64_t H_in = input_dims[dim]; - const int64_t H_out = output_dims[dim]; - const int64_t end = std::min(start + split_size, input_dims[dim]); - memcpy(static_cast(output->DataPtr()) + n * H_out * W, - static_cast(input->DataPtr()) + n * H_in * W + start * W, - (end - start) * W * sizeof(float)); - } - outputs.push_back(std::move(output)); - } - return outputs; -} - -std::shared_ptr SplitBackward(const std::vector &input_dims, int64_t split_size, int dim, - const std::vector> &grad_outputs) { - CHECK_GT(split_size, 0); - CHECK_GE(dim, 0) << "Currently we do not support negative dimension"; - CHECK_LT(dim, input_dims.size()); - CHECK_EQ(grad_outputs.size(), (input_dims[dim] + split_size - 1) / split_size); - - auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32); - grad_input->Fill(0.0f); - for (int64_t start = 0, idx = 0; start < input_dims[dim]; start += split_size, ++idx) { - auto output_dims = input_dims; - output_dims[dim] = std::min(split_size, input_dims[dim] - start); - const auto &grad_output = grad_outputs[idx]; - for (int dim_idx = 0; dim_idx < grad_input->Dims().size(); ++dim_idx) { - CHECK_EQ(output_dims[dim_idx], grad_output->Dims()[dim_idx]); - } - const int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); - const int64_t W - = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); - for (int64_t n = 0; n < N; ++n) { - // grad_input[0:N, start:end, 0:W] = grad_output[0:N, :, 0:W] - const int64_t H_in = input_dims[dim]; - const int64_t H_out = output_dims[dim]; - const int64_t end = std::min(start + split_size, input_dims[dim]); - memcpy(static_cast(grad_input->DataPtr()) + n * H_in * W + start * W, - static_cast(grad_output->DataPtr()) + n * H_out * W, - (end - start) * W * sizeof(float)); - } - } - return grad_input; -} -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_SPLIT_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_SPLIT_KERNEL(SplitForward) -REGISTER_CPU_SPLIT_KERNEL(SplitBackward) - -#undef REGISTER_CPU_SPLIT_KERNEL +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::vector> SplitForward(const std::shared_ptr &input, int64_t split_size, int dim) { + CHECK_GT(split_size, 0); + CHECK_GE(dim, 0) << "Currently we do not support negative dimension"; + const auto &input_dims = input->Dims(); + CHECK_LT(dim, input_dims.size()); + + std::vector> outputs; + for (int64_t start = 0; start < input_dims[dim]; start += split_size) { + auto output_dims = input_dims; + output_dims[dim] = std::min(split_size, input_dims[dim] - start); + auto output = std::make_shared(output_dims, DataType::kFLOAT32); + const int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + const int64_t W + = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); + for (int64_t n = 0; n < N; ++n) { + // out[0:N, :, 0:W] = in[0:N, start:end, 0:W] + const int64_t H_in = input_dims[dim]; + const int64_t H_out = output_dims[dim]; + const int64_t end = std::min(start + split_size, input_dims[dim]); + memcpy(static_cast(output->DataPtr()) + n * H_out * W, + static_cast(input->DataPtr()) + n * H_in * W + start * W, + (end - start) * W * sizeof(float)); + } + outputs.push_back(std::move(output)); + } + return outputs; +} + +std::shared_ptr SplitBackward(const std::vector &input_dims, int64_t split_size, int dim, + const std::vector> &grad_outputs) { + CHECK_GT(split_size, 0); + CHECK_GE(dim, 0) << "Currently we do not support negative dimension"; + CHECK_LT(dim, input_dims.size()); + CHECK_EQ(grad_outputs.size(), (input_dims[dim] + split_size - 1) / split_size); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32); + grad_input->Fill(0.0f); + for (int64_t start = 0, idx = 0; start < input_dims[dim]; start += split_size, ++idx) { + auto output_dims = input_dims; + output_dims[dim] = std::min(split_size, input_dims[dim] - start); + const auto &grad_output = grad_outputs[idx]; + for (int dim_idx = 0; dim_idx < grad_input->Dims().size(); ++dim_idx) { + CHECK_EQ(output_dims[dim_idx], grad_output->Dims()[dim_idx]); + } + const int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + const int64_t W + = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); + for (int64_t n = 0; n < N; ++n) { + // grad_input[0:N, start:end, 0:W] = grad_output[0:N, :, 0:W] + const int64_t H_in = input_dims[dim]; + const int64_t H_out = output_dims[dim]; + const int64_t end = std::min(start + split_size, input_dims[dim]); + memcpy(static_cast(grad_input->DataPtr()) + n * H_in * W + start * W, + static_cast(grad_output->DataPtr()) + n * H_out * W, + (end - start) * W * sizeof(float)); + } + } + return grad_input; +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_SPLIT_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_SPLIT_KERNEL(SplitForward) +REGISTER_CPU_SPLIT_KERNEL(SplitBackward) + +#undef REGISTER_CPU_SPLIT_KERNEL diff --git a/infini_train/src/kernels/cpu/stack.cc b/infini_train/src/kernels/cpu/stack.cc index e672e6b..3f30d0d 100644 --- a/infini_train/src/kernels/cpu/stack.cc +++ b/infini_train/src/kernels/cpu/stack.cc @@ -1,91 +1,91 @@ -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { -std::shared_ptr StackForward(const std::vector> &inputs, int64_t dim) { - CHECK(!inputs.empty()); - - const auto &base_dims = inputs[0]->Dims(); - if (dim < 0) { - dim += base_dims.size() + 1; - } - CHECK_GE(dim, 0); - CHECK_LE(dim, base_dims.size()); - - for (const auto &input : inputs) { CHECK(input->Dims() == base_dims); } - - std::vector output_dims = base_dims; - output_dims.insert(output_dims.begin() + dim, inputs.size()); - - auto output = std::make_shared(output_dims, DataType::kFLOAT32); - - const int64_t outer_size - = std::accumulate(output_dims.begin(), output_dims.begin() + dim, 1, std::multiplies()); - const int64_t inner_size - = std::accumulate(output_dims.begin() + dim + 1, output_dims.end(), 1, std::multiplies()); - const size_t elem_size = sizeof(float); - - for (size_t i = 0; i < inputs.size(); ++i) { - const float *src_ptr = static_cast(inputs[i]->DataPtr()); - float *dst_ptr = static_cast(output->DataPtr()); - for (int64_t n = 0; n < outer_size; ++n) { - float *dst_block = dst_ptr + (n * inputs.size() + i) * inner_size; - const float *src_block = src_ptr + n * inner_size; - memcpy(dst_block, src_block, inner_size * elem_size); - } - } - - return output; -} - -std::vector> StackBackward(const std::vector &input_dims, int64_t dim, - const std::shared_ptr &grad_output) { - const auto &grad_dims = grad_output->Dims(); - int64_t actual_dim = dim < 0 ? dim + grad_dims.size() : dim; - CHECK_GE(actual_dim, 0); - CHECK_LT(actual_dim, grad_dims.size()); - - const int64_t num_inputs = grad_dims[actual_dim]; - std::vector> grads; - - std::vector out_dims = grad_dims; - out_dims.erase(out_dims.begin() + actual_dim); // remove stack dim - - const int64_t outer_size - = std::accumulate(out_dims.begin(), out_dims.begin() + actual_dim, 1, std::multiplies()); - const int64_t inner_size - = std::accumulate(out_dims.begin() + actual_dim, out_dims.end(), 1, std::multiplies()); - - const float *src_base = static_cast(grad_output->DataPtr()); - - for (int i = 0; i < num_inputs; ++i) { - auto grad = std::make_shared(input_dims, DataType::kFLOAT32); - float *dst_ptr = static_cast(grad->DataPtr()); - - for (int64_t n = 0; n < outer_size; ++n) { - const float *src_ptr = src_base + (n * num_inputs + i) * inner_size; - memcpy(dst_ptr + n * inner_size, src_ptr, inner_size * sizeof(float)); - } - grads.push_back(grad); - } - - return grads; -} - -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_STACK_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_STACK_KERNEL(StackForward) -REGISTER_CPU_STACK_KERNEL(StackBackward) - -#undef REGISTER_CPU_STACK_KERNEL +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr StackForward(const std::vector> &inputs, int64_t dim) { + CHECK(!inputs.empty()); + + const auto &base_dims = inputs[0]->Dims(); + if (dim < 0) { + dim += base_dims.size() + 1; + } + CHECK_GE(dim, 0); + CHECK_LE(dim, base_dims.size()); + + for (const auto &input : inputs) { CHECK(input->Dims() == base_dims); } + + std::vector output_dims = base_dims; + output_dims.insert(output_dims.begin() + dim, inputs.size()); + + auto output = std::make_shared(output_dims, DataType::kFLOAT32); + + const int64_t outer_size + = std::accumulate(output_dims.begin(), output_dims.begin() + dim, 1, std::multiplies()); + const int64_t inner_size + = std::accumulate(output_dims.begin() + dim + 1, output_dims.end(), 1, std::multiplies()); + const size_t elem_size = sizeof(float); + + for (size_t i = 0; i < inputs.size(); ++i) { + const float *src_ptr = static_cast(inputs[i]->DataPtr()); + float *dst_ptr = static_cast(output->DataPtr()); + for (int64_t n = 0; n < outer_size; ++n) { + float *dst_block = dst_ptr + (n * inputs.size() + i) * inner_size; + const float *src_block = src_ptr + n * inner_size; + memcpy(dst_block, src_block, inner_size * elem_size); + } + } + + return output; +} + +std::vector> StackBackward(const std::vector &input_dims, int64_t dim, + const std::shared_ptr &grad_output) { + const auto &grad_dims = grad_output->Dims(); + int64_t actual_dim = dim < 0 ? dim + grad_dims.size() : dim; + CHECK_GE(actual_dim, 0); + CHECK_LT(actual_dim, grad_dims.size()); + + const int64_t num_inputs = grad_dims[actual_dim]; + std::vector> grads; + + std::vector out_dims = grad_dims; + out_dims.erase(out_dims.begin() + actual_dim); // remove stack dim + + const int64_t outer_size + = std::accumulate(out_dims.begin(), out_dims.begin() + actual_dim, 1, std::multiplies()); + const int64_t inner_size + = std::accumulate(out_dims.begin() + actual_dim, out_dims.end(), 1, std::multiplies()); + + const float *src_base = static_cast(grad_output->DataPtr()); + + for (int i = 0; i < num_inputs; ++i) { + auto grad = std::make_shared(input_dims, DataType::kFLOAT32); + float *dst_ptr = static_cast(grad->DataPtr()); + + for (int64_t n = 0; n < outer_size; ++n) { + const float *src_ptr = src_base + (n * num_inputs + i) * inner_size; + memcpy(dst_ptr + n * inner_size, src_ptr, inner_size * sizeof(float)); + } + grads.push_back(grad); + } + + return grads; +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_STACK_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_STACK_KERNEL(StackForward) +REGISTER_CPU_STACK_KERNEL(StackBackward) + +#undef REGISTER_CPU_STACK_KERNEL diff --git a/infini_train/src/kernels/cpu/transform.cc b/infini_train/src/kernels/cpu/transform.cc index 5fa6c25..283bad5 100644 --- a/infini_train/src/kernels/cpu/transform.cc +++ b/infini_train/src/kernels/cpu/transform.cc @@ -1,234 +1,234 @@ -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cpu { - -std::shared_ptr TrilForward(const std::shared_ptr &input, int64_t diagonal) { - CHECK_EQ(input->Dims().size(), 2); - - auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); - for (int i = 0; i < input->NumElements(); ++i) { - int64_t row = i / input->Dims()[1]; - int64_t col = i % input->Dims()[1]; - if (row - col + diagonal >= 0) { - static_cast(output->DataPtr())[i] = static_cast(input->DataPtr())[i]; - } else { - static_cast(output->DataPtr())[i] = 0.0; - } - } - return output; -} - -std::shared_ptr TrilBackward(const std::shared_ptr &grad_output, int64_t diagonal) { - auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); - for (int i = 0; i < grad_output->NumElements(); ++i) { - int64_t row = i / grad_output->Dims()[1]; - int64_t col = i % grad_output->Dims()[1]; - if (row - col + diagonal >= 0) { - static_cast(grad_input->DataPtr())[i] = static_cast(grad_output->DataPtr())[i]; - } else { - static_cast(grad_input->DataPtr())[i] = 0.0; - } - } - return grad_input; -} - -std::shared_ptr TriuForward(const std::shared_ptr &input, int64_t diagonal) { - CHECK_EQ(input->Dims().size(), 2); - - auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); - for (int i = 0; i < input->NumElements(); ++i) { - int64_t row = i / input->Dims()[1]; - int64_t col = i % input->Dims()[1]; - if (row - col + diagonal <= 0) { - static_cast(output->DataPtr())[i] = static_cast(input->DataPtr())[i]; - } else { - static_cast(output->DataPtr())[i] = 0.0f; - } - } - return output; -} - -std::shared_ptr TriuBackward(const std::shared_ptr &grad_output, int64_t diagonal) { - CHECK_EQ(grad_output->Dims().size(), 2); - - auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); - for (int i = 0; i < grad_output->NumElements(); ++i) { - int64_t row = i / grad_output->Dims()[1]; - int64_t col = i % grad_output->Dims()[1]; - if (row - col + diagonal <= 0) { - static_cast(grad_input->DataPtr())[i] = static_cast(grad_output->DataPtr())[i]; - } else { - static_cast(grad_input->DataPtr())[i] = 0.0f; - } - } - return grad_input; -} - -std::shared_ptr TransposeForward(const std::shared_ptr &input, int64_t dim0, int64_t dim1) { - dim0 = dim0 < 0 ? dim0 + input->Dims().size() : dim0; - dim1 = dim1 < 0 ? dim1 + input->Dims().size() : dim1; - CHECK(dim0 >= 0 && dim0 < input->Dims().size() && dim1 >= 0 && dim1 < input->Dims().size()); - - auto in_dims = input->Dims(); - std::vector out_dims = in_dims; - std::swap(out_dims[dim0], out_dims[dim1]); - - auto output = std::make_shared(out_dims, input->Dtype(), input->GetDevice()); - - const float *in_ptr = static_cast(input->DataPtr()); - float *out_ptr = static_cast(output->DataPtr()); - - // compute strides of in_dims and out_dims - std::vector in_strides(in_dims.size(), 1); - std::vector out_strides(out_dims.size(), 1); - for (int i = in_dims.size() - 2; i >= 0; --i) { - in_strides[i] = in_strides[i + 1] * in_dims[i + 1]; - out_strides[i] = out_strides[i + 1] * out_dims[i + 1]; - } - - for (int64_t idx = 0; idx < output->NumElements(); ++idx) { - // multi-dimensional indices from flat index of input - int64_t temp = idx; - std::vector in_index(in_dims.size()); - for (int i = 0; i < in_dims.size(); ++i) { - in_index[i] = temp / in_strides[i]; - temp %= in_strides[i]; - } - - // swap indices at dim0 and dim1 - std::swap(in_index[dim0], in_index[dim1]); - - // flat index of output - int64_t out_idx = 0; - for (int i = 0; i < out_dims.size(); ++i) { out_idx += in_index[i] * out_strides[i]; } - - out_ptr[out_idx] = in_ptr[idx]; - } - - return output; -} - -std::shared_ptr TransposeBackward(const std::shared_ptr &grad_output, int64_t dim0, int64_t dim1) { - return TransposeForward(grad_output, dim1, dim0); -} - -std::shared_ptr MaskForward(const std::shared_ptr &input, const std::shared_ptr &mask, - float value) { - CHECK_EQ(input->NumElements() % mask->NumElements(), 0); - CHECK_EQ(static_cast(input->Dtype()), static_cast(mask->Dtype())); - auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); - - const float *in_ptr = static_cast(input->DataPtr()); - - for (int i = 0; i < input->NumElements(); ++i) { - if ((std::abs(static_cast(mask->DataPtr())[i % mask->NumElements()] - 1.0f) < 1e-5)) { - static_cast(output->DataPtr())[i] = value; - } else { - static_cast(output->DataPtr())[i] = in_ptr[i]; - } - } - return output; -} - -std::shared_ptr MaskBackward(const std::shared_ptr &grad_output, const std::shared_ptr &mask) { - CHECK_EQ(grad_output->NumElements() % mask->NumElements(), 0); - CHECK_EQ(static_cast(grad_output->Dtype()), static_cast(mask->Dtype())); - auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); - - for (int i = 0; i < grad_output->NumElements(); ++i) { - if ((std::abs(static_cast(mask->DataPtr())[i % mask->NumElements()] - 1.0f) < 1e-5)) { - static_cast(grad_input->DataPtr())[i] = 0.0; - } else { - static_cast(grad_input->DataPtr())[i] = static_cast(grad_output->DataPtr())[i]; - } - } - return grad_input; -} - -std::shared_ptr RepeatInterleaveForward(const std::shared_ptr &input, int64_t repeat, int64_t dim) { - CHECK_GT(repeat, 0); - CHECK_GE(dim, 0); - CHECK_LT(dim, input->Dims().size()); - - const auto &input_dims = input->Dims(); - const int64_t outer = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); - const int64_t inner - = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); - const int64_t dim_size = input_dims[dim]; - - std::vector output_dims = input_dims; - output_dims[dim] = dim_size * repeat; - auto output = std::make_shared(output_dims, input->Dtype(), input->GetDevice()); - - const float *input_ptr = static_cast(input->DataPtr()); - float *output_ptr = static_cast(output->DataPtr()); - - for (int64_t o = 0; o < outer; ++o) { - for (int64_t i = 0; i < dim_size; ++i) { - for (int r = 0; r < repeat; ++r) { - std::memcpy(output_ptr + ((o * dim_size * repeat + i * repeat + r) * inner), - input_ptr + ((o * dim_size + i) * inner), sizeof(float) * inner); - } - } - } - - return output; -} - -std::shared_ptr RepeatInterleaveBackward(const std::shared_ptr &grad_output, - const std::vector &input_dims, int64_t dim) { - CHECK_GE(dim, 0); - CHECK_LT(dim, input_dims.size()); - - const int64_t outer = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); - const int64_t inner - = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); - const int64_t dim_size = input_dims[dim]; - - int repeat = grad_output->Dims()[dim] / dim_size; - CHECK_EQ(grad_output->Dims()[dim], dim_size * repeat); - - auto grad_input = std::make_shared(input_dims, grad_output->Dtype(), grad_output->GetDevice()); - grad_input->Fill(0.0f); - - const float *grad_out_ptr = static_cast(grad_output->DataPtr()); - float *grad_in_ptr = static_cast(grad_input->DataPtr()); - - for (int64_t o = 0; o < outer; ++o) { - for (int64_t i = 0; i < dim_size; ++i) { - for (int64_t j = 0; j < inner; ++j) { - float sum = 0.0f; - for (int r = 0; r < repeat; ++r) { - sum += grad_out_ptr[((o * dim_size * repeat + i * repeat + r) * inner) + j]; - } - grad_in_ptr[(o * dim_size + i) * inner + j] = sum; - } - } - } - - return grad_input; -} - -} // namespace infini_train::kernels::cpu - -#define REGISTER_CPU_TRANSFORM_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) - -REGISTER_CPU_TRANSFORM_KERNEL(TrilForward) -REGISTER_CPU_TRANSFORM_KERNEL(TrilBackward) -REGISTER_CPU_TRANSFORM_KERNEL(TriuForward) -REGISTER_CPU_TRANSFORM_KERNEL(TriuBackward) -REGISTER_CPU_TRANSFORM_KERNEL(TransposeForward) -REGISTER_CPU_TRANSFORM_KERNEL(TransposeBackward) -REGISTER_CPU_TRANSFORM_KERNEL(MaskForward) -REGISTER_CPU_TRANSFORM_KERNEL(MaskBackward) -REGISTER_CPU_TRANSFORM_KERNEL(RepeatInterleaveForward) -REGISTER_CPU_TRANSFORM_KERNEL(RepeatInterleaveBackward) - -#undef REGISTER_CPU_TRANSFORM_KERNEL +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { + +std::shared_ptr TrilForward(const std::shared_ptr &input, int64_t diagonal) { + CHECK_EQ(input->Dims().size(), 2); + + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + for (int i = 0; i < input->NumElements(); ++i) { + int64_t row = i / input->Dims()[1]; + int64_t col = i % input->Dims()[1]; + if (row - col + diagonal >= 0) { + static_cast(output->DataPtr())[i] = static_cast(input->DataPtr())[i]; + } else { + static_cast(output->DataPtr())[i] = 0.0; + } + } + return output; +} + +std::shared_ptr TrilBackward(const std::shared_ptr &grad_output, int64_t diagonal) { + auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); + for (int i = 0; i < grad_output->NumElements(); ++i) { + int64_t row = i / grad_output->Dims()[1]; + int64_t col = i % grad_output->Dims()[1]; + if (row - col + diagonal >= 0) { + static_cast(grad_input->DataPtr())[i] = static_cast(grad_output->DataPtr())[i]; + } else { + static_cast(grad_input->DataPtr())[i] = 0.0; + } + } + return grad_input; +} + +std::shared_ptr TriuForward(const std::shared_ptr &input, int64_t diagonal) { + CHECK_EQ(input->Dims().size(), 2); + + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + for (int i = 0; i < input->NumElements(); ++i) { + int64_t row = i / input->Dims()[1]; + int64_t col = i % input->Dims()[1]; + if (row - col + diagonal <= 0) { + static_cast(output->DataPtr())[i] = static_cast(input->DataPtr())[i]; + } else { + static_cast(output->DataPtr())[i] = 0.0f; + } + } + return output; +} + +std::shared_ptr TriuBackward(const std::shared_ptr &grad_output, int64_t diagonal) { + CHECK_EQ(grad_output->Dims().size(), 2); + + auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); + for (int i = 0; i < grad_output->NumElements(); ++i) { + int64_t row = i / grad_output->Dims()[1]; + int64_t col = i % grad_output->Dims()[1]; + if (row - col + diagonal <= 0) { + static_cast(grad_input->DataPtr())[i] = static_cast(grad_output->DataPtr())[i]; + } else { + static_cast(grad_input->DataPtr())[i] = 0.0f; + } + } + return grad_input; +} + +std::shared_ptr TransposeForward(const std::shared_ptr &input, int64_t dim0, int64_t dim1) { + dim0 = dim0 < 0 ? dim0 + input->Dims().size() : dim0; + dim1 = dim1 < 0 ? dim1 + input->Dims().size() : dim1; + CHECK(dim0 >= 0 && dim0 < input->Dims().size() && dim1 >= 0 && dim1 < input->Dims().size()); + + auto in_dims = input->Dims(); + std::vector out_dims = in_dims; + std::swap(out_dims[dim0], out_dims[dim1]); + + auto output = std::make_shared(out_dims, input->Dtype(), input->GetDevice()); + + const float *in_ptr = static_cast(input->DataPtr()); + float *out_ptr = static_cast(output->DataPtr()); + + // compute strides of in_dims and out_dims + std::vector in_strides(in_dims.size(), 1); + std::vector out_strides(out_dims.size(), 1); + for (int i = in_dims.size() - 2; i >= 0; --i) { + in_strides[i] = in_strides[i + 1] * in_dims[i + 1]; + out_strides[i] = out_strides[i + 1] * out_dims[i + 1]; + } + + for (int64_t idx = 0; idx < output->NumElements(); ++idx) { + // multi-dimensional indices from flat index of input + int64_t temp = idx; + std::vector in_index(in_dims.size()); + for (int i = 0; i < in_dims.size(); ++i) { + in_index[i] = temp / in_strides[i]; + temp %= in_strides[i]; + } + + // swap indices at dim0 and dim1 + std::swap(in_index[dim0], in_index[dim1]); + + // flat index of output + int64_t out_idx = 0; + for (int i = 0; i < out_dims.size(); ++i) { out_idx += in_index[i] * out_strides[i]; } + + out_ptr[out_idx] = in_ptr[idx]; + } + + return output; +} + +std::shared_ptr TransposeBackward(const std::shared_ptr &grad_output, int64_t dim0, int64_t dim1) { + return TransposeForward(grad_output, dim1, dim0); +} + +std::shared_ptr MaskForward(const std::shared_ptr &input, const std::shared_ptr &mask, + float value) { + CHECK_EQ(input->NumElements() % mask->NumElements(), 0); + CHECK_EQ(static_cast(input->Dtype()), static_cast(mask->Dtype())); + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + + const float *in_ptr = static_cast(input->DataPtr()); + + for (int i = 0; i < input->NumElements(); ++i) { + if ((std::abs(static_cast(mask->DataPtr())[i % mask->NumElements()] - 1.0f) < 1e-5)) { + static_cast(output->DataPtr())[i] = value; + } else { + static_cast(output->DataPtr())[i] = in_ptr[i]; + } + } + return output; +} + +std::shared_ptr MaskBackward(const std::shared_ptr &grad_output, const std::shared_ptr &mask) { + CHECK_EQ(grad_output->NumElements() % mask->NumElements(), 0); + CHECK_EQ(static_cast(grad_output->Dtype()), static_cast(mask->Dtype())); + auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); + + for (int i = 0; i < grad_output->NumElements(); ++i) { + if ((std::abs(static_cast(mask->DataPtr())[i % mask->NumElements()] - 1.0f) < 1e-5)) { + static_cast(grad_input->DataPtr())[i] = 0.0; + } else { + static_cast(grad_input->DataPtr())[i] = static_cast(grad_output->DataPtr())[i]; + } + } + return grad_input; +} + +std::shared_ptr RepeatInterleaveForward(const std::shared_ptr &input, int64_t repeat, int64_t dim) { + CHECK_GT(repeat, 0); + CHECK_GE(dim, 0); + CHECK_LT(dim, input->Dims().size()); + + const auto &input_dims = input->Dims(); + const int64_t outer = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + const int64_t inner + = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); + const int64_t dim_size = input_dims[dim]; + + std::vector output_dims = input_dims; + output_dims[dim] = dim_size * repeat; + auto output = std::make_shared(output_dims, input->Dtype(), input->GetDevice()); + + const float *input_ptr = static_cast(input->DataPtr()); + float *output_ptr = static_cast(output->DataPtr()); + + for (int64_t o = 0; o < outer; ++o) { + for (int64_t i = 0; i < dim_size; ++i) { + for (int r = 0; r < repeat; ++r) { + std::memcpy(output_ptr + ((o * dim_size * repeat + i * repeat + r) * inner), + input_ptr + ((o * dim_size + i) * inner), sizeof(float) * inner); + } + } + } + + return output; +} + +std::shared_ptr RepeatInterleaveBackward(const std::shared_ptr &grad_output, + const std::vector &input_dims, int64_t dim) { + CHECK_GE(dim, 0); + CHECK_LT(dim, input_dims.size()); + + const int64_t outer = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + const int64_t inner + = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); + const int64_t dim_size = input_dims[dim]; + + int repeat = grad_output->Dims()[dim] / dim_size; + CHECK_EQ(grad_output->Dims()[dim], dim_size * repeat); + + auto grad_input = std::make_shared(input_dims, grad_output->Dtype(), grad_output->GetDevice()); + grad_input->Fill(0.0f); + + const float *grad_out_ptr = static_cast(grad_output->DataPtr()); + float *grad_in_ptr = static_cast(grad_input->DataPtr()); + + for (int64_t o = 0; o < outer; ++o) { + for (int64_t i = 0; i < dim_size; ++i) { + for (int64_t j = 0; j < inner; ++j) { + float sum = 0.0f; + for (int r = 0; r < repeat; ++r) { + sum += grad_out_ptr[((o * dim_size * repeat + i * repeat + r) * inner) + j]; + } + grad_in_ptr[(o * dim_size + i) * inner + j] = sum; + } + } + } + + return grad_input; +} + +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_TRANSFORM_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_TRANSFORM_KERNEL(TrilForward) +REGISTER_CPU_TRANSFORM_KERNEL(TrilBackward) +REGISTER_CPU_TRANSFORM_KERNEL(TriuForward) +REGISTER_CPU_TRANSFORM_KERNEL(TriuBackward) +REGISTER_CPU_TRANSFORM_KERNEL(TransposeForward) +REGISTER_CPU_TRANSFORM_KERNEL(TransposeBackward) +REGISTER_CPU_TRANSFORM_KERNEL(MaskForward) +REGISTER_CPU_TRANSFORM_KERNEL(MaskBackward) +REGISTER_CPU_TRANSFORM_KERNEL(RepeatInterleaveForward) +REGISTER_CPU_TRANSFORM_KERNEL(RepeatInterleaveBackward) + +#undef REGISTER_CPU_TRANSFORM_KERNEL diff --git a/infini_train/src/kernels/cuda/accumulate_grad.cu b/infini_train/src/kernels/cuda/accumulate_grad.cu index 5f977c3..bb6093b 100644 --- a/infini_train/src/kernels/cuda/accumulate_grad.cu +++ b/infini_train/src/kernels/cuda/accumulate_grad.cu @@ -1,41 +1,70 @@ -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { - -__global__ void AccumulateGradKernel(const float *grad_ptr, float rate, float *tensor_ptr, size_t num_elements) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx < num_elements) { - tensor_ptr[idx] += rate * grad_ptr[idx]; - } -} - -void AccumulateGrad(const std::shared_ptr &gradient, float rate, const std::shared_ptr &tensor) { - size_t num_elements = gradient->NumElements(); - - const float *grad_ptr = static_cast(gradient->DataPtr()); - float *tensor_ptr = static_cast(tensor->DataPtr()); - - int threads_per_block = 256; - int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; - - AccumulateGradKernel<<>>(grad_ptr, rate, tensor_ptr, num_elements); -} - -void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_ptr ¶m, - const std::shared_ptr &m, const std::shared_ptr &v, float learning_rate, - float beta1, float beta2, float eps, int64_t t) { - // =================================== 作业 =================================== - // TODO：实现Adam优化器的梯度累积和参数更新 - // REF: - // =================================== 作业 =================================== -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AccumulateGrad) -REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AdamAccumulateGrad) - -#undef REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { + +__global__ void AccumulateGradKernel(const float *grad_ptr, float rate, float *tensor_ptr, size_t num_elements) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + tensor_ptr[idx] += rate * grad_ptr[idx]; + } +} + +void AccumulateGrad(const std::shared_ptr &gradient, float rate, const std::shared_ptr &tensor) { + size_t num_elements = gradient->NumElements(); + + const float *grad_ptr = static_cast(gradient->DataPtr()); + float *tensor_ptr = static_cast(tensor->DataPtr()); + + int threads_per_block = 256; + int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; + + AccumulateGradKernel<<>>(grad_ptr, rate, tensor_ptr, num_elements); +} + +__global__ void AdamAccumulateGradKernel(const float *grad_ptr, float *param_ptr, float *m_ptr, float *v_ptr, + float learning_rate, float beta1, float beta2, float eps, float alpha, + size_t num_elements) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + float g = grad_ptr[idx]; + m_ptr[idx] = beta1 * m_ptr[idx] + (1.0f - beta1) * g; + v_ptr[idx] = beta2 * v_ptr[idx] + (1.0f - beta2) * g * g; + param_ptr[idx] -= alpha * m_ptr[idx] / (sqrtf(v_ptr[idx]) + eps); + } +} + +void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_ptr ¶m, + const std::shared_ptr &m, const std::shared_ptr &v, float learning_rate, + float beta1, float beta2, float eps, int64_t t) { + // =================================== 作业 =================================== + // TODO：实现Adam优化器的梯度累积和参数更新 + // REF: + // =================================== 作业 =================================== + + size_t num_elements = grad->NumElements(); + + const float *grad_ptr = static_cast(grad->DataPtr()); + float *param_ptr = static_cast(param->DataPtr()); + float *m_ptr = static_cast(m->DataPtr()); + float *v_ptr = static_cast(v->DataPtr()); + + float beta1_t = powf(beta1, static_cast(t)); + float beta2_t = powf(beta2, static_cast(t)); + float alpha = learning_rate * sqrtf(1.0f - beta2_t) / (1.0f - beta1_t); + + int threads_per_block = 256; + int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; + + AdamAccumulateGradKernel<<>>(grad_ptr, param_ptr, m_ptr, v_ptr, learning_rate, beta1, + beta2, eps, alpha, num_elements); +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AccumulateGrad) +REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AdamAccumulateGrad) + +#undef REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL diff --git a/infini_train/src/kernels/cuda/cross_entropy.cu b/infini_train/src/kernels/cuda/cross_entropy.cu index 7938329..8288a8c 100644 --- a/infini_train/src/kernels/cuda/cross_entropy.cu +++ b/infini_train/src/kernels/cuda/cross_entropy.cu @@ -1,209 +1,209 @@ -#include -#include -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { -namespace { -constexpr float kNegativeInfinity = -std::numeric_limits::infinity(); -} - -#if defined(CUB_VERSION) && CUB_VERSION >= 200800 - #include - using CubSumOp = ::cuda::std::plus<>; - using CubMaxOp = ::cuda::maximum<>; - using CubMinOp = ::cuda::minimum<>; -#else - using CubSumOp = cub::Sum; - using CubMaxOp = cub::Max; - using CubMinOp = cub::Min; -#endif - -template -__global__ void CrossEntropyForwardKernel(const float *__restrict__ input_ptr, - const TargetType *__restrict__ target_ptr, float *__restrict__ loss_ptr, - int bs, int num_classes) { - __shared__ struct { - float max_logit; - float sum_exp; - TargetType target_class; - typename cub::BlockReduce::TempStorage reduce; - } shared; - - const int sample_idx = blockIdx.x; - if (sample_idx >= bs) { - return; - } - - const int tid = threadIdx.x; - const size_t base = sample_idx * num_classes; - - if (tid == 0) { - shared.target_class = target_ptr[sample_idx]; - } - __syncthreads(); - - // calculate the max - float thread_max = kNegativeInfinity; - for (int i = tid; i < num_classes; i += BLOCK_SIZE) { thread_max = fmaxf(thread_max, input_ptr[base + i]); } - shared.max_logit = cub::BlockReduce(shared.reduce).Reduce(thread_max, CubMaxOp()); - __syncthreads(); - - // calculate the sum of exponents - float thread_sum = 0.0f; - for (int i = tid; i < num_classes; i += BLOCK_SIZE) { thread_sum += expf(input_ptr[base + i] - shared.max_logit); } - shared.sum_exp = cub::BlockReduce(shared.reduce).Sum(thread_sum); - __syncthreads(); - - // calculate the loss - if (tid == 0) { - const float target_val = input_ptr[base + shared.target_class] - shared.max_logit; - loss_ptr[sample_idx] = logf(shared.sum_exp) - target_val; - } -} - -std::shared_ptr CrossEntropyForward(const std::shared_ptr &input, - const std::shared_ptr &target) { - const auto &input_dims = input->Dims(); - CHECK_GE(input_dims.size(), 2); - const int bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); - const int num_classes = *input_dims.rbegin(); - - auto batched_output = std::make_shared(std::vector{bs}, DataType::kFLOAT32, input->GetDevice()); - const float *input_ptr = static_cast(input->DataPtr()); - float *batched_loss_ptr = static_cast(batched_output->DataPtr()); - - constexpr int threads_per_block = 256; - int num_blocks = bs; - - switch (target->Dtype()) { - case DataType::kUINT8: { - const uint8_t *target_ptr = static_cast(target->DataPtr()); - CrossEntropyForwardKernel - <<>>(input_ptr, target_ptr, batched_loss_ptr, bs, num_classes); - break; - } - case DataType::kINT64: { - const int64_t *target_ptr = static_cast(target->DataPtr()); - CrossEntropyForwardKernel - <<>>(input_ptr, target_ptr, batched_loss_ptr, bs, num_classes); - break; - } - default: - LOG(FATAL) << "Unsupported target data type: " << static_cast(target->Dtype()); - } - cudaDeviceSynchronize(); - - auto loss_cpu = batched_output->To(Device()); - auto loss = std::make_shared(std::vector{}, DataType::kFLOAT32, Device()); - static_cast(loss->DataPtr())[0] - = std::accumulate(static_cast(loss_cpu.DataPtr()), - static_cast(loss_cpu.DataPtr()) + bs, 0.0f) - / bs; - - return {std::make_shared(loss->To(input->GetDevice()))}; -} - -template -__global__ void CrossEntropyBackwardKernel(const float *__restrict__ input_ptr, float *__restrict__ input_grad_ptr, - const TargetType *__restrict__ target_ptr, int bs, int num_classes) { - __shared__ struct { - float max_logit; - float sum_exp; - int target_class; - typename cub::BlockReduce::TempStorage reduce; - } shared; - - const int tid = threadIdx.x; - const int idx = blockIdx.x; - - if (idx >= bs) { - return; - } - - const size_t idx_base = idx * num_classes; - - if (tid == 0) { - shared.target_class = static_cast(target_ptr[idx]); - } - __syncthreads(); - - // calculate the max - float thread_max = kNegativeInfinity; - for (int i = tid; i < num_classes; i += BLOCK_SIZE) { thread_max = fmaxf(thread_max, input_ptr[idx_base + i]); } - shared.max_logit = cub::BlockReduce(shared.reduce).Reduce(thread_max, CubMaxOp()); - __syncthreads(); - - // calculate the sum - float thread_sum = 0.0f; - for (int i = tid; i < num_classes; i += BLOCK_SIZE) { - thread_sum += expf(input_ptr[idx_base + i] - shared.max_logit); - } - shared.sum_exp = cub::BlockReduce(shared.reduce).Sum(thread_sum); - __syncthreads(); - - // calculate the gradient - const float inv_bs = 1.0f / bs; - const float scale = 1.0f / shared.sum_exp; - const int target = shared.target_class; - - for (int i = tid; i < num_classes; i += BLOCK_SIZE) { - const int global_idx = idx_base + i; - const float exp_val = expf(input_ptr[global_idx] - shared.max_logit); - input_grad_ptr[global_idx] = (exp_val * scale - (i == target)) * inv_bs; - } -} - -std::shared_ptr CrossEntropyBackward(const std::shared_ptr &input, - const std::shared_ptr &target, - const std::shared_ptr &grad_output) { - const auto &input_dims = input->Dims(); - CHECK_GE(input_dims.size(), 2); - const int bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); - const int num_classes = *input_dims.rbegin(); - - CHECK_EQ(grad_output->Dims().size(), 0); - auto grad_input = std::make_shared(input->Dims(), DataType::kFLOAT32, grad_output->GetDevice()); - grad_input->Fill(0.0f); - const float *input_ptr = static_cast(input->DataPtr()); - float *input_grad_ptr = static_cast(grad_input->DataPtr()); - - constexpr int threads_per_block = 256; - int num_blocks = bs; - - switch (target->Dtype()) { - case DataType::kUINT8: { - const uint8_t *target_ptr = static_cast(target->DataPtr()); - CrossEntropyBackwardKernel - <<>>(input_ptr, input_grad_ptr, target_ptr, bs, num_classes); - break; - } - case DataType::kINT64: { - const int64_t *target_ptr = static_cast(target->DataPtr()); - CrossEntropyBackwardKernel - <<>>(input_ptr, input_grad_ptr, target_ptr, bs, num_classes); - break; - } - default: - LOG(FATAL) << "Unsupported target data type: " << static_cast(target->Dtype()); - } - - return {grad_input}; -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_CROSS_ENTROPY_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_CROSS_ENTROPY_KERNEL(CrossEntropyForward) -REGISTER_CUDA_CROSS_ENTROPY_KERNEL(CrossEntropyBackward) - -#undef REGISTER_CUDA_CROSS_ENTROPY_KERNEL +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { +namespace { +constexpr float kNegativeInfinity = -std::numeric_limits::infinity(); +} + +#if defined(CUB_VERSION) && CUB_VERSION >= 200800 + #include + using CubSumOp = ::cuda::std::plus<>; + using CubMaxOp = ::cuda::maximum<>; + using CubMinOp = ::cuda::minimum<>; +#else + using CubSumOp = cub::Sum; + using CubMaxOp = cub::Max; + using CubMinOp = cub::Min; +#endif + +template +__global__ void CrossEntropyForwardKernel(const float *__restrict__ input_ptr, + const TargetType *__restrict__ target_ptr, float *__restrict__ loss_ptr, + int bs, int num_classes) { + __shared__ struct { + float max_logit; + float sum_exp; + TargetType target_class; + typename cub::BlockReduce::TempStorage reduce; + } shared; + + const int sample_idx = blockIdx.x; + if (sample_idx >= bs) { + return; + } + + const int tid = threadIdx.x; + const size_t base = sample_idx * num_classes; + + if (tid == 0) { + shared.target_class = target_ptr[sample_idx]; + } + __syncthreads(); + + // calculate the max + float thread_max = kNegativeInfinity; + for (int i = tid; i < num_classes; i += BLOCK_SIZE) { thread_max = fmaxf(thread_max, input_ptr[base + i]); } + shared.max_logit = cub::BlockReduce(shared.reduce).Reduce(thread_max, CubMaxOp()); + __syncthreads(); + + // calculate the sum of exponents + float thread_sum = 0.0f; + for (int i = tid; i < num_classes; i += BLOCK_SIZE) { thread_sum += expf(input_ptr[base + i] - shared.max_logit); } + shared.sum_exp = cub::BlockReduce(shared.reduce).Sum(thread_sum); + __syncthreads(); + + // calculate the loss + if (tid == 0) { + const float target_val = input_ptr[base + shared.target_class] - shared.max_logit; + loss_ptr[sample_idx] = logf(shared.sum_exp) - target_val; + } +} + +std::shared_ptr CrossEntropyForward(const std::shared_ptr &input, + const std::shared_ptr &target) { + const auto &input_dims = input->Dims(); + CHECK_GE(input_dims.size(), 2); + const int bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int num_classes = *input_dims.rbegin(); + + auto batched_output = std::make_shared(std::vector{bs}, DataType::kFLOAT32, input->GetDevice()); + const float *input_ptr = static_cast(input->DataPtr()); + float *batched_loss_ptr = static_cast(batched_output->DataPtr()); + + constexpr int threads_per_block = 256; + int num_blocks = bs; + + switch (target->Dtype()) { + case DataType::kUINT8: { + const uint8_t *target_ptr = static_cast(target->DataPtr()); + CrossEntropyForwardKernel + <<>>(input_ptr, target_ptr, batched_loss_ptr, bs, num_classes); + break; + } + case DataType::kINT64: { + const int64_t *target_ptr = static_cast(target->DataPtr()); + CrossEntropyForwardKernel + <<>>(input_ptr, target_ptr, batched_loss_ptr, bs, num_classes); + break; + } + default: + LOG(FATAL) << "Unsupported target data type: " << static_cast(target->Dtype()); + } + cudaDeviceSynchronize(); + + auto loss_cpu = batched_output->To(Device()); + auto loss = std::make_shared(std::vector{}, DataType::kFLOAT32, Device()); + static_cast(loss->DataPtr())[0] + = std::accumulate(static_cast(loss_cpu.DataPtr()), + static_cast(loss_cpu.DataPtr()) + bs, 0.0f) + / bs; + + return {std::make_shared(loss->To(input->GetDevice()))}; +} + +template +__global__ void CrossEntropyBackwardKernel(const float *__restrict__ input_ptr, float *__restrict__ input_grad_ptr, + const TargetType *__restrict__ target_ptr, int bs, int num_classes) { + __shared__ struct { + float max_logit; + float sum_exp; + int target_class; + typename cub::BlockReduce::TempStorage reduce; + } shared; + + const int tid = threadIdx.x; + const int idx = blockIdx.x; + + if (idx >= bs) { + return; + } + + const size_t idx_base = idx * num_classes; + + if (tid == 0) { + shared.target_class = static_cast(target_ptr[idx]); + } + __syncthreads(); + + // calculate the max + float thread_max = kNegativeInfinity; + for (int i = tid; i < num_classes; i += BLOCK_SIZE) { thread_max = fmaxf(thread_max, input_ptr[idx_base + i]); } + shared.max_logit = cub::BlockReduce(shared.reduce).Reduce(thread_max, CubMaxOp()); + __syncthreads(); + + // calculate the sum + float thread_sum = 0.0f; + for (int i = tid; i < num_classes; i += BLOCK_SIZE) { + thread_sum += expf(input_ptr[idx_base + i] - shared.max_logit); + } + shared.sum_exp = cub::BlockReduce(shared.reduce).Sum(thread_sum); + __syncthreads(); + + // calculate the gradient + const float inv_bs = 1.0f / bs; + const float scale = 1.0f / shared.sum_exp; + const int target = shared.target_class; + + for (int i = tid; i < num_classes; i += BLOCK_SIZE) { + const int global_idx = idx_base + i; + const float exp_val = expf(input_ptr[global_idx] - shared.max_logit); + input_grad_ptr[global_idx] = (exp_val * scale - (i == target)) * inv_bs; + } +} + +std::shared_ptr CrossEntropyBackward(const std::shared_ptr &input, + const std::shared_ptr &target, + const std::shared_ptr &grad_output) { + const auto &input_dims = input->Dims(); + CHECK_GE(input_dims.size(), 2); + const int bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int num_classes = *input_dims.rbegin(); + + CHECK_EQ(grad_output->Dims().size(), 0); + auto grad_input = std::make_shared(input->Dims(), DataType::kFLOAT32, grad_output->GetDevice()); + grad_input->Fill(0.0f); + const float *input_ptr = static_cast(input->DataPtr()); + float *input_grad_ptr = static_cast(grad_input->DataPtr()); + + constexpr int threads_per_block = 256; + int num_blocks = bs; + + switch (target->Dtype()) { + case DataType::kUINT8: { + const uint8_t *target_ptr = static_cast(target->DataPtr()); + CrossEntropyBackwardKernel + <<>>(input_ptr, input_grad_ptr, target_ptr, bs, num_classes); + break; + } + case DataType::kINT64: { + const int64_t *target_ptr = static_cast(target->DataPtr()); + CrossEntropyBackwardKernel + <<>>(input_ptr, input_grad_ptr, target_ptr, bs, num_classes); + break; + } + default: + LOG(FATAL) << "Unsupported target data type: " << static_cast(target->Dtype()); + } + + return {grad_input}; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_CROSS_ENTROPY_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_CROSS_ENTROPY_KERNEL(CrossEntropyForward) +REGISTER_CUDA_CROSS_ENTROPY_KERNEL(CrossEntropyBackward) + +#undef REGISTER_CUDA_CROSS_ENTROPY_KERNEL diff --git a/infini_train/src/kernels/cuda/elementwise.cu b/infini_train/src/kernels/cuda/elementwise.cu index e1fa344..d52781c 100644 --- a/infini_train/src/kernels/cuda/elementwise.cu +++ b/infini_train/src/kernels/cuda/elementwise.cu @@ -1,500 +1,500 @@ -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { - -#define CEIL_DIV(x, y) (((x) + (y)-1) / (y)) - -namespace { - -template -__global__ void UnaryForwardKernel(T *output, Func fn, size_t num_elements, size_t offset, const T *input) { - size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; - - if (idx < num_elements) { - output[idx] = fn(input[idx]); - } -} - -// Helper for broadcast indexing -__device__ inline int64_t CalcOffset(int64_t idx, int ndim, const int64_t *strides, const int64_t *shape, - const int64_t *out_strides) { - int64_t offset = 0; - for (int i = 0; i < ndim; ++i) { - int64_t out_index = (idx / out_strides[i]) % shape[i]; - int64_t index = shape[i] == 1 ? 0 : out_index; - offset += index * strides[i]; - } - return offset; -} - -template -__global__ void BinaryForwardKernel(T *output, Func fn, int ndim, const int64_t *a_strides, const int64_t *a_shape, - const int64_t *b_strides, const int64_t *b_shape, const int64_t *out_strides, - const T *a, const T *b, size_t num_elements) { - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= num_elements) { - return; - } - - int64_t a_offset = CalcOffset(idx, ndim, a_strides, a_shape, out_strides); - int64_t b_offset = CalcOffset(idx, ndim, b_strides, b_shape, out_strides); - - output[idx] = fn(a[a_offset], b[b_offset]); -} - -// launch the given kernel function with the given output and inputs -template -void LaunchKernel(Kernel &&kernel, const std::shared_ptr &output, const Inputs &...inputs) { - auto extract_ptrs - = [](const auto &...ts) { return std::make_tuple(static_cast(ts ? ts->DataPtr() : nullptr)...); }; - auto input_ptrs = extract_ptrs(inputs...); - - const size_t num_elements = output->NumElements(); - dim3 block_dims(std::min(BLOCK_SIZE, static_cast(1024))); - dim3 grid_dims(CEIL_DIV(num_elements, block_dims.x)); - const size_t step = grid_dims.x * block_dims.x; - - for (size_t offset = 0; offset < num_elements; offset += step) { - std::apply([&](auto... ptrs) { kernel(grid_dims, block_dims, offset, ptrs...); }, input_ptrs); - } -} - -// Helper for stride calculation -std::vector ComputeStride(const std::vector &dims) { - std::vector strides(dims.size(), 1); - for (int i = dims.size() - 2; i >= 0; --i) { strides[i] = strides[i + 1] * dims[i + 1]; } - return strides; -} - -// launch a forward elementwise operation given the calculation function, output, and the inputs -// Note: currently only support unary and binary operations -template -void LaunchForward(Func func, const std::shared_ptr &output, const Inputs &...inputs) { - T *output_ptr = static_cast(output->DataPtr()); - - if constexpr (sizeof...(inputs) == 1) { - // Unary case - LaunchKernel( - [&](dim3 grid, dim3 block, size_t offset, auto... ptrs) { - UnaryForwardKernel<<>>(output_ptr, func, output->NumElements(), offset, ptrs...); - }, - output, inputs...); - } else if constexpr (sizeof...(inputs) == 2) { - // Binary case - auto input_tuple = std::make_tuple(inputs...); - const auto &input_a = std::get<0>(input_tuple); - const auto &input_b = std::get<1>(input_tuple); - - const auto &a_dims = input_a->Dims(); - const auto &b_dims = input_b->Dims(); - const auto &out_dims = output->Dims(); - int ndim = out_dims.size(); - - std::vector a_shape(ndim, 1), b_shape(ndim, 1), out_shape(ndim, 1); - std::copy_backward(a_dims.begin(), a_dims.end(), a_shape.end()); - std::copy_backward(b_dims.begin(), b_dims.end(), b_shape.end()); - std::copy_backward(out_dims.begin(), out_dims.end(), out_shape.end()); - - auto a_stride_host = ComputeStride(a_shape); - auto b_stride_host = ComputeStride(b_shape); - auto out_stride_host = ComputeStride(out_shape); - - int64_t *device_buffer; - cudaMallocAsync(&device_buffer, 5 * ndim * sizeof(int64_t), 0); - - int64_t *device_a_strides, *device_b_strides, *device_out_strides, *device_a_shape, *device_b_shape; - device_a_strides = device_buffer + ndim * 0; - device_b_strides = device_buffer + ndim * 1; - device_out_strides = device_buffer + ndim * 2; - device_a_shape = device_buffer + ndim * 3; - device_b_shape = device_buffer + ndim * 4; - - std::vector host_buffer; - host_buffer.insert(host_buffer.end(), a_stride_host.begin(), a_stride_host.end()); - host_buffer.insert(host_buffer.end(), b_stride_host.begin(), b_stride_host.end()); - host_buffer.insert(host_buffer.end(), out_stride_host.begin(), out_stride_host.end()); - host_buffer.insert(host_buffer.end(), a_shape.begin(), a_shape.end()); - host_buffer.insert(host_buffer.end(), b_shape.begin(), b_shape.end()); - - cudaMemcpyAsync(device_buffer, host_buffer.data(), 5 * ndim * sizeof(int64_t), cudaMemcpyHostToDevice, 0); - - LaunchKernel( - [&](dim3 grid, dim3 block, size_t offset, const T *a_ptr, const T *b_ptr) { - BinaryForwardKernel<<>>(output_ptr, func, ndim, device_a_strides, device_a_shape, - device_b_strides, device_b_shape, device_out_strides, a_ptr, b_ptr, - output->NumElements()); - }, - output, inputs...); - - cudaFreeAsync(device_buffer, 0); - } else { - static_assert(sizeof...(inputs) == 1 || sizeof...(inputs) == 2, - "LaunchForward currently only supports unary and binary operations."); - } -} - -// Backward kernel for unary operators -template -__global__ void UnaryBackwardKernel(T *output, Func fn, size_t num_elements, size_t offset, const T *grad_output, - const T *input) { - size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; - - if (idx < num_elements) { - output[idx] = grad_output[idx] * fn(input ? input[idx] : T(0)); - } -} - -// Backward kernel for binary operators -template -__global__ void BinaryBackwardKernel(T *output_a, T *output_b, FuncA fn_a, FuncB fn_b, int ndim, size_t num_elements, - const int64_t *a_strides, const int64_t *a_shape, const int64_t *b_strides, - const int64_t *b_shape, const int64_t *out_strides, const T *grad_output, - const T *input_a, const T *input_b) { - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= num_elements) { - return; - } - - int64_t a_offset = CalcOffset(idx, ndim, a_strides, a_shape, out_strides); - int64_t b_offset = CalcOffset(idx, ndim, b_strides, b_shape, out_strides); - - const T a_val = input_a ? input_a[a_offset] : T(0); - const T b_val = input_b ? input_b[b_offset] : T(0); - - output_a[a_offset] = grad_output[idx] * fn_a(a_val, b_val); - atomicAdd(&output_b[b_offset], grad_output[idx] * fn_b(a_val, b_val)); -} - -// launch unary operator's backward kernel -template -void LaunchBackward(Func func, const std::shared_ptr &output, const std::shared_ptr &grad_output, - const Inputs &...inputs) { - T *output_ptr = static_cast(output->DataPtr()); - const T *grad_ptr = static_cast(grad_output->DataPtr()); - - LaunchKernel( - [=](dim3 grid, dim3 block, size_t offset, auto... ptrs) { - UnaryBackwardKernel<<>>(output_ptr, func, output->NumElements(), offset, grad_ptr, ptrs...); - }, - output, inputs...); -} - -// launch binary operator's backward kernel -template -void LaunchBackward(FuncA fun_a, FuncB fun_b, const std::shared_ptr &output_a, - const std::shared_ptr &output_b, const std::vector &a_dims, - const std::vector &b_dims, const std::shared_ptr &grad_output, - const Inputs &...inputs) { - T *output_a_ptr = static_cast(output_a->DataPtr()); - T *output_b_ptr = static_cast(output_b->DataPtr()); - const T *grad_output_ptr = static_cast(grad_output->DataPtr()); - - const auto &out_dims = grad_output->Dims(); - int ndim = out_dims.size(); - - std::vector a_shape(ndim, 1), b_shape(ndim, 1), out_shape(ndim, 1); - std::copy_backward(a_dims.begin(), a_dims.end(), a_shape.end()); - std::copy_backward(b_dims.begin(), b_dims.end(), b_shape.end()); - std::copy_backward(out_dims.begin(), out_dims.end(), out_shape.end()); - - auto a_stride_host = ComputeStride(a_shape); - auto b_stride_host = ComputeStride(b_shape); - auto out_stride_host = ComputeStride(out_shape); - - int64_t *device_buffer; - cudaMallocAsync(&device_buffer, 5 * ndim * sizeof(int64_t), 0); - - int64_t *device_a_strides, *device_b_strides, *device_out_strides, *device_a_shape, *device_b_shape; - device_a_strides = device_buffer + ndim * 0; - device_b_strides = device_buffer + ndim * 1; - device_out_strides = device_buffer + ndim * 2; - device_a_shape = device_buffer + ndim * 3; - device_b_shape = device_buffer + ndim * 4; - - std::vector host_buffer; - host_buffer.insert(host_buffer.end(), a_stride_host.begin(), a_stride_host.end()); - host_buffer.insert(host_buffer.end(), b_stride_host.begin(), b_stride_host.end()); - host_buffer.insert(host_buffer.end(), out_stride_host.begin(), out_stride_host.end()); - host_buffer.insert(host_buffer.end(), a_shape.begin(), a_shape.end()); - host_buffer.insert(host_buffer.end(), b_shape.begin(), b_shape.end()); - - cudaMemcpyAsync(device_buffer, host_buffer.data(), 5 * ndim * sizeof(int64_t), cudaMemcpyHostToDevice, 0); - - const size_t num_elements = grad_output->NumElements(); - LaunchKernel( - [=](dim3 grid, dim3 block, size_t offset, auto... ptrs) { - BinaryBackwardKernel<<>>(output_a_ptr, output_b_ptr, fun_a, fun_b, ndim, num_elements, - device_a_strides, device_a_shape, device_b_strides, device_b_shape, - device_out_strides, grad_output_ptr, ptrs...); - }, - output_a, inputs...); - - cudaFreeAsync(device_buffer, 0); -} - -template std::shared_ptr UnaryForward(const std::shared_ptr &input, Func unary_fn) { - auto dtype = input->Dtype(); - auto output = std::make_shared(input->Dims(), dtype, input->GetDevice()); - - switch (dtype) { - case DataType::kFLOAT32: - LaunchForward<256, float>(unary_fn, output, input); - break; - default: - LOG(FATAL) << "CUDA unary forward: 'Unsupported data type' at " << __FILE__ << ":" << __LINE__; - } - - return output; -} - -template -std::shared_ptr UnaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr &a, - Func unary_fn) { - auto dtype = grad_output->Dtype(); - auto output = std::make_shared(grad_output->Dims(), dtype, grad_output->GetDevice()); - output->Fill(0.0f); - switch (dtype) { - case DataType::kFLOAT32: - LaunchBackward<256, float>(unary_fn, output, grad_output, a); - break; - default: - LOG(FATAL) << "CUDA unary backward: 'Unsupported data type' at " << __FILE__ << ":" << __LINE__; - } - - return output; -} - -template -std::shared_ptr BinaryForward(const std::shared_ptr &a, const std::shared_ptr &b, - Func binary_fn) { - auto dtype = a->Dtype(); - // Currently a and b should have the same data type and only one-way broadcasting from b to a is assumed by default - CHECK(dtype == b->Dtype() && a->NumElements() >= b->NumElements() && a->NumElements() % b->NumElements() == 0); - - auto output = std::make_shared(a->Dims(), dtype, a->GetDevice()); - - switch (dtype) { - case DataType::kFLOAT32: - LaunchForward<256, float>(binary_fn, output, a, b); - break; - default: - LOG(FATAL) << "CUDA binary forward: 'Unsupported data type' at " << __FILE__ << ":" << __LINE__; - } - - return output; -} - -template -std::pair, std::shared_ptr> -BinaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr &a, - const std::shared_ptr &b, const std::vector &a_dims, const std::vector &b_dims, - FuncA fn_a, FuncB fn_b) { - const auto a_num_elements = std::accumulate(a_dims.begin(), a_dims.end(), 1, std::multiplies()); - const auto b_num_elements = std::accumulate(b_dims.begin(), b_dims.end(), 1, std::multiplies()); - - CHECK(a_num_elements >= b_num_elements && a_num_elements % b_num_elements == 0); - if (a) { - CHECK(a_num_elements == a->NumElements()); - } - if (b) { - CHECK(b_num_elements == b->NumElements()); - } - auto dtype = grad_output->Dtype(); - auto device = grad_output->GetDevice(); - - // Currently a and b should have the same data type - if (a && b) { - CHECK(a->Dtype() == b->Dtype()); - } - auto grad_a = std::make_shared(a_dims, dtype, device); - auto grad_b = std::make_shared(b_dims, dtype, device); - grad_a->Fill(0.0f); - grad_b->Fill(0.0f); - switch (dtype) { - case DataType::kFLOAT32: - LaunchBackward<256, float>(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output, a, b); - break; - default: - LOG(FATAL) << "CUDA binary backward: 'Unsupported data type' at " << __FILE__ << ":" << __LINE__; - } - - return {grad_a, grad_b}; -} -} // namespace - -std::shared_ptr NegForward(const std::shared_ptr &input) { - return UnaryForward(input, [] __device__(float x) { return -x; }); -} - -std::shared_ptr NegBackward(const std::shared_ptr &grad_output) { - return UnaryBackward(grad_output, nullptr, [] __device__(float) { return -1.0f; }); -} - -std::shared_ptr ReciprocalForward(const std::shared_ptr &input) { - return UnaryForward(input, [] __device__(float x) { return 1.0f / x; }); -} - -std::shared_ptr ReciprocalBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &input) { - return UnaryBackward(grad_output, input, [] __device__(float x) { return -1.0f / (x * x); }); -} - -std::shared_ptr SinForward(const std::shared_ptr &input) { - return UnaryForward(input, [] __device__(float x) { return sinf(x); }); -} - -std::shared_ptr SinBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { - return UnaryBackward(grad_output, input, [] __device__(float x) { return cosf(x); }); -} - -std::shared_ptr CosForward(const std::shared_ptr &input) { - return UnaryForward(input, [] __device__(float x) { return cosf(x); }); -} - -std::shared_ptr CosBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { - return UnaryBackward(grad_output, input, [] __device__(float x) { return -sinf(x); }); -} - -std::shared_ptr TanhForward(const std::shared_ptr &input) { - return UnaryForward(input, [] __device__(float x) { return tanhf(x); }); -} - -std::shared_ptr TanhBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &output) { - return UnaryBackward(grad_output, output, [] __device__(float x) { return 1.0 - x * x; }); -} - -std::shared_ptr PowForward(const std::shared_ptr &input, float scalar, bool scalar_is_base) { - if (scalar_is_base) { - return UnaryForward(input, [scalar] __device__(float x) { return powf(scalar, x); }); - } else { - return UnaryForward(input, [scalar] __device__(float x) { return powf(x, scalar); }); - } -} - -std::shared_ptr PowBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, - float scalar, bool scalar_is_base) { - if (scalar_is_base) { - return UnaryBackward(grad_output, input, - [scalar] __device__(float x) { return logf(scalar) * powf(scalar, x); }); - } else { - return UnaryBackward(grad_output, input, - [scalar] __device__(float x) { return scalar * powf(x, scalar - 1.0f); }); - } -} - -std::shared_ptr RsqrtForward(const std::shared_ptr &input) { - return UnaryForward(input, [] __device__(float x) { return 1.0f / sqrtf(x); }); -} - -std::shared_ptr RsqrtBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &input) { - return UnaryBackward(grad_output, input, [] __device__(float x) { return -0.5f / (x * sqrtf(x)); }); -} - -std::shared_ptr EqualsScalarForward(const std::shared_ptr &a, float scalar) { - return UnaryForward(a, [scalar] __device__(float x) { return x == scalar ? 1.0f : 0.0f; }); -} - -std::shared_ptr AddForward(const std::shared_ptr &a, const std::shared_ptr &b) { - return BinaryForward(a, b, [] __device__(float x, float y) { return x + y; }); -} - -std::pair, std::shared_ptr> AddBackward(const std::shared_ptr &grad_output, - const std::vector &a_dims, - const std::vector &b_dims) { - return BinaryBackward( - grad_output, nullptr, nullptr, a_dims, b_dims, [] __device__(float, float) { return 1.f; }, - [] __device__(float, float) { return 1.f; }); -} - -std::shared_ptr AddScalarForward(const std::shared_ptr &a, float scalar) { - return UnaryForward(a, [scalar] __device__(float x) { return x + scalar; }); -} - -std::shared_ptr AddScalarBackward(const std::shared_ptr &grad_output) { - return UnaryBackward(grad_output, nullptr, [] __device__(float) { return 1.0f; }); -} - -std::shared_ptr SubForward(const std::shared_ptr &a, const std::shared_ptr &b) { - return BinaryForward(a, b, [] __device__(float x, float y) { return x - y; }); -} - -std::pair, std::shared_ptr> SubBackward(const std::shared_ptr &grad_output, - const std::vector &a_dims, - const std::vector &b_dims) { - return BinaryBackward( - grad_output, nullptr, nullptr, a_dims, b_dims, [] __device__(float, float) { return 1.f; }, - [] __device__(float, float) { return -1.f; }); -} - -std::shared_ptr MulForward(const std::shared_ptr &a, const std::shared_ptr &b) { - return BinaryForward(a, b, [] __device__(float x, float y) { return x * y; }); -} - -std::pair, std::shared_ptr> MulBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &a, - const std::shared_ptr &b) { - return BinaryBackward( - grad_output, a, b, a->Dims(), b->Dims(), [] __device__(float, float y) { return y; }, - [] __device__(float x, float) { return x; }); -} - -std::shared_ptr MulScalarForward(const std::shared_ptr &a, float scalar) { - return UnaryForward(a, [scalar] __device__(float x) { return x * scalar; }); -} - -std::shared_ptr MulScalarBackward(const std::shared_ptr &grad_output, float scalar) { - return UnaryBackward(grad_output, nullptr, [scalar] __device__(float) { return scalar; }); -} - -std::shared_ptr DivForward(const std::shared_ptr &a, const std::shared_ptr &b) { - return BinaryForward(a, b, [] __device__(float x, float y) { return x / y; }); -} - -std::pair, std::shared_ptr> DivBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &a, - const std::shared_ptr &b) { - return BinaryBackward( - grad_output, a, b, a->Dims(), b->Dims(), [] __device__(float, float y) { return 1 / y; }, - [] __device__(float x, float y) { return -x / (y * y); }); -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_ELEMENTWISE_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_ELEMENTWISE_KERNEL(NegForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(NegBackward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(ReciprocalForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(ReciprocalBackward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(SinForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(SinBackward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(CosForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(CosBackward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(TanhForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(TanhBackward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(PowForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(PowBackward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(RsqrtForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(RsqrtBackward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(EqualsScalarForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(AddForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(AddBackward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(AddScalarForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(AddScalarBackward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(SubForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(SubBackward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(MulForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(MulBackward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(MulScalarForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(MulScalarBackward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(DivForward) -REGISTER_CUDA_ELEMENTWISE_KERNEL(DivBackward) - -#undef REGISTER_CUDA_ELEMENTWISE_KERNEL +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { + +#define CEIL_DIV(x, y) (((x) + (y)-1) / (y)) + +namespace { + +template +__global__ void UnaryForwardKernel(T *output, Func fn, size_t num_elements, size_t offset, const T *input) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; + + if (idx < num_elements) { + output[idx] = fn(input[idx]); + } +} + +// Helper for broadcast indexing +__device__ inline int64_t CalcOffset(int64_t idx, int ndim, const int64_t *strides, const int64_t *shape, + const int64_t *out_strides) { + int64_t offset = 0; + for (int i = 0; i < ndim; ++i) { + int64_t out_index = (idx / out_strides[i]) % shape[i]; + int64_t index = shape[i] == 1 ? 0 : out_index; + offset += index * strides[i]; + } + return offset; +} + +template +__global__ void BinaryForwardKernel(T *output, Func fn, int ndim, const int64_t *a_strides, const int64_t *a_shape, + const int64_t *b_strides, const int64_t *b_shape, const int64_t *out_strides, + const T *a, const T *b, size_t num_elements) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_elements) { + return; + } + + int64_t a_offset = CalcOffset(idx, ndim, a_strides, a_shape, out_strides); + int64_t b_offset = CalcOffset(idx, ndim, b_strides, b_shape, out_strides); + + output[idx] = fn(a[a_offset], b[b_offset]); +} + +// launch the given kernel function with the given output and inputs +template +void LaunchKernel(Kernel &&kernel, const std::shared_ptr &output, const Inputs &...inputs) { + auto extract_ptrs + = [](const auto &...ts) { return std::make_tuple(static_cast(ts ? ts->DataPtr() : nullptr)...); }; + auto input_ptrs = extract_ptrs(inputs...); + + const size_t num_elements = output->NumElements(); + dim3 block_dims(std::min(BLOCK_SIZE, static_cast(1024))); + dim3 grid_dims(CEIL_DIV(num_elements, block_dims.x)); + const size_t step = grid_dims.x * block_dims.x; + + for (size_t offset = 0; offset < num_elements; offset += step) { + std::apply([&](auto... ptrs) { kernel(grid_dims, block_dims, offset, ptrs...); }, input_ptrs); + } +} + +// Helper for stride calculation +std::vector ComputeStride(const std::vector &dims) { + std::vector strides(dims.size(), 1); + for (int i = dims.size() - 2; i >= 0; --i) { strides[i] = strides[i + 1] * dims[i + 1]; } + return strides; +} + +// launch a forward elementwise operation given the calculation function, output, and the inputs +// Note: currently only support unary and binary operations +template +void LaunchForward(Func func, const std::shared_ptr &output, const Inputs &...inputs) { + T *output_ptr = static_cast(output->DataPtr()); + + if constexpr (sizeof...(inputs) == 1) { + // Unary case + LaunchKernel( + [&](dim3 grid, dim3 block, size_t offset, auto... ptrs) { + UnaryForwardKernel<<>>(output_ptr, func, output->NumElements(), offset, ptrs...); + }, + output, inputs...); + } else if constexpr (sizeof...(inputs) == 2) { + // Binary case + auto input_tuple = std::make_tuple(inputs...); + const auto &input_a = std::get<0>(input_tuple); + const auto &input_b = std::get<1>(input_tuple); + + const auto &a_dims = input_a->Dims(); + const auto &b_dims = input_b->Dims(); + const auto &out_dims = output->Dims(); + int ndim = out_dims.size(); + + std::vector a_shape(ndim, 1), b_shape(ndim, 1), out_shape(ndim, 1); + std::copy_backward(a_dims.begin(), a_dims.end(), a_shape.end()); + std::copy_backward(b_dims.begin(), b_dims.end(), b_shape.end()); + std::copy_backward(out_dims.begin(), out_dims.end(), out_shape.end()); + + auto a_stride_host = ComputeStride(a_shape); + auto b_stride_host = ComputeStride(b_shape); + auto out_stride_host = ComputeStride(out_shape); + + int64_t *device_buffer; + cudaMallocAsync(&device_buffer, 5 * ndim * sizeof(int64_t), 0); + + int64_t *device_a_strides, *device_b_strides, *device_out_strides, *device_a_shape, *device_b_shape; + device_a_strides = device_buffer + ndim * 0; + device_b_strides = device_buffer + ndim * 1; + device_out_strides = device_buffer + ndim * 2; + device_a_shape = device_buffer + ndim * 3; + device_b_shape = device_buffer + ndim * 4; + + std::vector host_buffer; + host_buffer.insert(host_buffer.end(), a_stride_host.begin(), a_stride_host.end()); + host_buffer.insert(host_buffer.end(), b_stride_host.begin(), b_stride_host.end()); + host_buffer.insert(host_buffer.end(), out_stride_host.begin(), out_stride_host.end()); + host_buffer.insert(host_buffer.end(), a_shape.begin(), a_shape.end()); + host_buffer.insert(host_buffer.end(), b_shape.begin(), b_shape.end()); + + cudaMemcpyAsync(device_buffer, host_buffer.data(), 5 * ndim * sizeof(int64_t), cudaMemcpyHostToDevice, 0); + + LaunchKernel( + [&](dim3 grid, dim3 block, size_t offset, const T *a_ptr, const T *b_ptr) { + BinaryForwardKernel<<>>(output_ptr, func, ndim, device_a_strides, device_a_shape, + device_b_strides, device_b_shape, device_out_strides, a_ptr, b_ptr, + output->NumElements()); + }, + output, inputs...); + + cudaFreeAsync(device_buffer, 0); + } else { + static_assert(sizeof...(inputs) == 1 || sizeof...(inputs) == 2, + "LaunchForward currently only supports unary and binary operations."); + } +} + +// Backward kernel for unary operators +template +__global__ void UnaryBackwardKernel(T *output, Func fn, size_t num_elements, size_t offset, const T *grad_output, + const T *input) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; + + if (idx < num_elements) { + output[idx] = grad_output[idx] * fn(input ? input[idx] : T(0)); + } +} + +// Backward kernel for binary operators +template +__global__ void BinaryBackwardKernel(T *output_a, T *output_b, FuncA fn_a, FuncB fn_b, int ndim, size_t num_elements, + const int64_t *a_strides, const int64_t *a_shape, const int64_t *b_strides, + const int64_t *b_shape, const int64_t *out_strides, const T *grad_output, + const T *input_a, const T *input_b) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_elements) { + return; + } + + int64_t a_offset = CalcOffset(idx, ndim, a_strides, a_shape, out_strides); + int64_t b_offset = CalcOffset(idx, ndim, b_strides, b_shape, out_strides); + + const T a_val = input_a ? input_a[a_offset] : T(0); + const T b_val = input_b ? input_b[b_offset] : T(0); + + output_a[a_offset] = grad_output[idx] * fn_a(a_val, b_val); + atomicAdd(&output_b[b_offset], grad_output[idx] * fn_b(a_val, b_val)); +} + +// launch unary operator's backward kernel +template +void LaunchBackward(Func func, const std::shared_ptr &output, const std::shared_ptr &grad_output, + const Inputs &...inputs) { + T *output_ptr = static_cast(output->DataPtr()); + const T *grad_ptr = static_cast(grad_output->DataPtr()); + + LaunchKernel( + [=](dim3 grid, dim3 block, size_t offset, auto... ptrs) { + UnaryBackwardKernel<<>>(output_ptr, func, output->NumElements(), offset, grad_ptr, ptrs...); + }, + output, inputs...); +} + +// launch binary operator's backward kernel +template +void LaunchBackward(FuncA fun_a, FuncB fun_b, const std::shared_ptr &output_a, + const std::shared_ptr &output_b, const std::vector &a_dims, + const std::vector &b_dims, const std::shared_ptr &grad_output, + const Inputs &...inputs) { + T *output_a_ptr = static_cast(output_a->DataPtr()); + T *output_b_ptr = static_cast(output_b->DataPtr()); + const T *grad_output_ptr = static_cast(grad_output->DataPtr()); + + const auto &out_dims = grad_output->Dims(); + int ndim = out_dims.size(); + + std::vector a_shape(ndim, 1), b_shape(ndim, 1), out_shape(ndim, 1); + std::copy_backward(a_dims.begin(), a_dims.end(), a_shape.end()); + std::copy_backward(b_dims.begin(), b_dims.end(), b_shape.end()); + std::copy_backward(out_dims.begin(), out_dims.end(), out_shape.end()); + + auto a_stride_host = ComputeStride(a_shape); + auto b_stride_host = ComputeStride(b_shape); + auto out_stride_host = ComputeStride(out_shape); + + int64_t *device_buffer; + cudaMallocAsync(&device_buffer, 5 * ndim * sizeof(int64_t), 0); + + int64_t *device_a_strides, *device_b_strides, *device_out_strides, *device_a_shape, *device_b_shape; + device_a_strides = device_buffer + ndim * 0; + device_b_strides = device_buffer + ndim * 1; + device_out_strides = device_buffer + ndim * 2; + device_a_shape = device_buffer + ndim * 3; + device_b_shape = device_buffer + ndim * 4; + + std::vector host_buffer; + host_buffer.insert(host_buffer.end(), a_stride_host.begin(), a_stride_host.end()); + host_buffer.insert(host_buffer.end(), b_stride_host.begin(), b_stride_host.end()); + host_buffer.insert(host_buffer.end(), out_stride_host.begin(), out_stride_host.end()); + host_buffer.insert(host_buffer.end(), a_shape.begin(), a_shape.end()); + host_buffer.insert(host_buffer.end(), b_shape.begin(), b_shape.end()); + + cudaMemcpyAsync(device_buffer, host_buffer.data(), 5 * ndim * sizeof(int64_t), cudaMemcpyHostToDevice, 0); + + const size_t num_elements = grad_output->NumElements(); + LaunchKernel( + [=](dim3 grid, dim3 block, size_t offset, auto... ptrs) { + BinaryBackwardKernel<<>>(output_a_ptr, output_b_ptr, fun_a, fun_b, ndim, num_elements, + device_a_strides, device_a_shape, device_b_strides, device_b_shape, + device_out_strides, grad_output_ptr, ptrs...); + }, + output_a, inputs...); + + cudaFreeAsync(device_buffer, 0); +} + +template std::shared_ptr UnaryForward(const std::shared_ptr &input, Func unary_fn) { + auto dtype = input->Dtype(); + auto output = std::make_shared(input->Dims(), dtype, input->GetDevice()); + + switch (dtype) { + case DataType::kFLOAT32: + LaunchForward<256, float>(unary_fn, output, input); + break; + default: + LOG(FATAL) << "CUDA unary forward: 'Unsupported data type' at " << __FILE__ << ":" << __LINE__; + } + + return output; +} + +template +std::shared_ptr UnaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr &a, + Func unary_fn) { + auto dtype = grad_output->Dtype(); + auto output = std::make_shared(grad_output->Dims(), dtype, grad_output->GetDevice()); + output->Fill(0.0f); + switch (dtype) { + case DataType::kFLOAT32: + LaunchBackward<256, float>(unary_fn, output, grad_output, a); + break; + default: + LOG(FATAL) << "CUDA unary backward: 'Unsupported data type' at " << __FILE__ << ":" << __LINE__; + } + + return output; +} + +template +std::shared_ptr BinaryForward(const std::shared_ptr &a, const std::shared_ptr &b, + Func binary_fn) { + auto dtype = a->Dtype(); + // Currently a and b should have the same data type and only one-way broadcasting from b to a is assumed by default + CHECK(dtype == b->Dtype() && a->NumElements() >= b->NumElements() && a->NumElements() % b->NumElements() == 0); + + auto output = std::make_shared(a->Dims(), dtype, a->GetDevice()); + + switch (dtype) { + case DataType::kFLOAT32: + LaunchForward<256, float>(binary_fn, output, a, b); + break; + default: + LOG(FATAL) << "CUDA binary forward: 'Unsupported data type' at " << __FILE__ << ":" << __LINE__; + } + + return output; +} + +template +std::pair, std::shared_ptr> +BinaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr &a, + const std::shared_ptr &b, const std::vector &a_dims, const std::vector &b_dims, + FuncA fn_a, FuncB fn_b) { + const auto a_num_elements = std::accumulate(a_dims.begin(), a_dims.end(), 1, std::multiplies()); + const auto b_num_elements = std::accumulate(b_dims.begin(), b_dims.end(), 1, std::multiplies()); + + CHECK(a_num_elements >= b_num_elements && a_num_elements % b_num_elements == 0); + if (a) { + CHECK(a_num_elements == a->NumElements()); + } + if (b) { + CHECK(b_num_elements == b->NumElements()); + } + auto dtype = grad_output->Dtype(); + auto device = grad_output->GetDevice(); + + // Currently a and b should have the same data type + if (a && b) { + CHECK(a->Dtype() == b->Dtype()); + } + auto grad_a = std::make_shared(a_dims, dtype, device); + auto grad_b = std::make_shared(b_dims, dtype, device); + grad_a->Fill(0.0f); + grad_b->Fill(0.0f); + switch (dtype) { + case DataType::kFLOAT32: + LaunchBackward<256, float>(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output, a, b); + break; + default: + LOG(FATAL) << "CUDA binary backward: 'Unsupported data type' at " << __FILE__ << ":" << __LINE__; + } + + return {grad_a, grad_b}; +} +} // namespace + +std::shared_ptr NegForward(const std::shared_ptr &input) { + return UnaryForward(input, [] __device__(float x) { return -x; }); +} + +std::shared_ptr NegBackward(const std::shared_ptr &grad_output) { + return UnaryBackward(grad_output, nullptr, [] __device__(float) { return -1.0f; }); +} + +std::shared_ptr ReciprocalForward(const std::shared_ptr &input) { + return UnaryForward(input, [] __device__(float x) { return 1.0f / x; }); +} + +std::shared_ptr ReciprocalBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &input) { + return UnaryBackward(grad_output, input, [] __device__(float x) { return -1.0f / (x * x); }); +} + +std::shared_ptr SinForward(const std::shared_ptr &input) { + return UnaryForward(input, [] __device__(float x) { return sinf(x); }); +} + +std::shared_ptr SinBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { + return UnaryBackward(grad_output, input, [] __device__(float x) { return cosf(x); }); +} + +std::shared_ptr CosForward(const std::shared_ptr &input) { + return UnaryForward(input, [] __device__(float x) { return cosf(x); }); +} + +std::shared_ptr CosBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { + return UnaryBackward(grad_output, input, [] __device__(float x) { return -sinf(x); }); +} + +std::shared_ptr TanhForward(const std::shared_ptr &input) { + return UnaryForward(input, [] __device__(float x) { return tanhf(x); }); +} + +std::shared_ptr TanhBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &output) { + return UnaryBackward(grad_output, output, [] __device__(float x) { return 1.0 - x * x; }); +} + +std::shared_ptr PowForward(const std::shared_ptr &input, float scalar, bool scalar_is_base) { + if (scalar_is_base) { + return UnaryForward(input, [scalar] __device__(float x) { return powf(scalar, x); }); + } else { + return UnaryForward(input, [scalar] __device__(float x) { return powf(x, scalar); }); + } +} + +std::shared_ptr PowBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + float scalar, bool scalar_is_base) { + if (scalar_is_base) { + return UnaryBackward(grad_output, input, + [scalar] __device__(float x) { return logf(scalar) * powf(scalar, x); }); + } else { + return UnaryBackward(grad_output, input, + [scalar] __device__(float x) { return scalar * powf(x, scalar - 1.0f); }); + } +} + +std::shared_ptr RsqrtForward(const std::shared_ptr &input) { + return UnaryForward(input, [] __device__(float x) { return 1.0f / sqrtf(x); }); +} + +std::shared_ptr RsqrtBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &input) { + return UnaryBackward(grad_output, input, [] __device__(float x) { return -0.5f / (x * sqrtf(x)); }); +} + +std::shared_ptr EqualsScalarForward(const std::shared_ptr &a, float scalar) { + return UnaryForward(a, [scalar] __device__(float x) { return x == scalar ? 1.0f : 0.0f; }); +} + +std::shared_ptr AddForward(const std::shared_ptr &a, const std::shared_ptr &b) { + return BinaryForward(a, b, [] __device__(float x, float y) { return x + y; }); +} + +std::pair, std::shared_ptr> AddBackward(const std::shared_ptr &grad_output, + const std::vector &a_dims, + const std::vector &b_dims) { + return BinaryBackward( + grad_output, nullptr, nullptr, a_dims, b_dims, [] __device__(float, float) { return 1.f; }, + [] __device__(float, float) { return 1.f; }); +} + +std::shared_ptr AddScalarForward(const std::shared_ptr &a, float scalar) { + return UnaryForward(a, [scalar] __device__(float x) { return x + scalar; }); +} + +std::shared_ptr AddScalarBackward(const std::shared_ptr &grad_output) { + return UnaryBackward(grad_output, nullptr, [] __device__(float) { return 1.0f; }); +} + +std::shared_ptr SubForward(const std::shared_ptr &a, const std::shared_ptr &b) { + return BinaryForward(a, b, [] __device__(float x, float y) { return x - y; }); +} + +std::pair, std::shared_ptr> SubBackward(const std::shared_ptr &grad_output, + const std::vector &a_dims, + const std::vector &b_dims) { + return BinaryBackward( + grad_output, nullptr, nullptr, a_dims, b_dims, [] __device__(float, float) { return 1.f; }, + [] __device__(float, float) { return -1.f; }); +} + +std::shared_ptr MulForward(const std::shared_ptr &a, const std::shared_ptr &b) { + return BinaryForward(a, b, [] __device__(float x, float y) { return x * y; }); +} + +std::pair, std::shared_ptr> MulBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &a, + const std::shared_ptr &b) { + return BinaryBackward( + grad_output, a, b, a->Dims(), b->Dims(), [] __device__(float, float y) { return y; }, + [] __device__(float x, float) { return x; }); +} + +std::shared_ptr MulScalarForward(const std::shared_ptr &a, float scalar) { + return UnaryForward(a, [scalar] __device__(float x) { return x * scalar; }); +} + +std::shared_ptr MulScalarBackward(const std::shared_ptr &grad_output, float scalar) { + return UnaryBackward(grad_output, nullptr, [scalar] __device__(float) { return scalar; }); +} + +std::shared_ptr DivForward(const std::shared_ptr &a, const std::shared_ptr &b) { + return BinaryForward(a, b, [] __device__(float x, float y) { return x / y; }); +} + +std::pair, std::shared_ptr> DivBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &a, + const std::shared_ptr &b) { + return BinaryBackward( + grad_output, a, b, a->Dims(), b->Dims(), [] __device__(float, float y) { return 1 / y; }, + [] __device__(float x, float y) { return -x / (y * y); }); +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_ELEMENTWISE_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_ELEMENTWISE_KERNEL(NegForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(NegBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(ReciprocalForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(ReciprocalBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(SinForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(SinBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(CosForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(CosBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(TanhForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(TanhBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(PowForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(PowBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(RsqrtForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(RsqrtBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(EqualsScalarForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(AddForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(AddBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(AddScalarForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(AddScalarBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(SubForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(SubBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(MulForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(MulBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(MulScalarForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(MulScalarBackward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(DivForward) +REGISTER_CUDA_ELEMENTWISE_KERNEL(DivBackward) + +#undef REGISTER_CUDA_ELEMENTWISE_KERNEL diff --git a/infini_train/src/kernels/cuda/embedding.cu b/infini_train/src/kernels/cuda/embedding.cu index 3ac0a8d..437b960 100644 --- a/infini_train/src/kernels/cuda/embedding.cu +++ b/infini_train/src/kernels/cuda/embedding.cu @@ -1,97 +1,97 @@ -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { - -#define CUDA_CHECK(call) \ - do { \ - cudaError_t status = call; \ - if (status != cudaSuccess) { \ - LOG(FATAL) << "CUDA Error: " << cudaGetErrorString(status) << " at " << __FILE__ << ":" << __LINE__; \ - } \ - } while (0) - -__global__ void EmbeddingForwardKernel(const int64_t *input, float *output, const float *weight, int batch_size, - int max_seqlen, int embed_dim) { - int idx = (blockIdx.x * blockDim.x + threadIdx.x); - if (idx >= batch_size * max_seqlen * embed_dim) { - return; - } - - int bt = idx / embed_dim; - int b = bt / max_seqlen; - int t = bt % max_seqlen; - int c = idx % embed_dim; - - int ix = static_cast(input[b * max_seqlen + t]); - - output[b * max_seqlen * embed_dim + t * embed_dim + c] = weight[ix * embed_dim + c]; -} - -std::shared_ptr EmbeddingForward(const std::shared_ptr &input, const std::shared_ptr &weight) { - CHECK(input->Dtype() == DataType::kINT64); - CHECK_EQ(weight->Dims().size(), 2); - - const int batch_size = input->Dims().size() == 2 ? input->Dims()[0] : 1; - const int max_seqlen = input->Dims().size() == 2 ? input->Dims()[1] : input->Dims()[0]; - const int embed_dim = weight->Dims()[1]; - auto output_dims = input->Dims(); - output_dims.push_back(embed_dim); - - auto output = std::make_shared(output_dims, DataType::kFLOAT32, input->GetDevice()); - int threads_per_block = 256; - int num_blocks = (batch_size * max_seqlen * embed_dim + threads_per_block - 1) / threads_per_block; - EmbeddingForwardKernel<<>>( - static_cast(input->DataPtr()), static_cast(output->DataPtr()), - static_cast(weight->DataPtr()), batch_size, max_seqlen, embed_dim); - return output; -} - -__global__ void EmbeddingBackwardKernel(const int64_t *input_ptr, const float *grad_output_ptr, float *grad_weight_ptr, - int num_tokens, int embedding_dim) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= num_tokens) { - return; - } - - int token_id = static_cast(input_ptr[idx]); - if (token_id < 0) { - return; - } - - for (int j = 0; j < embedding_dim; ++j) { - atomicAdd(&grad_weight_ptr[token_id * embedding_dim + j], grad_output_ptr[idx * embedding_dim + j]); - } -} - -std::shared_ptr EmbeddingBackward(const std::shared_ptr &input, const std::vector &weight_dims, - const std::shared_ptr &grad_output) { - CHECK(input->Dtype() == DataType::kINT64); - CHECK_EQ(weight_dims.size(), 2); - const int embedding_dim = weight_dims[1]; - CHECK_EQ(input->Dims().size() + 1, grad_output->Dims().size()); - for (int idx = 0; idx < input->Dims().size(); ++idx) { CHECK_EQ(input->Dims()[idx], grad_output->Dims()[idx]); } - CHECK_EQ(*grad_output->Dims().rbegin(), embedding_dim); - - auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32, grad_output->GetDevice()); - grad_weight->Fill(0.0f); - const int num_tokens = input->NumElements(); - const int threads_per_block = 256; - const int num_blocks = (num_tokens + threads_per_block - 1) / threads_per_block; - - EmbeddingBackwardKernel<<>>( - static_cast(input->DataPtr()), static_cast(grad_output->DataPtr()), - static_cast(grad_weight->DataPtr()), num_tokens, embedding_dim); - return grad_weight; -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_EMBEDDING_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_EMBEDDING_KERNEL(EmbeddingForward) -REGISTER_CUDA_EMBEDDING_KERNEL(EmbeddingBackward) - -#undef REGISTER_CUDA_EMBEDDING_KERNEL +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { + +#define CUDA_CHECK(call) \ + do { \ + cudaError_t status = call; \ + if (status != cudaSuccess) { \ + LOG(FATAL) << "CUDA Error: " << cudaGetErrorString(status) << " at " << __FILE__ << ":" << __LINE__; \ + } \ + } while (0) + +__global__ void EmbeddingForwardKernel(const int64_t *input, float *output, const float *weight, int batch_size, + int max_seqlen, int embed_dim) { + int idx = (blockIdx.x * blockDim.x + threadIdx.x); + if (idx >= batch_size * max_seqlen * embed_dim) { + return; + } + + int bt = idx / embed_dim; + int b = bt / max_seqlen; + int t = bt % max_seqlen; + int c = idx % embed_dim; + + int ix = static_cast(input[b * max_seqlen + t]); + + output[b * max_seqlen * embed_dim + t * embed_dim + c] = weight[ix * embed_dim + c]; +} + +std::shared_ptr EmbeddingForward(const std::shared_ptr &input, const std::shared_ptr &weight) { + CHECK(input->Dtype() == DataType::kINT64); + CHECK_EQ(weight->Dims().size(), 2); + + const int batch_size = input->Dims().size() == 2 ? input->Dims()[0] : 1; + const int max_seqlen = input->Dims().size() == 2 ? input->Dims()[1] : input->Dims()[0]; + const int embed_dim = weight->Dims()[1]; + auto output_dims = input->Dims(); + output_dims.push_back(embed_dim); + + auto output = std::make_shared(output_dims, DataType::kFLOAT32, input->GetDevice()); + int threads_per_block = 256; + int num_blocks = (batch_size * max_seqlen * embed_dim + threads_per_block - 1) / threads_per_block; + EmbeddingForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), + static_cast(weight->DataPtr()), batch_size, max_seqlen, embed_dim); + return output; +} + +__global__ void EmbeddingBackwardKernel(const int64_t *input_ptr, const float *grad_output_ptr, float *grad_weight_ptr, + int num_tokens, int embedding_dim) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_tokens) { + return; + } + + int token_id = static_cast(input_ptr[idx]); + if (token_id < 0) { + return; + } + + for (int j = 0; j < embedding_dim; ++j) { + atomicAdd(&grad_weight_ptr[token_id * embedding_dim + j], grad_output_ptr[idx * embedding_dim + j]); + } +} + +std::shared_ptr EmbeddingBackward(const std::shared_ptr &input, const std::vector &weight_dims, + const std::shared_ptr &grad_output) { + CHECK(input->Dtype() == DataType::kINT64); + CHECK_EQ(weight_dims.size(), 2); + const int embedding_dim = weight_dims[1]; + CHECK_EQ(input->Dims().size() + 1, grad_output->Dims().size()); + for (int idx = 0; idx < input->Dims().size(); ++idx) { CHECK_EQ(input->Dims()[idx], grad_output->Dims()[idx]); } + CHECK_EQ(*grad_output->Dims().rbegin(), embedding_dim); + + auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32, grad_output->GetDevice()); + grad_weight->Fill(0.0f); + const int num_tokens = input->NumElements(); + const int threads_per_block = 256; + const int num_blocks = (num_tokens + threads_per_block - 1) / threads_per_block; + + EmbeddingBackwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(grad_output->DataPtr()), + static_cast(grad_weight->DataPtr()), num_tokens, embedding_dim); + return grad_weight; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_EMBEDDING_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_EMBEDDING_KERNEL(EmbeddingForward) +REGISTER_CUDA_EMBEDDING_KERNEL(EmbeddingBackward) + +#undef REGISTER_CUDA_EMBEDDING_KERNEL diff --git a/infini_train/src/kernels/cuda/fill.cu b/infini_train/src/kernels/cuda/fill.cu index 470e10c..a302d54 100644 --- a/infini_train/src/kernels/cuda/fill.cu +++ b/infini_train/src/kernels/cuda/fill.cu @@ -1,23 +1,23 @@ -#include "glog/logging.h" -#include -#include -#include -#include - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { -void Fill(std::shared_ptr tensor, void *value_ptr) { - // FIXME(zbl): support other data types - thrust::device_ptr dev_ptr(reinterpret_cast(tensor->DataPtr())); - thrust::fill(thrust::cuda::par.on(0), dev_ptr, dev_ptr + tensor->NumElements(), *(static_cast(value_ptr))); -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_FILL_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_FILL_KERNEL(Fill) - -#undef REGISTER_CUDA_FILL_KERNEL +#include "glog/logging.h" +#include +#include +#include +#include + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { +void Fill(std::shared_ptr tensor, void *value_ptr) { + // FIXME(zbl): support other data types + thrust::device_ptr dev_ptr(reinterpret_cast(tensor->DataPtr())); + thrust::fill(thrust::cuda::par.on(0), dev_ptr, dev_ptr + tensor->NumElements(), *(static_cast(value_ptr))); +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_FILL_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_FILL_KERNEL(Fill) + +#undef REGISTER_CUDA_FILL_KERNEL diff --git a/infini_train/src/kernels/cuda/layernorm.cu b/infini_train/src/kernels/cuda/layernorm.cu index 4b6dd0a..7b94f63 100644 --- a/infini_train/src/kernels/cuda/layernorm.cu +++ b/infini_train/src/kernels/cuda/layernorm.cu @@ -1,174 +1,174 @@ -#include "glog/logging.h" -#include - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { - -template -__global__ void LayerNormForwardKernel(const float *input, const float *weight, const float *bias, float *mean_out, - float *rstd_out, float *output, float eps, int embed_dim) { - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage_mean; - __shared__ typename BlockReduce::TempStorage temp_storage_rstd; - __shared__ float shared_mean; - __shared__ float shared_rstd; - - const int token_idx = blockIdx.x; - const float *x = input + token_idx * embed_dim; - float *y = output + token_idx * embed_dim; - - float sum = 0.0f; - float sqsum = 0.0f; - for (int i = threadIdx.x; i < embed_dim; i += BLOCK_SIZE) { - float val = x[i]; - sum += val; - sqsum += val * val; - } - - float total_sum = BlockReduce(temp_storage_mean).Sum(sum); - __syncthreads(); - - float total_sqsum = BlockReduce(temp_storage_rstd).Sum(sqsum); - __syncthreads(); - - if (threadIdx.x == 0) { - float mean = total_sum / embed_dim; - float var = total_sqsum / embed_dim - mean * mean; - float rstd = rsqrtf(var + eps); - shared_mean = mean; - shared_rstd = rstd; - if (mean_out) { - mean_out[token_idx] = mean; - } - if (rstd_out) { - rstd_out[token_idx] = rstd; - } - } - __syncthreads(); - - for (int i = threadIdx.x; i < embed_dim; i += BLOCK_SIZE) { - float norm = (x[i] - shared_mean) * shared_rstd; - y[i] = norm * weight[i] + bias[i]; - } -} - -std::tuple, std::shared_ptr, std::shared_ptr> -LayerNormForward(const std::shared_ptr &input, const std::shared_ptr &weight, - const std::shared_ptr &bias, const float eps) { - CHECK_EQ(input->Dims().size(), 3); - CHECK_LE(input->Dims()[2], weight->Dims()[0]); - CHECK_LE(input->Dims()[2], bias->Dims()[0]); - - const int batch_size = input->Dims()[0]; - const int max_seqlen = input->Dims()[1]; - const int embed_dim = input->Dims()[2]; - - auto output = std::make_shared(input->Dims(), DataType::kFLOAT32, input->GetDevice()); - auto mean = std::make_shared(std::vector{batch_size, max_seqlen}, DataType::kFLOAT32, - input->GetDevice()); - auto rstd = std::make_shared(std::vector{batch_size, max_seqlen}, DataType::kFLOAT32, - input->GetDevice()); - mean->Fill(0.0f); - rstd->Fill(0.0f); - - constexpr int BLOCK_SIZE = 256; - int threads_per_block = BLOCK_SIZE; - int num_blocks = batch_size * max_seqlen; - - LayerNormForwardKernel<<>>( - static_cast(input->DataPtr()), static_cast(weight->DataPtr()), - static_cast(bias->DataPtr()), static_cast(mean->DataPtr()), - static_cast(rstd->DataPtr()), static_cast(output->DataPtr()), eps, embed_dim); - return {output, mean, rstd}; -} - -template -__global__ void LayerNormBackwardKernel(const float *__restrict__ input, const float *__restrict__ grad_output, - const float *__restrict__ mean, const float *__restrict__ rstd, - const float *__restrict__ weight, float *__restrict__ grad_input, - float *__restrict__ grad_weight, float *__restrict__ grad_bias, int embed_dim) { - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage_mean; - __shared__ typename BlockReduce::TempStorage temp_storage_norm; - __shared__ float shared_mean; - __shared__ float shared_norm; - - int tid = threadIdx.x; - int token_idx = blockIdx.x; - - const float *input_ptr = input + token_idx * embed_dim; - const float *grad_output_ptr = grad_output + token_idx * embed_dim; - float *grad_input_ptr = grad_input + token_idx * embed_dim; - - float mean_val = mean[token_idx]; - float rstd_val = rstd[token_idx]; - - float dnorm_mean = 0.f; - float dnorm_norm_mean = 0.f; - - for (int i = tid; i < embed_dim; i += BLOCK_SIZE) { - float dnorm = weight[i] * grad_output_ptr[i]; - dnorm_mean += dnorm; - dnorm_norm_mean += dnorm * (input_ptr[i] - mean_val); - } - - dnorm_mean = BlockReduce(temp_storage_mean).Sum(dnorm_mean); - __syncthreads(); - dnorm_norm_mean = BlockReduce(temp_storage_norm).Sum(dnorm_norm_mean); - __syncthreads(); - - if (tid == 0) { - float mean_d = dnorm_mean / embed_dim; - float norm_d = (dnorm_norm_mean / embed_dim) * rstd_val - mean_d * mean_val * rstd_val; - shared_mean = mean_d; - shared_norm = norm_d; - } - __syncthreads(); - - for (int i = tid; i < embed_dim; i += BLOCK_SIZE) { - float norm = (input_ptr[i] - mean_val) * rstd_val; - - grad_input_ptr[i] = (weight[i] * grad_output_ptr[i] - shared_mean - norm * shared_norm) * rstd_val; - - atomicAdd(&grad_weight[i], grad_output_ptr[i] * norm); - atomicAdd(&grad_bias[i], grad_output_ptr[i]); - } -} - -std::tuple, std::shared_ptr, std::shared_ptr> -LayerNormBackward(const std::shared_ptr &input, const std::shared_ptr &weight, - const std::shared_ptr &bias, const std::shared_ptr &mean, - const std::shared_ptr &rstd, const std::shared_ptr &grad_output) { - const int batch_size = input->Dims()[0]; - const int max_seqlen = input->Dims()[1]; - const int embed_dim = input->Dims()[2]; - - auto grad_input = std::make_shared(input->Dims(), DataType::kFLOAT32, grad_output->GetDevice()); - auto grad_weight = std::make_shared(weight->Dims(), DataType::kFLOAT32, grad_output->GetDevice()); - auto grad_bias = std::make_shared(bias->Dims(), DataType::kFLOAT32, grad_output->GetDevice()); - grad_input->Fill(0.0f); - grad_weight->Fill(0.0f); - grad_bias->Fill(0.0f); - - constexpr int BLOCK_SIZE = 256; - int threads_per_block = BLOCK_SIZE; - int num_blocks = batch_size * max_seqlen; - - LayerNormBackwardKernel<<>>( - static_cast(input->DataPtr()), static_cast(grad_output->DataPtr()), - static_cast(mean->DataPtr()), static_cast(rstd->DataPtr()), - static_cast(weight->DataPtr()), static_cast(grad_input->DataPtr()), - static_cast(grad_weight->DataPtr()), static_cast(grad_bias->DataPtr()), embed_dim); - return {grad_input, grad_weight, grad_bias}; -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_LAYERNORM_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_LAYERNORM_KERNEL(LayerNormForward) -REGISTER_CUDA_LAYERNORM_KERNEL(LayerNormBackward) - -#undef REGISTER_CUDA_LAYERNORM_KERNEL +#include "glog/logging.h" +#include + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { + +template +__global__ void LayerNormForwardKernel(const float *input, const float *weight, const float *bias, float *mean_out, + float *rstd_out, float *output, float eps, int embed_dim) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage_mean; + __shared__ typename BlockReduce::TempStorage temp_storage_rstd; + __shared__ float shared_mean; + __shared__ float shared_rstd; + + const int token_idx = blockIdx.x; + const float *x = input + token_idx * embed_dim; + float *y = output + token_idx * embed_dim; + + float sum = 0.0f; + float sqsum = 0.0f; + for (int i = threadIdx.x; i < embed_dim; i += BLOCK_SIZE) { + float val = x[i]; + sum += val; + sqsum += val * val; + } + + float total_sum = BlockReduce(temp_storage_mean).Sum(sum); + __syncthreads(); + + float total_sqsum = BlockReduce(temp_storage_rstd).Sum(sqsum); + __syncthreads(); + + if (threadIdx.x == 0) { + float mean = total_sum / embed_dim; + float var = total_sqsum / embed_dim - mean * mean; + float rstd = rsqrtf(var + eps); + shared_mean = mean; + shared_rstd = rstd; + if (mean_out) { + mean_out[token_idx] = mean; + } + if (rstd_out) { + rstd_out[token_idx] = rstd; + } + } + __syncthreads(); + + for (int i = threadIdx.x; i < embed_dim; i += BLOCK_SIZE) { + float norm = (x[i] - shared_mean) * shared_rstd; + y[i] = norm * weight[i] + bias[i]; + } +} + +std::tuple, std::shared_ptr, std::shared_ptr> +LayerNormForward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias, const float eps) { + CHECK_EQ(input->Dims().size(), 3); + CHECK_LE(input->Dims()[2], weight->Dims()[0]); + CHECK_LE(input->Dims()[2], bias->Dims()[0]); + + const int batch_size = input->Dims()[0]; + const int max_seqlen = input->Dims()[1]; + const int embed_dim = input->Dims()[2]; + + auto output = std::make_shared(input->Dims(), DataType::kFLOAT32, input->GetDevice()); + auto mean = std::make_shared(std::vector{batch_size, max_seqlen}, DataType::kFLOAT32, + input->GetDevice()); + auto rstd = std::make_shared(std::vector{batch_size, max_seqlen}, DataType::kFLOAT32, + input->GetDevice()); + mean->Fill(0.0f); + rstd->Fill(0.0f); + + constexpr int BLOCK_SIZE = 256; + int threads_per_block = BLOCK_SIZE; + int num_blocks = batch_size * max_seqlen; + + LayerNormForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(weight->DataPtr()), + static_cast(bias->DataPtr()), static_cast(mean->DataPtr()), + static_cast(rstd->DataPtr()), static_cast(output->DataPtr()), eps, embed_dim); + return {output, mean, rstd}; +} + +template +__global__ void LayerNormBackwardKernel(const float *__restrict__ input, const float *__restrict__ grad_output, + const float *__restrict__ mean, const float *__restrict__ rstd, + const float *__restrict__ weight, float *__restrict__ grad_input, + float *__restrict__ grad_weight, float *__restrict__ grad_bias, int embed_dim) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage_mean; + __shared__ typename BlockReduce::TempStorage temp_storage_norm; + __shared__ float shared_mean; + __shared__ float shared_norm; + + int tid = threadIdx.x; + int token_idx = blockIdx.x; + + const float *input_ptr = input + token_idx * embed_dim; + const float *grad_output_ptr = grad_output + token_idx * embed_dim; + float *grad_input_ptr = grad_input + token_idx * embed_dim; + + float mean_val = mean[token_idx]; + float rstd_val = rstd[token_idx]; + + float dnorm_mean = 0.f; + float dnorm_norm_mean = 0.f; + + for (int i = tid; i < embed_dim; i += BLOCK_SIZE) { + float dnorm = weight[i] * grad_output_ptr[i]; + dnorm_mean += dnorm; + dnorm_norm_mean += dnorm * (input_ptr[i] - mean_val); + } + + dnorm_mean = BlockReduce(temp_storage_mean).Sum(dnorm_mean); + __syncthreads(); + dnorm_norm_mean = BlockReduce(temp_storage_norm).Sum(dnorm_norm_mean); + __syncthreads(); + + if (tid == 0) { + float mean_d = dnorm_mean / embed_dim; + float norm_d = (dnorm_norm_mean / embed_dim) * rstd_val - mean_d * mean_val * rstd_val; + shared_mean = mean_d; + shared_norm = norm_d; + } + __syncthreads(); + + for (int i = tid; i < embed_dim; i += BLOCK_SIZE) { + float norm = (input_ptr[i] - mean_val) * rstd_val; + + grad_input_ptr[i] = (weight[i] * grad_output_ptr[i] - shared_mean - norm * shared_norm) * rstd_val; + + atomicAdd(&grad_weight[i], grad_output_ptr[i] * norm); + atomicAdd(&grad_bias[i], grad_output_ptr[i]); + } +} + +std::tuple, std::shared_ptr, std::shared_ptr> +LayerNormBackward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias, const std::shared_ptr &mean, + const std::shared_ptr &rstd, const std::shared_ptr &grad_output) { + const int batch_size = input->Dims()[0]; + const int max_seqlen = input->Dims()[1]; + const int embed_dim = input->Dims()[2]; + + auto grad_input = std::make_shared(input->Dims(), DataType::kFLOAT32, grad_output->GetDevice()); + auto grad_weight = std::make_shared(weight->Dims(), DataType::kFLOAT32, grad_output->GetDevice()); + auto grad_bias = std::make_shared(bias->Dims(), DataType::kFLOAT32, grad_output->GetDevice()); + grad_input->Fill(0.0f); + grad_weight->Fill(0.0f); + grad_bias->Fill(0.0f); + + constexpr int BLOCK_SIZE = 256; + int threads_per_block = BLOCK_SIZE; + int num_blocks = batch_size * max_seqlen; + + LayerNormBackwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(grad_output->DataPtr()), + static_cast(mean->DataPtr()), static_cast(rstd->DataPtr()), + static_cast(weight->DataPtr()), static_cast(grad_input->DataPtr()), + static_cast(grad_weight->DataPtr()), static_cast(grad_bias->DataPtr()), embed_dim); + return {grad_input, grad_weight, grad_bias}; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_LAYERNORM_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_LAYERNORM_KERNEL(LayerNormForward) +REGISTER_CUDA_LAYERNORM_KERNEL(LayerNormBackward) + +#undef REGISTER_CUDA_LAYERNORM_KERNEL diff --git a/infini_train/src/kernels/cuda/linear.cu b/infini_train/src/kernels/cuda/linear.cu index efaaaa6..bed5108 100644 --- a/infini_train/src/kernels/cuda/linear.cu +++ b/infini_train/src/kernels/cuda/linear.cu @@ -1,243 +1,357 @@ -#include "cublas_v2.h" -#include "glog/logging.h" -#include - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { - -#define CUDA_CHECK(call) \ - do { \ - cudaError_t status = call; \ - if (status != cudaSuccess) { \ - LOG(FATAL) << "CUDA Error: " << cudaGetErrorString(status) << " at " << __FILE__ << ":" << __LINE__; \ - } \ - } while (0) - -#define CUBLAS_CHECK(call) \ - do { \ - cublasStatus_t status = call; \ - if (status != CUBLAS_STATUS_SUCCESS) { \ - LOG(FATAL) << "CUBLAS Error: " << cublasGetStatusString(status) << " at " << __FILE__ << ":" << __LINE__; \ - } \ - } while (0) - -std::shared_ptr MatmulForward(const std::shared_ptr &input, const std::shared_ptr &other) { - // =================================== 作业 =================================== - // TODO：实现CUDA上的矩阵乘法前向计算 - // REF: - // =================================== 作业 =================================== - - auto output = std::make_shared(); - return output; -} - -std::tuple, std::shared_ptr> -MatmulBackward(const std::shared_ptr &input, const std::shared_ptr &other, - const std::shared_ptr &grad_output) { - // =================================== 作业 =================================== - // TODO：实现CUDA上的矩阵乘法反向传播 - // REF: - // =================================== 作业 =================================== - - auto grad_input = std::make_shared(); - auto grad_other = std::make_shared(); - return {grad_input, grad_other}; -} - -__global__ void BiasCopyKernel(float *output, const float *bias, int bs, int out_features) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= bs * out_features) { - return; - } - int j = idx % out_features; - output[idx] = bias[j]; -} - -std::shared_ptr LinearForward(const std::shared_ptr &input, const std::shared_ptr &weight, - bool transpose, const std::shared_ptr &bias) { - - /* - !transpose: output = input * weight + bias - output[*, out_features] = input[*, in_features] * weight[in_features, out_features] + bias[out_features] - - transpose: output = input * weight^T + bias - output[*, out_features] = input[*, in_features] * weight[out_features, in_features]^T + bias[out_features] - */ - - const auto &input_dims = input->Dims(); - CHECK_GE(input_dims.size(), 2); - const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); - const int64_t in_features = *input_dims.rbegin(); - - const auto &weight_dims = weight->Dims(); - CHECK_EQ(weight_dims.size(), 2); - CHECK_EQ(in_features, weight_dims[transpose ? 1 : 0]); - - // As for cublas: - // C = alpha * op(B) * op(A) + beta * C - // Dimensions: - // input: (bs, in_features) - // weight: (in_features, out_features) or (out_features, in_features) if transposed - // output: (bs, out_features) - const int64_t out_features = weight_dims[transpose ? 0 : 1]; - - auto output_dims = input_dims; - *output_dims.rbegin() = out_features; - auto output = std::make_shared(output_dims, DataType::kFLOAT32, input->GetDevice()); - - if (bias) { - CHECK_EQ(bias->Dims().size(), 1); - CHECK_EQ(bias->Dims()[0], out_features); - int threads_per_block = 256; - int num_blocks = (bs * out_features + threads_per_block - 1) / threads_per_block; - BiasCopyKernel<<>>( - static_cast(output->DataPtr()), static_cast(bias->DataPtr()), bs, out_features); - } else { - output->Fill(0.0f); - } - - const float alpha = 1.0f; - const float beta = 1.0f; - cublasHandle_t handle; - CUBLAS_CHECK(cublasCreate(&handle)); - if (transpose) { - // weight is [out_features, in_features] here - - // output = input * weight.T --> output.T = weight * input.T - // C = output.T[out_features, bs] - // A = weight.T[in_features, out_features] - // B = input.T[in_features, bs] - CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, out_features, bs, in_features, &alpha, - static_cast(weight->DataPtr()), in_features, - static_cast(input->DataPtr()), in_features, &beta, - static_cast(output->DataPtr()), out_features)); - } else { - // output = input * weight --> output.T = weight.T * input.T - // C = output.T[out_features, bs] - // A = weight.T[out_features, in_features] - // B = input.T[in_features, bs] - CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, out_features, bs, in_features, &alpha, - static_cast(weight->DataPtr()), out_features, - static_cast(input->DataPtr()), in_features, &beta, - static_cast(output->DataPtr()), out_features)); - } - CUBLAS_CHECK(cublasDestroy(handle)); - return output; -} - -template -__global__ void ReduceColumnsKernel(const float *__restrict__ input, float *__restrict__ output, int num_rows, - int num_cols) { - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - - int row = blockIdx.x; - float sum = 0.0f; - - for (int col = threadIdx.x; col < num_cols; col += blockDim.x) { sum += input[row * num_cols + col]; } - - float reduced = BlockReduce(temp_storage).Sum(sum); - - if (threadIdx.x == 0) { - output[row] = reduced; - } -} - -std::tuple, std::shared_ptr, std::shared_ptr> -LinearBackward(const std::shared_ptr &input, const std::shared_ptr &weight, bool transpose, - int64_t out_features, const std::shared_ptr &grad_output, const bool bias) { - const auto &input_dims = input->Dims(); - CHECK_GE(input_dims.size(), 2); - const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); - const int64_t in_features = *input_dims.rbegin(); - - const auto &weight_dims = weight->Dims(); - CHECK_EQ(weight_dims.size(), 2); - CHECK_EQ(in_features, weight_dims[transpose ? 1 : 0]); - CHECK_EQ(out_features, weight_dims[transpose ? 0 : 1]); - - auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, grad_output->GetDevice()); - auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32, grad_output->GetDevice()); - grad_input->Fill(0.0f); - grad_weight->Fill(0.0f); - std::shared_ptr grad_bias = nullptr; - if (bias) { - grad_bias = std::make_shared(std::vector{out_features}, DataType::kFLOAT32, - grad_output->GetDevice()); - grad_bias->Fill(0.0f); - } - - float alpha = 1.0f; - float beta = 0.0f; - cublasHandle_t handle; - CUBLAS_CHECK(cublasCreate(&handle)); - - if (transpose) { - // weight is [out_features, in_features] here - - // d_input = d_output * weight --> d_input.T = weight.T * d_output.T - // C = d_input.T[in_features, bs] - // A = weight.T[in_features, out_features] - // B = d_output.T[out_features, bs] - CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, in_features, bs, out_features, &alpha, - static_cast(weight->DataPtr()), in_features, - static_cast(grad_output->DataPtr()), out_features, &beta, - static_cast(grad_input->DataPtr()), in_features)); - - // d_weight = d_output.T * input --> d_weight.T = input.T * d_output - // C = d_weight.T[in_features, out_features] - // A = input.T[in_features, bs] - // B = d_output.T[out_features, bs] - CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, in_features, out_features, bs, &alpha, - static_cast(input->DataPtr()), in_features, - static_cast(grad_output->DataPtr()), out_features, &beta, - static_cast(grad_weight->DataPtr()), in_features)); - } else { - // weight is [in_features, out_features] here - - // d_input = d_output * weight.T --> d_input.T = weight * d_output.T - // C = d_input.T[in_features, bs] - // A = weight.T[out_features, in_features] - // B = d_output.T[out_features, bs] - CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, in_features, bs, out_features, &alpha, - static_cast(weight->DataPtr()), out_features, - static_cast(grad_output->DataPtr()), out_features, &beta, - static_cast(grad_input->DataPtr()), in_features)); - - // d_weight = input.T * d_output --> d_weight.T = d_output.T * input - // C = d_weight.T[out_features, in_features] - // A = d_output.T[out_features, bs] - // B = input.T[in_features, bs] - CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, out_features, in_features, bs, &alpha, - static_cast(grad_output->DataPtr()), out_features, - static_cast(input->DataPtr()), in_features, &beta, - static_cast(grad_weight->DataPtr()), out_features)); - } - - // d_bias = \sum_i(i=0, bs-1) d_output[i] - if (bias) { - constexpr int BLOCK_SIZE = 256; - int threads_per_block = BLOCK_SIZE; - int num_blocks = out_features; - ReduceColumnsKernel - <<>>(static_cast(grad_output->DataPtr()), - static_cast(grad_bias->DataPtr()), out_features, bs); - } - - CUBLAS_CHECK(cublasDestroy(handle)); - - return {grad_input, grad_weight, grad_bias}; -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_LINEAR_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_LINEAR_KERNEL(MatmulForward) -REGISTER_CUDA_LINEAR_KERNEL(MatmulBackward) -REGISTER_CUDA_LINEAR_KERNEL(LinearForward) -REGISTER_CUDA_LINEAR_KERNEL(LinearBackward) - -#undef REGISTER_CUDA_LINEAR_KERNEL +#include "cublas_v2.h" +#include "glog/logging.h" +#include + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { + +#define CUDA_CHECK(call) \ + do { \ + cudaError_t status = call; \ + if (status != cudaSuccess) { \ + LOG(FATAL) << "CUDA Error: " << cudaGetErrorString(status) << " at " << __FILE__ << ":" << __LINE__; \ + } \ + } while (0) + +#define CUBLAS_CHECK(call) \ + do { \ + cublasStatus_t status = call; \ + if (status != CUBLAS_STATUS_SUCCESS) { \ + LOG(FATAL) << "CUBLAS Error: " << cublasGetStatusString(status) << " at " << __FILE__ << ":" << __LINE__; \ + } \ + } while (0) + +std::shared_ptr MatmulForward(const std::shared_ptr &input, const std::shared_ptr &other) { + // =================================== 作业 =================================== + // TODO：实现CUDA上的矩阵乘法前向计算 + // REF: + // =================================== 作业 =================================== + + const auto &input_dims = input->Dims(); + const auto &other_dims = other->Dims(); + CHECK_GE(input_dims.size(), 2); + CHECK_GE(other_dims.size(), 2); + const int64_t M = input_dims[input_dims.size() - 2]; + const int64_t K = input_dims[input_dims.size() - 1]; + const int64_t N = other_dims[other_dims.size() - 1]; + CHECK_EQ(K, other_dims[other_dims.size() - 2]); + + int64_t batch_input = 1; + for (size_t i = 0; i < input_dims.size() - 2; ++i) { + batch_input *= input_dims[i]; + } + int64_t batch_other = 1; + for (size_t i = 0; i < other_dims.size() - 2; ++i) { + batch_other *= other_dims[i]; + } + int64_t batch = std::max(batch_input, batch_other); + + std::vector output_dims; + if (batch > 1) { + output_dims = (input_dims.size() >= other_dims.size()) ? input_dims : other_dims; + output_dims[output_dims.size() - 2] = M; + output_dims[output_dims.size() - 1] = N; + } else { + output_dims = {M, N}; + } + + auto output = std::make_shared(output_dims, DataType::kFLOAT32, input->GetDevice()); + output->Fill(0.0f); + + const float alpha = 1.0f; + const float beta = 0.0f; + cublasHandle_t handle; + CUBLAS_CHECK(cublasCreate(&handle)); + + if (batch <= 1) { + // C = output^T[N, M] = other^T[N, K] * input^T[K, M] + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, + static_cast(other->DataPtr()), N, + static_cast(input->DataPtr()), K, &beta, + static_cast(output->DataPtr()), N)); + } else { + // Use strided batched GEMM + int64_t strideA = (batch_other == 1) ? 0 : K * N; + int64_t strideB = (batch_input == 1) ? 0 : M * K; + int64_t strideC = M * N; + CUBLAS_CHECK(cublasSgemmStridedBatched(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, + static_cast(other->DataPtr()), N, strideA, + static_cast(input->DataPtr()), K, strideB, &beta, + static_cast(output->DataPtr()), N, strideC, batch)); + } + + CUBLAS_CHECK(cublasDestroy(handle)); + return output; +} + +std::tuple, std::shared_ptr> +MatmulBackward(const std::shared_ptr &input, const std::shared_ptr &other, + const std::shared_ptr &grad_output) { + // =================================== 作业 =================================== + // TODO：实现CUDA上的矩阵乘法反向传播 + // REF: + // =================================== 作业 =================================== + + const auto &input_dims = input->Dims(); + const auto &other_dims = other->Dims(); + const int64_t M = input_dims[input_dims.size() - 2]; + const int64_t K = input_dims[input_dims.size() - 1]; + const int64_t N = other_dims[other_dims.size() - 1]; + + int64_t batch_input = 1; + for (size_t i = 0; i < input_dims.size() - 2; ++i) { + batch_input *= input_dims[i]; + } + int64_t batch_other = 1; + for (size_t i = 0; i < other_dims.size() - 2; ++i) { + batch_other *= other_dims[i]; + } + int64_t batch = std::max(batch_input, batch_other); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, grad_output->GetDevice()); + auto grad_other = std::make_shared(other_dims, DataType::kFLOAT32, grad_output->GetDevice()); + grad_input->Fill(0.0f); + grad_other->Fill(0.0f); + + const float alpha = 1.0f; + const float beta = 0.0f; + cublasHandle_t handle; + CUBLAS_CHECK(cublasCreate(&handle)); + + if (batch <= 1) { + // grad_input = grad_output * other^T + // grad_input^T[K, M] = other[K, N] * grad_output^T[N, M] + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, K, M, N, &alpha, + static_cast(other->DataPtr()), N, + static_cast(grad_output->DataPtr()), N, &beta, + static_cast(grad_input->DataPtr()), K)); + + // grad_other = input^T * grad_output + // grad_other^T[N, K] = grad_output^T[N, M] * input[M, K] + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, N, K, M, &alpha, + static_cast(grad_output->DataPtr()), N, + static_cast(input->DataPtr()), K, &beta, + static_cast(grad_other->DataPtr()), N)); + } else { + // Strided batched for grad_input + int64_t strideA_gi = (batch_other == 1) ? 0 : K * N; + int64_t strideB_gi = M * N; + int64_t strideC_gi = (batch_input == 1) ? 0 : M * K; + CUBLAS_CHECK(cublasSgemmStridedBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, K, M, N, &alpha, + static_cast(other->DataPtr()), N, strideA_gi, + static_cast(grad_output->DataPtr()), N, strideB_gi, &beta, + static_cast(grad_input->DataPtr()), K, strideC_gi, batch)); + + // Strided batched for grad_other + int64_t strideA_go = M * N; + int64_t strideB_go = (batch_input == 1) ? 0 : M * K; + int64_t strideC_go = (batch_other == 1) ? 0 : K * N; + CUBLAS_CHECK(cublasSgemmStridedBatched(handle, CUBLAS_OP_N, CUBLAS_OP_T, N, K, M, &alpha, + static_cast(grad_output->DataPtr()), N, strideA_go, + static_cast(input->DataPtr()), K, strideB_go, &beta, + static_cast(grad_other->DataPtr()), N, strideC_go, batch)); + } + + CUBLAS_CHECK(cublasDestroy(handle)); + return {grad_input, grad_other}; +} + +__global__ void BiasCopyKernel(float *output, const float *bias, int bs, int out_features) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= bs * out_features) { + return; + } + int j = idx % out_features; + output[idx] = bias[j]; +} + +std::shared_ptr LinearForward(const std::shared_ptr &input, const std::shared_ptr &weight, + bool transpose, const std::shared_ptr &bias) { + + /* + !transpose: output = input * weight + bias + output[*, out_features] = input[*, in_features] * weight[in_features, out_features] + bias[out_features] + + transpose: output = input * weight^T + bias + output[*, out_features] = input[*, in_features] * weight[out_features, in_features]^T + bias[out_features] + */ + + const auto &input_dims = input->Dims(); + CHECK_GE(input_dims.size(), 2); + const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int64_t in_features = *input_dims.rbegin(); + + const auto &weight_dims = weight->Dims(); + CHECK_EQ(weight_dims.size(), 2); + CHECK_EQ(in_features, weight_dims[transpose ? 1 : 0]); + + // As for cublas: + // C = alpha * op(B) * op(A) + beta * C + // Dimensions: + // input: (bs, in_features) + // weight: (in_features, out_features) or (out_features, in_features) if transposed + // output: (bs, out_features) + const int64_t out_features = weight_dims[transpose ? 0 : 1]; + + auto output_dims = input_dims; + *output_dims.rbegin() = out_features; + auto output = std::make_shared(output_dims, DataType::kFLOAT32, input->GetDevice()); + + if (bias) { + CHECK_EQ(bias->Dims().size(), 1); + CHECK_EQ(bias->Dims()[0], out_features); + int threads_per_block = 256; + int num_blocks = (bs * out_features + threads_per_block - 1) / threads_per_block; + BiasCopyKernel<<>>( + static_cast(output->DataPtr()), static_cast(bias->DataPtr()), bs, out_features); + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaDeviceSynchronize()); + } else { + output->Fill(0.0f); + } + + const float alpha = 1.0f; + const float beta = 1.0f; + cublasHandle_t handle; + CUBLAS_CHECK(cublasCreate(&handle)); + if (transpose) { + // weight is [out_features, in_features] here + + // output = input * weight.T --> output.T = weight * input.T + // C = output.T[out_features, bs] + // A = weight.T[in_features, out_features] + // B = input.T[in_features, bs] + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, out_features, bs, in_features, &alpha, + static_cast(weight->DataPtr()), in_features, + static_cast(input->DataPtr()), in_features, &beta, + static_cast(output->DataPtr()), out_features)); + } else { + // output = input * weight --> output.T = weight.T * input.T + // C = output.T[out_features, bs] + // A = weight.T[out_features, in_features] + // B = input.T[in_features, bs] + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, out_features, bs, in_features, &alpha, + static_cast(weight->DataPtr()), out_features, + static_cast(input->DataPtr()), in_features, &beta, + static_cast(output->DataPtr()), out_features)); + } + CUBLAS_CHECK(cublasDestroy(handle)); + return output; +} + +template +__global__ void ReduceColumnsKernel(const float *__restrict__ input, float *__restrict__ output, int num_rows, + int num_cols) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + int row = blockIdx.x; + float sum = 0.0f; + + for (int col = threadIdx.x; col < num_cols; col += blockDim.x) { sum += input[row * num_cols + col]; } + + float reduced = BlockReduce(temp_storage).Sum(sum); + + if (threadIdx.x == 0) { + output[row] = reduced; + } +} + +std::tuple, std::shared_ptr, std::shared_ptr> +LinearBackward(const std::shared_ptr &input, const std::shared_ptr &weight, bool transpose, + int64_t out_features, const std::shared_ptr &grad_output, const bool bias) { + const auto &input_dims = input->Dims(); + CHECK_GE(input_dims.size(), 2); + const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int64_t in_features = *input_dims.rbegin(); + + const auto &weight_dims = weight->Dims(); + CHECK_EQ(weight_dims.size(), 2); + CHECK_EQ(in_features, weight_dims[transpose ? 1 : 0]); + CHECK_EQ(out_features, weight_dims[transpose ? 0 : 1]); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, grad_output->GetDevice()); + auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32, grad_output->GetDevice()); + grad_input->Fill(0.0f); + grad_weight->Fill(0.0f); + std::shared_ptr grad_bias = nullptr; + if (bias) { + grad_bias = std::make_shared(std::vector{out_features}, DataType::kFLOAT32, + grad_output->GetDevice()); + grad_bias->Fill(0.0f); + } + + float alpha = 1.0f; + float beta = 0.0f; + cublasHandle_t handle; + CUBLAS_CHECK(cublasCreate(&handle)); + + if (transpose) { + // weight is [out_features, in_features] here + + // d_input = d_output * weight --> d_input.T = weight.T * d_output.T + // C = d_input.T[in_features, bs] + // A = weight.T[in_features, out_features] + // B = d_output.T[out_features, bs] + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, in_features, bs, out_features, &alpha, + static_cast(weight->DataPtr()), in_features, + static_cast(grad_output->DataPtr()), out_features, &beta, + static_cast(grad_input->DataPtr()), in_features)); + + // d_weight = d_output.T * input --> d_weight.T = input.T * d_output + // C = d_weight.T[in_features, out_features] + // A = input.T[in_features, bs] + // B = d_output.T[out_features, bs] + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, in_features, out_features, bs, &alpha, + static_cast(input->DataPtr()), in_features, + static_cast(grad_output->DataPtr()), out_features, &beta, + static_cast(grad_weight->DataPtr()), in_features)); + } else { + // weight is [in_features, out_features] here + + // d_input = d_output * weight.T --> d_input.T = weight * d_output.T + // C = d_input.T[in_features, bs] + // A = weight.T[out_features, in_features] + // B = d_output.T[out_features, bs] + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, in_features, bs, out_features, &alpha, + static_cast(weight->DataPtr()), out_features, + static_cast(grad_output->DataPtr()), out_features, &beta, + static_cast(grad_input->DataPtr()), in_features)); + + // d_weight = input.T * d_output --> d_weight.T = d_output.T * input + // C = d_weight.T[out_features, in_features] + // A = d_output.T[out_features, bs] + // B = input.T[in_features, bs] + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, out_features, in_features, bs, &alpha, + static_cast(grad_output->DataPtr()), out_features, + static_cast(input->DataPtr()), in_features, &beta, + static_cast(grad_weight->DataPtr()), out_features)); + } + + // d_bias = \sum_i(i=0, bs-1) d_output[i] + if (bias) { + constexpr int BLOCK_SIZE = 256; + int threads_per_block = BLOCK_SIZE; + int num_blocks = out_features; + ReduceColumnsKernel + <<>>(static_cast(grad_output->DataPtr()), + static_cast(grad_bias->DataPtr()), out_features, bs); + } + + CUBLAS_CHECK(cublasDestroy(handle)); + + return {grad_input, grad_weight, grad_bias}; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_LINEAR_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_LINEAR_KERNEL(MatmulForward) +REGISTER_CUDA_LINEAR_KERNEL(MatmulBackward) +REGISTER_CUDA_LINEAR_KERNEL(LinearForward) +REGISTER_CUDA_LINEAR_KERNEL(LinearBackward) + +#undef REGISTER_CUDA_LINEAR_KERNEL diff --git a/infini_train/src/kernels/cuda/no_op.cu b/infini_train/src/kernels/cuda/no_op.cu index eba88c5..428acb6 100644 --- a/infini_train/src/kernels/cuda/no_op.cu +++ b/infini_train/src/kernels/cuda/no_op.cu @@ -1,30 +1,30 @@ -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { -std::shared_ptr NoOpForward(const std::shared_ptr &input, const std::vector &dims) { - const int64_t num_elements = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); - CHECK_EQ(input->NumElements(), num_elements); - - auto output = std::make_shared(*input, 0, dims); - return output; -} - -std::shared_ptr NoOpBackward(const std::vector &dims, const std::shared_ptr &grad_output) { - auto num_elements = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); - CHECK_EQ(num_elements, grad_output->NumElements()); - - auto grad_input = std::make_shared(*grad_output, 0, dims); - return grad_input; -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_NO_OP_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_NO_OP_KERNEL(NoOpForward) -REGISTER_CUDA_NO_OP_KERNEL(NoOpBackward) - -#undef REGISTER_CUDA_NO_OP_KERNEL +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { +std::shared_ptr NoOpForward(const std::shared_ptr &input, const std::vector &dims) { + const int64_t num_elements = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); + CHECK_EQ(input->NumElements(), num_elements); + + auto output = std::make_shared(*input, 0, dims); + return output; +} + +std::shared_ptr NoOpBackward(const std::vector &dims, const std::shared_ptr &grad_output) { + auto num_elements = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); + CHECK_EQ(num_elements, grad_output->NumElements()); + + auto grad_input = std::make_shared(*grad_output, 0, dims); + return grad_input; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_NO_OP_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_NO_OP_KERNEL(NoOpForward) +REGISTER_CUDA_NO_OP_KERNEL(NoOpBackward) + +#undef REGISTER_CUDA_NO_OP_KERNEL diff --git a/infini_train/src/kernels/cuda/outer.cu b/infini_train/src/kernels/cuda/outer.cu index 2d330bd..5a22c17 100644 --- a/infini_train/src/kernels/cuda/outer.cu +++ b/infini_train/src/kernels/cuda/outer.cu @@ -1,112 +1,112 @@ -#include -#include - -#include "cublas_v2.h" -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { - -#define CUDA_CHECK(call) \ - do { \ - cudaError_t status = call; \ - if (status != cudaSuccess) { \ - LOG(FATAL) << "CUDA Error: " << cudaGetErrorString(status) << " at " << __FILE__ << ":" << __LINE__; \ - } \ - } while (0) - -#define CUBLAS_CHECK(call) \ - do { \ - cublasStatus_t status = call; \ - if (status != CUBLAS_STATUS_SUCCESS) { \ - LOG(FATAL) << "CUBLAS Error: " << cublasGetStatusString(status) << " at " << __FILE__ << ":" << __LINE__; \ - } \ - } while (0) - -std::shared_ptr OuterForward(const std::shared_ptr &input, const std::shared_ptr &other) { - /* - Computes outer product: output[i, j] = input[i] * other[j] - Equivalent to: input: [M, 1], other: [1, N] → output: [M, N] - */ - - const auto &in_dims = input->Dims(); - const auto &ot_dims = other->Dims(); - CHECK_EQ(in_dims.size(), 1); - CHECK_EQ(ot_dims.size(), 1); - - const int64_t M = in_dims[0]; - const int64_t N = ot_dims[0]; - - auto output = std::make_shared(std::vector{M, N}, DataType::kFLOAT32, input->GetDevice()); - - // reinterpret input: [M] as column vector [M, 1] - // reinterpret other: [N] as row vector [1, N] - // output[M, N] = input[M, 1] * other.T[1, N] - // output.T[N, M] = other[N, 1] * input.T[1, M] - float alpha = 1.0f; - float beta = 0.0f; - cublasHandle_t handle; - CUBLAS_CHECK(cublasCreate(&handle)); - - CUBLAS_CHECK(cublasSgemm( - handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, 1, &alpha, static_cast(other->DataPtr()), N, - static_cast(input->DataPtr()), 1, &beta, static_cast(output->DataPtr()), N)); - - CUBLAS_CHECK(cublasDestroy(handle)); - return output; -} - -std::tuple, std::shared_ptr> OuterBackward(const std::shared_ptr &input, - const std::shared_ptr &other, - const std::shared_ptr &grad_output) { - /* - grad_input: [M] = grad_output: [M, N] × other: [N] - grad_other: [N] = grad_output.T: [N, M] × input: [M] - */ - const int64_t M = input->Dims()[0]; - const int64_t N = other->Dims()[0]; - CHECK_EQ(grad_output->Dims().size(), 2); - CHECK_EQ(grad_output->Dims()[0], M); - CHECK_EQ(grad_output->Dims()[1], N); - - auto grad_input = std::make_shared(std::vector{M}, DataType::kFLOAT32, grad_output->GetDevice()); - auto grad_other = std::make_shared(std::vector{N}, DataType::kFLOAT32, grad_output->GetDevice()); - grad_input->Fill(0.0f); - grad_other->Fill(0.0f); - - float alpha = 1.0f; - float beta = 0.0f; - cublasHandle_t handle; - CUBLAS_CHECK(cublasCreate(&handle)); - - // grad_input[M, 1] = grad_output[M, N] × other[N, 1] - // y = grad_input[M] - // A = grad_output.T[N, M] - // x = other[N] - CUBLAS_CHECK(cublasSgemv(handle, CUBLAS_OP_T, N, M, &alpha, static_cast(grad_output->DataPtr()), N, - static_cast(other->DataPtr()), 1, &beta, - static_cast(grad_input->DataPtr()), 1)); - - // grad_other[N, 1] = grad_output.T[N, M] × input[M, 1] - // y = grad_other[N] - // A = grad_output.T[N, M] - // x = input[M] - CUBLAS_CHECK(cublasSgemv(handle, CUBLAS_OP_N, N, M, &alpha, static_cast(grad_output->DataPtr()), N, - static_cast(input->DataPtr()), 1, &beta, - static_cast(grad_other->DataPtr()), 1)); - - CUBLAS_CHECK(cublasDestroy(handle)); - return {grad_input, grad_other}; -} - -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_OUTER_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_OUTER_KERNEL(OuterForward) -REGISTER_CUDA_OUTER_KERNEL(OuterBackward) - -#undef REGISTER_CUDA_OUTER_KERNEL +#include +#include + +#include "cublas_v2.h" +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { + +#define CUDA_CHECK(call) \ + do { \ + cudaError_t status = call; \ + if (status != cudaSuccess) { \ + LOG(FATAL) << "CUDA Error: " << cudaGetErrorString(status) << " at " << __FILE__ << ":" << __LINE__; \ + } \ + } while (0) + +#define CUBLAS_CHECK(call) \ + do { \ + cublasStatus_t status = call; \ + if (status != CUBLAS_STATUS_SUCCESS) { \ + LOG(FATAL) << "CUBLAS Error: " << cublasGetStatusString(status) << " at " << __FILE__ << ":" << __LINE__; \ + } \ + } while (0) + +std::shared_ptr OuterForward(const std::shared_ptr &input, const std::shared_ptr &other) { + /* + Computes outer product: output[i, j] = input[i] * other[j] + Equivalent to: input: [M, 1], other: [1, N] → output: [M, N] + */ + + const auto &in_dims = input->Dims(); + const auto &ot_dims = other->Dims(); + CHECK_EQ(in_dims.size(), 1); + CHECK_EQ(ot_dims.size(), 1); + + const int64_t M = in_dims[0]; + const int64_t N = ot_dims[0]; + + auto output = std::make_shared(std::vector{M, N}, DataType::kFLOAT32, input->GetDevice()); + + // reinterpret input: [M] as column vector [M, 1] + // reinterpret other: [N] as row vector [1, N] + // output[M, N] = input[M, 1] * other.T[1, N] + // output.T[N, M] = other[N, 1] * input.T[1, M] + float alpha = 1.0f; + float beta = 0.0f; + cublasHandle_t handle; + CUBLAS_CHECK(cublasCreate(&handle)); + + CUBLAS_CHECK(cublasSgemm( + handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, 1, &alpha, static_cast(other->DataPtr()), N, + static_cast(input->DataPtr()), 1, &beta, static_cast(output->DataPtr()), N)); + + CUBLAS_CHECK(cublasDestroy(handle)); + return output; +} + +std::tuple, std::shared_ptr> OuterBackward(const std::shared_ptr &input, + const std::shared_ptr &other, + const std::shared_ptr &grad_output) { + /* + grad_input: [M] = grad_output: [M, N] × other: [N] + grad_other: [N] = grad_output.T: [N, M] × input: [M] + */ + const int64_t M = input->Dims()[0]; + const int64_t N = other->Dims()[0]; + CHECK_EQ(grad_output->Dims().size(), 2); + CHECK_EQ(grad_output->Dims()[0], M); + CHECK_EQ(grad_output->Dims()[1], N); + + auto grad_input = std::make_shared(std::vector{M}, DataType::kFLOAT32, grad_output->GetDevice()); + auto grad_other = std::make_shared(std::vector{N}, DataType::kFLOAT32, grad_output->GetDevice()); + grad_input->Fill(0.0f); + grad_other->Fill(0.0f); + + float alpha = 1.0f; + float beta = 0.0f; + cublasHandle_t handle; + CUBLAS_CHECK(cublasCreate(&handle)); + + // grad_input[M, 1] = grad_output[M, N] × other[N, 1] + // y = grad_input[M] + // A = grad_output.T[N, M] + // x = other[N] + CUBLAS_CHECK(cublasSgemv(handle, CUBLAS_OP_T, N, M, &alpha, static_cast(grad_output->DataPtr()), N, + static_cast(other->DataPtr()), 1, &beta, + static_cast(grad_input->DataPtr()), 1)); + + // grad_other[N, 1] = grad_output.T[N, M] × input[M, 1] + // y = grad_other[N] + // A = grad_output.T[N, M] + // x = input[M] + CUBLAS_CHECK(cublasSgemv(handle, CUBLAS_OP_N, N, M, &alpha, static_cast(grad_output->DataPtr()), N, + static_cast(input->DataPtr()), 1, &beta, + static_cast(grad_other->DataPtr()), 1)); + + CUBLAS_CHECK(cublasDestroy(handle)); + return {grad_input, grad_other}; +} + +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_OUTER_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_OUTER_KERNEL(OuterForward) +REGISTER_CUDA_OUTER_KERNEL(OuterBackward) + +#undef REGISTER_CUDA_OUTER_KERNEL diff --git a/infini_train/src/kernels/cuda/reduction.cu b/infini_train/src/kernels/cuda/reduction.cu index e85e113..4e61dd4 100644 --- a/infini_train/src/kernels/cuda/reduction.cu +++ b/infini_train/src/kernels/cuda/reduction.cu @@ -1,230 +1,230 @@ -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { -namespace { -constexpr float kInfinity = std::numeric_limits::infinity(); -} // namespace - -namespace { - -#if defined(CUB_VERSION) && CUB_VERSION >= 200800 - #include - using CubSumOp = ::cuda::std::plus<>; - using CubMaxOp = ::cuda::maximum<>; - using CubMinOp = ::cuda::minimum<>; -#else - using CubSumOp = cub::Sum; - using CubMaxOp = cub::Max; - using CubMinOp = cub::Min; -#endif - -// Reduction operators -template struct CubOp; - -template <> struct CubOp { - __device__ static float Init() { return 0.0f; } - __device__ static float Reduce(float a, float b) { return a + b; } - __device__ static CubSumOp Op() { return CubSumOp(); } -}; - -template <> struct CubOp { - __device__ static float Init() { return -kInfinity; } - __device__ static float Reduce(float a, float b) { return fmaxf(a, b); } - __device__ static CubMaxOp Op() { return CubMaxOp(); } -}; - -template <> struct CubOp { - __device__ static float Init() { return kInfinity; } - __device__ static float Reduce(float a, float b) { return fminf(a, b); } - __device__ static CubMinOp Op() { return CubMinOp(); } -}; - -// Finalization strategies -struct MeanFinalize { - __device__ __forceinline__ float operator()(float sum, int64_t count) const { - return sum / static_cast(count); - } -}; - -struct IdentityFinalize { - __device__ __forceinline__ float operator()(float val, int64_t) const { return val; } -}; - -// Generic reduction kernel -template -__global__ void GenericReduceKernel(const float *input, float *output, int64_t N, int64_t H, int64_t W, - FinalizeOp finalize_op) { - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - - int idx = blockIdx.x; - if (idx >= N * W) { - return; - } - - int n = idx / W; - int w = idx % W; - - float acc = CubOp::Init(); - for (int64_t h = threadIdx.x; h < H; h += blockDim.x) { - int input_idx = (n * H + h) * W + w; - acc = CubOp::Reduce(acc, input[input_idx]); - } - - float reduced = BlockReduce(temp_storage).Reduce(acc, CubOp::Op()); - - if (threadIdx.x == 0) { - output[idx] = finalize_op(reduced, H); - } -} - -// Unified backward kernel for Mean, Sum, Max, and Min -__global__ void GenericReduceBackwardKernel(float *grad_input, const float *grad_output, const float *input, - const float *reduced, int64_t N, int64_t H, int64_t W, bool is_mean, - bool is_masked) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= N * H * W) { - return; - } - - int n = idx / (H * W); - int hw = idx % (H * W); - int w = hw % W; - - int reduced_idx = n * W + w; - - if (is_masked) { - float selected = reduced[reduced_idx]; - float value = input[idx]; - grad_input[idx] = (value == selected) ? grad_output[reduced_idx] : 0.0f; - } else { - grad_input[idx] = grad_output[reduced_idx]; - if (is_mean) { - grad_input[idx] /= static_cast(H); - } - } -} -} // namespace - -// Common forward implementation for reduce ops -template -std::shared_ptr ReduceOpForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim, - FinalizeOp finalize_op) { - const auto &input_dims = input->Dims(); - int64_t actual_dim = dim < 0 ? dim + input_dims.size() : dim; - CHECK_GE(actual_dim, 0); - CHECK_LT(actual_dim, input_dims.size()); - - std::vector output_dims = input_dims; - if (keep_dim) { - output_dims[actual_dim] = 1; - } else { - output_dims.erase(output_dims.begin() + actual_dim); - } - - auto output = std::make_shared(output_dims, DataType::kFLOAT32, input->GetDevice()); - - int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + actual_dim, 1, std::multiplies()); - int64_t H = input_dims[actual_dim]; - int64_t W = std::accumulate(input_dims.begin() + actual_dim + 1, input_dims.end(), 1, std::multiplies()); - - const float *input_ptr = static_cast(input->DataPtr()); - float *output_ptr = static_cast(output->DataPtr()); - - constexpr int BLOCK_SIZE = 256; - int threads_per_block = BLOCK_SIZE; - int num_blocks = N * W; - - GenericReduceKernel - <<>>(input_ptr, output_ptr, N, H, W, finalize_op); - - return output; -} - -// Common backward implementation for reduce ops -std::shared_ptr ReduceOpBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &input, const std::shared_ptr &reduced, - const std::vector &input_dims, const int64_t dim, bool keep_dim, - bool is_mean, bool is_masked) { - int64_t actual_dim = dim < 0 ? dim + input_dims.size() : dim; - CHECK_GE(actual_dim, 0); - CHECK_LT(actual_dim, input_dims.size()); - - auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, grad_output->GetDevice()); - grad_input->Fill(0.0f); - - int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + actual_dim, 1, std::multiplies()); - int64_t H = input_dims[actual_dim]; - int64_t W = std::accumulate(input_dims.begin() + actual_dim + 1, input_dims.end(), 1, std::multiplies()); - - float *grad_input_ptr = static_cast(grad_input->DataPtr()); - const float *grad_output_ptr = static_cast(grad_output->DataPtr()); - const float *input_ptr = input ? static_cast(input->DataPtr()) : nullptr; - const float *reduced_ptr = reduced ? static_cast(reduced->DataPtr()) : nullptr; - - int threads_per_block = 256; - int num_blocks = (N * H * W + threads_per_block - 1) / threads_per_block; - - GenericReduceBackwardKernel<<>>(grad_input_ptr, grad_output_ptr, input_ptr, - reduced_ptr, N, H, W, is_mean, is_masked); - - return grad_input; -} - -std::shared_ptr MeanForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { - return ReduceOpForward(input, dim, keep_dim, MeanFinalize{}); -} - -std::shared_ptr SumForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { - return ReduceOpForward(input, dim, keep_dim, IdentityFinalize{}); -} - -std::shared_ptr MaxForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { - return ReduceOpForward(input, dim, keep_dim, IdentityFinalize{}); -} - -std::shared_ptr MinForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { - return ReduceOpForward(input, dim, keep_dim, IdentityFinalize{}); -} - -std::shared_ptr MeanBackward(const std::shared_ptr &grad_output, const std::vector &input_dims, - const int64_t dim, bool keep_dim) { - return ReduceOpBackward(grad_output, nullptr, nullptr, input_dims, dim, keep_dim, true, false); -} - -std::shared_ptr SumBackward(const std::shared_ptr &grad_output, const std::vector &input_dims, - const int64_t dim, bool keep_dim) { - return ReduceOpBackward(grad_output, nullptr, nullptr, input_dims, dim, keep_dim, false, false); -} - -std::shared_ptr MaxBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, - const std::shared_ptr &reduced, const int64_t dim, bool keep_dim) { - return ReduceOpBackward(grad_output, input, reduced, input->Dims(), dim, keep_dim, false, true); -} - -std::shared_ptr MinBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, - const std::shared_ptr &reduced, const int64_t dim, bool keep_dim) { - return ReduceOpBackward(grad_output, input, reduced, input->Dims(), dim, keep_dim, false, true); -} - -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_REDUCTION_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_REDUCTION_KERNEL(MeanForward) -REGISTER_CUDA_REDUCTION_KERNEL(SumForward) -REGISTER_CUDA_REDUCTION_KERNEL(MaxForward) -REGISTER_CUDA_REDUCTION_KERNEL(MinForward) -REGISTER_CUDA_REDUCTION_KERNEL(MeanBackward) -REGISTER_CUDA_REDUCTION_KERNEL(SumBackward) -REGISTER_CUDA_REDUCTION_KERNEL(MaxBackward) -REGISTER_CUDA_REDUCTION_KERNEL(MinBackward) - -#undef REGISTER_CUDA_REDUCTION_KERNEL +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { +namespace { +constexpr float kInfinity = std::numeric_limits::infinity(); +} // namespace + +namespace { + +#if defined(CUB_VERSION) && CUB_VERSION >= 200800 + #include + using CubSumOp = ::cuda::std::plus<>; + using CubMaxOp = ::cuda::maximum<>; + using CubMinOp = ::cuda::minimum<>; +#else + using CubSumOp = cub::Sum; + using CubMaxOp = cub::Max; + using CubMinOp = cub::Min; +#endif + +// Reduction operators +template struct CubOp; + +template <> struct CubOp { + __device__ static float Init() { return 0.0f; } + __device__ static float Reduce(float a, float b) { return a + b; } + __device__ static CubSumOp Op() { return CubSumOp(); } +}; + +template <> struct CubOp { + __device__ static float Init() { return -kInfinity; } + __device__ static float Reduce(float a, float b) { return fmaxf(a, b); } + __device__ static CubMaxOp Op() { return CubMaxOp(); } +}; + +template <> struct CubOp { + __device__ static float Init() { return kInfinity; } + __device__ static float Reduce(float a, float b) { return fminf(a, b); } + __device__ static CubMinOp Op() { return CubMinOp(); } +}; + +// Finalization strategies +struct MeanFinalize { + __device__ __forceinline__ float operator()(float sum, int64_t count) const { + return sum / static_cast(count); + } +}; + +struct IdentityFinalize { + __device__ __forceinline__ float operator()(float val, int64_t) const { return val; } +}; + +// Generic reduction kernel +template +__global__ void GenericReduceKernel(const float *input, float *output, int64_t N, int64_t H, int64_t W, + FinalizeOp finalize_op) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + int idx = blockIdx.x; + if (idx >= N * W) { + return; + } + + int n = idx / W; + int w = idx % W; + + float acc = CubOp::Init(); + for (int64_t h = threadIdx.x; h < H; h += blockDim.x) { + int input_idx = (n * H + h) * W + w; + acc = CubOp::Reduce(acc, input[input_idx]); + } + + float reduced = BlockReduce(temp_storage).Reduce(acc, CubOp::Op()); + + if (threadIdx.x == 0) { + output[idx] = finalize_op(reduced, H); + } +} + +// Unified backward kernel for Mean, Sum, Max, and Min +__global__ void GenericReduceBackwardKernel(float *grad_input, const float *grad_output, const float *input, + const float *reduced, int64_t N, int64_t H, int64_t W, bool is_mean, + bool is_masked) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * H * W) { + return; + } + + int n = idx / (H * W); + int hw = idx % (H * W); + int w = hw % W; + + int reduced_idx = n * W + w; + + if (is_masked) { + float selected = reduced[reduced_idx]; + float value = input[idx]; + grad_input[idx] = (value == selected) ? grad_output[reduced_idx] : 0.0f; + } else { + grad_input[idx] = grad_output[reduced_idx]; + if (is_mean) { + grad_input[idx] /= static_cast(H); + } + } +} +} // namespace + +// Common forward implementation for reduce ops +template +std::shared_ptr ReduceOpForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim, + FinalizeOp finalize_op) { + const auto &input_dims = input->Dims(); + int64_t actual_dim = dim < 0 ? dim + input_dims.size() : dim; + CHECK_GE(actual_dim, 0); + CHECK_LT(actual_dim, input_dims.size()); + + std::vector output_dims = input_dims; + if (keep_dim) { + output_dims[actual_dim] = 1; + } else { + output_dims.erase(output_dims.begin() + actual_dim); + } + + auto output = std::make_shared(output_dims, DataType::kFLOAT32, input->GetDevice()); + + int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + actual_dim, 1, std::multiplies()); + int64_t H = input_dims[actual_dim]; + int64_t W = std::accumulate(input_dims.begin() + actual_dim + 1, input_dims.end(), 1, std::multiplies()); + + const float *input_ptr = static_cast(input->DataPtr()); + float *output_ptr = static_cast(output->DataPtr()); + + constexpr int BLOCK_SIZE = 256; + int threads_per_block = BLOCK_SIZE; + int num_blocks = N * W; + + GenericReduceKernel + <<>>(input_ptr, output_ptr, N, H, W, finalize_op); + + return output; +} + +// Common backward implementation for reduce ops +std::shared_ptr ReduceOpBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &input, const std::shared_ptr &reduced, + const std::vector &input_dims, const int64_t dim, bool keep_dim, + bool is_mean, bool is_masked) { + int64_t actual_dim = dim < 0 ? dim + input_dims.size() : dim; + CHECK_GE(actual_dim, 0); + CHECK_LT(actual_dim, input_dims.size()); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, grad_output->GetDevice()); + grad_input->Fill(0.0f); + + int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + actual_dim, 1, std::multiplies()); + int64_t H = input_dims[actual_dim]; + int64_t W = std::accumulate(input_dims.begin() + actual_dim + 1, input_dims.end(), 1, std::multiplies()); + + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + const float *input_ptr = input ? static_cast(input->DataPtr()) : nullptr; + const float *reduced_ptr = reduced ? static_cast(reduced->DataPtr()) : nullptr; + + int threads_per_block = 256; + int num_blocks = (N * H * W + threads_per_block - 1) / threads_per_block; + + GenericReduceBackwardKernel<<>>(grad_input_ptr, grad_output_ptr, input_ptr, + reduced_ptr, N, H, W, is_mean, is_masked); + + return grad_input; +} + +std::shared_ptr MeanForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + return ReduceOpForward(input, dim, keep_dim, MeanFinalize{}); +} + +std::shared_ptr SumForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + return ReduceOpForward(input, dim, keep_dim, IdentityFinalize{}); +} + +std::shared_ptr MaxForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + return ReduceOpForward(input, dim, keep_dim, IdentityFinalize{}); +} + +std::shared_ptr MinForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + return ReduceOpForward(input, dim, keep_dim, IdentityFinalize{}); +} + +std::shared_ptr MeanBackward(const std::shared_ptr &grad_output, const std::vector &input_dims, + const int64_t dim, bool keep_dim) { + return ReduceOpBackward(grad_output, nullptr, nullptr, input_dims, dim, keep_dim, true, false); +} + +std::shared_ptr SumBackward(const std::shared_ptr &grad_output, const std::vector &input_dims, + const int64_t dim, bool keep_dim) { + return ReduceOpBackward(grad_output, nullptr, nullptr, input_dims, dim, keep_dim, false, false); +} + +std::shared_ptr MaxBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + const std::shared_ptr &reduced, const int64_t dim, bool keep_dim) { + return ReduceOpBackward(grad_output, input, reduced, input->Dims(), dim, keep_dim, false, true); +} + +std::shared_ptr MinBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + const std::shared_ptr &reduced, const int64_t dim, bool keep_dim) { + return ReduceOpBackward(grad_output, input, reduced, input->Dims(), dim, keep_dim, false, true); +} + +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_REDUCTION_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_REDUCTION_KERNEL(MeanForward) +REGISTER_CUDA_REDUCTION_KERNEL(SumForward) +REGISTER_CUDA_REDUCTION_KERNEL(MaxForward) +REGISTER_CUDA_REDUCTION_KERNEL(MinForward) +REGISTER_CUDA_REDUCTION_KERNEL(MeanBackward) +REGISTER_CUDA_REDUCTION_KERNEL(SumBackward) +REGISTER_CUDA_REDUCTION_KERNEL(MaxBackward) +REGISTER_CUDA_REDUCTION_KERNEL(MinBackward) + +#undef REGISTER_CUDA_REDUCTION_KERNEL diff --git a/infini_train/src/kernels/cuda/sigmoid.cu b/infini_train/src/kernels/cuda/sigmoid.cu index 3f32447..a77aa7d 100644 --- a/infini_train/src/kernels/cuda/sigmoid.cu +++ b/infini_train/src/kernels/cuda/sigmoid.cu @@ -1,66 +1,66 @@ -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { -__global__ void SigmoidForwardKernel(const float *input_ptr, float *output_ptr, size_t num_elements) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx < num_elements) { - output_ptr[idx] = 1.0f / (1.0f + exp(-input_ptr[idx])); - } -} - -std::shared_ptr SigmoidForward(const std::shared_ptr &input) { - size_t num_elements = input->NumElements(); - - auto output = std::make_shared(input->Dims(), DataType::kFLOAT32, input->GetDevice()); - - const float *input_ptr = static_cast(input->DataPtr()); - float *output_ptr = static_cast(output->DataPtr()); - - int threads_per_block = 256; - int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; - - SigmoidForwardKernel<<>>(input_ptr, output_ptr, num_elements); - - return output; -} - -__global__ void SigmoidBackwardKernel(const float *output_ptr, const float *grad_output_ptr, float *grad_input_ptr, - size_t num_elements) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx < num_elements) { - grad_input_ptr[idx] = grad_output_ptr[idx] * output_ptr[idx] * (1.0f - output_ptr[idx]); - } -} - -std::shared_ptr SigmoidBackward(const std::shared_ptr &output, - const std::shared_ptr &grad_output) { - size_t num_elements = output->NumElements(); - - auto grad_input = std::make_shared(output->Dims(), DataType::kFLOAT32, grad_output->GetDevice()); - grad_input->Fill(0.0f); - const float *output_ptr = static_cast(output->DataPtr()); - const float *grad_output_ptr = static_cast(grad_output->DataPtr()); - float *grad_input_ptr = static_cast(grad_input->DataPtr()); - - int threads_per_block = 256; - int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; - - SigmoidBackwardKernel<<>>(output_ptr, grad_output_ptr, grad_input_ptr, num_elements); - - return grad_input; -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_SIGMOID_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_SIGMOID_KERNEL(SigmoidForward) -REGISTER_CUDA_SIGMOID_KERNEL(SigmoidBackward) - -#undef REGISTER_CUDA_SIGMOID_KERNEL +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { +__global__ void SigmoidForwardKernel(const float *input_ptr, float *output_ptr, size_t num_elements) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + output_ptr[idx] = 1.0f / (1.0f + exp(-input_ptr[idx])); + } +} + +std::shared_ptr SigmoidForward(const std::shared_ptr &input) { + size_t num_elements = input->NumElements(); + + auto output = std::make_shared(input->Dims(), DataType::kFLOAT32, input->GetDevice()); + + const float *input_ptr = static_cast(input->DataPtr()); + float *output_ptr = static_cast(output->DataPtr()); + + int threads_per_block = 256; + int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; + + SigmoidForwardKernel<<>>(input_ptr, output_ptr, num_elements); + + return output; +} + +__global__ void SigmoidBackwardKernel(const float *output_ptr, const float *grad_output_ptr, float *grad_input_ptr, + size_t num_elements) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + grad_input_ptr[idx] = grad_output_ptr[idx] * output_ptr[idx] * (1.0f - output_ptr[idx]); + } +} + +std::shared_ptr SigmoidBackward(const std::shared_ptr &output, + const std::shared_ptr &grad_output) { + size_t num_elements = output->NumElements(); + + auto grad_input = std::make_shared(output->Dims(), DataType::kFLOAT32, grad_output->GetDevice()); + grad_input->Fill(0.0f); + const float *output_ptr = static_cast(output->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + + int threads_per_block = 256; + int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; + + SigmoidBackwardKernel<<>>(output_ptr, grad_output_ptr, grad_input_ptr, num_elements); + + return grad_input; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_SIGMOID_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_SIGMOID_KERNEL(SigmoidForward) +REGISTER_CUDA_SIGMOID_KERNEL(SigmoidBackward) + +#undef REGISTER_CUDA_SIGMOID_KERNEL diff --git a/infini_train/src/kernels/cuda/slice.cu b/infini_train/src/kernels/cuda/slice.cu index 61ea478..52d38bd 100644 --- a/infini_train/src/kernels/cuda/slice.cu +++ b/infini_train/src/kernels/cuda/slice.cu @@ -1,173 +1,173 @@ -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { -__global__ void SliceForwardKernel(const float *input, float *output, const int64_t *new_dims, const int64_t *starts, - const int64_t *steps, const int64_t *in_strides, const int64_t *out_strides, - int num_dims, int64_t total_elements) { - int64_t out_idx = blockIdx.x * blockDim.x + threadIdx.x; - if (out_idx >= total_elements) { - return; - } - - int64_t in_index = 0; - for (int i = 0; i < num_dims; ++i) { - int64_t idx = (out_idx / out_strides[i]) % new_dims[i]; - in_index += (starts[i] + idx * steps[i]) * in_strides[i]; - } - - output[out_idx] = input[in_index]; -} - -std::shared_ptr SliceForward(const std::shared_ptr &input, const std::vector &starts, - const std::vector &ends, const std::vector &steps) { - CHECK_EQ(starts.size(), ends.size()); - CHECK_EQ(starts.size(), steps.size()); - auto &dims = input->Dims(); - CHECK_EQ(starts.size(), dims.size()); - const int64_t num_dims = dims.size(); - - std::vector new_dims; - for (int i = 0; i < starts.size(); i++) { - CHECK_LE(starts[i], ends[i]); - CHECK_LE(0, steps[i]); - new_dims.push_back((ends[i] - starts[i] + steps[i] - 1) / steps[i]); - } - - auto new_tensor = std::make_shared(new_dims, input->Dtype(), input->GetDevice()); - // NOTE(zbl): must initialize with 0 - new_tensor->Fill(0.0f); - - std::vector src_strides(dims.size(), 0), dst_strides(new_dims.size(), 0); - int64_t stride = 1; - for (int i = dims.size() - 1; i >= 0; --i) { - src_strides[i] = stride; - stride *= dims[i]; - } - - stride = 1; - for (int i = new_dims.size() - 1; i >= 0; --i) { - dst_strides[i] = stride; - stride *= new_dims[i]; - } - - int64_t total_elements = stride; - - int64_t *new_dims_dev, *starts_dev, *steps_dev, *input_strides_dev, *output_strides_dev; - - cudaMallocAsync(&new_dims_dev, - (ends.size() + starts.size() + steps.size() + dims.size() + new_dims.size()) * sizeof(int64_t), 0); - starts_dev = new_dims_dev + ends.size(); - steps_dev = starts_dev + starts.size(); - input_strides_dev = steps_dev + steps.size(); - output_strides_dev = input_strides_dev + dims.size(); - - cudaMemcpyAsync(new_dims_dev, new_dims.data(), ends.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); - cudaMemcpyAsync(starts_dev, starts.data(), starts.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); - cudaMemcpyAsync(steps_dev, steps.data(), steps.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); - cudaMemcpyAsync(input_strides_dev, src_strides.data(), dims.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); - cudaMemcpyAsync(output_strides_dev, dst_strides.data(), new_dims.size() * sizeof(int64_t), cudaMemcpyHostToDevice, - 0); - - int threads_per_block = 256; - int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; - - SliceForwardKernel<<>>( - static_cast(input->DataPtr()), static_cast(new_tensor->DataPtr()), new_dims_dev, - starts_dev, steps_dev, input_strides_dev, output_strides_dev, num_dims, total_elements); - - cudaFreeAsync(new_dims_dev, 0); - - return new_tensor; -} - -__global__ void SliceBackwardKernel(const float *grad_output, float *grad_input, const int64_t *new_dims, - const int64_t *starts, const int64_t *steps, const int64_t *in_strides, - const int64_t *out_strides, int num_dims, int64_t total_elements) { - int64_t out_idx = blockIdx.x * blockDim.x + threadIdx.x; - if (out_idx >= total_elements) { - return; - } - - int64_t in_index = 0; - for (int i = 0; i < num_dims; ++i) { - int64_t idx = (out_idx / out_strides[i]) % new_dims[i]; - in_index += (starts[i] + idx * steps[i]) * in_strides[i]; - } - atomicAdd(&grad_input[in_index], grad_output[out_idx]); -} - -std::shared_ptr SliceBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, - const std::vector &starts, const std::vector &ends, - const std::vector &steps) { - CHECK_EQ(starts.size(), ends.size()); - CHECK_EQ(starts.size(), steps.size()); - auto &dims = input->Dims(); - CHECK_EQ(starts.size(), dims.size()); - const int64_t num_dims = dims.size(); - - std::vector new_dims; - for (int i = 0; i < starts.size(); i++) { - CHECK_LE(starts[i], ends[i]); - CHECK_LE(0, steps[i]); - new_dims.push_back((ends[i] - starts[i] + steps[i] - 1) / steps[i]); - } - - auto grad_input = std::make_shared(input->Dims(), input->Dtype(), grad_output->GetDevice()); - grad_input->Fill(0.0); - - std::vector src_strides(dims.size()); - int64_t stride = 1; - for (int i = src_strides.size() - 1; i >= 0; --i) { - src_strides[i] = stride; - stride *= dims[i]; - } - - std::vector dst_strides(new_dims.size()); - stride = 1; - for (int i = dst_strides.size() - 1; i >= 0; --i) { - dst_strides[i] = stride; - stride *= new_dims[i]; - } - - int64_t total_elements = stride; - - int dims_size = dims.size(); - int64_t *new_dims_dev, *starts_dev, *steps_dev, *input_strides_dev, *output_strides_dev; - - cudaMallocAsync(&new_dims_dev, - (ends.size() + starts.size() + steps.size() + dims.size() + new_dims.size()) * sizeof(int64_t), 0); - starts_dev = new_dims_dev + ends.size(); - steps_dev = starts_dev + starts.size(); - input_strides_dev = steps_dev + steps.size(); - output_strides_dev = input_strides_dev + dims.size(); - - cudaMemcpyAsync(new_dims_dev, new_dims.data(), ends.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); - cudaMemcpyAsync(starts_dev, starts.data(), starts.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); - cudaMemcpyAsync(steps_dev, steps.data(), steps.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); - cudaMemcpyAsync(input_strides_dev, src_strides.data(), dims.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); - cudaMemcpyAsync(output_strides_dev, dst_strides.data(), new_dims.size() * sizeof(int64_t), cudaMemcpyHostToDevice, - 0); - - int threads_per_block = 256; - int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; - - SliceBackwardKernel<<>>( - static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), new_dims_dev, - starts_dev, steps_dev, input_strides_dev, output_strides_dev, num_dims, total_elements); - - cudaFreeAsync(new_dims_dev, 0); - - return grad_input; -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_SLICE_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_SLICE_KERNEL(SliceForward) -REGISTER_CUDA_SLICE_KERNEL(SliceBackward) - -#undef REGISTER_CUDA_SLICE_KERNEL +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { +__global__ void SliceForwardKernel(const float *input, float *output, const int64_t *new_dims, const int64_t *starts, + const int64_t *steps, const int64_t *in_strides, const int64_t *out_strides, + int num_dims, int64_t total_elements) { + int64_t out_idx = blockIdx.x * blockDim.x + threadIdx.x; + if (out_idx >= total_elements) { + return; + } + + int64_t in_index = 0; + for (int i = 0; i < num_dims; ++i) { + int64_t idx = (out_idx / out_strides[i]) % new_dims[i]; + in_index += (starts[i] + idx * steps[i]) * in_strides[i]; + } + + output[out_idx] = input[in_index]; +} + +std::shared_ptr SliceForward(const std::shared_ptr &input, const std::vector &starts, + const std::vector &ends, const std::vector &steps) { + CHECK_EQ(starts.size(), ends.size()); + CHECK_EQ(starts.size(), steps.size()); + auto &dims = input->Dims(); + CHECK_EQ(starts.size(), dims.size()); + const int64_t num_dims = dims.size(); + + std::vector new_dims; + for (int i = 0; i < starts.size(); i++) { + CHECK_LE(starts[i], ends[i]); + CHECK_LE(0, steps[i]); + new_dims.push_back((ends[i] - starts[i] + steps[i] - 1) / steps[i]); + } + + auto new_tensor = std::make_shared(new_dims, input->Dtype(), input->GetDevice()); + // NOTE(zbl): must initialize with 0 + new_tensor->Fill(0.0f); + + std::vector src_strides(dims.size(), 0), dst_strides(new_dims.size(), 0); + int64_t stride = 1; + for (int i = dims.size() - 1; i >= 0; --i) { + src_strides[i] = stride; + stride *= dims[i]; + } + + stride = 1; + for (int i = new_dims.size() - 1; i >= 0; --i) { + dst_strides[i] = stride; + stride *= new_dims[i]; + } + + int64_t total_elements = stride; + + int64_t *new_dims_dev, *starts_dev, *steps_dev, *input_strides_dev, *output_strides_dev; + + cudaMallocAsync(&new_dims_dev, + (ends.size() + starts.size() + steps.size() + dims.size() + new_dims.size()) * sizeof(int64_t), 0); + starts_dev = new_dims_dev + ends.size(); + steps_dev = starts_dev + starts.size(); + input_strides_dev = steps_dev + steps.size(); + output_strides_dev = input_strides_dev + dims.size(); + + cudaMemcpyAsync(new_dims_dev, new_dims.data(), ends.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); + cudaMemcpyAsync(starts_dev, starts.data(), starts.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); + cudaMemcpyAsync(steps_dev, steps.data(), steps.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); + cudaMemcpyAsync(input_strides_dev, src_strides.data(), dims.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); + cudaMemcpyAsync(output_strides_dev, dst_strides.data(), new_dims.size() * sizeof(int64_t), cudaMemcpyHostToDevice, + 0); + + int threads_per_block = 256; + int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; + + SliceForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(new_tensor->DataPtr()), new_dims_dev, + starts_dev, steps_dev, input_strides_dev, output_strides_dev, num_dims, total_elements); + + cudaFreeAsync(new_dims_dev, 0); + + return new_tensor; +} + +__global__ void SliceBackwardKernel(const float *grad_output, float *grad_input, const int64_t *new_dims, + const int64_t *starts, const int64_t *steps, const int64_t *in_strides, + const int64_t *out_strides, int num_dims, int64_t total_elements) { + int64_t out_idx = blockIdx.x * blockDim.x + threadIdx.x; + if (out_idx >= total_elements) { + return; + } + + int64_t in_index = 0; + for (int i = 0; i < num_dims; ++i) { + int64_t idx = (out_idx / out_strides[i]) % new_dims[i]; + in_index += (starts[i] + idx * steps[i]) * in_strides[i]; + } + atomicAdd(&grad_input[in_index], grad_output[out_idx]); +} + +std::shared_ptr SliceBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + const std::vector &starts, const std::vector &ends, + const std::vector &steps) { + CHECK_EQ(starts.size(), ends.size()); + CHECK_EQ(starts.size(), steps.size()); + auto &dims = input->Dims(); + CHECK_EQ(starts.size(), dims.size()); + const int64_t num_dims = dims.size(); + + std::vector new_dims; + for (int i = 0; i < starts.size(); i++) { + CHECK_LE(starts[i], ends[i]); + CHECK_LE(0, steps[i]); + new_dims.push_back((ends[i] - starts[i] + steps[i] - 1) / steps[i]); + } + + auto grad_input = std::make_shared(input->Dims(), input->Dtype(), grad_output->GetDevice()); + grad_input->Fill(0.0); + + std::vector src_strides(dims.size()); + int64_t stride = 1; + for (int i = src_strides.size() - 1; i >= 0; --i) { + src_strides[i] = stride; + stride *= dims[i]; + } + + std::vector dst_strides(new_dims.size()); + stride = 1; + for (int i = dst_strides.size() - 1; i >= 0; --i) { + dst_strides[i] = stride; + stride *= new_dims[i]; + } + + int64_t total_elements = stride; + + int dims_size = dims.size(); + int64_t *new_dims_dev, *starts_dev, *steps_dev, *input_strides_dev, *output_strides_dev; + + cudaMallocAsync(&new_dims_dev, + (ends.size() + starts.size() + steps.size() + dims.size() + new_dims.size()) * sizeof(int64_t), 0); + starts_dev = new_dims_dev + ends.size(); + steps_dev = starts_dev + starts.size(); + input_strides_dev = steps_dev + steps.size(); + output_strides_dev = input_strides_dev + dims.size(); + + cudaMemcpyAsync(new_dims_dev, new_dims.data(), ends.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); + cudaMemcpyAsync(starts_dev, starts.data(), starts.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); + cudaMemcpyAsync(steps_dev, steps.data(), steps.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); + cudaMemcpyAsync(input_strides_dev, src_strides.data(), dims.size() * sizeof(int64_t), cudaMemcpyHostToDevice, 0); + cudaMemcpyAsync(output_strides_dev, dst_strides.data(), new_dims.size() * sizeof(int64_t), cudaMemcpyHostToDevice, + 0); + + int threads_per_block = 256; + int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; + + SliceBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), new_dims_dev, + starts_dev, steps_dev, input_strides_dev, output_strides_dev, num_dims, total_elements); + + cudaFreeAsync(new_dims_dev, 0); + + return grad_input; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_SLICE_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_SLICE_KERNEL(SliceForward) +REGISTER_CUDA_SLICE_KERNEL(SliceBackward) + +#undef REGISTER_CUDA_SLICE_KERNEL diff --git a/infini_train/src/kernels/cuda/softmax.cu b/infini_train/src/kernels/cuda/softmax.cu index 0f266d0..1151fd9 100644 --- a/infini_train/src/kernels/cuda/softmax.cu +++ b/infini_train/src/kernels/cuda/softmax.cu @@ -1,215 +1,215 @@ -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { - -#if defined(CUB_VERSION) && CUB_VERSION >= 200800 - #include - using CubSumOp = ::cuda::std::plus<>; - using CubMaxOp = ::cuda::maximum<>; - using CubMinOp = ::cuda::minimum<>; -#else - using CubSumOp = cub::Sum; - using CubMaxOp = cub::Max; - using CubMinOp = cub::Min; -#endif - -template -__global__ void SoftmaxForwardKernel(T *output, const T *input, int64_t outer_size, int64_t axis_size, - int64_t inner_size) { - using BlockReduce = cub::BlockReduce; - - __shared__ typename BlockReduce::TempStorage temp_storage_max; - __shared__ typename BlockReduce::TempStorage temp_storage_sum; - __shared__ T row_max; - __shared__ T row_sum; - - const int64_t group = blockIdx.x; // row of the grid - const int64_t inner_idx = blockIdx.y; // column of the grid - const int tid = threadIdx.x; - - // calculate the maximum for each group - T thread_max = -INFINITY; - for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { - int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; - thread_max = max(thread_max, input[idx]); - } - T block_max = BlockReduce(temp_storage_max).Reduce(thread_max, CubMaxOp()); - - if (tid == 0) { - row_max = block_max; - } - __syncthreads(); - - // calculate the sum of exponents - T thread_sum = 0; - for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { - int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; - T exp_val = exp(input[idx] - row_max); - output[idx] = exp_val; - thread_sum += exp_val; - } - T block_sum = BlockReduce(temp_storage_sum).Sum(thread_sum); - - if (tid == 0) { - row_sum = block_sum; - } - __syncthreads(); - - // normalize - for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { - int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; - output[idx] /= row_sum; - } -} - -template -void LaunchForward(const std::shared_ptr &output, const std::shared_ptr &input, int64_t dim) { - const auto &input_dims = input->Dims(); - int64_t outer_size = 1; - int64_t axis_size = input_dims[dim]; - int64_t inner_size = 1; - - for (int i = 0; i < dim; ++i) { outer_size *= input_dims[i]; }; - for (int i = dim + 1; i < input_dims.size(); ++i) { inner_size *= input_dims[i]; }; - if (axis_size == 0) { - LOG(INFO) << "CUDA softmax forward: 'input_dims[dim] == 0' at " << __FILE__ << ":" << __LINE__; - return; - } - if (outer_size == 0) { - return; - } - - T *output_ptr = static_cast(output->DataPtr()); - const T *input_ptr = static_cast(input->DataPtr()); - - if (BLOCK_SIZE > 1024) { - LOG(FATAL) << "CUDA softmax forward: 'BLOCK_SIZE used is larger than the max number of thread per block' at " - << __FILE__ << ":" << __LINE__; - } - dim3 block_dims(BLOCK_SIZE); - dim3 grid_dims(outer_size, inner_size); - - SoftmaxForwardKernel - <<>>(output_ptr, input_ptr, outer_size, axis_size, inner_size); -} - -std::shared_ptr SoftmaxForward(const std::shared_ptr &input, int64_t dim) { - auto dtype = input->Dtype(); - const auto &input_dims = input->Dims(); - dim = dim < 0 ? dim + input_dims.size() : dim; - CHECK(dim >= 0 && dim < input_dims.size()); - auto output = std::make_shared(input_dims, dtype, input->GetDevice()); - - switch (dtype) { - case DataType::kFLOAT32: - LaunchForward<256, float>(output, input, dim); - break; - default: - LOG(FATAL) << "CUDA softmax forward: 'Unsupported data type' at " << __FILE__ << ":" << __LINE__; - } - return output; -} - -template -__global__ void SoftmaxBackwardKernel(T *grad_input, const T *grad_output, const T *output, int64_t outer_size, - int64_t axis_size, int64_t inner_size) { - using BlockReduce = cub::BlockReduce; - - __shared__ typename BlockReduce::TempStorage temp_storage_sum; - __shared__ T row_sum; - - const int64_t group = blockIdx.x; - const int64_t inner_idx = blockIdx.y; - const int tid = threadIdx.x; - - // calculate the sum of the dot product of gradients - T thread_sum = 0; - for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { - const int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; - thread_sum += grad_output[idx] * output[idx]; - } - T block_sum = BlockReduce(temp_storage_sum).Sum(thread_sum); - - if (tid == 0) { - row_sum = block_sum; - } - __syncthreads(); - - // update the input gradient - for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { - const int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; - grad_input[idx] = output[idx] * (grad_output[idx] - row_sum); - } -} - -template -void LaunchBackward(const std::shared_ptr &grad_input, const std::shared_ptr &grad_output, - const std::shared_ptr &output, int64_t dim) { - const auto &output_dims = output->Dims(); - int64_t outer_size = 1; - int64_t axis_size = output_dims[dim]; - int64_t inner_size = 1; - - for (int i = 0; i < dim; ++i) { outer_size *= output_dims[i]; }; - for (int i = dim + 1; i < output_dims.size(); ++i) { inner_size *= output_dims[i]; }; - if (axis_size == 0) { - LOG(INFO) << "CUDA softmax backward: 'output_dims[dim] == 0' at " << __FILE__ << ":" << __LINE__; - return; - } - if (outer_size == 0) { - return; - } - - T *grad_input_ptr = static_cast(grad_input->DataPtr()); - const T *grad_output_ptr = static_cast(grad_output->DataPtr()); - const T *output_ptr = static_cast(output->DataPtr()); - - if (BLOCK_SIZE > 1024) { - LOG(FATAL) << "CUDA softmax backward: 'BLOCK_SIZE used is larger than the max number of thread per block' at " - << __FILE__ << ":" << __LINE__; - } - dim3 block(BLOCK_SIZE); - dim3 grid(outer_size, inner_size); - - SoftmaxBackwardKernel - <<>>(grad_input_ptr, grad_output_ptr, output_ptr, outer_size, axis_size, inner_size); -} - -std::shared_ptr SoftmaxBackward(const std::shared_ptr &grad_output, - const std::shared_ptr &output, int64_t dim) { - auto dtype = output->Dtype(); - const auto &output_dims = output->Dims(); - dim = dim < 0 ? dim + output->Dims().size() : dim; - CHECK(dim >= 0 && dim < output->Dims().size()); - - auto grad_input = std::make_shared(output_dims, dtype, output->GetDevice()); - grad_input->Fill(0.0f); - - switch (dtype) { - case DataType::kFLOAT32: - LaunchBackward<256, float>(grad_input, grad_output, output, dim); - break; - default: - LOG(FATAL) << "CUDA softmax backward: 'Unsupported data type' at " << __FILE__ << ":" << __LINE__; - } - - return grad_input; -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_SOFTMAX_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_SOFTMAX_KERNEL(SoftmaxForward) -REGISTER_CUDA_SOFTMAX_KERNEL(SoftmaxBackward) - -#undef REGISTER_CUDA_SOFTMAX_KERNEL +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { + +#if defined(CUB_VERSION) && CUB_VERSION >= 200800 + #include + using CubSumOp = ::cuda::std::plus<>; + using CubMaxOp = ::cuda::maximum<>; + using CubMinOp = ::cuda::minimum<>; +#else + using CubSumOp = cub::Sum; + using CubMaxOp = cub::Max; + using CubMinOp = cub::Min; +#endif + +template +__global__ void SoftmaxForwardKernel(T *output, const T *input, int64_t outer_size, int64_t axis_size, + int64_t inner_size) { + using BlockReduce = cub::BlockReduce; + + __shared__ typename BlockReduce::TempStorage temp_storage_max; + __shared__ typename BlockReduce::TempStorage temp_storage_sum; + __shared__ T row_max; + __shared__ T row_sum; + + const int64_t group = blockIdx.x; // row of the grid + const int64_t inner_idx = blockIdx.y; // column of the grid + const int tid = threadIdx.x; + + // calculate the maximum for each group + T thread_max = -INFINITY; + for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { + int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; + thread_max = max(thread_max, input[idx]); + } + T block_max = BlockReduce(temp_storage_max).Reduce(thread_max, CubMaxOp()); + + if (tid == 0) { + row_max = block_max; + } + __syncthreads(); + + // calculate the sum of exponents + T thread_sum = 0; + for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { + int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; + T exp_val = exp(input[idx] - row_max); + output[idx] = exp_val; + thread_sum += exp_val; + } + T block_sum = BlockReduce(temp_storage_sum).Sum(thread_sum); + + if (tid == 0) { + row_sum = block_sum; + } + __syncthreads(); + + // normalize + for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { + int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; + output[idx] /= row_sum; + } +} + +template +void LaunchForward(const std::shared_ptr &output, const std::shared_ptr &input, int64_t dim) { + const auto &input_dims = input->Dims(); + int64_t outer_size = 1; + int64_t axis_size = input_dims[dim]; + int64_t inner_size = 1; + + for (int i = 0; i < dim; ++i) { outer_size *= input_dims[i]; }; + for (int i = dim + 1; i < input_dims.size(); ++i) { inner_size *= input_dims[i]; }; + if (axis_size == 0) { + LOG(INFO) << "CUDA softmax forward: 'input_dims[dim] == 0' at " << __FILE__ << ":" << __LINE__; + return; + } + if (outer_size == 0) { + return; + } + + T *output_ptr = static_cast(output->DataPtr()); + const T *input_ptr = static_cast(input->DataPtr()); + + if (BLOCK_SIZE > 1024) { + LOG(FATAL) << "CUDA softmax forward: 'BLOCK_SIZE used is larger than the max number of thread per block' at " + << __FILE__ << ":" << __LINE__; + } + dim3 block_dims(BLOCK_SIZE); + dim3 grid_dims(outer_size, inner_size); + + SoftmaxForwardKernel + <<>>(output_ptr, input_ptr, outer_size, axis_size, inner_size); +} + +std::shared_ptr SoftmaxForward(const std::shared_ptr &input, int64_t dim) { + auto dtype = input->Dtype(); + const auto &input_dims = input->Dims(); + dim = dim < 0 ? dim + input_dims.size() : dim; + CHECK(dim >= 0 && dim < input_dims.size()); + auto output = std::make_shared(input_dims, dtype, input->GetDevice()); + + switch (dtype) { + case DataType::kFLOAT32: + LaunchForward<256, float>(output, input, dim); + break; + default: + LOG(FATAL) << "CUDA softmax forward: 'Unsupported data type' at " << __FILE__ << ":" << __LINE__; + } + return output; +} + +template +__global__ void SoftmaxBackwardKernel(T *grad_input, const T *grad_output, const T *output, int64_t outer_size, + int64_t axis_size, int64_t inner_size) { + using BlockReduce = cub::BlockReduce; + + __shared__ typename BlockReduce::TempStorage temp_storage_sum; + __shared__ T row_sum; + + const int64_t group = blockIdx.x; + const int64_t inner_idx = blockIdx.y; + const int tid = threadIdx.x; + + // calculate the sum of the dot product of gradients + T thread_sum = 0; + for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { + const int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; + thread_sum += grad_output[idx] * output[idx]; + } + T block_sum = BlockReduce(temp_storage_sum).Sum(thread_sum); + + if (tid == 0) { + row_sum = block_sum; + } + __syncthreads(); + + // update the input gradient + for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { + const int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; + grad_input[idx] = output[idx] * (grad_output[idx] - row_sum); + } +} + +template +void LaunchBackward(const std::shared_ptr &grad_input, const std::shared_ptr &grad_output, + const std::shared_ptr &output, int64_t dim) { + const auto &output_dims = output->Dims(); + int64_t outer_size = 1; + int64_t axis_size = output_dims[dim]; + int64_t inner_size = 1; + + for (int i = 0; i < dim; ++i) { outer_size *= output_dims[i]; }; + for (int i = dim + 1; i < output_dims.size(); ++i) { inner_size *= output_dims[i]; }; + if (axis_size == 0) { + LOG(INFO) << "CUDA softmax backward: 'output_dims[dim] == 0' at " << __FILE__ << ":" << __LINE__; + return; + } + if (outer_size == 0) { + return; + } + + T *grad_input_ptr = static_cast(grad_input->DataPtr()); + const T *grad_output_ptr = static_cast(grad_output->DataPtr()); + const T *output_ptr = static_cast(output->DataPtr()); + + if (BLOCK_SIZE > 1024) { + LOG(FATAL) << "CUDA softmax backward: 'BLOCK_SIZE used is larger than the max number of thread per block' at " + << __FILE__ << ":" << __LINE__; + } + dim3 block(BLOCK_SIZE); + dim3 grid(outer_size, inner_size); + + SoftmaxBackwardKernel + <<>>(grad_input_ptr, grad_output_ptr, output_ptr, outer_size, axis_size, inner_size); +} + +std::shared_ptr SoftmaxBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &output, int64_t dim) { + auto dtype = output->Dtype(); + const auto &output_dims = output->Dims(); + dim = dim < 0 ? dim + output->Dims().size() : dim; + CHECK(dim >= 0 && dim < output->Dims().size()); + + auto grad_input = std::make_shared(output_dims, dtype, output->GetDevice()); + grad_input->Fill(0.0f); + + switch (dtype) { + case DataType::kFLOAT32: + LaunchBackward<256, float>(grad_input, grad_output, output, dim); + break; + default: + LOG(FATAL) << "CUDA softmax backward: 'Unsupported data type' at " << __FILE__ << ":" << __LINE__; + } + + return grad_input; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_SOFTMAX_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_SOFTMAX_KERNEL(SoftmaxForward) +REGISTER_CUDA_SOFTMAX_KERNEL(SoftmaxBackward) + +#undef REGISTER_CUDA_SOFTMAX_KERNEL diff --git a/infini_train/src/kernels/cuda/split.cu b/infini_train/src/kernels/cuda/split.cu index 9474229..760c5bf 100644 --- a/infini_train/src/kernels/cuda/split.cu +++ b/infini_train/src/kernels/cuda/split.cu @@ -1,144 +1,144 @@ -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { -__global__ void SplitForwardKernel(const float *input, float *output, int64_t N, int64_t H_in, int64_t H_out, int64_t W, - int64_t start_idx) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total = N * H_out * W; - - if (idx < total) { - int w = idx % W; - int h = (idx / W) % H_out; - int n = idx / (H_out * W); - - int input_h = h + start_idx; - int input_idx = n * H_in * W + input_h * W + w; - int output_idx = n * H_out * W + h * W + w; - - output[output_idx] = input[input_idx]; - } -} - -std::vector> SplitForward(const std::shared_ptr &input, int64_t split_size, int dim) { - CHECK_GT(split_size, 0); - CHECK_GE(dim, 0) << "Currently we do not support negative dimension"; - const auto &input_dims = input->Dims(); - CHECK_LT(dim, input_dims.size()); - - std::vector> outputs; - - const int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); - const int64_t W = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); - const int64_t H_in = input_dims[dim]; - - for (int64_t start = 0; start < H_in; start += split_size) { - auto output_dims = input_dims; - const int64_t H_out = std::min(split_size, H_in - start); - output_dims[dim] = H_out; - - auto output = std::make_shared(output_dims, DataType::kFLOAT32, input->GetDevice()); - - int64_t total = N * H_out * W; - int threads_per_block = 256; - int num_blocks = (total + threads_per_block - 1) / threads_per_block; - - SplitForwardKernel<<>>(static_cast(input->DataPtr()), - static_cast(output->DataPtr()), N, H_in, H_out, - W, start); - outputs.push_back(std::move(output)); - } - - return outputs; -} - -__global__ void SplitBackwardKernel(const float *const *grad_outputs, float *grad_input, int64_t N, int64_t H_in, - int64_t W, int64_t split_size, int64_t num_splits, const int64_t *H_outs) { - int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; - int64_t total = N * H_in * W; - if (idx >= total) { - return; - } - - int64_t w = idx % W; - int64_t h = (idx / W) % H_in; - int64_t n = idx / (H_in * W); - - int64_t split_idx = h / split_size; - if (split_idx >= num_splits) { - return; - } - - int64_t H_out = H_outs[split_idx]; - int64_t local_h = h - split_idx * split_size; - - if (local_h >= H_out) { - return; - } - - const float *grad_output = grad_outputs[split_idx]; - float value = grad_output[(n * H_out + local_h) * W + w]; - grad_input[(n * H_in + h) * W + w] = value; -} - -std::shared_ptr SplitBackward(const std::vector &input_dims, int64_t split_size, int dim, - const std::vector> &grad_outputs) { - CHECK_GT(split_size, 0); - CHECK_GE(dim, 0) << "Currently we do not support negative dimension"; - CHECK_LT(dim, input_dims.size()); - - auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, grad_outputs[0]->GetDevice()); - grad_input->Fill(0.0f); - int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); - int64_t W = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); - int64_t H_in = input_dims[dim]; - int64_t num_splits = grad_outputs.size(); - - // init the array of grad_output ptrs - std::vector host_grad_output_ptrs; - for (const auto &grad_output : grad_outputs) { - host_grad_output_ptrs.push_back(static_cast(grad_output->DataPtr())); - } - - void *device_ptr; - const float **device_grad_output_ptrs; - int64_t *device_H_outs; - cudaMallocAsync(&device_ptr, (sizeof(float *) + sizeof(int64_t)) * num_splits, 0); - device_grad_output_ptrs = (const float **)(device_ptr); - device_H_outs = reinterpret_cast(device_grad_output_ptrs + num_splits); - - cudaMemcpyAsync(device_grad_output_ptrs, host_grad_output_ptrs.data(), sizeof(float *) * num_splits, - cudaMemcpyHostToDevice, 0); - - // init H_out for each split - std::vector H_outs(num_splits); - for (int i = 0; i < num_splits; ++i) { H_outs[i] = std::min(split_size, H_in - i * split_size); } - - cudaMemcpyAsync(device_H_outs, H_outs.data(), sizeof(int64_t) * num_splits, cudaMemcpyHostToDevice, 0); - - int64_t total_elements = N * H_in * W; - int threads_per_block = 256; - int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; - - SplitBackwardKernel<<>>(device_grad_output_ptrs, - static_cast(grad_input->DataPtr()), N, H_in, W, - split_size, num_splits, device_H_outs); - - cudaFreeAsync(device_ptr, 0); - return grad_input; -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_SPLIT_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_SPLIT_KERNEL(SplitForward) -REGISTER_CUDA_SPLIT_KERNEL(SplitBackward) - -#undef REGISTER_CUDA_SPLIT_KERNEL +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { +__global__ void SplitForwardKernel(const float *input, float *output, int64_t N, int64_t H_in, int64_t H_out, int64_t W, + int64_t start_idx) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = N * H_out * W; + + if (idx < total) { + int w = idx % W; + int h = (idx / W) % H_out; + int n = idx / (H_out * W); + + int input_h = h + start_idx; + int input_idx = n * H_in * W + input_h * W + w; + int output_idx = n * H_out * W + h * W + w; + + output[output_idx] = input[input_idx]; + } +} + +std::vector> SplitForward(const std::shared_ptr &input, int64_t split_size, int dim) { + CHECK_GT(split_size, 0); + CHECK_GE(dim, 0) << "Currently we do not support negative dimension"; + const auto &input_dims = input->Dims(); + CHECK_LT(dim, input_dims.size()); + + std::vector> outputs; + + const int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + const int64_t W = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); + const int64_t H_in = input_dims[dim]; + + for (int64_t start = 0; start < H_in; start += split_size) { + auto output_dims = input_dims; + const int64_t H_out = std::min(split_size, H_in - start); + output_dims[dim] = H_out; + + auto output = std::make_shared(output_dims, DataType::kFLOAT32, input->GetDevice()); + + int64_t total = N * H_out * W; + int threads_per_block = 256; + int num_blocks = (total + threads_per_block - 1) / threads_per_block; + + SplitForwardKernel<<>>(static_cast(input->DataPtr()), + static_cast(output->DataPtr()), N, H_in, H_out, + W, start); + outputs.push_back(std::move(output)); + } + + return outputs; +} + +__global__ void SplitBackwardKernel(const float *const *grad_outputs, float *grad_input, int64_t N, int64_t H_in, + int64_t W, int64_t split_size, int64_t num_splits, const int64_t *H_outs) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t total = N * H_in * W; + if (idx >= total) { + return; + } + + int64_t w = idx % W; + int64_t h = (idx / W) % H_in; + int64_t n = idx / (H_in * W); + + int64_t split_idx = h / split_size; + if (split_idx >= num_splits) { + return; + } + + int64_t H_out = H_outs[split_idx]; + int64_t local_h = h - split_idx * split_size; + + if (local_h >= H_out) { + return; + } + + const float *grad_output = grad_outputs[split_idx]; + float value = grad_output[(n * H_out + local_h) * W + w]; + grad_input[(n * H_in + h) * W + w] = value; +} + +std::shared_ptr SplitBackward(const std::vector &input_dims, int64_t split_size, int dim, + const std::vector> &grad_outputs) { + CHECK_GT(split_size, 0); + CHECK_GE(dim, 0) << "Currently we do not support negative dimension"; + CHECK_LT(dim, input_dims.size()); + + auto grad_input = std::make_shared(input_dims, DataType::kFLOAT32, grad_outputs[0]->GetDevice()); + grad_input->Fill(0.0f); + int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + int64_t W = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); + int64_t H_in = input_dims[dim]; + int64_t num_splits = grad_outputs.size(); + + // init the array of grad_output ptrs + std::vector host_grad_output_ptrs; + for (const auto &grad_output : grad_outputs) { + host_grad_output_ptrs.push_back(static_cast(grad_output->DataPtr())); + } + + void *device_ptr; + const float **device_grad_output_ptrs; + int64_t *device_H_outs; + cudaMallocAsync(&device_ptr, (sizeof(float *) + sizeof(int64_t)) * num_splits, 0); + device_grad_output_ptrs = (const float **)(device_ptr); + device_H_outs = reinterpret_cast(device_grad_output_ptrs + num_splits); + + cudaMemcpyAsync(device_grad_output_ptrs, host_grad_output_ptrs.data(), sizeof(float *) * num_splits, + cudaMemcpyHostToDevice, 0); + + // init H_out for each split + std::vector H_outs(num_splits); + for (int i = 0; i < num_splits; ++i) { H_outs[i] = std::min(split_size, H_in - i * split_size); } + + cudaMemcpyAsync(device_H_outs, H_outs.data(), sizeof(int64_t) * num_splits, cudaMemcpyHostToDevice, 0); + + int64_t total_elements = N * H_in * W; + int threads_per_block = 256; + int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; + + SplitBackwardKernel<<>>(device_grad_output_ptrs, + static_cast(grad_input->DataPtr()), N, H_in, W, + split_size, num_splits, device_H_outs); + + cudaFreeAsync(device_ptr, 0); + return grad_input; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_SPLIT_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_SPLIT_KERNEL(SplitForward) +REGISTER_CUDA_SPLIT_KERNEL(SplitBackward) + +#undef REGISTER_CUDA_SPLIT_KERNEL diff --git a/infini_train/src/kernels/cuda/stack.cu b/infini_train/src/kernels/cuda/stack.cu index e2c8d25..90e61c0 100644 --- a/infini_train/src/kernels/cuda/stack.cu +++ b/infini_train/src/kernels/cuda/stack.cu @@ -1,138 +1,138 @@ -#include -#include -#include -#include -#include - -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { - -#define CUDA_CHECK(call) \ - do { \ - cudaError_t status = call; \ - if (status != cudaSuccess) { \ - LOG(FATAL) << "CUDA Error: " << cudaGetErrorString(status) << " at " << __FILE__ << ":" << __LINE__; \ - } \ - } while (0) - -__global__ void StackForwardKernel(const float **inputs, float *output, int64_t N, int64_t D, int64_t num_inputs) { - int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; - int64_t total = N * num_inputs * D; - - if (idx >= total) { - return; - } - - int64_t d = idx % D; - int64_t s = (idx / D) % num_inputs; - int64_t n = idx / (D * num_inputs); - - const float *input = inputs[s]; - output[idx] = input[n * D + d]; -} - -std::shared_ptr StackForward(const std::vector> &inputs, int64_t dim) { - CHECK(!inputs.empty()); - - const auto &base_dims = inputs[0]->Dims(); - if (dim < 0) { - dim += base_dims.size() + 1; - } - CHECK_GE(dim, 0); - CHECK_LE(dim, base_dims.size()); - for (const auto &input : inputs) { CHECK(input->Dims() == base_dims); } - - std::vector out_dims = base_dims; - out_dims.insert(out_dims.begin() + dim, inputs.size()); - auto output = std::make_shared(out_dims, DataType::kFLOAT32, inputs[0]->GetDevice()); - - const int64_t N = std::accumulate(base_dims.begin(), base_dims.begin() + dim, 1, std::multiplies()); - const int64_t D = std::accumulate(base_dims.begin() + dim, base_dims.end(), 1, std::multiplies()); - const int64_t num_inputs = inputs.size(); - - std::vector host_input_ptrs; - for (const auto &t : inputs) { host_input_ptrs.push_back(static_cast(t->DataPtr())); } - - const float **device_input_ptrs; - cudaMallocAsync(&device_input_ptrs, sizeof(float *) * num_inputs, 0); - cudaMemcpyAsync(device_input_ptrs, host_input_ptrs.data(), sizeof(float *) * num_inputs, cudaMemcpyHostToDevice, 0); - - int64_t total = N * num_inputs * D; - int threads_per_block = 256; - int num_blocks = (total + threads_per_block - 1) / threads_per_block; - StackForwardKernel<<>>(device_input_ptrs, static_cast(output->DataPtr()), N, - D, num_inputs); - - cudaFreeAsync(device_input_ptrs, 0); - return output; -} - -__global__ void StackBackwardKernel(const float *grad_output, float **grad_inputs, int64_t N, int64_t D, - int64_t num_inputs) { - int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; - int64_t total = N * num_inputs * D; - - if (idx >= total) { - return; - } - - int64_t d = idx % D; - int64_t s = (idx / D) % num_inputs; - int64_t n = idx / (D * num_inputs); - - if (s < num_inputs) { - grad_inputs[s][n * D + d] = grad_output[idx]; - } -} - -std::vector> StackBackward(const std::vector &input_dims, int64_t dim, - const std::shared_ptr &grad_output) { - if (dim < 0) { - dim += input_dims.size() + 1; - } - const int64_t num_inputs = grad_output->Dims()[dim]; - std::vector base_dims = grad_output->Dims(); - base_dims.erase(base_dims.begin() + dim); - - std::vector> grads; - for (int i = 0; i < num_inputs; ++i) { - auto t = std::make_shared(base_dims, DataType::kFLOAT32, grad_output->GetDevice()); - t->Fill(0.0f); - grads.push_back(t); - } - - int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); - int64_t D = std::accumulate(input_dims.begin() + dim, input_dims.end(), 1, std::multiplies()); - - std::vector host_ptrs; - for (auto &t : grads) { host_ptrs.push_back(static_cast(t->DataPtr())); } - - float **device_ptrs; - cudaMallocAsync(&device_ptrs, sizeof(float *) * num_inputs, 0); - cudaMemcpyAsync(device_ptrs, host_ptrs.data(), sizeof(float *) * num_inputs, cudaMemcpyHostToDevice, 0); - - int64_t total = N * num_inputs * D; - int threads_per_block = 256; - int num_blocks = (total + threads_per_block - 1) / threads_per_block; - - StackBackwardKernel<<>>(static_cast(grad_output->DataPtr()), - device_ptrs, N, D, num_inputs); - - CUDA_CHECK(cudaGetLastError()); - cudaFreeAsync(device_ptrs, 0); - return grads; -} - -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_STACK_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_STACK_KERNEL(StackForward) -REGISTER_CUDA_STACK_KERNEL(StackBackward) - -#undef REGISTER_CUDA_STACK_KERNEL +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { + +#define CUDA_CHECK(call) \ + do { \ + cudaError_t status = call; \ + if (status != cudaSuccess) { \ + LOG(FATAL) << "CUDA Error: " << cudaGetErrorString(status) << " at " << __FILE__ << ":" << __LINE__; \ + } \ + } while (0) + +__global__ void StackForwardKernel(const float **inputs, float *output, int64_t N, int64_t D, int64_t num_inputs) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t total = N * num_inputs * D; + + if (idx >= total) { + return; + } + + int64_t d = idx % D; + int64_t s = (idx / D) % num_inputs; + int64_t n = idx / (D * num_inputs); + + const float *input = inputs[s]; + output[idx] = input[n * D + d]; +} + +std::shared_ptr StackForward(const std::vector> &inputs, int64_t dim) { + CHECK(!inputs.empty()); + + const auto &base_dims = inputs[0]->Dims(); + if (dim < 0) { + dim += base_dims.size() + 1; + } + CHECK_GE(dim, 0); + CHECK_LE(dim, base_dims.size()); + for (const auto &input : inputs) { CHECK(input->Dims() == base_dims); } + + std::vector out_dims = base_dims; + out_dims.insert(out_dims.begin() + dim, inputs.size()); + auto output = std::make_shared(out_dims, DataType::kFLOAT32, inputs[0]->GetDevice()); + + const int64_t N = std::accumulate(base_dims.begin(), base_dims.begin() + dim, 1, std::multiplies()); + const int64_t D = std::accumulate(base_dims.begin() + dim, base_dims.end(), 1, std::multiplies()); + const int64_t num_inputs = inputs.size(); + + std::vector host_input_ptrs; + for (const auto &t : inputs) { host_input_ptrs.push_back(static_cast(t->DataPtr())); } + + const float **device_input_ptrs; + cudaMallocAsync(&device_input_ptrs, sizeof(float *) * num_inputs, 0); + cudaMemcpyAsync(device_input_ptrs, host_input_ptrs.data(), sizeof(float *) * num_inputs, cudaMemcpyHostToDevice, 0); + + int64_t total = N * num_inputs * D; + int threads_per_block = 256; + int num_blocks = (total + threads_per_block - 1) / threads_per_block; + StackForwardKernel<<>>(device_input_ptrs, static_cast(output->DataPtr()), N, + D, num_inputs); + + cudaFreeAsync(device_input_ptrs, 0); + return output; +} + +__global__ void StackBackwardKernel(const float *grad_output, float **grad_inputs, int64_t N, int64_t D, + int64_t num_inputs) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t total = N * num_inputs * D; + + if (idx >= total) { + return; + } + + int64_t d = idx % D; + int64_t s = (idx / D) % num_inputs; + int64_t n = idx / (D * num_inputs); + + if (s < num_inputs) { + grad_inputs[s][n * D + d] = grad_output[idx]; + } +} + +std::vector> StackBackward(const std::vector &input_dims, int64_t dim, + const std::shared_ptr &grad_output) { + if (dim < 0) { + dim += input_dims.size() + 1; + } + const int64_t num_inputs = grad_output->Dims()[dim]; + std::vector base_dims = grad_output->Dims(); + base_dims.erase(base_dims.begin() + dim); + + std::vector> grads; + for (int i = 0; i < num_inputs; ++i) { + auto t = std::make_shared(base_dims, DataType::kFLOAT32, grad_output->GetDevice()); + t->Fill(0.0f); + grads.push_back(t); + } + + int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + int64_t D = std::accumulate(input_dims.begin() + dim, input_dims.end(), 1, std::multiplies()); + + std::vector host_ptrs; + for (auto &t : grads) { host_ptrs.push_back(static_cast(t->DataPtr())); } + + float **device_ptrs; + cudaMallocAsync(&device_ptrs, sizeof(float *) * num_inputs, 0); + cudaMemcpyAsync(device_ptrs, host_ptrs.data(), sizeof(float *) * num_inputs, cudaMemcpyHostToDevice, 0); + + int64_t total = N * num_inputs * D; + int threads_per_block = 256; + int num_blocks = (total + threads_per_block - 1) / threads_per_block; + + StackBackwardKernel<<>>(static_cast(grad_output->DataPtr()), + device_ptrs, N, D, num_inputs); + + CUDA_CHECK(cudaGetLastError()); + cudaFreeAsync(device_ptrs, 0); + return grads; +} + +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_STACK_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_STACK_KERNEL(StackForward) +REGISTER_CUDA_STACK_KERNEL(StackBackward) + +#undef REGISTER_CUDA_STACK_KERNEL diff --git a/infini_train/src/kernels/cuda/transform.cu b/infini_train/src/kernels/cuda/transform.cu index 88467ee..45e9c67 100644 --- a/infini_train/src/kernels/cuda/transform.cu +++ b/infini_train/src/kernels/cuda/transform.cu @@ -1,402 +1,402 @@ -#include "cuda_runtime.h" -#include "glog/logging.h" - -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::kernels::cuda { - -__global__ void TrilForwardKernel(const float *input, float *output, int rows, int cols, int64_t diagonal) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= rows * cols) { - return; - } - - int row = idx / cols; - int col = idx % cols; - - if (row - col + diagonal >= 0) { - output[idx] = input[idx]; - } else { - output[idx] = 0.0f; - } -} - -std::shared_ptr TrilForward(const std::shared_ptr &input, int64_t diagonal) { - CHECK_EQ(input->Dims().size(), 2); - int64_t rows = input->Dims()[0]; - int64_t cols = input->Dims()[1]; - - auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); - - int threads_per_block = 256; - int num_blocks = (rows * cols + threads_per_block - 1) / threads_per_block; - - TrilForwardKernel<<>>(static_cast(input->DataPtr()), - static_cast(output->DataPtr()), rows, cols, diagonal); - return output; -} - -__global__ void TrilBackwardKernel(const float *grad_output, float *grad_input, int rows, int cols, int64_t diagonal) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= rows * cols) { - return; - } - - int row = idx / cols; - int col = idx % cols; - - if (row - col + diagonal >= 0) { - grad_input[idx] = grad_output[idx]; - } else { - grad_input[idx] = 0.0f; - } -} - -std::shared_ptr TrilBackward(const std::shared_ptr &grad_output, int64_t diagonal) { - int rows = grad_output->Dims()[0]; - int cols = grad_output->Dims()[1]; - - auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); - grad_input->Fill(0.0f); - - int threads_per_block = 256; - int num_blocks = (rows * cols + threads_per_block - 1) / threads_per_block; - - TrilBackwardKernel<<>>(static_cast(grad_output->DataPtr()), - static_cast(grad_input->DataPtr()), rows, cols, - diagonal); - - return grad_input; -} - -__global__ void TriuForwardKernel(const float *input, float *output, int rows, int cols, int64_t diagonal) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= rows * cols) { - return; - } - - int row = idx / cols; - int col = idx % cols; - - if (row - col + diagonal <= 0) { - output[idx] = input[idx]; - } else { - output[idx] = 0.0f; - } -} - -std::shared_ptr TriuForward(const std::shared_ptr &input, int64_t diagonal) { - CHECK_EQ(input->Dims().size(), 2); - int64_t rows = input->Dims()[0]; - int64_t cols = input->Dims()[1]; - - auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); - - int threads_per_block = 256; - int num_blocks = (rows * cols + threads_per_block - 1) / threads_per_block; - - TriuForwardKernel<<>>(static_cast(input->DataPtr()), - static_cast(output->DataPtr()), rows, cols, diagonal); - - return output; -} - -__global__ void TriuBackwardKernel(const float *grad_output, float *grad_input, int rows, int cols, int64_t diagonal) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= rows * cols) { - return; - } - - int row = idx / cols; - int col = idx % cols; - - if (row - col + diagonal <= 0) { - grad_input[idx] = grad_output[idx]; - } else { - grad_input[idx] = 0.0f; - } -} - -std::shared_ptr TriuBackward(const std::shared_ptr &grad_output, int64_t diagonal) { - int rows = grad_output->Dims()[0]; - int cols = grad_output->Dims()[1]; - - auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); - grad_input->Fill(0.0f); - - int threads_per_block = 256; - int num_blocks = (rows * cols + threads_per_block - 1) / threads_per_block; - - TriuBackwardKernel<<>>(static_cast(grad_output->DataPtr()), - static_cast(grad_input->DataPtr()), rows, cols, - diagonal); - - return grad_input; -} - -__global__ void TransposeForwardKernel(const float *input, float *output, const int64_t *in_dims, - const int64_t *in_strides, const int64_t *out_strides, int64_t ndim, - int64_t dim0, int64_t dim1, int64_t num_elements) { - int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= num_elements) { - return; - } - - int64_t remaining = idx; - int64_t coords[8]; - - // 1. decode coord from output index - for (int i = 0; i < ndim; ++i) { - coords[i] = remaining / out_strides[i]; - remaining %= out_strides[i]; - } - - // 2. swap the coordinates - int64_t tmp = coords[dim0]; - coords[dim0] = coords[dim1]; - coords[dim1] = tmp; - - // 3. compute input flat index - int64_t in_flat_idx = 0; - for (int i = 0; i < ndim; ++i) { in_flat_idx += coords[i] * in_strides[i]; } - - output[idx] = input[in_flat_idx]; -} - -std::shared_ptr TransposeForward(const std::shared_ptr &input, int64_t dim0, int64_t dim1) { - CHECK_LE(input->Dims().size(), 8); - dim0 = dim0 < 0 ? dim0 + input->Dims().size() : dim0; - dim1 = dim1 < 0 ? dim1 + input->Dims().size() : dim1; - CHECK(dim0 >= 0 && dim0 < input->Dims().size() && dim1 >= 0 && dim1 < input->Dims().size()); - - auto in_dims = input->Dims(); - std::vector out_dims = in_dims; - std::swap(out_dims[dim0], out_dims[dim1]); - - auto output = std::make_shared(out_dims, input->Dtype(), input->GetDevice()); - output->Fill(0.0f); - int64_t ndim = in_dims.size(); - int64_t num_elements = output->NumElements(); - - // compute strides of in_dims and out_dims - std::vector in_strides(ndim, 1); - std::vector out_strides(ndim, 1); - for (int i = ndim - 2; i >= 0; --i) { - in_strides[i] = in_strides[i + 1] * in_dims[i + 1]; - out_strides[i] = out_strides[i + 1] * out_dims[i + 1]; - } - - // Allocate device memory for dims and strides - int64_t *device_buffer; - cudaMallocAsync(&device_buffer, 3 * ndim * sizeof(int64_t), 0); - - int64_t *in_dims_dev = device_buffer; - int64_t *in_strides_dev = device_buffer + ndim; - int64_t *out_strides_dev = device_buffer + 2 * ndim; - - std::vector host_buffer; - host_buffer.insert(host_buffer.end(), in_dims.begin(), in_dims.end()); - host_buffer.insert(host_buffer.end(), in_strides.begin(), in_strides.end()); - host_buffer.insert(host_buffer.end(), out_strides.begin(), out_strides.end()); - - cudaMemcpyAsync(device_buffer, host_buffer.data(), 3 * ndim * sizeof(int64_t), cudaMemcpyHostToDevice, 0); - - int threads_per_block = 256; - int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; - - TransposeForwardKernel<<>>( - static_cast(input->DataPtr()), static_cast(output->DataPtr()), in_dims_dev, - in_strides_dev, out_strides_dev, ndim, dim0, dim1, num_elements); - - cudaFreeAsync(device_buffer, 0); - - return output; -} - -std::shared_ptr TransposeBackward(const std::shared_ptr &grad_output, int64_t dim0, int64_t dim1) { - return TransposeForward(grad_output, dim1, dim0); -} - -__global__ void MaskForwardKernel(const float *input, const float *mask, float *output, float value, int batch_size, - int mask_size) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < batch_size * mask_size) { - output[i] = (mask[i % mask_size] == 1.0f) ? value : input[i]; - } -} - -std::shared_ptr MaskForward(const std::shared_ptr &input, const std::shared_ptr &mask, - float value) { - auto input_shape = input->Dims(); - auto mask_shape = mask->Dims(); - CHECK_EQ(static_cast(input->Dtype()), static_cast(mask->Dtype())); - - int64_t input_dims = input_shape.size(); - int64_t mask_dims = mask_shape.size(); - for (int i = 0; i < mask_dims; ++i) { - int input_dim = input_shape[input_dims - mask_dims + i]; - int mask_dim = mask_shape[i]; - CHECK(input_dim == mask_dim || mask_dim == 1); - } - - int64_t mask_size = mask->NumElements(); - int64_t batch_size = input->NumElements() / mask_size; - - auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); - - int threads_per_block = 256; - int num_blocks = (input->NumElements() + threads_per_block - 1) / threads_per_block; - - MaskForwardKernel<<>>( - static_cast(input->DataPtr()), static_cast(mask->DataPtr()), - static_cast(output->DataPtr()), value, batch_size, mask_size); - return output; -} - -__global__ void MaskBackwardKernel(const float *grad_output, const float *mask, float *grad_input, int batch_size, - int mask_size) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < batch_size * mask_size) { - grad_input[i] = (mask[i % mask_size] == 1.0f) ? 0.0f : grad_output[i]; - } -} - -std::shared_ptr MaskBackward(const std::shared_ptr &grad_output, const std::shared_ptr &mask) { - auto output_shape = grad_output->Dims(); - auto mask_shape = mask->Dims(); - CHECK_EQ(static_cast(grad_output->Dtype()), static_cast(mask->Dtype())); - - int64_t output_dims = output_shape.size(); - int64_t mask_dims = mask_shape.size(); - for (int i = 0; i < mask_dims; ++i) { - int out_dim = output_shape[output_dims - mask_dims + i]; - int mask_dim = mask_shape[i]; - CHECK(out_dim == mask_dim || mask_dim == 1); - } - - int64_t mask_size = mask->NumElements(); - int64_t batch_size = grad_output->NumElements() / mask_size; - - auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); - grad_input->Fill(0.0f); - - int threads_per_block = 256; - int num_blocks = (grad_output->NumElements() + threads_per_block - 1) / threads_per_block; - - MaskBackwardKernel<<>>( - static_cast(grad_output->DataPtr()), static_cast(mask->DataPtr()), - static_cast(grad_input->DataPtr()), batch_size, mask_size); - return grad_input; -} - -__global__ void RepeatInterleaveForwardKernel(const float *input, float *output, int64_t outer, int64_t dim_size, - int64_t inner, int64_t repeat) { - int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; - int64_t total = outer * dim_size * repeat * inner; - if (idx >= total) { - return; - } - - int64_t i = idx / inner; - int64_t j = idx % inner; - - int64_t o = i / (dim_size * repeat); - int64_t di = (i / repeat) % dim_size; - - output[idx] = input[(o * dim_size + di) * inner + j]; -} - -std::shared_ptr RepeatInterleaveForward(const std::shared_ptr &input, int64_t repeat, int64_t dim) { - CHECK_GT(repeat, 0); - CHECK_GE(dim, 0); - CHECK_LT(dim, input->Dims().size()); - - const auto &input_dims = input->Dims(); - const int64_t outer = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); - const int64_t inner - = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); - const int64_t dim_size = input_dims[dim]; - - std::vector output_dims = input_dims; - output_dims[dim] = dim_size * repeat; - auto output = std::make_shared(output_dims, input->Dtype(), input->GetDevice()); - - const float *input_ptr = static_cast(input->DataPtr()); - float *output_ptr = static_cast(output->DataPtr()); - - int64_t total_elements = outer * dim_size * repeat * inner; - int threads_per_block = 256; - int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; - RepeatInterleaveForwardKernel<<>>(input_ptr, output_ptr, outer, dim_size, inner, - repeat); - - return output; -} - -__global__ void RepeatInterleaveBackwardKernel(const float *grad_output, float *grad_input, int64_t outer, - int64_t dim_size, int64_t inner, int64_t repeat) { - int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; - int64_t total = outer * dim_size * inner; - if (idx >= total) { - return; - } - - int64_t i = idx / inner; - int64_t j = idx % inner; - - int64_t o = i / dim_size; - int64_t di = i % dim_size; - - float sum = 0.0f; - for (int64_t r = 0; r < repeat; ++r) { - int64_t out_idx = ((o * dim_size * repeat + di * repeat + r) * inner) + j; - sum += grad_output[out_idx]; - } - grad_input[idx] = sum; -} - -std::shared_ptr RepeatInterleaveBackward(const std::shared_ptr &grad_output, - const std::vector &input_dims, int64_t dim) { - CHECK_GE(dim, 0); - CHECK_LT(dim, input_dims.size()); - - const int64_t outer = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); - const int64_t inner - = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); - const int64_t dim_size = input_dims[dim]; - - int64_t repeat = grad_output->Dims()[dim] / dim_size; - CHECK_EQ(grad_output->Dims()[dim], dim_size * repeat); - - auto grad_input = std::make_shared(input_dims, grad_output->Dtype(), grad_output->GetDevice()); - grad_input->Fill(0.0f); - - const float *grad_out_ptr = static_cast(grad_output->DataPtr()); - float *grad_in_ptr = static_cast(grad_input->DataPtr()); - - int64_t total_elements = outer * dim_size * inner; - int threads_per_block = 256; - int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; - RepeatInterleaveBackwardKernel<<>>(grad_out_ptr, grad_in_ptr, outer, dim_size, inner, - repeat); - - return grad_input; -} -} // namespace infini_train::kernels::cuda - -#define REGISTER_CUDA_TRANSFORM_KERNEL(kernel_name) \ - REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) - -REGISTER_CUDA_TRANSFORM_KERNEL(TrilForward) -REGISTER_CUDA_TRANSFORM_KERNEL(TrilBackward) -REGISTER_CUDA_TRANSFORM_KERNEL(TriuForward) -REGISTER_CUDA_TRANSFORM_KERNEL(TriuBackward) -REGISTER_CUDA_TRANSFORM_KERNEL(TransposeForward) -REGISTER_CUDA_TRANSFORM_KERNEL(TransposeBackward) -REGISTER_CUDA_TRANSFORM_KERNEL(MaskForward) -REGISTER_CUDA_TRANSFORM_KERNEL(MaskBackward) -REGISTER_CUDA_TRANSFORM_KERNEL(RepeatInterleaveForward) -REGISTER_CUDA_TRANSFORM_KERNEL(RepeatInterleaveBackward) - -#undef REGISTER_CUDA_TRANSFORM_KERNEL +#include "cuda_runtime.h" +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cuda { + +__global__ void TrilForwardKernel(const float *input, float *output, int rows, int cols, int64_t diagonal) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= rows * cols) { + return; + } + + int row = idx / cols; + int col = idx % cols; + + if (row - col + diagonal >= 0) { + output[idx] = input[idx]; + } else { + output[idx] = 0.0f; + } +} + +std::shared_ptr TrilForward(const std::shared_ptr &input, int64_t diagonal) { + CHECK_EQ(input->Dims().size(), 2); + int64_t rows = input->Dims()[0]; + int64_t cols = input->Dims()[1]; + + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + + int threads_per_block = 256; + int num_blocks = (rows * cols + threads_per_block - 1) / threads_per_block; + + TrilForwardKernel<<>>(static_cast(input->DataPtr()), + static_cast(output->DataPtr()), rows, cols, diagonal); + return output; +} + +__global__ void TrilBackwardKernel(const float *grad_output, float *grad_input, int rows, int cols, int64_t diagonal) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= rows * cols) { + return; + } + + int row = idx / cols; + int col = idx % cols; + + if (row - col + diagonal >= 0) { + grad_input[idx] = grad_output[idx]; + } else { + grad_input[idx] = 0.0f; + } +} + +std::shared_ptr TrilBackward(const std::shared_ptr &grad_output, int64_t diagonal) { + int rows = grad_output->Dims()[0]; + int cols = grad_output->Dims()[1]; + + auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); + grad_input->Fill(0.0f); + + int threads_per_block = 256; + int num_blocks = (rows * cols + threads_per_block - 1) / threads_per_block; + + TrilBackwardKernel<<>>(static_cast(grad_output->DataPtr()), + static_cast(grad_input->DataPtr()), rows, cols, + diagonal); + + return grad_input; +} + +__global__ void TriuForwardKernel(const float *input, float *output, int rows, int cols, int64_t diagonal) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= rows * cols) { + return; + } + + int row = idx / cols; + int col = idx % cols; + + if (row - col + diagonal <= 0) { + output[idx] = input[idx]; + } else { + output[idx] = 0.0f; + } +} + +std::shared_ptr TriuForward(const std::shared_ptr &input, int64_t diagonal) { + CHECK_EQ(input->Dims().size(), 2); + int64_t rows = input->Dims()[0]; + int64_t cols = input->Dims()[1]; + + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + + int threads_per_block = 256; + int num_blocks = (rows * cols + threads_per_block - 1) / threads_per_block; + + TriuForwardKernel<<>>(static_cast(input->DataPtr()), + static_cast(output->DataPtr()), rows, cols, diagonal); + + return output; +} + +__global__ void TriuBackwardKernel(const float *grad_output, float *grad_input, int rows, int cols, int64_t diagonal) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= rows * cols) { + return; + } + + int row = idx / cols; + int col = idx % cols; + + if (row - col + diagonal <= 0) { + grad_input[idx] = grad_output[idx]; + } else { + grad_input[idx] = 0.0f; + } +} + +std::shared_ptr TriuBackward(const std::shared_ptr &grad_output, int64_t diagonal) { + int rows = grad_output->Dims()[0]; + int cols = grad_output->Dims()[1]; + + auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); + grad_input->Fill(0.0f); + + int threads_per_block = 256; + int num_blocks = (rows * cols + threads_per_block - 1) / threads_per_block; + + TriuBackwardKernel<<>>(static_cast(grad_output->DataPtr()), + static_cast(grad_input->DataPtr()), rows, cols, + diagonal); + + return grad_input; +} + +__global__ void TransposeForwardKernel(const float *input, float *output, const int64_t *in_dims, + const int64_t *in_strides, const int64_t *out_strides, int64_t ndim, + int64_t dim0, int64_t dim1, int64_t num_elements) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_elements) { + return; + } + + int64_t remaining = idx; + int64_t coords[8]; + + // 1. decode coord from output index + for (int i = 0; i < ndim; ++i) { + coords[i] = remaining / out_strides[i]; + remaining %= out_strides[i]; + } + + // 2. swap the coordinates + int64_t tmp = coords[dim0]; + coords[dim0] = coords[dim1]; + coords[dim1] = tmp; + + // 3. compute input flat index + int64_t in_flat_idx = 0; + for (int i = 0; i < ndim; ++i) { in_flat_idx += coords[i] * in_strides[i]; } + + output[idx] = input[in_flat_idx]; +} + +std::shared_ptr TransposeForward(const std::shared_ptr &input, int64_t dim0, int64_t dim1) { + CHECK_LE(input->Dims().size(), 8); + dim0 = dim0 < 0 ? dim0 + input->Dims().size() : dim0; + dim1 = dim1 < 0 ? dim1 + input->Dims().size() : dim1; + CHECK(dim0 >= 0 && dim0 < input->Dims().size() && dim1 >= 0 && dim1 < input->Dims().size()); + + auto in_dims = input->Dims(); + std::vector out_dims = in_dims; + std::swap(out_dims[dim0], out_dims[dim1]); + + auto output = std::make_shared(out_dims, input->Dtype(), input->GetDevice()); + output->Fill(0.0f); + int64_t ndim = in_dims.size(); + int64_t num_elements = output->NumElements(); + + // compute strides of in_dims and out_dims + std::vector in_strides(ndim, 1); + std::vector out_strides(ndim, 1); + for (int i = ndim - 2; i >= 0; --i) { + in_strides[i] = in_strides[i + 1] * in_dims[i + 1]; + out_strides[i] = out_strides[i + 1] * out_dims[i + 1]; + } + + // Allocate device memory for dims and strides + int64_t *device_buffer; + cudaMallocAsync(&device_buffer, 3 * ndim * sizeof(int64_t), 0); + + int64_t *in_dims_dev = device_buffer; + int64_t *in_strides_dev = device_buffer + ndim; + int64_t *out_strides_dev = device_buffer + 2 * ndim; + + std::vector host_buffer; + host_buffer.insert(host_buffer.end(), in_dims.begin(), in_dims.end()); + host_buffer.insert(host_buffer.end(), in_strides.begin(), in_strides.end()); + host_buffer.insert(host_buffer.end(), out_strides.begin(), out_strides.end()); + + cudaMemcpyAsync(device_buffer, host_buffer.data(), 3 * ndim * sizeof(int64_t), cudaMemcpyHostToDevice, 0); + + int threads_per_block = 256; + int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; + + TransposeForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), in_dims_dev, + in_strides_dev, out_strides_dev, ndim, dim0, dim1, num_elements); + + cudaFreeAsync(device_buffer, 0); + + return output; +} + +std::shared_ptr TransposeBackward(const std::shared_ptr &grad_output, int64_t dim0, int64_t dim1) { + return TransposeForward(grad_output, dim1, dim0); +} + +__global__ void MaskForwardKernel(const float *input, const float *mask, float *output, float value, int batch_size, + int mask_size) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < batch_size * mask_size) { + output[i] = (mask[i % mask_size] == 1.0f) ? value : input[i]; + } +} + +std::shared_ptr MaskForward(const std::shared_ptr &input, const std::shared_ptr &mask, + float value) { + auto input_shape = input->Dims(); + auto mask_shape = mask->Dims(); + CHECK_EQ(static_cast(input->Dtype()), static_cast(mask->Dtype())); + + int64_t input_dims = input_shape.size(); + int64_t mask_dims = mask_shape.size(); + for (int i = 0; i < mask_dims; ++i) { + int input_dim = input_shape[input_dims - mask_dims + i]; + int mask_dim = mask_shape[i]; + CHECK(input_dim == mask_dim || mask_dim == 1); + } + + int64_t mask_size = mask->NumElements(); + int64_t batch_size = input->NumElements() / mask_size; + + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + + int threads_per_block = 256; + int num_blocks = (input->NumElements() + threads_per_block - 1) / threads_per_block; + + MaskForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(mask->DataPtr()), + static_cast(output->DataPtr()), value, batch_size, mask_size); + return output; +} + +__global__ void MaskBackwardKernel(const float *grad_output, const float *mask, float *grad_input, int batch_size, + int mask_size) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < batch_size * mask_size) { + grad_input[i] = (mask[i % mask_size] == 1.0f) ? 0.0f : grad_output[i]; + } +} + +std::shared_ptr MaskBackward(const std::shared_ptr &grad_output, const std::shared_ptr &mask) { + auto output_shape = grad_output->Dims(); + auto mask_shape = mask->Dims(); + CHECK_EQ(static_cast(grad_output->Dtype()), static_cast(mask->Dtype())); + + int64_t output_dims = output_shape.size(); + int64_t mask_dims = mask_shape.size(); + for (int i = 0; i < mask_dims; ++i) { + int out_dim = output_shape[output_dims - mask_dims + i]; + int mask_dim = mask_shape[i]; + CHECK(out_dim == mask_dim || mask_dim == 1); + } + + int64_t mask_size = mask->NumElements(); + int64_t batch_size = grad_output->NumElements() / mask_size; + + auto grad_input = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); + grad_input->Fill(0.0f); + + int threads_per_block = 256; + int num_blocks = (grad_output->NumElements() + threads_per_block - 1) / threads_per_block; + + MaskBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(mask->DataPtr()), + static_cast(grad_input->DataPtr()), batch_size, mask_size); + return grad_input; +} + +__global__ void RepeatInterleaveForwardKernel(const float *input, float *output, int64_t outer, int64_t dim_size, + int64_t inner, int64_t repeat) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t total = outer * dim_size * repeat * inner; + if (idx >= total) { + return; + } + + int64_t i = idx / inner; + int64_t j = idx % inner; + + int64_t o = i / (dim_size * repeat); + int64_t di = (i / repeat) % dim_size; + + output[idx] = input[(o * dim_size + di) * inner + j]; +} + +std::shared_ptr RepeatInterleaveForward(const std::shared_ptr &input, int64_t repeat, int64_t dim) { + CHECK_GT(repeat, 0); + CHECK_GE(dim, 0); + CHECK_LT(dim, input->Dims().size()); + + const auto &input_dims = input->Dims(); + const int64_t outer = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + const int64_t inner + = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); + const int64_t dim_size = input_dims[dim]; + + std::vector output_dims = input_dims; + output_dims[dim] = dim_size * repeat; + auto output = std::make_shared(output_dims, input->Dtype(), input->GetDevice()); + + const float *input_ptr = static_cast(input->DataPtr()); + float *output_ptr = static_cast(output->DataPtr()); + + int64_t total_elements = outer * dim_size * repeat * inner; + int threads_per_block = 256; + int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; + RepeatInterleaveForwardKernel<<>>(input_ptr, output_ptr, outer, dim_size, inner, + repeat); + + return output; +} + +__global__ void RepeatInterleaveBackwardKernel(const float *grad_output, float *grad_input, int64_t outer, + int64_t dim_size, int64_t inner, int64_t repeat) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t total = outer * dim_size * inner; + if (idx >= total) { + return; + } + + int64_t i = idx / inner; + int64_t j = idx % inner; + + int64_t o = i / dim_size; + int64_t di = i % dim_size; + + float sum = 0.0f; + for (int64_t r = 0; r < repeat; ++r) { + int64_t out_idx = ((o * dim_size * repeat + di * repeat + r) * inner) + j; + sum += grad_output[out_idx]; + } + grad_input[idx] = sum; +} + +std::shared_ptr RepeatInterleaveBackward(const std::shared_ptr &grad_output, + const std::vector &input_dims, int64_t dim) { + CHECK_GE(dim, 0); + CHECK_LT(dim, input_dims.size()); + + const int64_t outer = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + const int64_t inner + = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); + const int64_t dim_size = input_dims[dim]; + + int64_t repeat = grad_output->Dims()[dim] / dim_size; + CHECK_EQ(grad_output->Dims()[dim], dim_size * repeat); + + auto grad_input = std::make_shared(input_dims, grad_output->Dtype(), grad_output->GetDevice()); + grad_input->Fill(0.0f); + + const float *grad_out_ptr = static_cast(grad_output->DataPtr()); + float *grad_in_ptr = static_cast(grad_input->DataPtr()); + + int64_t total_elements = outer * dim_size * inner; + int threads_per_block = 256; + int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; + RepeatInterleaveBackwardKernel<<>>(grad_out_ptr, grad_in_ptr, outer, dim_size, inner, + repeat); + + return grad_input; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_TRANSFORM_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_TRANSFORM_KERNEL(TrilForward) +REGISTER_CUDA_TRANSFORM_KERNEL(TrilBackward) +REGISTER_CUDA_TRANSFORM_KERNEL(TriuForward) +REGISTER_CUDA_TRANSFORM_KERNEL(TriuBackward) +REGISTER_CUDA_TRANSFORM_KERNEL(TransposeForward) +REGISTER_CUDA_TRANSFORM_KERNEL(TransposeBackward) +REGISTER_CUDA_TRANSFORM_KERNEL(MaskForward) +REGISTER_CUDA_TRANSFORM_KERNEL(MaskBackward) +REGISTER_CUDA_TRANSFORM_KERNEL(RepeatInterleaveForward) +REGISTER_CUDA_TRANSFORM_KERNEL(RepeatInterleaveBackward) + +#undef REGISTER_CUDA_TRANSFORM_KERNEL diff --git a/infini_train/src/nn/functional.cc b/infini_train/src/nn/functional.cc index b01e88b..c493e2e 100644 --- a/infini_train/src/nn/functional.cc +++ b/infini_train/src/nn/functional.cc @@ -1,65 +1,65 @@ -#include "infini_train/include/nn/functional.h" - -#include -#include -#include - -#include "infini_train/include/autograd/activations.h" -#include "infini_train/include/autograd/elementwise.h" -#include "infini_train/include/autograd/misc.h" -#include "infini_train/include/autograd/reduction.h" -#include "infini_train/include/autograd/softmax.h" -#include "infini_train/include/autograd/transform.h" -#include "infini_train/include/nn/init.h" - -namespace infini_train::nn::function { -std::shared_ptr Tril(const std::shared_ptr &input, int64_t diagonal) { - return std::make_shared(diagonal)->Apply({input})[0]; -} - -std::shared_ptr Triu(const std::shared_ptr &input, int64_t diagonal) { - return std::make_shared(diagonal)->Apply({input})[0]; -} - -std::shared_ptr Ones(const std::vector size) { - auto ones = std::make_shared(size, DataType::kFLOAT32); - return init::Ones(ones); -} - -std::shared_ptr Reciprocal(const std::shared_ptr &input) { return input->Reciprocal(); } - -std::shared_ptr Sin(const std::shared_ptr &input) { return input->Sin(); } - -std::shared_ptr Cos(const std::shared_ptr &input) { return input->Cos(); } - -std::shared_ptr Tanh(const std::shared_ptr &input) { return input->Tanh(); } - -std::shared_ptr Pow(const std::shared_ptr &input, float exponent) { return input->Pow(exponent); } - -std::shared_ptr Pow(float base, const std::shared_ptr &input) { - return std::make_shared(base, true)->Apply({input})[0]; -} - -std::shared_ptr Rsqrt(const std::shared_ptr &input) { return input->Rsqrt(); } - -std::shared_ptr Mean(const std::shared_ptr &input, int64_t dim, bool keep_dim) { - return std::make_shared(dim, keep_dim)->Apply({input})[0]; -} - -std::shared_ptr Slice(const std::shared_ptr &input, const std::vector &starts, - const std::vector &ends, const std::vector &steps) { - return input->Slice(starts, ends, steps); -} - -std::shared_ptr Stack(const std::vector> &inputs, int64_t dim) { - return std::make_shared(dim)->Apply(inputs)[0]; -} - -std::shared_ptr Softmax(const std::shared_ptr &input, int64_t dim) { - return std::make_shared(dim)->Apply({input})[0]; -} - -std::shared_ptr Sigmoid(const std::shared_ptr &input) { - return std::make_shared()->Apply({input})[0]; -} -} // namespace infini_train::nn::function +#include "infini_train/include/nn/functional.h" + +#include +#include +#include + +#include "infini_train/include/autograd/activations.h" +#include "infini_train/include/autograd/elementwise.h" +#include "infini_train/include/autograd/misc.h" +#include "infini_train/include/autograd/reduction.h" +#include "infini_train/include/autograd/softmax.h" +#include "infini_train/include/autograd/transform.h" +#include "infini_train/include/nn/init.h" + +namespace infini_train::nn::function { +std::shared_ptr Tril(const std::shared_ptr &input, int64_t diagonal) { + return std::make_shared(diagonal)->Apply({input})[0]; +} + +std::shared_ptr Triu(const std::shared_ptr &input, int64_t diagonal) { + return std::make_shared(diagonal)->Apply({input})[0]; +} + +std::shared_ptr Ones(const std::vector size) { + auto ones = std::make_shared(size, DataType::kFLOAT32); + return init::Ones(ones); +} + +std::shared_ptr Reciprocal(const std::shared_ptr &input) { return input->Reciprocal(); } + +std::shared_ptr Sin(const std::shared_ptr &input) { return input->Sin(); } + +std::shared_ptr Cos(const std::shared_ptr &input) { return input->Cos(); } + +std::shared_ptr Tanh(const std::shared_ptr &input) { return input->Tanh(); } + +std::shared_ptr Pow(const std::shared_ptr &input, float exponent) { return input->Pow(exponent); } + +std::shared_ptr Pow(float base, const std::shared_ptr &input) { + return std::make_shared(base, true)->Apply({input})[0]; +} + +std::shared_ptr Rsqrt(const std::shared_ptr &input) { return input->Rsqrt(); } + +std::shared_ptr Mean(const std::shared_ptr &input, int64_t dim, bool keep_dim) { + return std::make_shared(dim, keep_dim)->Apply({input})[0]; +} + +std::shared_ptr Slice(const std::shared_ptr &input, const std::vector &starts, + const std::vector &ends, const std::vector &steps) { + return input->Slice(starts, ends, steps); +} + +std::shared_ptr Stack(const std::vector> &inputs, int64_t dim) { + return std::make_shared(dim)->Apply(inputs)[0]; +} + +std::shared_ptr Softmax(const std::shared_ptr &input, int64_t dim) { + return std::make_shared(dim)->Apply({input})[0]; +} + +std::shared_ptr Sigmoid(const std::shared_ptr &input) { + return std::make_shared()->Apply({input})[0]; +} +} // namespace infini_train::nn::function diff --git a/infini_train/src/nn/init.cc b/infini_train/src/nn/init.cc index 4c1a05b..f1aa5a6 100644 --- a/infini_train/src/nn/init.cc +++ b/infini_train/src/nn/init.cc @@ -1,249 +1,249 @@ -#include "infini_train/include/nn/init.h" - -#include -#include -#include -#include -#include -#include -#include - -#include "glog/logging.h" -#ifdef USE_CUDA -#include "cuda_runtime_api.h" -#endif - -#include "infini_train/include/tensor.h" - -namespace infini_train::nn::init { -namespace { -static std::random_device rd; -static std::mt19937 gen(rd()); -} // namespace - -std::shared_ptr Normal(const std::shared_ptr &tensor, float mean, float std, - std::optional generator) { - const int64_t num_elements = tensor->NumElements(); - std::vector buffer(num_elements); - std::normal_distribution dis(mean, std); - std::generate(buffer.begin(), buffer.end(), [&]() { return generator ? dis(generator.value()) : dis(gen); }); - - switch (tensor->GetDevice().Type()) { - case DeviceType::kCPU: { - memcpy(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float)); - break; - } -#ifdef USE_CUDA - case DeviceType::kCUDA: { - cudaMemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float), cudaMemcpyHostToDevice, 0); - break; - } -#endif - default: { - LOG(FATAL) << "Unsupported device type: " << static_cast(tensor->GetDevice().Type()); - break; - } - } - return tensor; -} - -std::pair CalculateFanInAndFanOut(const std::shared_ptr &tensor) { - if (tensor->Dims().size() < 2) { - LOG(FATAL) << "Fan in and fan out can not be computed for tensor with less than 2 dimensions"; - } - const auto num_input_fmaps = tensor->Dims()[1]; - const auto num_output_fmaps = tensor->Dims()[0]; - int64_t receptive_field_size = 1; - if (tensor->Dims().size() > 2) { - receptive_field_size - *= std::accumulate(tensor->Dims().begin() + 2, tensor->Dims().end(), 1, std::multiplies()); - } - const auto fan_in = num_input_fmaps * receptive_field_size; - const auto fan_out = num_output_fmaps * receptive_field_size; - return {fan_in, fan_out}; -} - -namespace { -int64_t CalculateCorrectFan(const std::shared_ptr &tensor, KaimingMode mode) { - const auto [fan_in, fan_out] = CalculateFanInAndFanOut(tensor); - return mode == KaimingMode::kFanIn ? fan_in : fan_out; -} - -float CalculateGain(NonLinearityType nonlinearity, std::optional param = std::nullopt) { - static std::unordered_set kLinearFns = { - NonLinearityType::kLinear, NonLinearityType::kConv1D, NonLinearityType::kConv2D, - NonLinearityType::kConv3D, NonLinearityType::kConvTransposed1d, NonLinearityType::kConvTransposed2d, - NonLinearityType::kConvTransposed3d, - }; - if (kLinearFns.contains(nonlinearity) || nonlinearity == NonLinearityType::kSigmoid) { - return 1.0f; - } else if (nonlinearity == NonLinearityType::kTanh) { - return 5.0f / 3; - } else if (nonlinearity == NonLinearityType::kReLU) { - return sqrt(2.0f); - } else if (nonlinearity == NonLinearityType::kLeakyReLU) { - const float negative_slope = param ? *param : 0.01f; - return sqrt(2.0f / (1 + negative_slope * negative_slope)); - } else if (nonlinearity == NonLinearityType::kSELU) { - return 3.0f / 4; // Value found empirically (https://github.com/pytorch/pytorch/pull/50664) - } else { - LOG(FATAL) << "Unsupported non-linearity type: " << static_cast(nonlinearity); - } - return -1.0f; -} -} // namespace - -std::shared_ptr KaimingUniform(const std::shared_ptr &tensor, float a, KaimingMode mode, - NonLinearityType nonlinearity, std::optional generator) { - for (const auto dim : tensor->Dims()) { - if (dim == 0) { - LOG(WARNING) << "Initializing zero-element tensors is a no-op"; - return tensor; - } - } - const auto fan = CalculateCorrectFan(tensor, mode); - const auto gain = CalculateGain(nonlinearity, a); - const float std = gain / sqrt(fan); - const float bound = sqrt(3.0f) * std; // Calculate uniform bounds from standard deviation - return tensor->Uniform(-bound, bound, generator); -} - -std::shared_ptr Uniform(const std::shared_ptr &tensor, float a, float b, - std::optional generator) { - const int64_t num_elements = tensor->NumElements(); - std::vector buffer(num_elements); - std::uniform_real_distribution dis(a, b); - std::generate(buffer.begin(), buffer.end(), [&]() { return generator ? dis(generator.value()) : dis(gen); }); - - switch (tensor->GetDevice().Type()) { - case DeviceType::kCPU: { - memcpy(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float)); - break; - } -#ifdef USE_CUDA - case DeviceType::kCUDA: { - cudaMemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float), cudaMemcpyHostToDevice, 0); - break; - } -#endif - default: { - LOG(FATAL) << "Unsupported device type: " << static_cast(tensor->GetDevice().Type()); - break; - } - } - return tensor; -} - -std::shared_ptr Ones(const std::shared_ptr &tensor) { - CHECK_EQ(static_cast(tensor->Dtype()), static_cast(DataType::kFLOAT32)); - const int64_t num_elements = tensor->NumElements(); - std::vector buffer(num_elements, 1.0f); - - switch (tensor->GetDevice().Type()) { - case DeviceType::kCPU: { - memcpy(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float)); - break; - } -#ifdef USE_CUDA - case DeviceType::kCUDA: { - cudaMemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float), cudaMemcpyHostToDevice, 0); - break; - } -#endif - default: { - LOG(FATAL) << "Unsupported device type: " << static_cast(tensor->GetDevice().Type()); - break; - } - } - return tensor; -} - -std::shared_ptr Zeros(const std::shared_ptr &tensor) { - CHECK_EQ(static_cast(tensor->Dtype()), static_cast(DataType::kFLOAT32)); - const int64_t num_elements = tensor->NumElements(); - std::vector buffer(num_elements, 0.0f); - - switch (tensor->GetDevice().Type()) { - case DeviceType::kCPU: { - memcpy(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float)); - break; - } -#ifdef USE_CUDA - case DeviceType::kCUDA: { - cudaMemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float), cudaMemcpyHostToDevice, 0); - break; - } -#endif - default: { - LOG(FATAL) << "Unsupported device type: " << static_cast(tensor->GetDevice().Type()); - break; - } - } - return tensor; -} - -#define CASE(DATA_TYPE, TYPE) \ - case DATA_TYPE: { \ - std::vector buffer(num_elements); \ - std::iota(buffer.begin(), buffer.end(), static_cast(start)); \ - memcpy(tensor->DataPtr(), buffer.data(), num_elements * sizeof(TYPE)); \ - break; \ - } -#define CUDA_CASE(DATA_TYPE, TYPE) \ - case DATA_TYPE: { \ - std::vector buffer(num_elements); \ - std::iota(buffer.begin(), buffer.end(), static_cast(start)); \ - cudaMemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(TYPE), cudaMemcpyHostToDevice, 0); \ - break; \ - } - -std::shared_ptr Arange(int64_t start, int64_t end, DataType dtype, Device device) { - int64_t num_elements = end - start; - auto tensor = std::make_shared(std::vector{num_elements}, dtype, device); - if (device == Device()) { - switch (dtype) { - CASE(DataType::kUINT8, uint8_t) - CASE(DataType::kINT8, int8_t) - CASE(DataType::kUINT16, uint16_t) - CASE(DataType::kINT16, int16_t) - CASE(DataType::kUINT32, uint32_t) - CASE(DataType::kINT32, int32_t) - CASE(DataType::kUINT64, uint64_t) - CASE(DataType::kINT64, int64_t) - // CASE(DataType::kBFLOAT16, bf16) - // CASE(DataType::kFLOAT16, fp16) - CASE(DataType::kFLOAT32, float) - CASE(DataType::kFLOAT64, double) - default: - LOG(FATAL) << "Unsupported data type: " << static_cast(dtype); - break; - } - } else { -#ifdef USE_CUDA - switch (dtype) { - CUDA_CASE(DataType::kUINT8, uint8_t) - CUDA_CASE(DataType::kINT8, int8_t) - CUDA_CASE(DataType::kUINT16, uint16_t) - CUDA_CASE(DataType::kINT16, int16_t) - CUDA_CASE(DataType::kUINT32, uint32_t) - CUDA_CASE(DataType::kINT32, int32_t) - CUDA_CASE(DataType::kUINT64, uint64_t) - CUDA_CASE(DataType::kINT64, int64_t) - // CUDA_CASE(DataType::kBFLOAT16, bf16) - // CUDA_CASE(DataType::kFLOAT16, fp16) - CUDA_CASE(DataType::kFLOAT32, float) - CUDA_CASE(DataType::kFLOAT64, double) - default: - LOG(FATAL) << "Unsupported data type: " << static_cast(dtype); - break; - } -#else - LOG(FATAL) << "Unsupported device type: " << static_cast(device.Type()); -#endif - } - return tensor; -} - -#undef CASE -#undef CUDA_CASE -} // namespace infini_train::nn::init +#include "infini_train/include/nn/init.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" +#ifdef USE_CUDA +#include "cuda_runtime_api.h" +#endif + +#include "infini_train/include/tensor.h" + +namespace infini_train::nn::init { +namespace { +static std::random_device rd; +static std::mt19937 gen(rd()); +} // namespace + +std::shared_ptr Normal(const std::shared_ptr &tensor, float mean, float std, + std::optional generator) { + const int64_t num_elements = tensor->NumElements(); + std::vector buffer(num_elements); + std::normal_distribution dis(mean, std); + std::generate(buffer.begin(), buffer.end(), [&]() { return generator ? dis(generator.value()) : dis(gen); }); + + switch (tensor->GetDevice().Type()) { + case DeviceType::kCPU: { + memcpy(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float)); + break; + } +#ifdef USE_CUDA + case DeviceType::kCUDA: { + cudaMemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float), cudaMemcpyHostToDevice, 0); + break; + } +#endif + default: { + LOG(FATAL) << "Unsupported device type: " << static_cast(tensor->GetDevice().Type()); + break; + } + } + return tensor; +} + +std::pair CalculateFanInAndFanOut(const std::shared_ptr &tensor) { + if (tensor->Dims().size() < 2) { + LOG(FATAL) << "Fan in and fan out can not be computed for tensor with less than 2 dimensions"; + } + const auto num_input_fmaps = tensor->Dims()[1]; + const auto num_output_fmaps = tensor->Dims()[0]; + int64_t receptive_field_size = 1; + if (tensor->Dims().size() > 2) { + receptive_field_size + *= std::accumulate(tensor->Dims().begin() + 2, tensor->Dims().end(), 1, std::multiplies()); + } + const auto fan_in = num_input_fmaps * receptive_field_size; + const auto fan_out = num_output_fmaps * receptive_field_size; + return {fan_in, fan_out}; +} + +namespace { +int64_t CalculateCorrectFan(const std::shared_ptr &tensor, KaimingMode mode) { + const auto [fan_in, fan_out] = CalculateFanInAndFanOut(tensor); + return mode == KaimingMode::kFanIn ? fan_in : fan_out; +} + +float CalculateGain(NonLinearityType nonlinearity, std::optional param = std::nullopt) { + static std::unordered_set kLinearFns = { + NonLinearityType::kLinear, NonLinearityType::kConv1D, NonLinearityType::kConv2D, + NonLinearityType::kConv3D, NonLinearityType::kConvTransposed1d, NonLinearityType::kConvTransposed2d, + NonLinearityType::kConvTransposed3d, + }; + if (kLinearFns.contains(nonlinearity) || nonlinearity == NonLinearityType::kSigmoid) { + return 1.0f; + } else if (nonlinearity == NonLinearityType::kTanh) { + return 5.0f / 3; + } else if (nonlinearity == NonLinearityType::kReLU) { + return sqrt(2.0f); + } else if (nonlinearity == NonLinearityType::kLeakyReLU) { + const float negative_slope = param ? *param : 0.01f; + return sqrt(2.0f / (1 + negative_slope * negative_slope)); + } else if (nonlinearity == NonLinearityType::kSELU) { + return 3.0f / 4; // Value found empirically (https://github.com/pytorch/pytorch/pull/50664) + } else { + LOG(FATAL) << "Unsupported non-linearity type: " << static_cast(nonlinearity); + } + return -1.0f; +} +} // namespace + +std::shared_ptr KaimingUniform(const std::shared_ptr &tensor, float a, KaimingMode mode, + NonLinearityType nonlinearity, std::optional generator) { + for (const auto dim : tensor->Dims()) { + if (dim == 0) { + LOG(WARNING) << "Initializing zero-element tensors is a no-op"; + return tensor; + } + } + const auto fan = CalculateCorrectFan(tensor, mode); + const auto gain = CalculateGain(nonlinearity, a); + const float std = gain / sqrt(fan); + const float bound = sqrt(3.0f) * std; // Calculate uniform bounds from standard deviation + return tensor->Uniform(-bound, bound, generator); +} + +std::shared_ptr Uniform(const std::shared_ptr &tensor, float a, float b, + std::optional generator) { + const int64_t num_elements = tensor->NumElements(); + std::vector buffer(num_elements); + std::uniform_real_distribution dis(a, b); + std::generate(buffer.begin(), buffer.end(), [&]() { return generator ? dis(generator.value()) : dis(gen); }); + + switch (tensor->GetDevice().Type()) { + case DeviceType::kCPU: { + memcpy(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float)); + break; + } +#ifdef USE_CUDA + case DeviceType::kCUDA: { + cudaMemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float), cudaMemcpyHostToDevice, 0); + break; + } +#endif + default: { + LOG(FATAL) << "Unsupported device type: " << static_cast(tensor->GetDevice().Type()); + break; + } + } + return tensor; +} + +std::shared_ptr Ones(const std::shared_ptr &tensor) { + CHECK_EQ(static_cast(tensor->Dtype()), static_cast(DataType::kFLOAT32)); + const int64_t num_elements = tensor->NumElements(); + std::vector buffer(num_elements, 1.0f); + + switch (tensor->GetDevice().Type()) { + case DeviceType::kCPU: { + memcpy(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float)); + break; + } +#ifdef USE_CUDA + case DeviceType::kCUDA: { + cudaMemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float), cudaMemcpyHostToDevice, 0); + break; + } +#endif + default: { + LOG(FATAL) << "Unsupported device type: " << static_cast(tensor->GetDevice().Type()); + break; + } + } + return tensor; +} + +std::shared_ptr Zeros(const std::shared_ptr &tensor) { + CHECK_EQ(static_cast(tensor->Dtype()), static_cast(DataType::kFLOAT32)); + const int64_t num_elements = tensor->NumElements(); + std::vector buffer(num_elements, 0.0f); + + switch (tensor->GetDevice().Type()) { + case DeviceType::kCPU: { + memcpy(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float)); + break; + } +#ifdef USE_CUDA + case DeviceType::kCUDA: { + cudaMemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(float), cudaMemcpyHostToDevice, 0); + break; + } +#endif + default: { + LOG(FATAL) << "Unsupported device type: " << static_cast(tensor->GetDevice().Type()); + break; + } + } + return tensor; +} + +#define CASE(DATA_TYPE, TYPE) \ + case DATA_TYPE: { \ + std::vector buffer(num_elements); \ + std::iota(buffer.begin(), buffer.end(), static_cast(start)); \ + memcpy(tensor->DataPtr(), buffer.data(), num_elements * sizeof(TYPE)); \ + break; \ + } +#define CUDA_CASE(DATA_TYPE, TYPE) \ + case DATA_TYPE: { \ + std::vector buffer(num_elements); \ + std::iota(buffer.begin(), buffer.end(), static_cast(start)); \ + cudaMemcpyAsync(tensor->DataPtr(), buffer.data(), num_elements * sizeof(TYPE), cudaMemcpyHostToDevice, 0); \ + break; \ + } + +std::shared_ptr Arange(int64_t start, int64_t end, DataType dtype, Device device) { + int64_t num_elements = end - start; + auto tensor = std::make_shared(std::vector{num_elements}, dtype, device); + if (device == Device()) { + switch (dtype) { + CASE(DataType::kUINT8, uint8_t) + CASE(DataType::kINT8, int8_t) + CASE(DataType::kUINT16, uint16_t) + CASE(DataType::kINT16, int16_t) + CASE(DataType::kUINT32, uint32_t) + CASE(DataType::kINT32, int32_t) + CASE(DataType::kUINT64, uint64_t) + CASE(DataType::kINT64, int64_t) + // CASE(DataType::kBFLOAT16, bf16) + // CASE(DataType::kFLOAT16, fp16) + CASE(DataType::kFLOAT32, float) + CASE(DataType::kFLOAT64, double) + default: + LOG(FATAL) << "Unsupported data type: " << static_cast(dtype); + break; + } + } else { +#ifdef USE_CUDA + switch (dtype) { + CUDA_CASE(DataType::kUINT8, uint8_t) + CUDA_CASE(DataType::kINT8, int8_t) + CUDA_CASE(DataType::kUINT16, uint16_t) + CUDA_CASE(DataType::kINT16, int16_t) + CUDA_CASE(DataType::kUINT32, uint32_t) + CUDA_CASE(DataType::kINT32, int32_t) + CUDA_CASE(DataType::kUINT64, uint64_t) + CUDA_CASE(DataType::kINT64, int64_t) + // CUDA_CASE(DataType::kBFLOAT16, bf16) + // CUDA_CASE(DataType::kFLOAT16, fp16) + CUDA_CASE(DataType::kFLOAT32, float) + CUDA_CASE(DataType::kFLOAT64, double) + default: + LOG(FATAL) << "Unsupported data type: " << static_cast(dtype); + break; + } +#else + LOG(FATAL) << "Unsupported device type: " << static_cast(device.Type()); +#endif + } + return tensor; +} + +#undef CASE +#undef CUDA_CASE +} // namespace infini_train::nn::init diff --git a/infini_train/src/nn/modules/activations.cc b/infini_train/src/nn/modules/activations.cc index 77b7ad3..80d8891 100644 --- a/infini_train/src/nn/modules/activations.cc +++ b/infini_train/src/nn/modules/activations.cc @@ -1,13 +1,13 @@ -#include "infini_train/include/nn/modules/activations.h" - -#include -#include - -#include "infini_train/include/autograd/activations.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -std::vector> Sigmoid::Forward(const std::vector> &input_tensors) { - return std::make_shared()->Apply(input_tensors); -} -} // namespace infini_train::nn +#include "infini_train/include/nn/modules/activations.h" + +#include +#include + +#include "infini_train/include/autograd/activations.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +std::vector> Sigmoid::Forward(const std::vector> &input_tensors) { + return std::make_shared()->Apply(input_tensors); +} +} // namespace infini_train::nn diff --git a/infini_train/src/nn/modules/container.cc b/infini_train/src/nn/modules/container.cc index 8e2cb3d..19fc05b 100644 --- a/infini_train/src/nn/modules/container.cc +++ b/infini_train/src/nn/modules/container.cc @@ -1,31 +1,31 @@ -#include "infini_train/include/nn/modules/container.h" - -#include -#include -#include - -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -Sequential::Sequential(std::vector> &&layers) { - int idx = 0; - for (auto &layer : layers) { - modules_[std::to_string(idx)] = std::move(layer); - ++idx; - } -} - -std::vector> Sequential::Forward(const std::vector> &input_tensors) { - auto &x = const_cast> &>(input_tensors); - for (int idx = 0; idx < modules_.size(); ++idx) { x = modules_[std::to_string(idx)]->Forward(x); } - return x; -} - -ModuleDict::ModuleDict(std::unordered_map> &&modules) { - for (auto &[name, layer] : modules) { modules_[name] = std::move(layer); } -} - -std::vector> ModuleDict::Forward(const std::vector> &input_tensors) { - LOG(FATAL) << "Not implemented"; -} -} // namespace infini_train::nn +#include "infini_train/include/nn/modules/container.h" + +#include +#include +#include + +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +Sequential::Sequential(std::vector> &&layers) { + int idx = 0; + for (auto &layer : layers) { + modules_[std::to_string(idx)] = std::move(layer); + ++idx; + } +} + +std::vector> Sequential::Forward(const std::vector> &input_tensors) { + auto &x = const_cast> &>(input_tensors); + for (int idx = 0; idx < modules_.size(); ++idx) { x = modules_[std::to_string(idx)]->Forward(x); } + return x; +} + +ModuleDict::ModuleDict(std::unordered_map> &&modules) { + for (auto &[name, layer] : modules) { modules_[name] = std::move(layer); } +} + +std::vector> ModuleDict::Forward(const std::vector> &input_tensors) { + LOG(FATAL) << "Not implemented"; +} +} // namespace infini_train::nn diff --git a/infini_train/src/nn/modules/linear.cc b/infini_train/src/nn/modules/linear.cc index b385f07..e2ace45 100644 --- a/infini_train/src/nn/modules/linear.cc +++ b/infini_train/src/nn/modules/linear.cc @@ -1,41 +1,41 @@ -#include "infini_train/include/nn/modules/linear.h" - -#include -#include -#include - -#include "infini_train/include/autograd/linear.h" -#include "infini_train/include/device.h" -#include "infini_train/include/nn/init.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -Linear::Linear(int64_t in_features, int64_t out_features, bool bias, Device device) : Module(kType), bias_(bias) { - device_ = device; - - parameters_[kParamWeightName] - = std::make_shared(std::vector{out_features, in_features}, DataType::kFLOAT32, device) - ->RequiresGrad(); - if (bias) { - parameters_[kParamBiasName] - = std::make_shared(std::vector{out_features}, DataType::kFLOAT32, device)->RequiresGrad(); - } - ResetParameters(); -} - -std::vector> Linear::Forward(const std::vector> &input_tensors) { - return std::make_shared()->Apply( - bias_ ? std::vector>{input_tensors[0], parameters_[kParamWeightName], - parameters_[kParamBiasName]} - : std::vector>{input_tensors[0], parameters_[kParamWeightName]}); -} - -void Linear::ResetParameters() { - init::KaimingUniform(parameters_[kParamWeightName], sqrt(5.0f)); - if (bias_) { - const auto [fan_in, _] = init::CalculateFanInAndFanOut(parameters_[kParamWeightName]); - const float bound = fan_in > 0 ? 1.0 / sqrt(fan_in) : 0.0; - init::Uniform(parameters_[kParamBiasName], -bound, bound); - } -} -} // namespace infini_train::nn +#include "infini_train/include/nn/modules/linear.h" + +#include +#include +#include + +#include "infini_train/include/autograd/linear.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +Linear::Linear(int64_t in_features, int64_t out_features, bool bias, Device device) : Module(kType), bias_(bias) { + device_ = device; + + parameters_[kParamWeightName] + = std::make_shared(std::vector{out_features, in_features}, DataType::kFLOAT32, device) + ->RequiresGrad(); + if (bias) { + parameters_[kParamBiasName] + = std::make_shared(std::vector{out_features}, DataType::kFLOAT32, device)->RequiresGrad(); + } + ResetParameters(); +} + +std::vector> Linear::Forward(const std::vector> &input_tensors) { + return std::make_shared()->Apply( + bias_ ? std::vector>{input_tensors[0], parameters_[kParamWeightName], + parameters_[kParamBiasName]} + : std::vector>{input_tensors[0], parameters_[kParamWeightName]}); +} + +void Linear::ResetParameters() { + init::KaimingUniform(parameters_[kParamWeightName], sqrt(5.0f)); + if (bias_) { + const auto [fan_in, _] = init::CalculateFanInAndFanOut(parameters_[kParamWeightName]); + const float bound = fan_in > 0 ? 1.0 / sqrt(fan_in) : 0.0; + init::Uniform(parameters_[kParamBiasName], -bound, bound); + } +} +} // namespace infini_train::nn diff --git a/infini_train/src/nn/modules/loss.cc b/infini_train/src/nn/modules/loss.cc index a3f9763..ee07bd8 100644 --- a/infini_train/src/nn/modules/loss.cc +++ b/infini_train/src/nn/modules/loss.cc @@ -1,14 +1,14 @@ -#include "infini_train/include/nn/modules/loss.h" - -#include -#include - -#include "infini_train/include/autograd/loss.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -std::vector> -CrossEntropyLoss::Forward(const std::vector> &input_tensors) { - return std::make_shared()->Apply(input_tensors); -} -} // namespace infini_train::nn +#include "infini_train/include/nn/modules/loss.h" + +#include +#include + +#include "infini_train/include/autograd/loss.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +std::vector> +CrossEntropyLoss::Forward(const std::vector> &input_tensors) { + return std::make_shared()->Apply(input_tensors); +} +} // namespace infini_train::nn diff --git a/infini_train/src/nn/modules/module.cc b/infini_train/src/nn/modules/module.cc index b315075..b7d3d90 100644 --- a/infini_train/src/nn/modules/module.cc +++ b/infini_train/src/nn/modules/module.cc @@ -1,77 +1,77 @@ -#include "infini_train/include/nn/modules/module.h" - -#include -#include - -#include "infini_train/include/device.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -const std::string &Module::type() const { return type_; } - -std::vector> Module::Parameters() const { - std::vector> params; - for (auto &[_, param] : parameters_) { params.push_back(param); } - for (auto &[_, layer] : modules_) { - for (auto ¶m : layer->Parameters()) { params.push_back(param); } - } - return params; -} - -bool Module::has_parameter(const std::string &name) const { return parameters_.find(name) != parameters_.end(); } - -std::shared_ptr *Module::mutable_parameter(const std::string &name) { - CHECK(parameters_.find(name) != parameters_.end()); - return ¶meters_.at(name); -} - -const std::shared_ptr &Module::parameter(const std::string &name) const { - CHECK(parameters_.find(name) != parameters_.end()); - return parameters_.at(name); -} - -std::vector Module::modules() const { - std::vector modules; - for (auto &[_, module] : modules_) { modules.push_back(module.get()); } - return modules; -} - -Module *Module::mutable_module(const std::string &name) { - CHECK(modules_.find(name) != modules_.end()); - return modules_.at(name).get(); -} - -const Module &Module::module(const std::string &name) const { - CHECK(modules_.find(name) != modules_.end()); - return *modules_.at(name).get(); -} - -std::unordered_map> Module::StateDict() const { - std::unordered_map> state; - for (auto &[name, param] : parameters_) { state.emplace(name, param); } - for (auto &[name, layer] : modules_) { - for (auto &[sub_name, param] : layer->StateDict()) { state.emplace(name + "." + sub_name, param); } - } - return state; -} - -void Module::To(Device device) { - if (device == device_) { - return; - } - - std::unordered_map> new_parameters; - for (auto &[name, param] : parameters_) { - new_parameters.emplace(name, std::make_shared(param->To(device))); - } - parameters_ = std::move(new_parameters); - device_ = device; - - for (auto &[_, layer] : modules_) { layer->To(device); } -} - -void Module::Apply(std::function fn) { - for (auto *module : modules()) { module->Apply(fn); } - fn(this); -} -} // namespace infini_train::nn +#include "infini_train/include/nn/modules/module.h" + +#include +#include + +#include "infini_train/include/device.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +const std::string &Module::type() const { return type_; } + +std::vector> Module::Parameters() const { + std::vector> params; + for (auto &[_, param] : parameters_) { params.push_back(param); } + for (auto &[_, layer] : modules_) { + for (auto ¶m : layer->Parameters()) { params.push_back(param); } + } + return params; +} + +bool Module::has_parameter(const std::string &name) const { return parameters_.find(name) != parameters_.end(); } + +std::shared_ptr *Module::mutable_parameter(const std::string &name) { + CHECK(parameters_.find(name) != parameters_.end()); + return ¶meters_.at(name); +} + +const std::shared_ptr &Module::parameter(const std::string &name) const { + CHECK(parameters_.find(name) != parameters_.end()); + return parameters_.at(name); +} + +std::vector Module::modules() const { + std::vector modules; + for (auto &[_, module] : modules_) { modules.push_back(module.get()); } + return modules; +} + +Module *Module::mutable_module(const std::string &name) { + CHECK(modules_.find(name) != modules_.end()); + return modules_.at(name).get(); +} + +const Module &Module::module(const std::string &name) const { + CHECK(modules_.find(name) != modules_.end()); + return *modules_.at(name).get(); +} + +std::unordered_map> Module::StateDict() const { + std::unordered_map> state; + for (auto &[name, param] : parameters_) { state.emplace(name, param); } + for (auto &[name, layer] : modules_) { + for (auto &[sub_name, param] : layer->StateDict()) { state.emplace(name + "." + sub_name, param); } + } + return state; +} + +void Module::To(Device device) { + if (device == device_) { + return; + } + + std::unordered_map> new_parameters; + for (auto &[name, param] : parameters_) { + new_parameters.emplace(name, std::make_shared(param->To(device))); + } + parameters_ = std::move(new_parameters); + device_ = device; + + for (auto &[_, layer] : modules_) { layer->To(device); } +} + +void Module::Apply(std::function fn) { + for (auto *module : modules()) { module->Apply(fn); } + fn(this); +} +} // namespace infini_train::nn diff --git a/infini_train/src/nn/modules/normalization.cc b/infini_train/src/nn/modules/normalization.cc index 11db0af..c6946cc 100644 --- a/infini_train/src/nn/modules/normalization.cc +++ b/infini_train/src/nn/modules/normalization.cc @@ -1,31 +1,31 @@ -#include "infini_train/include/nn/modules/normalization.h" - -#include -#include - -#include "infini_train/include/autograd/normalization.h" -#include "infini_train/include/device.h" -#include "infini_train/include/nn/init.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { -LayerNorm::LayerNorm(const std::vector &normalized_shape, float eps, Device device) : eps_(eps) { - device_ = device; - - parameters_[kParamWeightName] - = std::make_shared(normalized_shape, DataType::kFLOAT32, device)->RequiresGrad(); - parameters_[kParamBiasName] - = std::make_shared(normalized_shape, DataType::kFLOAT32, device)->RequiresGrad(); - ResetParameters(); -} - -std::vector> LayerNorm::Forward(const std::vector> &input_tensors) { - return std::make_shared(eps_)->Apply( - {input_tensors[0], parameters_[kParamWeightName], parameters_[kParamBiasName]}); -} - -void LayerNorm::ResetParameters() { - init::Ones(parameters_[kParamWeightName]); - init::Zeros(parameters_[kParamBiasName]); -} -} // namespace infini_train::nn +#include "infini_train/include/nn/modules/normalization.h" + +#include +#include + +#include "infini_train/include/autograd/normalization.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { +LayerNorm::LayerNorm(const std::vector &normalized_shape, float eps, Device device) : eps_(eps) { + device_ = device; + + parameters_[kParamWeightName] + = std::make_shared(normalized_shape, DataType::kFLOAT32, device)->RequiresGrad(); + parameters_[kParamBiasName] + = std::make_shared(normalized_shape, DataType::kFLOAT32, device)->RequiresGrad(); + ResetParameters(); +} + +std::vector> LayerNorm::Forward(const std::vector> &input_tensors) { + return std::make_shared(eps_)->Apply( + {input_tensors[0], parameters_[kParamWeightName], parameters_[kParamBiasName]}); +} + +void LayerNorm::ResetParameters() { + init::Ones(parameters_[kParamWeightName]); + init::Zeros(parameters_[kParamBiasName]); +} +} // namespace infini_train::nn diff --git a/infini_train/src/nn/modules/sparse.cc b/infini_train/src/nn/modules/sparse.cc index 57223ca..4fc9930 100644 --- a/infini_train/src/nn/modules/sparse.cc +++ b/infini_train/src/nn/modules/sparse.cc @@ -1,27 +1,27 @@ -#include "infini_train/include/nn/modules/sparse.h" - -#include -#include - -#include "infini_train/include/autograd/sparse.h" -#include "infini_train/include/device.h" -#include "infini_train/include/nn/init.h" -#include "infini_train/include/tensor.h" - -namespace infini_train::nn { - -Embedding::Embedding(int num_embeddings, int embedding_dim, Device device) : Module(kType) { - device_ = device; - - parameters_[kParamWeightName] - = std::make_shared(std::vector{num_embeddings, embedding_dim}, DataType::kFLOAT32, device) - ->RequiresGrad(); - ResetParameters(); -} - -std::vector> Embedding::Forward(const std::vector> &input_tensors) { - return std::make_shared()->Apply({input_tensors[0], parameters_[kParamWeightName]}); -} - -void Embedding::ResetParameters() { init::Normal(parameters_[kParamWeightName]); } -} // namespace infini_train::nn +#include "infini_train/include/nn/modules/sparse.h" + +#include +#include + +#include "infini_train/include/autograd/sparse.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/init.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::nn { + +Embedding::Embedding(int num_embeddings, int embedding_dim, Device device) : Module(kType) { + device_ = device; + + parameters_[kParamWeightName] + = std::make_shared(std::vector{num_embeddings, embedding_dim}, DataType::kFLOAT32, device) + ->RequiresGrad(); + ResetParameters(); +} + +std::vector> Embedding::Forward(const std::vector> &input_tensors) { + return std::make_shared()->Apply({input_tensors[0], parameters_[kParamWeightName]}); +} + +void Embedding::ResetParameters() { init::Normal(parameters_[kParamWeightName]); } +} // namespace infini_train::nn diff --git a/infini_train/src/optimizer.cc b/infini_train/src/optimizer.cc index 9ff2b5c..b89a987 100644 --- a/infini_train/src/optimizer.cc +++ b/infini_train/src/optimizer.cc @@ -1,55 +1,55 @@ -#include "infini_train/include/optimizer.h" - -#include - -#include "infini_train/include/device.h" -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/tensor.h" - -namespace infini_train { -Optimizer::Optimizer(const std::vector> ¶ms) : params_(params) {} - -void Optimizer::ZeroGrad() { - for (auto param : params_) { param->ZeroGrad(); } -} - -namespace optimizers { - -SGD::SGD(const std::vector> ¶ms, float learning_rate) - : Optimizer(params), learning_rate_(learning_rate) {} - -void SGD::Step() { - for (auto param : params_) { - auto device = param->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "AccumulateGrad"}); - kernel.Call(param->grad(), -learning_rate_, param); - } -} - -Adam::Adam(const std::vector> ¶ms, float learning_rate, float beta1, float beta2, float eps) - : Optimizer(params), t_(0), learning_rate_(learning_rate), beta1_(beta1), beta2_(beta2), eps_(eps) { - - for (const auto ¶m : params_) { - m_.emplace_back(std::make_shared(param->Dims(), param->Dtype(), param->GetDevice())); - v_.emplace_back(std::make_shared(param->Dims(), param->Dtype(), param->GetDevice())); - m_.back()->Fill(0.0f); - v_.back()->Fill(0.0f); - } -} - -void Adam::Step() { - ++t_; - - for (size_t i = 0; i < params_.size(); ++i) { - auto ¶m = params_[i]; - const auto &grad = param->grad(); - auto &m = m_[i]; - auto &v = v_[i]; - - auto device = param->GetDevice().Type(); - auto kernel = Dispatcher::Instance().GetKernel({device, "AdamAccumulateGrad"}); - kernel.Call(grad, param, m, v, learning_rate_, beta1_, beta2_, eps_, t_); - } -} -} // namespace optimizers -} // namespace infini_train +#include "infini_train/include/optimizer.h" + +#include + +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train { +Optimizer::Optimizer(const std::vector> ¶ms) : params_(params) {} + +void Optimizer::ZeroGrad() { + for (auto param : params_) { param->ZeroGrad(); } +} + +namespace optimizers { + +SGD::SGD(const std::vector> ¶ms, float learning_rate) + : Optimizer(params), learning_rate_(learning_rate) {} + +void SGD::Step() { + for (auto param : params_) { + auto device = param->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "AccumulateGrad"}); + kernel.Call(param->grad(), -learning_rate_, param); + } +} + +Adam::Adam(const std::vector> ¶ms, float learning_rate, float beta1, float beta2, float eps) + : Optimizer(params), t_(0), learning_rate_(learning_rate), beta1_(beta1), beta2_(beta2), eps_(eps) { + + for (const auto ¶m : params_) { + m_.emplace_back(std::make_shared(param->Dims(), param->Dtype(), param->GetDevice())); + v_.emplace_back(std::make_shared(param->Dims(), param->Dtype(), param->GetDevice())); + m_.back()->Fill(0.0f); + v_.back()->Fill(0.0f); + } +} + +void Adam::Step() { + ++t_; + + for (size_t i = 0; i < params_.size(); ++i) { + auto ¶m = params_[i]; + const auto &grad = param->grad(); + auto &m = m_[i]; + auto &v = v_[i]; + + auto device = param->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "AdamAccumulateGrad"}); + kernel.Call(grad, param, m, v, learning_rate_, beta1_, beta2_, eps_, t_); + } +} +} // namespace optimizers +} // namespace infini_train diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc index 8f8c744..d475c8c 100644 --- a/infini_train/src/tensor.cc +++ b/infini_train/src/tensor.cc @@ -1,717 +1,761 @@ -#include "infini_train/include/tensor.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -#ifdef USE_CUDA -#include "cuda_runtime_api.h" -#endif -#include "Eigen/Dense" -#include "glog/logging.h" - -#include "infini_train/include/autograd/elementwise.h" -#include "infini_train/include/autograd/matmul.h" -#include "infini_train/include/autograd/misc.h" -#include "infini_train/include/autograd/outer.h" -#include "infini_train/include/autograd/transform.h" -#include "infini_train/include/device.h" -#include "infini_train/include/dispatcher.h" -#include "infini_train/include/nn/init.h" - -namespace infini_train { -namespace { -const std::unordered_map kDataTypeToSize = { - {DataType::kUINT8, 1}, {DataType::kINT8, 1}, {DataType::kUINT16, 2}, {DataType::kINT16, 2}, - {DataType::kUINT32, 4}, {DataType::kINT32, 4}, {DataType::kUINT64, 8}, {DataType::kINT64, 8}, - {DataType::kBFLOAT16, 2}, {DataType::kFLOAT16, 2}, {DataType::kFLOAT32, 4}, {DataType::kFLOAT64, 8}, -}; - -const std::unordered_map kDataTypeToDesc = { - {DataType::kUINT8, "uint8"}, {DataType::kINT8, "int8"}, {DataType::kUINT16, "uint16"}, - {DataType::kINT16, "int16"}, {DataType::kUINT32, "uint32"}, {DataType::kINT32, "int32"}, - {DataType::kUINT64, "uint64"}, {DataType::kINT64, "int64"}, {DataType::kBFLOAT16, "bf16"}, - {DataType::kFLOAT16, "fp16"}, {DataType::kFLOAT32, "fp32"}, {DataType::kFLOAT64, "fp64"}, -}; - -template struct TypeMap; - -template <> struct TypeMap { - using type = float; -}; -template <> struct TypeMap { - using type = double; -}; -template <> struct TypeMap { - using type = int32_t; -}; -template <> struct TypeMap { - using type = int64_t; -}; -} // namespace - -TensorBuffer::TensorBuffer(Device device, size_t size) : device_(device), size_(size) { - switch (device_.Type()) { - case DeviceType::kCPU: - data_ = malloc(size); - break; -#ifdef USE_CUDA - case DeviceType::kCUDA: - cudaMallocAsync(&data_, size, 0); - break; -#endif - default: - LOG(FATAL) << "Unsupported device type: " << static_cast(device_.Type()); - break; - } -} - -TensorBuffer::~TensorBuffer() { - switch (device_.Type()) { - case DeviceType::kCPU: - free(data_); - break; -#ifdef USE_CUDA - case DeviceType::kCUDA: - cudaFreeAsync(data_, 0); - break; -#endif - default: - LOG(FATAL) << "Unsupported device type: " << static_cast(device_.Type()); - break; - } -} - -void *TensorBuffer::DataPtr() { return data_; } - -const void *TensorBuffer::DataPtr() const { return data_; } - -Device TensorBuffer::GetDevice() const { return device_; } - -size_t TensorBuffer::Size() const { return size_; } - -// Tensor implementation -Tensor::Tensor(const std::vector &dims, DataType dtype, Device device) : dims_(dims), dtype_(dtype) { - num_elements_ = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); - buffer_ = std::make_shared(device, kDataTypeToSize.at(dtype) * num_elements_); -} - -Tensor::Tensor(const Tensor &tensor, size_t offset, const std::vector &dims) - : buffer_(tensor.buffer_), offset_(offset), dims_(dims), - num_elements_(std::accumulate(dims.begin(), dims.end(), 1, std::multiplies())), dtype_(tensor.dtype_) { - CHECK_LE(offset_ + kDataTypeToSize.at(dtype_) * num_elements_, buffer_->Size()); -} - -Device Tensor::GetDevice() const { return buffer_->GetDevice(); } - -void *Tensor::DataPtr() { return reinterpret_cast(buffer_->DataPtr()) + offset_; } - -const void *Tensor::DataPtr() const { return reinterpret_cast(buffer_->DataPtr()) + offset_; } - -size_t Tensor::SizeInBytes() const { return kDataTypeToSize.at(dtype_) * num_elements_; } - -const std::vector &Tensor::Dims() const { return dims_; } - -size_t Tensor::NumElements() const { return num_elements_; } - -DataType Tensor::Dtype() const { return dtype_; } - -template void Tensor::Fill(T value) { - DataType dtype = Dtype(); - - uint64_t storage = 0; - - switch (dtype) { - case DataType::kFLOAT32: { - using TargetT = typename TypeMap::type; - TargetT casted_value = static_cast(value); - std::memcpy(&storage, &casted_value, sizeof(TargetT)); - break; - } - case DataType::kFLOAT64: { - using TargetT = typename TypeMap::type; - TargetT casted_value = static_cast(value); - std::memcpy(&storage, &casted_value, sizeof(TargetT)); - break; - } - case DataType::kINT32: { - using TargetT = typename TypeMap::type; - TargetT casted_value = static_cast(value); - std::memcpy(&storage, &casted_value, sizeof(TargetT)); - break; - } - case DataType::kINT64: { - using TargetT = typename TypeMap::type; - TargetT casted_value = static_cast(value); - std::memcpy(&storage, &casted_value, sizeof(TargetT)); - break; - } - default: - throw std::runtime_error("Unsupported data type in Tensor::Fill()"); - } - - auto kernel = Dispatcher::Instance().GetKernel({GetDevice().Type(), "Fill"}); - kernel.Call(shared_from_this(), static_cast(&storage)); -} - -template void Tensor::Fill(float); - -Eigen::Map> Tensor::EigenMatrix() { - const int64_t bs = std::accumulate(dims_.rbegin() + 1, dims_.rend(), 1, std::multiplies()); - return Eigen::Map>( - reinterpret_cast(DataPtr()), bs, *dims_.rbegin()); -} - -Eigen::Map> Tensor::EigenVector() { - CHECK_EQ(dims_.size(), 1); - return Eigen::Map>(reinterpret_cast(DataPtr()), 1, - dims_[0]); -} - -Tensor Tensor::To(Device device) { - if (device == buffer_->GetDevice()) { - auto new_tensor = Tensor(*this, offset_, dims_); - if (grad_) { - new_tensor.grad_ = std::make_unique(*grad_.get(), grad_->offset_, grad_->dims_); - } - return new_tensor; - } - - Tensor new_tensor; - switch (device.Type()) { -#ifdef USE_CUDA - case DeviceType::kCPU: - // CUDA -> CPU - new_tensor = Tensor(dims_, dtype_, Device(DeviceType::kCPU, 0)); - cudaMemcpyAsync(new_tensor.DataPtr(), DataPtr(), SizeInBytes(), cudaMemcpyDeviceToHost, 0); - break; - case DeviceType::kCUDA: - // CPU -> CUDA - new_tensor = Tensor(dims_, dtype_, Device(DeviceType::kCUDA, 0)); - cudaMemcpyAsync(new_tensor.DataPtr(), DataPtr(), SizeInBytes(), cudaMemcpyHostToDevice, 0); - break; -#endif - default: - LOG(FATAL) << "Unsupported device type: " << static_cast(device.Type()); - } - - if (grad_) { - new_tensor.grad_ = std::make_unique(grad_->To(device)); - } - - new_tensor.requires_grad_ = requires_grad_; - - return new_tensor; -} - -// operator overloading -std::shared_ptr Tensor::Equals(float scalar) { - return std::make_shared(scalar)->Apply({shared_from_this()})[0]; -} - -std::shared_ptr Tensor::Add(const std::shared_ptr &other) { - CHECK_EQ(static_cast(GetDevice().Type()), static_cast(other->GetDevice().Type())); - return std::make_shared()->Apply({shared_from_this(), other})[0]; -} - -std::shared_ptr Tensor::Add(float scalar) { - return std::make_shared(scalar)->Apply({shared_from_this()})[0]; -} - -std::shared_ptr Tensor::Sub(const std::shared_ptr &other) { - CHECK_EQ(static_cast(GetDevice().Type()), static_cast(other->GetDevice().Type())); - return std::make_shared()->Apply({shared_from_this(), other})[0]; -} - -std::shared_ptr Tensor::Mul(const std::shared_ptr &other) { - CHECK_EQ(static_cast(GetDevice().Type()), static_cast(other->GetDevice().Type())); - return std::make_shared()->Apply({shared_from_this(), other})[0]; -} - -std::shared_ptr Tensor::Mul(float scalar) { - return std::make_shared(scalar)->Apply({shared_from_this()})[0]; -} - -std::shared_ptr Tensor::Div(const std::shared_ptr &other) { - CHECK_EQ(static_cast(GetDevice().Type()), static_cast(other->GetDevice().Type())); - return std::make_shared()->Apply({shared_from_this(), other})[0]; -} - -std::shared_ptr Tensor::Neg() { return std::make_shared()->Apply({shared_from_this()})[0]; } - -std::shared_ptr Tensor::Reciprocal() { - return std::make_shared()->Apply({shared_from_this()})[0]; -} - -std::shared_ptr Tensor::Sin() { return std::make_shared()->Apply({shared_from_this()})[0]; } - -std::shared_ptr Tensor::Cos() { return std::make_shared()->Apply({shared_from_this()})[0]; } - -std::shared_ptr Tensor::Tanh() { return std::make_shared()->Apply({shared_from_this()})[0]; } - -std::shared_ptr Tensor::Pow(float exponent) { - return std::make_shared(exponent)->Apply({shared_from_this()})[0]; -} - -std::shared_ptr Tensor::Rsqrt() { return std::make_shared()->Apply({shared_from_this()})[0]; } - -std::vector> Tensor::Split(int split_size, int dim) { - return std::make_shared(split_size, dim)->Apply({shared_from_this()}); -} - -std::shared_ptr Tensor::RepeatInterleave(int64_t repeat, int64_t dim) { - return std::make_shared(repeat, dim)->Apply({shared_from_this()})[0]; -} - -std::shared_ptr Tensor::View(const std::vector &dims) { - return std::make_shared(dims)->Apply({shared_from_this()})[0]; -} - -std::shared_ptr Tensor::Contiguous() { - return std::make_shared(dims_)->Apply({shared_from_this()})[0]; -} - -std::shared_ptr Tensor::Flatten(int64_t start, int64_t end) { - // return Contiguous()->View(new_shape); - // =================================== 作业 =================================== - // TODO：实现张量扁平化操作，将指定维度范围[start, end]内的所有维度合并为一个维度 - // HINT: - // =================================== 作业 =================================== - - return std::make_shared(); -} - -std::shared_ptr Tensor::Squeeze(int64_t dim) { - std::vector new_shape = dims_; - if (dim < 0) { - dim += new_shape.size(); - } - CHECK_GE(dim, 0); - CHECK_LT(dim, new_shape.size()); - CHECK_EQ(new_shape[dim], 1) << "Cannot squeeze dim " << dim << " because size (" << new_shape[dim] << ") != 1."; - - new_shape.erase(new_shape.begin() + dim); - - return Contiguous()->View(new_shape); -} - -std::shared_ptr Tensor::Slice(const std::vector &starts, const std::vector &ends, - const std::vector &steps) { - return std::make_shared(starts, ends, steps)->Apply({shared_from_this()})[0]; -} - -std::shared_ptr Tensor::Slice(int64_t dim, int64_t start, int64_t end, int64_t step) { - // Slice only on one dimension - if (dim < 0) { - dim += dims_.size(); - } - CHECK_GE(dim, 0); - CHECK_LT(dim, dims_.size()); - std::vector starts(dims_.size(), 0); - std::vector ends = dims_; - std::vector steps(dims_.size(), 1); - - starts[dim] = start; - ends[dim] = end; - steps[dim] = step; - return Slice(starts, ends, steps); -} - -std::shared_ptr Tensor::Transpose(int dim0, int dim1) { - return std::make_shared(dim0, dim1)->Apply({shared_from_this()})[0]; -} - -std::shared_ptr Tensor::MaskedFill(const std::shared_ptr &mask, float value) { - return std::make_shared(mask, value)->Apply({shared_from_this()})[0]; -} - -std::shared_ptr Tensor::Matmul(const std::shared_ptr &other) { - return std::make_shared()->Apply({shared_from_this(), other})[0]; -} - -std::shared_ptr Tensor::Outer(const std::shared_ptr &other) { - return std::make_shared()->Apply({shared_from_this(), other})[0]; -} - -// distribution -std::shared_ptr Tensor::Uniform(float from, float to, std::optional generator) { - return nn::init::Uniform(shared_from_this(), from, to, generator); -} - -// autograd related -std::shared_ptr Tensor::RequiresGrad() { - requires_grad_ = true; - if (!grad_) { - grad_ = std::make_unique(dims_, dtype_, GetDevice()); - grad_->Fill(0.0f); - } - return shared_from_this(); -} - -void Tensor::Backward(std::shared_ptr gradient, bool retain_graph, bool create_graph) const { - // =================================== 作业 =================================== - // TODO：实现自动微分反向传播 - // 功能描述：1. 计算当前张量对叶子节点的梯度 2. 支持多输出场景的梯度累加 - // =================================== 作业 =================================== -} - -void Tensor::ZeroGrad() { - if (grad_) { - grad_->Fill(0.0f); - } -} - -std::ostream &operator<<(std::ostream &os, const Tensor &tensor) { - os << "Tensor(data_ptr=" << static_cast(tensor.DataPtr()) << ", dims=["; - for (const auto &dim : tensor.Dims()) { os << dim << ", "; } - os << "], dtype=" << kDataTypeToDesc.at(tensor.Dtype()) << ")"; - return os; -} - -std::shared_ptr operator==(const std::shared_ptr &t, float scalar) { return t->Equals(scalar); } - -std::shared_ptr operator+(const std::shared_ptr &t1, const std::shared_ptr &t2) { - return t1->Add(t2); -} - -std::shared_ptr operator+(float scalar, const std::shared_ptr &t) { return t->Add(scalar); } - -std::shared_ptr operator+(const std::shared_ptr &t, float scalar) { return t->Add(scalar); } - -std::shared_ptr operator-(const std::shared_ptr &t1, const std::shared_ptr &t2) { - return t1->Sub(t2); -} - -std::shared_ptr operator-(float scalar, const std::shared_ptr &t) { return t->Neg()->Add(scalar); } - -std::shared_ptr operator-(const std::shared_ptr &t, float scalar) { return t->Add(-scalar); } - -std::shared_ptr operator-(const std::shared_ptr &t) { return t->Neg(); } - -std::shared_ptr operator*(const std::shared_ptr &t1, const std::shared_ptr &t2) { - return t1->Mul(t2); -} - -std::shared_ptr operator*(float scalar, const std::shared_ptr &t) { return t->Mul(scalar); } - -std::shared_ptr operator*(const std::shared_ptr &t, float scalar) { return t->Mul(scalar); } - -std::shared_ptr operator/(const std::shared_ptr &t1, const std::shared_ptr &t2) { - return t1->Div(t2); -} - -std::shared_ptr operator/(float scalar, const std::shared_ptr &t) { - return t->Reciprocal()->Mul(scalar); -} - -std::shared_ptr operator/(const std::shared_ptr &t, float scalar) { return t->Mul(1.0f / scalar); } - -void Tensor::SaveAsNpy(const std::string &path) const { - CHECK(dtype_ == DataType::kFLOAT32); - - const size_t num_elements = NumElements(); - const size_t num_bytes = num_elements * sizeof(float); - - // Prepare host buffer - std::vector host_buffer(num_elements); - - if (GetDevice().Type() == DeviceType::kCPU) { - // If on CPU, direct copy - std::memcpy(host_buffer.data(), DataPtr(), num_bytes); - } -#ifdef USE_CUDA - else if (GetDevice().Type() == DeviceType::kCUDA) { - // If on CUDA, copy back to host - cudaDeviceSynchronize(); - cudaError_t err = cudaMemcpy(host_buffer.data(), DataPtr(), num_bytes, cudaMemcpyDeviceToHost); - CHECK_EQ(err, cudaSuccess) << "cudaMemcpy failed: " << cudaGetErrorString(err); - } -#endif - else { - LOG(FATAL) << "Unsupported device type for SaveAsNpy."; - } - - // Write .npy file - std::ofstream file(path, std::ios::binary); - CHECK(file.is_open()) << "Failed to open file for writing: " << path; - - // Write magic string - file.write("\x93NUMPY", 6); - - // Write version - uint8_t major = 1; - uint8_t minor = 0; - file.put(major); - file.put(minor); - - // Construct header - std::ostringstream header_ss; - header_ss << "{'descr': '(header.size()); - file.write(reinterpret_cast(&header_size), sizeof(header_size)); - file.write(header.c_str(), header.size()); - - // Write data - file.write(reinterpret_cast(host_buffer.data()), num_bytes); - - file.close(); -} - -void Tensor::SetPrintOptions(std::optional precision, std::optional threshold, - std::optional edge_items, std::optional linewidth, - std::optional profile, std::optional sci_mode) { - PrintOptions &opts = PrintOptions::Get(); - if (profile) { - // ref: https://github.com/pytorch/pytorch/blob/main/torch/_tensor_str.py - std::string &profile_name = *profile; - std::transform(profile_name.begin(), profile_name.end(), profile_name.begin(), ::tolower); - if (profile_name == "default") { - opts.precision = 4; - opts.threshold = 1000; - opts.edge_items = 3; - opts.linewidth = 80; - opts.sci_mode = std::nullopt; - } else if (profile_name == "short") { - opts.precision = 4; - opts.threshold = 100; - opts.edge_items = 2; - opts.linewidth = 80; - opts.sci_mode = std::nullopt; - } else if (profile_name == "full") { - opts.precision = 4; - opts.threshold = std::numeric_limits::max(); - opts.edge_items = 3; - opts.linewidth = 80; - opts.sci_mode = std::nullopt; - } else { - LOG(WARNING) << "Undefined profile name: " << profile_name; - } - } - - if (precision) { - opts.precision = *precision; - } - if (threshold) { - opts.threshold = *threshold; - } - if (edge_items) { - opts.edge_items = *edge_items; - } - if (linewidth) { - opts.linewidth = *linewidth; - } - if (sci_mode) { - opts.sci_mode = *sci_mode; - } -} - -void Tensor::Print(std::ostream &os) const { - /* - Print tensor in torch.tensor/np.array style. - */ - CHECK(dtype_ == DataType::kFLOAT32); - - const size_t num_elements = NumElements(); - const size_t num_bytes = num_elements * sizeof(float); - - std::vector host_buffer(num_elements); - - if (GetDevice().Type() == DeviceType::kCPU) { - std::memcpy(host_buffer.data(), DataPtr(), num_bytes); - } -#ifdef USE_CUDA - else if (GetDevice().Type() == DeviceType::kCUDA) { - cudaDeviceSynchronize(); - cudaError_t err = cudaMemcpy(host_buffer.data(), DataPtr(), num_bytes, cudaMemcpyDeviceToHost); - CHECK_EQ(err, cudaSuccess) << "cudaMemcpy failed: " << cudaGetErrorString(err); - } -#endif - else { - LOG(FATAL) << "Unsupported device type for Print."; - } - - const PrintOptions &opts = PrintOptions::Get(); - const int64_t precision = opts.precision; - const int64_t threshold = opts.threshold; - const int64_t edge_items = opts.edge_items; - const int64_t linewidth = opts.linewidth; - const int64_t base_indent = 8; // length of "tensor([" - - bool use_sci = opts.sci_mode.value_or(false); - if (!opts.sci_mode.has_value()) { - for (float v : host_buffer) { - float abs_v = std::fabs(v); - if ((abs_v > 0.0f && abs_v < 1e-4f) || abs_v >= 1e+4f) { - use_sci = true; - break; - } - } - } - - auto format_float = [&](float val) -> std::string { - std::ostringstream ss; - if (use_sci) { - ss << std::scientific << std::setprecision(precision); - } else { - ss << std::fixed << std::setprecision(precision); - } - ss << val; - return ss.str(); - }; - - std::vector str_vals(num_elements); - size_t max_width = 0; - for (size_t i = 0; i < num_elements; ++i) { - str_vals[i] = format_float(host_buffer[i]); - max_width = std::max(max_width, str_vals[i].length()); - } - - const int ndim = dims_.size(); - - std::function print_rec; - print_rec = [&](int dim, size_t offset, int indent) { - os << "["; - size_t step = 1; - for (int d = dim + 1; d < ndim; ++d) { step *= dims_[d]; } - int n = dims_[dim]; - - if (dim == ndim - 1) { - if (n <= 2 * edge_items || num_elements <= threshold) { - int line_len = base_indent + indent + 1; - for (int i = 0; i < n; ++i) { - if (i > 0) { - os << ", "; - line_len += 2; - } - std::string item = str_vals[offset + i]; - if (linewidth > 0 && line_len + max_width > linewidth) { - os << "\n" << std::string(base_indent + indent + 1, ' '); - line_len = base_indent + indent + 1; - } - os << std::setw(max_width) << item; - line_len += max_width; - } - } else { - int line_len = base_indent + indent + 1; - for (int i = 0; i < edge_items; ++i) { - if (i > 0) { - os << ", "; - line_len += 2; - } - std::string item = str_vals[offset + i]; - if (linewidth > 0 && line_len + max_width > linewidth) { - os << "\n" << std::string(base_indent + indent + 1, ' '); - line_len = base_indent + indent + 1; - } - os << std::setw(max_width) << item; - line_len += max_width; - } - os << ", ..."; - line_len += 5; // length of ", ..." - if (linewidth > 0 && line_len + max_width > linewidth) { - os << "\n" << std::string(base_indent + indent + 1, ' '); - line_len = base_indent + indent + 1; - } else { - os << ", "; - line_len += 2; - } - for (int i = n - edge_items; i < n; ++i) { - if (i > n - edge_items) { - os << ", "; - line_len += 2; - } - std::string item = str_vals[offset + i]; - if (linewidth > 0 && line_len + max_width > linewidth) { - os << "\n" << std::string(base_indent + indent + 1, ' '); - line_len = base_indent + indent + 1; - } - os << std::setw(max_width) << item; - line_len += max_width; - } - } - } else { - if (n <= 2 * edge_items || num_elements <= threshold) { - for (int i = 0; i < n; ++i) { - if (i > 0) { - if (dim < ndim - 2) { - os << ",\n\n" << std::string(base_indent + indent, ' '); - } else { - os << ",\n" << std::string(base_indent + indent, ' '); - } - } - print_rec(dim + 1, offset + i * step, indent + 1); - } - } else { - for (int i = 0; i < edge_items; ++i) { - if (i > 0) { - if (dim < ndim - 2) { - os << ",\n\n" << std::string(base_indent + indent, ' '); - } else { - os << ",\n" << std::string(base_indent + indent, ' '); - } - } - print_rec(dim + 1, offset + i * step, indent + 1); - } - os << ",\n" - << std::string(base_indent + indent, ' ') << "...\n" - << std::string(base_indent + indent, ' '); - for (int i = n - edge_items; i < n; ++i) { - if (i > n - edge_items) { - if (dim < ndim - 2) { - os << ",\n\n" << std::string(base_indent + indent, ' '); - } else { - os << ",\n" << std::string(base_indent + indent, ' '); - } - } - print_rec(dim + 1, offset + i * step, indent + 1); - } - } - } - os << "]"; - }; - - os << "Tensor("; - if (num_elements == 0) { - os << "[], "; - } else { - print_rec(0, 0, 0); - os << ", \n"; - } - - os << std::string(base_indent - 1, ' ') << "dtype=float32, shape=("; - for (size_t i = 0; i < dims_.size(); ++i) { - if (i > 0) { - os << ", "; - } - os << dims_[i]; - } - if (dims_.size() == 1) { - os << ","; - } - os << "))\n"; -} -} // namespace infini_train +#include "infini_train/include/tensor.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef USE_CUDA +#include "cuda_runtime_api.h" +#endif +#include "Eigen/Dense" +#include "glog/logging.h" + +#include "infini_train/include/autograd/elementwise.h" +#include "infini_train/include/autograd/matmul.h" +#include "infini_train/include/autograd/misc.h" +#include "infini_train/include/autograd/outer.h" +#include "infini_train/include/autograd/transform.h" +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/nn/init.h" + +namespace infini_train { +namespace { +const std::unordered_map kDataTypeToSize = { + {DataType::kUINT8, 1}, {DataType::kINT8, 1}, {DataType::kUINT16, 2}, {DataType::kINT16, 2}, + {DataType::kUINT32, 4}, {DataType::kINT32, 4}, {DataType::kUINT64, 8}, {DataType::kINT64, 8}, + {DataType::kBFLOAT16, 2}, {DataType::kFLOAT16, 2}, {DataType::kFLOAT32, 4}, {DataType::kFLOAT64, 8}, +}; + +const std::unordered_map kDataTypeToDesc = { + {DataType::kUINT8, "uint8"}, {DataType::kINT8, "int8"}, {DataType::kUINT16, "uint16"}, + {DataType::kINT16, "int16"}, {DataType::kUINT32, "uint32"}, {DataType::kINT32, "int32"}, + {DataType::kUINT64, "uint64"}, {DataType::kINT64, "int64"}, {DataType::kBFLOAT16, "bf16"}, + {DataType::kFLOAT16, "fp16"}, {DataType::kFLOAT32, "fp32"}, {DataType::kFLOAT64, "fp64"}, +}; + +template struct TypeMap; + +template <> struct TypeMap { + using type = float; +}; +template <> struct TypeMap { + using type = double; +}; +template <> struct TypeMap { + using type = int32_t; +}; +template <> struct TypeMap { + using type = int64_t; +}; +} // namespace + +TensorBuffer::TensorBuffer(Device device, size_t size) : device_(device), size_(size) { + switch (device_.Type()) { + case DeviceType::kCPU: + data_ = malloc(size); + break; +#ifdef USE_CUDA + case DeviceType::kCUDA: + cudaMallocAsync(&data_, size, 0); + break; +#endif + default: + LOG(FATAL) << "Unsupported device type: " << static_cast(device_.Type()); + break; + } +} + +TensorBuffer::~TensorBuffer() { + switch (device_.Type()) { + case DeviceType::kCPU: + free(data_); + break; +#ifdef USE_CUDA + case DeviceType::kCUDA: + cudaFreeAsync(data_, 0); + break; +#endif + default: + LOG(FATAL) << "Unsupported device type: " << static_cast(device_.Type()); + break; + } +} + +void *TensorBuffer::DataPtr() { return data_; } + +const void *TensorBuffer::DataPtr() const { return data_; } + +Device TensorBuffer::GetDevice() const { return device_; } + +size_t TensorBuffer::Size() const { return size_; } + +// Tensor implementation +Tensor::Tensor(const std::vector &dims, DataType dtype, Device device) : dims_(dims), dtype_(dtype) { + num_elements_ = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); + buffer_ = std::make_shared(device, kDataTypeToSize.at(dtype) * num_elements_); +} + +Tensor::Tensor(const Tensor &tensor, size_t offset, const std::vector &dims) + : buffer_(tensor.buffer_), offset_(offset), dims_(dims), + num_elements_(std::accumulate(dims.begin(), dims.end(), 1, std::multiplies())), dtype_(tensor.dtype_) { + CHECK_LE(offset_ + kDataTypeToSize.at(dtype_) * num_elements_, buffer_->Size()); +} + +Device Tensor::GetDevice() const { return buffer_->GetDevice(); } + +void *Tensor::DataPtr() { return reinterpret_cast(buffer_->DataPtr()) + offset_; } + +const void *Tensor::DataPtr() const { return reinterpret_cast(buffer_->DataPtr()) + offset_; } + +size_t Tensor::SizeInBytes() const { return kDataTypeToSize.at(dtype_) * num_elements_; } + +const std::vector &Tensor::Dims() const { return dims_; } + +size_t Tensor::NumElements() const { return num_elements_; } + +DataType Tensor::Dtype() const { return dtype_; } + +template void Tensor::Fill(T value) { + DataType dtype = Dtype(); + + uint64_t storage = 0; + + switch (dtype) { + case DataType::kFLOAT32: { + using TargetT = typename TypeMap::type; + TargetT casted_value = static_cast(value); + std::memcpy(&storage, &casted_value, sizeof(TargetT)); + break; + } + case DataType::kFLOAT64: { + using TargetT = typename TypeMap::type; + TargetT casted_value = static_cast(value); + std::memcpy(&storage, &casted_value, sizeof(TargetT)); + break; + } + case DataType::kINT32: { + using TargetT = typename TypeMap::type; + TargetT casted_value = static_cast(value); + std::memcpy(&storage, &casted_value, sizeof(TargetT)); + break; + } + case DataType::kINT64: { + using TargetT = typename TypeMap::type; + TargetT casted_value = static_cast(value); + std::memcpy(&storage, &casted_value, sizeof(TargetT)); + break; + } + default: + throw std::runtime_error("Unsupported data type in Tensor::Fill()"); + } + + auto kernel = Dispatcher::Instance().GetKernel({GetDevice().Type(), "Fill"}); + kernel.Call(shared_from_this(), static_cast(&storage)); +} + +template void Tensor::Fill(float); + +Eigen::Map> Tensor::EigenMatrix() { + const int64_t bs = std::accumulate(dims_.rbegin() + 1, dims_.rend(), 1, std::multiplies()); + return Eigen::Map>( + reinterpret_cast(DataPtr()), bs, *dims_.rbegin()); +} + +Eigen::Map> Tensor::EigenVector() { + CHECK_EQ(dims_.size(), 1); + return Eigen::Map>(reinterpret_cast(DataPtr()), 1, + dims_[0]); +} + +Tensor Tensor::To(Device device) { + if (device == buffer_->GetDevice()) { + auto new_tensor = Tensor(*this, offset_, dims_); + if (grad_) { + new_tensor.grad_ = std::make_unique(*grad_.get(), grad_->offset_, grad_->dims_); + } + return new_tensor; + } + + Tensor new_tensor; + switch (device.Type()) { +#ifdef USE_CUDA + case DeviceType::kCPU: + // CUDA -> CPU + new_tensor = Tensor(dims_, dtype_, Device(DeviceType::kCPU, 0)); + cudaMemcpyAsync(new_tensor.DataPtr(), DataPtr(), SizeInBytes(), cudaMemcpyDeviceToHost, 0); + break; + case DeviceType::kCUDA: + // CPU -> CUDA + new_tensor = Tensor(dims_, dtype_, Device(DeviceType::kCUDA, 0)); + cudaMemcpyAsync(new_tensor.DataPtr(), DataPtr(), SizeInBytes(), cudaMemcpyHostToDevice, 0); + break; +#else + case DeviceType::kCPU: + // CPU -> CPU copy + new_tensor = Tensor(dims_, dtype_, Device(DeviceType::kCPU, 0)); + memcpy(new_tensor.DataPtr(), DataPtr(), SizeInBytes()); + break; +#endif + default: + LOG(FATAL) << "Unsupported device type: " << static_cast(device.Type()); + } + + if (grad_) { + new_tensor.grad_ = std::make_unique(grad_->To(device)); + } + + new_tensor.requires_grad_ = requires_grad_; + + return new_tensor; +} + +// operator overloading +std::shared_ptr Tensor::Equals(float scalar) { + return std::make_shared(scalar)->Apply({shared_from_this()})[0]; +} + +std::shared_ptr Tensor::Add(const std::shared_ptr &other) { + CHECK_EQ(static_cast(GetDevice().Type()), static_cast(other->GetDevice().Type())); + return std::make_shared()->Apply({shared_from_this(), other})[0]; +} + +std::shared_ptr Tensor::Add(float scalar) { + return std::make_shared(scalar)->Apply({shared_from_this()})[0]; +} + +std::shared_ptr Tensor::Sub(const std::shared_ptr &other) { + CHECK_EQ(static_cast(GetDevice().Type()), static_cast(other->GetDevice().Type())); + return std::make_shared()->Apply({shared_from_this(), other})[0]; +} + +std::shared_ptr Tensor::Mul(const std::shared_ptr &other) { + CHECK_EQ(static_cast(GetDevice().Type()), static_cast(other->GetDevice().Type())); + return std::make_shared()->Apply({shared_from_this(), other})[0]; +} + +std::shared_ptr Tensor::Mul(float scalar) { + return std::make_shared(scalar)->Apply({shared_from_this()})[0]; +} + +std::shared_ptr Tensor::Div(const std::shared_ptr &other) { + CHECK_EQ(static_cast(GetDevice().Type()), static_cast(other->GetDevice().Type())); + return std::make_shared()->Apply({shared_from_this(), other})[0]; +} + +std::shared_ptr Tensor::Neg() { return std::make_shared()->Apply({shared_from_this()})[0]; } + +std::shared_ptr Tensor::Reciprocal() { + return std::make_shared()->Apply({shared_from_this()})[0]; +} + +std::shared_ptr Tensor::Sin() { return std::make_shared()->Apply({shared_from_this()})[0]; } + +std::shared_ptr Tensor::Cos() { return std::make_shared()->Apply({shared_from_this()})[0]; } + +std::shared_ptr Tensor::Tanh() { return std::make_shared()->Apply({shared_from_this()})[0]; } + +std::shared_ptr Tensor::Pow(float exponent) { + return std::make_shared(exponent)->Apply({shared_from_this()})[0]; +} + +std::shared_ptr Tensor::Rsqrt() { return std::make_shared()->Apply({shared_from_this()})[0]; } + +std::vector> Tensor::Split(int split_size, int dim) { + return std::make_shared(split_size, dim)->Apply({shared_from_this()}); +} + +std::shared_ptr Tensor::RepeatInterleave(int64_t repeat, int64_t dim) { + return std::make_shared(repeat, dim)->Apply({shared_from_this()})[0]; +} + +std::shared_ptr Tensor::View(const std::vector &dims) { + return std::make_shared(dims)->Apply({shared_from_this()})[0]; +} + +std::shared_ptr Tensor::Contiguous() { + return std::make_shared(dims_)->Apply({shared_from_this()})[0]; +} + +std::shared_ptr Tensor::Flatten(int64_t start, int64_t end) { + // return Contiguous()->View(new_shape); + // =================================== 作业 =================================== + // TODO：实现张量扁平化操作，将指定维度范围[start, end]内的所有维度合并为一个维度 + // HINT: + // =================================== 作业 =================================== + + int64_t ndim = static_cast(dims_.size()); + if (end < 0) { + end += ndim; + } + CHECK_GE(start, 0); + CHECK_LT(start, ndim); + CHECK_GE(end, start); + CHECK_LT(end, ndim); + + std::vector new_shape; + for (int64_t i = 0; i < start; ++i) { + new_shape.push_back(dims_[i]); + } + int64_t flattened_size = 1; + for (int64_t i = start; i <= end; ++i) { + flattened_size *= dims_[i]; + } + new_shape.push_back(flattened_size); + for (int64_t i = end + 1; i < ndim; ++i) { + new_shape.push_back(dims_[i]); + } + + return Contiguous()->View(new_shape); +} + +std::shared_ptr Tensor::Squeeze(int64_t dim) { + std::vector new_shape = dims_; + if (dim < 0) { + dim += new_shape.size(); + } + CHECK_GE(dim, 0); + CHECK_LT(dim, new_shape.size()); + CHECK_EQ(new_shape[dim], 1) << "Cannot squeeze dim " << dim << " because size (" << new_shape[dim] << ") != 1."; + + new_shape.erase(new_shape.begin() + dim); + + return Contiguous()->View(new_shape); +} + +std::shared_ptr Tensor::Slice(const std::vector &starts, const std::vector &ends, + const std::vector &steps) { + return std::make_shared(starts, ends, steps)->Apply({shared_from_this()})[0]; +} + +std::shared_ptr Tensor::Slice(int64_t dim, int64_t start, int64_t end, int64_t step) { + // Slice only on one dimension + if (dim < 0) { + dim += dims_.size(); + } + CHECK_GE(dim, 0); + CHECK_LT(dim, dims_.size()); + std::vector starts(dims_.size(), 0); + std::vector ends = dims_; + std::vector steps(dims_.size(), 1); + + starts[dim] = start; + ends[dim] = end; + steps[dim] = step; + return Slice(starts, ends, steps); +} + +std::shared_ptr Tensor::Transpose(int dim0, int dim1) { + return std::make_shared(dim0, dim1)->Apply({shared_from_this()})[0]; +} + +std::shared_ptr Tensor::MaskedFill(const std::shared_ptr &mask, float value) { + return std::make_shared(mask, value)->Apply({shared_from_this()})[0]; +} + +std::shared_ptr Tensor::Matmul(const std::shared_ptr &other) { + return std::make_shared()->Apply({shared_from_this(), other})[0]; +} + +std::shared_ptr Tensor::Outer(const std::shared_ptr &other) { + return std::make_shared()->Apply({shared_from_this(), other})[0]; +} + +// distribution +std::shared_ptr Tensor::Uniform(float from, float to, std::optional generator) { + return nn::init::Uniform(shared_from_this(), from, to, generator); +} + +// autograd related +std::shared_ptr Tensor::RequiresGrad() { + requires_grad_ = true; + if (!grad_) { + grad_ = std::make_unique(dims_, dtype_, GetDevice()); + grad_->Fill(0.0f); + } + return shared_from_this(); +} + +void Tensor::Backward(std::shared_ptr gradient, bool retain_graph, bool create_graph) const { + // =================================== 作业 =================================== + // TODO：实现自动微分反向传播 + // 功能描述：1. 计算当前张量对叶子节点的梯度 2. 支持多输出场景的梯度累加 + // =================================== 作业 =================================== + + if (!gradient) { + gradient = std::make_shared(dims_, dtype_, GetDevice()); + gradient->Fill(1.0f); + } + + if (is_leaf_) { + if (requires_grad_ && grad_) { + auto device = grad_->GetDevice().Type(); + auto kernel = Dispatcher::Instance().GetKernel({device, "AccumulateGrad"}); + kernel.Call(gradient, 1.0f, grad_); + } + } else if (grad_fn_) { + grad_fn_->BackwardPartial(gradient, output_idx_); + } +} + +void Tensor::ZeroGrad() { + if (grad_) { + grad_->Fill(0.0f); + } +} + +std::ostream &operator<<(std::ostream &os, const Tensor &tensor) { + os << "Tensor(data_ptr=" << static_cast(tensor.DataPtr()) << ", dims=["; + for (const auto &dim : tensor.Dims()) { os << dim << ", "; } + os << "], dtype=" << kDataTypeToDesc.at(tensor.Dtype()) << ")"; + return os; +} + +std::shared_ptr operator==(const std::shared_ptr &t, float scalar) { return t->Equals(scalar); } + +std::shared_ptr operator+(const std::shared_ptr &t1, const std::shared_ptr &t2) { + return t1->Add(t2); +} + +std::shared_ptr operator+(float scalar, const std::shared_ptr &t) { return t->Add(scalar); } + +std::shared_ptr operator+(const std::shared_ptr &t, float scalar) { return t->Add(scalar); } + +std::shared_ptr operator-(const std::shared_ptr &t1, const std::shared_ptr &t2) { + return t1->Sub(t2); +} + +std::shared_ptr operator-(float scalar, const std::shared_ptr &t) { return t->Neg()->Add(scalar); } + +std::shared_ptr operator-(const std::shared_ptr &t, float scalar) { return t->Add(-scalar); } + +std::shared_ptr operator-(const std::shared_ptr &t) { return t->Neg(); } + +std::shared_ptr operator*(const std::shared_ptr &t1, const std::shared_ptr &t2) { + return t1->Mul(t2); +} + +std::shared_ptr operator*(float scalar, const std::shared_ptr &t) { return t->Mul(scalar); } + +std::shared_ptr operator*(const std::shared_ptr &t, float scalar) { return t->Mul(scalar); } + +std::shared_ptr operator/(const std::shared_ptr &t1, const std::shared_ptr &t2) { + return t1->Div(t2); +} + +std::shared_ptr operator/(float scalar, const std::shared_ptr &t) { + return t->Reciprocal()->Mul(scalar); +} + +std::shared_ptr operator/(const std::shared_ptr &t, float scalar) { return t->Mul(1.0f / scalar); } + +void Tensor::SaveAsNpy(const std::string &path) const { + CHECK(dtype_ == DataType::kFLOAT32); + + const size_t num_elements = NumElements(); + const size_t num_bytes = num_elements * sizeof(float); + + // Prepare host buffer + std::vector host_buffer(num_elements); + + if (GetDevice().Type() == DeviceType::kCPU) { + // If on CPU, direct copy + std::memcpy(host_buffer.data(), DataPtr(), num_bytes); + } +#ifdef USE_CUDA + else if (GetDevice().Type() == DeviceType::kCUDA) { + // If on CUDA, copy back to host + cudaDeviceSynchronize(); + cudaError_t err = cudaMemcpy(host_buffer.data(), DataPtr(), num_bytes, cudaMemcpyDeviceToHost); + CHECK_EQ(err, cudaSuccess) << "cudaMemcpy failed: " << cudaGetErrorString(err); + } +#endif + else { + LOG(FATAL) << "Unsupported device type for SaveAsNpy."; + } + + // Write .npy file + std::ofstream file(path, std::ios::binary); + CHECK(file.is_open()) << "Failed to open file for writing: " << path; + + // Write magic string + file.write("\x93NUMPY", 6); + + // Write version + uint8_t major = 1; + uint8_t minor = 0; + file.put(major); + file.put(minor); + + // Construct header + std::ostringstream header_ss; + header_ss << "{'descr': '(header.size()); + file.write(reinterpret_cast(&header_size), sizeof(header_size)); + file.write(header.c_str(), header.size()); + + // Write data + file.write(reinterpret_cast(host_buffer.data()), num_bytes); + + file.close(); +} + +void Tensor::SetPrintOptions(std::optional precision, std::optional threshold, + std::optional edge_items, std::optional linewidth, + std::optional profile, std::optional sci_mode) { + PrintOptions &opts = PrintOptions::Get(); + if (profile) { + // ref: https://github.com/pytorch/pytorch/blob/main/torch/_tensor_str.py + std::string &profile_name = *profile; + std::transform(profile_name.begin(), profile_name.end(), profile_name.begin(), ::tolower); + if (profile_name == "default") { + opts.precision = 4; + opts.threshold = 1000; + opts.edge_items = 3; + opts.linewidth = 80; + opts.sci_mode = std::nullopt; + } else if (profile_name == "short") { + opts.precision = 4; + opts.threshold = 100; + opts.edge_items = 2; + opts.linewidth = 80; + opts.sci_mode = std::nullopt; + } else if (profile_name == "full") { + opts.precision = 4; + opts.threshold = std::numeric_limits::max(); + opts.edge_items = 3; + opts.linewidth = 80; + opts.sci_mode = std::nullopt; + } else { + LOG(WARNING) << "Undefined profile name: " << profile_name; + } + } + + if (precision) { + opts.precision = *precision; + } + if (threshold) { + opts.threshold = *threshold; + } + if (edge_items) { + opts.edge_items = *edge_items; + } + if (linewidth) { + opts.linewidth = *linewidth; + } + if (sci_mode) { + opts.sci_mode = *sci_mode; + } +} + +void Tensor::Print(std::ostream &os) const { + /* + Print tensor in torch.tensor/np.array style. + */ + CHECK(dtype_ == DataType::kFLOAT32); + + const size_t num_elements = NumElements(); + const size_t num_bytes = num_elements * sizeof(float); + + std::vector host_buffer(num_elements); + + if (GetDevice().Type() == DeviceType::kCPU) { + std::memcpy(host_buffer.data(), DataPtr(), num_bytes); + } +#ifdef USE_CUDA + else if (GetDevice().Type() == DeviceType::kCUDA) { + cudaDeviceSynchronize(); + cudaError_t err = cudaMemcpy(host_buffer.data(), DataPtr(), num_bytes, cudaMemcpyDeviceToHost); + CHECK_EQ(err, cudaSuccess) << "cudaMemcpy failed: " << cudaGetErrorString(err); + } +#endif + else { + LOG(FATAL) << "Unsupported device type for Print."; + } + + const PrintOptions &opts = PrintOptions::Get(); + const int64_t precision = opts.precision; + const int64_t threshold = opts.threshold; + const int64_t edge_items = opts.edge_items; + const int64_t linewidth = opts.linewidth; + const int64_t base_indent = 8; // length of "tensor([" + + bool use_sci = opts.sci_mode.value_or(false); + if (!opts.sci_mode.has_value()) { + for (float v : host_buffer) { + float abs_v = std::fabs(v); + if ((abs_v > 0.0f && abs_v < 1e-4f) || abs_v >= 1e+4f) { + use_sci = true; + break; + } + } + } + + auto format_float = [&](float val) -> std::string { + std::ostringstream ss; + if (use_sci) { + ss << std::scientific << std::setprecision(precision); + } else { + ss << std::fixed << std::setprecision(precision); + } + ss << val; + return ss.str(); + }; + + std::vector str_vals(num_elements); + size_t max_width = 0; + for (size_t i = 0; i < num_elements; ++i) { + str_vals[i] = format_float(host_buffer[i]); + max_width = std::max(max_width, str_vals[i].length()); + } + + const int ndim = dims_.size(); + + std::function print_rec; + print_rec = [&](int dim, size_t offset, int indent) { + os << "["; + size_t step = 1; + for (int d = dim + 1; d < ndim; ++d) { step *= dims_[d]; } + int n = dims_[dim]; + + if (dim == ndim - 1) { + if (n <= 2 * edge_items || num_elements <= threshold) { + int line_len = base_indent + indent + 1; + for (int i = 0; i < n; ++i) { + if (i > 0) { + os << ", "; + line_len += 2; + } + std::string item = str_vals[offset + i]; + if (linewidth > 0 && line_len + max_width > linewidth) { + os << "\n" << std::string(base_indent + indent + 1, ' '); + line_len = base_indent + indent + 1; + } + os << std::setw(max_width) << item; + line_len += max_width; + } + } else { + int line_len = base_indent + indent + 1; + for (int i = 0; i < edge_items; ++i) { + if (i > 0) { + os << ", "; + line_len += 2; + } + std::string item = str_vals[offset + i]; + if (linewidth > 0 && line_len + max_width > linewidth) { + os << "\n" << std::string(base_indent + indent + 1, ' '); + line_len = base_indent + indent + 1; + } + os << std::setw(max_width) << item; + line_len += max_width; + } + os << ", ..."; + line_len += 5; // length of ", ..." + if (linewidth > 0 && line_len + max_width > linewidth) { + os << "\n" << std::string(base_indent + indent + 1, ' '); + line_len = base_indent + indent + 1; + } else { + os << ", "; + line_len += 2; + } + for (int i = n - edge_items; i < n; ++i) { + if (i > n - edge_items) { + os << ", "; + line_len += 2; + } + std::string item = str_vals[offset + i]; + if (linewidth > 0 && line_len + max_width > linewidth) { + os << "\n" << std::string(base_indent + indent + 1, ' '); + line_len = base_indent + indent + 1; + } + os << std::setw(max_width) << item; + line_len += max_width; + } + } + } else { + if (n <= 2 * edge_items || num_elements <= threshold) { + for (int i = 0; i < n; ++i) { + if (i > 0) { + if (dim < ndim - 2) { + os << ",\n\n" << std::string(base_indent + indent, ' '); + } else { + os << ",\n" << std::string(base_indent + indent, ' '); + } + } + print_rec(dim + 1, offset + i * step, indent + 1); + } + } else { + for (int i = 0; i < edge_items; ++i) { + if (i > 0) { + if (dim < ndim - 2) { + os << ",\n\n" << std::string(base_indent + indent, ' '); + } else { + os << ",\n" << std::string(base_indent + indent, ' '); + } + } + print_rec(dim + 1, offset + i * step, indent + 1); + } + os << ",\n" + << std::string(base_indent + indent, ' ') << "...\n" + << std::string(base_indent + indent, ' '); + for (int i = n - edge_items; i < n; ++i) { + if (i > n - edge_items) { + if (dim < ndim - 2) { + os << ",\n\n" << std::string(base_indent + indent, ' '); + } else { + os << ",\n" << std::string(base_indent + indent, ' '); + } + } + print_rec(dim + 1, offset + i * step, indent + 1); + } + } + } + os << "]"; + }; + + os << "Tensor("; + if (num_elements == 0) { + os << "[], "; + } else { + print_rec(0, 0, 0); + os << ", \n"; + } + + os << std::string(base_indent - 1, ' ') << "dtype=float32, shape=("; + for (size_t i = 0; i < dims_.size(); ++i) { + if (i > 0) { + os << ", "; + } + os << dims_[i]; + } + if (dims_.size() == 1) { + os << ","; + } + os << "))\n"; +} +} // namespace infini_train