Implement all elementwise operations as maps and zips

Author: txus

include/tensor/ops.hpp

@@ -40,6 +40,15 @@ template <typename T, typename D>
 Tensor<std::remove_const_t<T>, D> pow(const TensorView<T, D>& tensor,
                                       std::remove_const_t<T> scalar);
 
+template <typename T, typename D>
+Tensor<std::remove_const_t<T>, D> cos(const TensorView<T, D>& tensor);
+
+template <typename T, typename D>
+Tensor<std::remove_const_t<T>, D> sin(const TensorView<T, D>& tensor);
+
+template <typename T, typename D>
+Tensor<std::remove_const_t<T>, D> exp(const TensorView<T, D>& tensor);
+
 template <typename T, typename D>
 Tensor<std::remove_const_t<T>, D> masked_fill(const TensorView<T, D>& input,
                                               const TensorView<int, D>& mask,

src/llama/rope.cpp

@@ -61,10 +61,10 @@ precompute_rope_values(size_t head_dim, float theta_base, size_t context_length)
 
   angles = cat(angles.view(), angles.view(), 1); // context length, head_Dim
 
-  auto sin = angles.view().sin();
-  auto cos = angles.view().cos();
+  auto sin_ = sin(angles.view());
+  auto cos_ = cos(angles.view());
 
-  return std::make_tuple(std::move(cos), std::move(sin));
+  return std::make_tuple(std::move(cos_), std::move(sin_));
 }
 
 template <typename T, typename D>

src/tensor/cpu/ops.cpp

@@ -152,6 +152,21 @@ Tensor<std::remove_const_t<T>, D> tril(const TensorView<T, D>& tensor, const boo
   return out;
 }
 
+template <typename T, typename D>
+Tensor<std::remove_const_t<T>, D> cos(const TensorView<T, D>& tensor) {
+  return tensor.template map<std::remove_const_t<T>>([](T val) { return std::cos(val); });
+}
+
+template <typename T, typename D>
+Tensor<std::remove_const_t<T>, D> sin(const TensorView<T, D>& tensor) {
+  return tensor.template map<std::remove_const_t<T>>([](T val) { return std::sin(val); });
+}
+
+template <typename T, typename D>
+Tensor<std::remove_const_t<T>, D> exp(const TensorView<T, D>& tensor) {
+  return tensor.template map<std::remove_const_t<T>>([](T val) { return std::exp(val); });
+}
+
 template <typename T, typename D>
 Tensor<std::remove_const_t<T>, D> pow(std::remove_const_t<T> scalar,
                                       const TensorView<T, D>& tensor) {
@@ -578,9 +593,15 @@ template Tensor<float, CPU> max(const TensorView<float, CPU>&, int, bool);
 template Tensor<bfloat16, CPU> masked_fill(const TensorView<bfloat16, CPU>&,
                                            const TensorView<int, CPU>&, bfloat16);
 
+// cat
 template Tensor<bfloat16, CPU> cat(const TensorView<bfloat16, CPU>&,
                                    const TensorView<bfloat16, CPU>&, int);
 template Tensor<float, CPU> cat(const TensorView<float, CPU>&, const TensorView<float, CPU>&, int);
+
+template Tensor<float, CPU> cos(const TensorView<float, CPU>& tensor);
+template Tensor<float, CPU> sin(const TensorView<float, CPU>& tensor);
+template Tensor<float, CPU> exp(const TensorView<float, CPU>& tensor);
+
 template Tensor<bfloat16, CPU> pow(bfloat16, const TensorView<bfloat16, CPU>&);
 template Tensor<float, CPU> pow(const TensorView<float, CPU>&, float);
 template Tensor<float, CPU> pow(float, const TensorView<float, CPU>&);

src/tensor/cuda/CMakeLists.txt

@@ -1,6 +1,6 @@
 file(GLOB HEADER_LIST CONFIGURE_DEPENDS "${PROJECT_SOURCE_DIR}/include/tensor/*.hpp")
 
-add_library(tensor_cuda STATIC storage.cu loader.cu ops.cu kernels/fill.cu kernels/argmax.cu kernels/arange.cu kernels/add.cu kernels/sub.cu kernels/div.cu kernels/mul.cu kernels/sum.cu kernels/max.cu kernels/masked_fill.cu kernels/cat.cu)
+add_library(tensor_cuda STATIC storage.cu loader.cu ops.cu kernels/fill.cu kernels/argmax.cu kernels/arange.cu kernels/sum.cu kernels/max.cu kernels/masked_fill.cu kernels/cat.cu kernels/map.cu kernels/zip.cu)
 
 set_target_properties(tensor_cuda PROPERTIES
   CUDA_SEPARABLE_COMPILATION ON