#pragma once #ifdef _OPENMP#include <omp.h>#endif #include <algorithm>#include <cassert>#include <cmath>#include <cstddef>#include <cstring>#include <functional>#include <cstdint>#include <cstdlib>#include <ctime>#include <iostream>#include <limits>#include <random>#include <ranges>#include <stdexcept>#include <type_traits>#include <unordered_map>#include <vector> #ifdef __ARM_NEON#include <arm_neon.h>#endif namespace tinytorch { const size_t kTensorMemAlign = 16; template <typename T>inline void AssertAligned(T ptr) { assert(((uintptr_t)(ptr)) % kTensorMemAlign == 0);} enum TensorType { kF32, kI32, kLEN // number of tensor types}; const size_t kTypeSize[kLEN] = { sizeof(float), sizeof(int32_t),}; template <typename T>inline bool IsTypeCompatible(TensorType type) { switch (type) { case kF32: return std::is_same<T, float>::value; case kI32: return std::is_same<T, int32_t>::value; default: throw std::runtime_error("Malformed tensor type"); }} struct Object { size_t offset; size_t size; Object *next; std::byte padding[8];}; const size_t kObjectSize = sizeof(struct Object); template <typename T>class TensorContextT { public: const size_t TENSOR_SIZE = sizeof(T); explicit TensorContextT(size_t mem_size) : mem_size_(mem_size) { mem_buffer_ = new std::byte[mem_size]; n_objects_ = 0; objects_begin_ = nullptr; objects_end_ = nullptr; } ~TensorContextT() { delete[] mem_buffer_; } T *NewTensor(const std::vector<int> &dims, float *data) { return NewTensor(dims, kF32, reinterpret_cast<std::byte *>(data)); } T *NewTensor(const std::vector<int> &dims, TensorType type = kF32, std::byte *data = nullptr) { const int n_dims = dims.size(); size_t size_needed = 0; if (data == nullptr) { size_t data_size = kTypeSize[type]; for (int i = 0; i < n_dims; i++) { data_size *= dims[i]; } size_needed += ((data_size + kTensorMemAlign - 1) / kTensorMemAlign) * kTensorMemAlign; } size_needed += TENSOR_SIZE; // layout // [Struct Object][Struct Tensor][data] std::byte *cur = mem_buffer_; if (objects_end_ != nullptr) { cur += objects_end_->offset + objects_end_->size; } if (cur + size_needed + kObjectSize > mem_buffer_ + mem_size_) { throw std::runtime_error("Out of tensor memory"); } Object *object = reinterpret_cast<Object *>(cur); *object = {.offset = (size_t)(cur - mem_buffer_) + kObjectSize, .size = size_needed, .next = nullptr}; AssertAligned(object); if (objects_end_ != nullptr) { objects_end_->next = object; } else { objects_begin_ = object; } objects_end_ = object; T *tensor = reinterpret_cast<T *>(cur + kObjectSize); AssertAligned(tensor); *tensor = T(this, dims, type, data == nullptr ? cur + kObjectSize + TENSOR_SIZE : data); AssertAligned(tensor->data_); n_objects_++; return tensor; } void PrintLayout(bool verbose = false) { std::cout << "TensorContext Layout" << std::endl; std::cout << "---------------------" << std::endl; std::cout << "Total memory size: " << mem_size_ << std::endl; std::cout << "Used memory size: " << (objects_end_ == nullptr ? 0 : (objects_end_->offset + objects_end_->size)) << std::endl; std::cout << "Number of objects: " << n_objects_ << std::endl; if (verbose) { std::cout << "Objects:" << std::endl; Object *cur = objects_begin_; while (cur != nullptr) { std::cout << " offset: " << cur->offset << ", size: " << cur->size << std::endl; cur = cur->next; } } } private: size_t mem_size_; std::byte *mem_buffer_; int n_objects_; Object *objects_begin_; Object *objects_end_;}; class Tensor;using TensorContext = TensorContextT<Tensor>; class NormalDist { public: NormalDist() { generator_.seed(std::random_device{}()); } float operator()() { return normal_dist_(generator_); } private: std::normal_distribution<float> normal_dist_; std::default_random_engine generator_;}; enum TensorOp { kOpNone, kOpAdd, kOpMul, kOpMatmul, kOpLookup, kOpNorm, kOpBroadcast, kOpView, kOpTranspose, kOpGelu, kOpSoftmax, kOpCrossEntropy, kOpLEN // number of tensor operations}; const std::string TENSOR_OP_NAMES[kOpLEN] = { // NOLINT "NONE", "ADD", "MUL", "MATMUL", "LOOKUP", "NORM", "BROADCAST", "VIEW", "TRANSPOSE", "GELU", "SOFTMAX", "CROSS_ENTROPY",}; const int kMaxTensorDims = 4;const int kMaxTensorOpParams = 2; // Tensor operations time profileclass Profile { public: void Reset() { times_.clear(); counts_.clear(); } void AddTime(TensorOp op, double time) { times_[op] += time; counts_[op]++; } void Print() { std::cout << "Profile" << std::endl; std::cout << "-------" << std::endl; for (int i = 0; i < kOpLEN; i++) { if (counts_[i] > 0) { std::cout << TENSOR_OP_NAMES[i] << ": " << times_[i] * 1000 << "ms (" << (times_[i] / counts_[i] * 1000) << "ms per op, " << counts_[i] << " times)" << std::endl; } } } private: std::unordered_map<int, double> times_; std::unordered_map<int, size_t> counts_;}; class Tensor { public: // Add Tensor &operator+(Tensor &other) { return operator2(other, kOpAdd); } // Mul Tensor &operator*(Tensor &other) { return operator2(other, kOpMul); } // Mul by scalar Tensor &operator*(float val) { assert(type_ == kF32); return *this * *(ctx_->NewTensor({1}, type_)->Fill(val)); } private: Tensor &operator2(Tensor &other_ref, TensorOp op) { auto other = &other_ref; assert(other != this); if (!SameShape(*other)) { assert(can_broadcast_to(*other, *this)); other = &broadcast_to(ctx_, *other, *this); } Tensor *dst = ctx_->NewTensor(Dims()); dst->op_ = op; dst->src0_ = this; dst->src1_ = other; return *dst; } public: // Lookup Tensor &operator[](Tensor &index) { assert(index.type_ == kI32); assert(n_dims_ + index.n_dims_ - 1 <= kMaxTensorDims); std::vector<int> ds; for (int i = kMaxTensorDims - 1; i > kMaxTensorDims - n_dims_; i--) { ds.push_back(dims_[i]); } for (int i = kMaxTensorDims - 1; i >= kMaxTensorDims - index.n_dims_; i--) { ds.push_back(index.dims_[i]); } reverse(ds.begin(), ds.end()); Tensor *dst = ctx_->NewTensor(ds, type_); dst->op_ = kOpLookup; dst->src0_ = this; dst->src1_ = &index; return *dst; } // Norm Tensor &Norm() { assert(type_ == kF32); Tensor *dst = ctx_->NewTensor(Dims(), type_); dst->op_ = kOpNorm; dst->src0_ = this; return *dst; } // Gelu Tensor &Gelu() { assert(type_ == kF32); Tensor *dst = ctx_->NewTensor(Dims(), type_); dst->op_ = kOpGelu; dst->src0_ = this; return *dst; } // Softmax Tensor &Softmax(bool is_causal = false, int vocab_size = 0) { assert(type_ == kF32); if (vocab_size > 0) { assert(vocab_size <= dims_[kMaxTensorDims - 1]); } Tensor *dst = ctx_->NewTensor(Dims(), type_); dst->op_ = kOpSoftmax; dst->src0_ = this; dst->op_params_[0] = is_causal; dst->op_params_[1] = vocab_size; return *dst; } // CrossEntropy Tensor &CrossEntropy(Tensor &target) { assert(type_ == kF32 && target.type_ == kI32); auto shape = Dims(); shape.pop_back(); assert(shape == target.Dims()); Tensor *dst = ctx_->NewTensor(shape, type_); dst->op_ = kOpCrossEntropy; dst->src0_ = this; dst->src1_ = ⌖ return *dst; } // Split std::vector<Tensor *> Split(int size, int axis) { assert(axis < n_dims_); auto dimi = kMaxTensorDims - n_dims_ + axis; assert(dims_[dimi] % size == 0); std::vector<Tensor *> tensors; if (dims_[dimi] == size) { tensors.push_back(this); return tensors; } std::vector<int> shape = Dims(); shape[axis] = size; for (int i = 0; i < dims_[dimi] / size; i++) { tensors.push_back(&View(shape, i, axis)); } return tensors; } // View // TODO(ysg): the view is actually a copy, we need to implement a real view Tensor &View(const std::vector<int> &shape, int split_no = 0, int split_axis = 0) { assert(NumElements() % num_of_elements(shape) == 0); int dimi = kMaxTensorDims - n_dims_ + split_axis; assert(dims_[dimi] % (NumElements() / num_of_elements(shape)) == 0); Tensor *dst = ctx_->NewTensor(shape, type_); dst->op_ = kOpView; dst->src0_ = this; dst->op_params_[0] = split_no; dst->op_params_[1] = split_axis; return *dst; } // Transpose Tensor &Transpose(int axis0, int axis1) { assert(axis0 < n_dims_ && axis1 < n_dims_); auto dimi0 = kMaxTensorDims - n_dims_ + axis0; auto dimi1 = kMaxTensorDims - n_dims_ + axis1; std::vector<int> shape = Dims(); std::swap(shape[axis0], shape[axis1]); Tensor *dst = ctx_->NewTensor(shape, type_); dst->op_ = kOpTranspose; dst->src0_ = this; dst->op_params_[0] = dimi0; dst->op_params_[1] = dimi1; return *dst; } // Matmul // (B, M, N) x (B, P, N) -> (B, M, P) // we assume that the input tensors are in the format (B, M, N) and (B, P, N) Tensor &MatMul(Tensor &other_ref) { auto other = &other_ref; assert(other != this); if (!can_matmul(*this, *other)) { assert(can_broadcast_to(*other, *this, 2)); other = &broadcast_to(ctx_, *other, *this, 2); assert(can_matmul(*this, *other)); } std::vector<int> dst_dims = {dims_[0], dims_[1], dims_[2], other->dims_[2]}; dst_dims.erase(dst_dims.begin(), dst_dims.begin() + dst_dims.size() - n_dims_); Tensor *dst = ctx_->NewTensor(dst_dims); dst->op_ = kOpMatmul; dst->src0_ = this; dst->src1_ = other; return *dst; } void Forward() { std::vector<Tensor *> sorted = topo_sort(this); struct timespec start, end; ForwardProfile.Reset(); for (auto *t : sorted) { clock_gettime(CLOCK_MONOTONIC, &start); switch (t->op_) { case kOpAdd: add_forward(t, t->src0_, t->src1_); break; case kOpMul: mul_forward(t, t->src0_, t->src1_); break; case kOpMatmul: matmul_forward(t, t->src0_, t->src1_); break; case kOpLookup: lookup_forward(t, t->src0_, t->src1_); break; case kOpNorm: norm_forward(t, t->src0_); break; case kOpTranspose: transpose_forward(t, t->src0_, t->op_params_[0], t->op_params_[1]); break; case kOpView: view_forward(t, t->src0_, t->op_params_[0], t->op_params_[1]); break; case kOpBroadcast: broadcast_forward(t, t->src0_); break; case kOpGelu: gelu_forward(t, t->src0_); break; case kOpSoftmax: softmax_forward(t, t->src0_, t->op_params_[0], t->op_params_[1]); break; case kOpCrossEntropy: cross_entropy_forward(t, t->src0_, t->src1_); break; case kOpNone: // no-op break; default: throw std::runtime_error("Forward(): Not implemented, " + TENSOR_OP_NAMES[t->op_]); } clock_gettime(CLOCK_MONOTONIC, &end); ForwardProfile.AddTime(t->op_, end.tv_sec - start.tv_sec + (end.tv_nsec - start.tv_nsec) / 1e9); } } void Backward(bool init_grad = true, float init_val = 1.0f) { std::vector<Tensor *> sorted = topo_sort(this); if (init_grad) { AllocGrad(false)->grad()->Fill(init_val); } struct timespec start, end; BackwardProfile.Reset(); for (auto *t : sorted | std::ranges::views::reverse) { clock_gettime(CLOCK_MONOTONIC, &start); switch (t->op_) { case kOpAdd: add_backward(t, t->src0_, t->src1_); break; case kOpMul: mul_backward(t, t->src0_, t->src1_); break; case kOpMatmul: matmul_backward(t, t->src0_, t->src1_); break; case kOpLookup: lookup_backward(t, t->src0_, t->src1_); break; case kOpNorm: norm_backward(t, t->src0_); break; case kOpTranspose: transpose_backward(t, t->src0_, t->op_params_[0], t->op_params_[1]); break; case kOpView: view_backward(t, t->src0_, t->op_params_[0], t->op_params_[1]); break; case kOpBroadcast: broadcast_backward(t, t->src0_); break; case kOpGelu: gelu_backward(t, t->src0_); break; case kOpSoftmax: softmax_backward(t, t->src0_, t->op_params_[0], t->op_params_[1]); break; case kOpCrossEntropy: cross_entropy_backward(t, t->src0_, t->src1_); break;Lines 1–500 of 1438. Load more
#pragma once #ifdef _OPENMP #include <omp.h> #endif #include <algorithm> #include <cassert> #include <cmath> #include <cstddef> #include <cstring> #include <functional> #include <cstdint> #include <cstdlib> #include <ctime> #include <iostream> #include <limits> #include <random> #include <ranges> #include <stdexcept> #include <type_traits> #include <unordered_map> #include <vector> #ifdef __ARM_NEON #include <arm_neon.h> #endif namespace tinytorch { const size_t kTensorMemAlign = 16; template <typename T> inline void AssertAligned(T ptr) { assert(((uintptr_t)(ptr)) % kTensorMemAlign == 0); } enum TensorType { kF32, kI32, kLEN // number of tensor types }; const size_t kTypeSize[kLEN] = { sizeof(float), sizeof(int32_t), }; template <typename T> inline bool IsTypeCompatible(TensorType type) { switch (type) { case kF32: return std::is_same<T, float>::value; case kI32: return std::is_same<T, int32_t>::value; default: throw std::runtime_error("Malformed tensor type"); } } struct Object { size_t offset; size_t size; Object *next; std::byte padding[8]; }; const size_t kObjectSize = sizeof(struct Object); template <typename T> class TensorContextT { public: const size_t TENSOR_SIZE = sizeof(T); explicit TensorContextT(size_t mem_size) : mem_size_(mem_size) { mem_buffer_ = new std::byte[mem_size]; n_objects_ = 0; objects_begin_ = nullptr; objects_end_ = nullptr; } ~TensorContextT() { delete[] mem_buffer_; } T *NewTensor(const std::vector<int> &dims, float *data) { return NewTensor(dims, kF32, reinterpret_cast<std::byte *>(data)); } T *NewTensor(const std::vector<int> &dims, TensorType type = kF32, std::byte *data = nullptr) { const int n_dims = dims.size(); size_t size_needed = 0; if (data == nullptr) { size_t data_size = kTypeSize[type]; for (int i = 0; i < n_dims; i++) { data_size *= dims[i]; } size_needed += ((data_size + kTensorMemAlign - 1) / kTensorMemAlign) * kTensorMemAlign; } size_needed += TENSOR_SIZE; // layout // [Struct Object][Struct Tensor][data] std::byte *cur = mem_buffer_; if (objects_end_ != nullptr) { cur += objects_end_->offset + objects_end_->size; } if (cur + size_needed + kObjectSize > mem_buffer_ + mem_size_) { throw std::runtime_error("Out of tensor memory"); } Object *object = reinterpret_cast<Object *>(cur); *object = {.offset = (size_t)(cur - mem_buffer_) + kObjectSize, .size = size_needed, .next = nullptr}; AssertAligned(object); if (objects_end_ != nullptr) { objects_end_->next = object; } else { objects_begin_ = object; } objects_end_ = object; T *tensor = reinterpret_cast<T *>(cur + kObjectSize); AssertAligned(tensor); *tensor = T(this, dims, type, data == nullptr ? cur + kObjectSize + TENSOR_SIZE : data); AssertAligned(tensor->data_); n_objects_++; return tensor; } void PrintLayout(bool verbose = false) { std::cout << "TensorContext Layout" << std::endl; std::cout << "---------------------" << std::endl; std::cout << "Total memory size: " << mem_size_ << std::endl; std::cout << "Used memory size: " << (objects_end_ == nullptr ? 0 : (objects_end_->offset + objects_end_->size)) << std::endl; std::cout << "Number of objects: " << n_objects_ << std::endl; if (verbose) { std::cout << "Objects:" << std::endl; Object *cur = objects_begin_; while (cur != nullptr) { std::cout << " offset: " << cur->offset << ", size: " << cur->size << std::endl; cur = cur->next; } } } private: size_t mem_size_; std::byte *mem_buffer_; int n_objects_; Object *objects_begin_; Object *objects_end_; }; class Tensor; using TensorContext = TensorContextT<Tensor>; class NormalDist { public: NormalDist() { generator_.seed(std::random_device{}()); } float operator()() { return normal_dist_(generator_); } private: std::normal_distribution<float> normal_dist_; std::default_random_engine generator_; }; enum TensorOp { kOpNone, kOpAdd, kOpMul, kOpMatmul, kOpLookup, kOpNorm, kOpBroadcast, kOpView, kOpTranspose, kOpGelu, kOpSoftmax, kOpCrossEntropy, kOpLEN // number of tensor operations }; const std::string TENSOR_OP_NAMES[kOpLEN] = { // NOLINT "NONE", "ADD", "MUL", "MATMUL", "LOOKUP", "NORM", "BROADCAST", "VIEW", "TRANSPOSE", "GELU", "SOFTMAX", "CROSS_ENTROPY", }; const int kMaxTensorDims = 4; const int kMaxTensorOpParams = 2; // Tensor operations time profile class Profile { public: void Reset() { times_.clear(); counts_.clear(); } void AddTime(TensorOp op, double time) { times_[op] += time; counts_[op]++; } void Print() { std::cout << "Profile" << std::endl; std::cout << "-------" << std::endl; for (int i = 0; i < kOpLEN; i++) { if (counts_[i] > 0) { std::cout << TENSOR_OP_NAMES[i] << ": " << times_[i] * 1000 << "ms (" << (times_[i] / counts_[i] * 1000) << "ms per op, " << counts_[i] << " times)" << std::endl; } } } private: std::unordered_map<int, double> times_; std::unordered_map<int, size_t> counts_; }; class Tensor { public: // Add Tensor &operator+(Tensor &other) { return operator2(other, kOpAdd); } // Mul Tensor &operator*(Tensor &other) { return operator2(other, kOpMul); } // Mul by scalar Tensor &operator*(float val) { assert(type_ == kF32); return *this * *(ctx_->NewTensor({1}, type_)->Fill(val)); } private: Tensor &operator2(Tensor &other_ref, TensorOp op) { auto other = &other_ref; assert(other != this); if (!SameShape(*other)) { assert(can_broadcast_to(*other, *this)); other = &broadcast_to(ctx_, *other, *this); } Tensor *dst = ctx_->NewTensor(Dims()); dst->op_ = op; dst->src0_ = this; dst->src1_ = other; return *dst; } public: // Lookup Tensor &operator[](Tensor &index) { assert(index.type_ == kI32); assert(n_dims_ + index.n_dims_ - 1 <= kMaxTensorDims); std::vector<int> ds; for (int i = kMaxTensorDims - 1; i > kMaxTensorDims - n_dims_; i--) { ds.push_back(dims_[i]); } for (int i = kMaxTensorDims - 1; i >= kMaxTensorDims - index.n_dims_; i--) { ds.push_back(index.dims_[i]); } reverse(ds.begin(), ds.end()); Tensor *dst = ctx_->NewTensor(ds, type_); dst->op_ = kOpLookup; dst->src0_ = this; dst->src1_ = &index; return *dst; } // Norm Tensor &Norm() { assert(type_ == kF32); Tensor *dst = ctx_->NewTensor(Dims(), type_); dst->op_ = kOpNorm; dst->src0_ = this; return *dst; } // Gelu Tensor &Gelu() { assert(type_ == kF32); Tensor *dst = ctx_->NewTensor(Dims(), type_); dst->op_ = kOpGelu; dst->src0_ = this; return *dst; } // Softmax Tensor &Softmax(bool is_causal = false, int vocab_size = 0) { assert(type_ == kF32); if (vocab_size > 0) { assert(vocab_size <= dims_[kMaxTensorDims - 1]); } Tensor *dst = ctx_->NewTensor(Dims(), type_); dst->op_ = kOpSoftmax; dst->src0_ = this; dst->op_params_[0] = is_causal; dst->op_params_[1] = vocab_size; return *dst; } // CrossEntropy Tensor &CrossEntropy(Tensor &target) { assert(type_ == kF32 && target.type_ == kI32); auto shape = Dims(); shape.pop_back(); assert(shape == target.Dims()); Tensor *dst = ctx_->NewTensor(shape, type_); dst->op_ = kOpCrossEntropy; dst->src0_ = this; dst->src1_ = ⌖ return *dst; } // Split std::vector<Tensor *> Split(int size, int axis) { assert(axis < n_dims_); auto dimi = kMaxTensorDims - n_dims_ + axis; assert(dims_[dimi] % size == 0); std::vector<Tensor *> tensors; if (dims_[dimi] == size) { tensors.push_back(this); return tensors; } std::vector<int> shape = Dims(); shape[axis] = size; for (int i = 0; i < dims_[dimi] / size; i++) { tensors.push_back(&View(shape, i, axis)); } return tensors; } // View // TODO(ysg): the view is actually a copy, we need to implement a real view Tensor &View(const std::vector<int> &shape, int split_no = 0, int split_axis = 0) { assert(NumElements() % num_of_elements(shape) == 0); int dimi = kMaxTensorDims - n_dims_ + split_axis; assert(dims_[dimi] % (NumElements() / num_of_elements(shape)) == 0); Tensor *dst = ctx_->NewTensor(shape, type_); dst->op_ = kOpView; dst->src0_ = this; dst->op_params_[0] = split_no; dst->op_params_[1] = split_axis; return *dst; } // Transpose Tensor &Transpose(int axis0, int axis1) { assert(axis0 < n_dims_ && axis1 < n_dims_); auto dimi0 = kMaxTensorDims - n_dims_ + axis0; auto dimi1 = kMaxTensorDims - n_dims_ + axis1; std::vector<int> shape = Dims(); std::swap(shape[axis0], shape[axis1]); Tensor *dst = ctx_->NewTensor(shape, type_); dst->op_ = kOpTranspose; dst->src0_ = this; dst->op_params_[0] = dimi0; dst->op_params_[1] = dimi1; return *dst; } // Matmul // (B, M, N) x (B, P, N) -> (B, M, P) // we assume that the input tensors are in the format (B, M, N) and (B, P, N) Tensor &MatMul(Tensor &other_ref) { auto other = &other_ref; assert(other != this); if (!can_matmul(*this, *other)) { assert(can_broadcast_to(*other, *this, 2)); other = &broadcast_to(ctx_, *other, *this, 2); assert(can_matmul(*this, *other)); } std::vector<int> dst_dims = {dims_[0], dims_[1], dims_[2], other->dims_[2]}; dst_dims.erase(dst_dims.begin(), dst_dims.begin() + dst_dims.size() - n_dims_); Tensor *dst = ctx_->NewTensor(dst_dims); dst->op_ = kOpMatmul; dst->src0_ = this; dst->src1_ = other; return *dst; } void Forward() { std::vector<Tensor *> sorted = topo_sort(this); struct timespec start, end; ForwardProfile.Reset(); for (auto *t : sorted) { clock_gettime(CLOCK_MONOTONIC, &start); switch (t->op_) { case kOpAdd: add_forward(t, t->src0_, t->src1_); break; case kOpMul: mul_forward(t, t->src0_, t->src1_); break; case kOpMatmul: matmul_forward(t, t->src0_, t->src1_); break; case kOpLookup: lookup_forward(t, t->src0_, t->src1_); break; case kOpNorm: norm_forward(t, t->src0_); break; case kOpTranspose: transpose_forward(t, t->src0_, t->op_params_[0], t->op_params_[1]); break; case kOpView: view_forward(t, t->src0_, t->op_params_[0], t->op_params_[1]); break; case kOpBroadcast: broadcast_forward(t, t->src0_); break; case kOpGelu: gelu_forward(t, t->src0_); break; case kOpSoftmax: softmax_forward(t, t->src0_, t->op_params_[0], t->op_params_[1]); break; case kOpCrossEntropy: cross_entropy_forward(t, t->src0_, t->src1_); break; case kOpNone: // no-op break; default: throw std::runtime_error("Forward(): Not implemented, " + TENSOR_OP_NAMES[t->op_]); } clock_gettime(CLOCK_MONOTONIC, &end); ForwardProfile.AddTime(t->op_, end.tv_sec - start.tv_sec + (end.tv_nsec - start.tv_nsec) / 1e9); } } void Backward(bool init_grad = true, float init_val = 1.0f) { std::vector<Tensor *> sorted = topo_sort(this); if (init_grad) { AllocGrad(false)->grad()->Fill(init_val); } struct timespec start, end; BackwardProfile.Reset(); for (auto *t : sorted | std::ranges::views::reverse) { clock_gettime(CLOCK_MONOTONIC, &start); switch (t->op_) { case kOpAdd: add_backward(t, t->src0_, t->src1_); break; case kOpMul: mul_backward(t, t->src0_, t->src1_); break; case kOpMatmul: matmul_backward(t, t->src0_, t->src1_); break; case kOpLookup: lookup_backward(t, t->src0_, t->src1_); break; case kOpNorm: norm_backward(t, t->src0_); break; case kOpTranspose: transpose_backward(t, t->src0_, t->op_params_[0], t->op_params_[1]); break; case kOpView: view_backward(t, t->src0_, t->op_params_[0], t->op_params_[1]); break; case kOpBroadcast: broadcast_backward(t, t->src0_); break; case kOpGelu: gelu_backward(t, t->src0_); break; case kOpSoftmax: softmax_backward(t, t->src0_, t->op_params_[0], t->op_params_[1]); break; case kOpCrossEntropy: cross_entropy_backward(t, t->src0_, t->src1_); break;