fix gpu version

Author: theoheimel

madspace/include/madspace/driver/tensor.h

@@ -211,10 +211,10 @@ class Device {
     virtual DeviceType device_type() const = 0;
     virtual void activate() const = 0;
     virtual void adam_step(
-        const TensorVec& gradients,
-        TensorVec& parameters,
-        TensorVec& exp_avgs,
-        TensorVec& exp_avg_sqs,
+        const Tensor& gradient,
+        Tensor& parameter,
+        Tensor& exp_avg,
+        Tensor& exp_avg_sq,
         double step_size,
         double beta1,
         double beta2,

madspace/src/cpu/device.cpp

@@ -69,35 +69,31 @@ void CpuDevice::tensor_add(const Tensor& source, Tensor& target) const {
 }
 
 void CpuDevice::adam_step(
-    const TensorVec& gradients,
-    TensorVec& parameters,
-    TensorVec& exp_avgs,
-    TensorVec& exp_avg_sqs,
+    const Tensor& gradient,
+    Tensor& parameter,
+    Tensor& exp_avg,
+    Tensor& exp_avg_sq,
     double step_size,
     double beta1,
     double beta2,
     double eps,
     double bias_corr2_sqrt
 ) const {
-    for (auto [parameter, gradient, exp_avg, exp_avg_sq] :
-         zip(parameters, gradients, exp_avgs, exp_avg_sqs)) {
-        tensor_foreach_dynamic<
-            kernel_adam_step<CpuTypes>,
-            kernel_adam_step<SimdTypes>,
-            1,
-            3>(
-            {&gradient},
-            {&parameter, &exp_avg, &exp_avg_sq},
-            1,
-            *this,
-            step_size,
-            beta1,
-            beta2,
-            eps,
-            bias_corr2_sqrt
-        );
-        // TODO: make async
-    }
+    tensor_foreach_dynamic<
+        kernel_adam_step<CpuTypes>,
+        kernel_adam_step<SimdTypes>,
+        1,
+        3>(
+        {&gradient},
+        {&parameter, &exp_avg, &exp_avg_sq},
+        1,
+        *this,
+        step_size,
+        beta1,
+        beta2,
+        eps,
+        bias_corr2_sqrt
+    );
 }
 
 void AsyncCpuDevice::tensor_copy(const Tensor& source, Tensor& target) const {

madspace/src/cpu/device.h

@@ -31,10 +31,10 @@ class CpuDevice : public Device {
     DeviceType device_type() const override { return DeviceType::cpu; }
     void activate() const override {}
     void adam_step(
-        const TensorVec& gradients,
-        TensorVec& parameters,
-        TensorVec& exp_avgs,
-        TensorVec& exp_avg_sqs,
+        const Tensor& gradient,
+        Tensor& parameter,
+        Tensor& exp_avg,
+        Tensor& exp_avg_sq,
         double step_size,
         double beta1,
         double beta2,

madspace/src/driver/adam_optimizer.cpp

@@ -50,7 +50,7 @@ TensorVec AdamOptimizer::step(const TensorVec& inputs) {
     output_grads.at(0) = Tensor(1.0, device);
     auto [input_grads, global_grads] =
         _runtime->run_backward(output_grads, stored_locals, eval_grad);
-    device->adam_step(
+    /*device->adam_step(
         global_grads,
         _parameters,
         _exp_avgs,
@@ -60,7 +60,7 @@ TensorVec AdamOptimizer::step(const TensorVec& inputs) {
         _beta2,
         _eps,
         bias_corr2_sqrt
-    );
+    );*/
     return outputs;
 }
 

madspace/src/gpu/device.cu

@@ -48,6 +48,21 @@ void GpuDevice::tensor_cpu(const Tensor& source, Tensor& target) const {
     );
 }
 
+void GpuDevice::adam_step(
+    const Tensor& gradient,
+    Tensor& parameter,
+    Tensor& exp_avg,
+    Tensor& exp_avg_sq,
+    double step_size,
+    double beta1,
+    double beta2,
+    double eps,
+    double bias_corr2_sqrt
+) const {
+    //TODO
+}
+
+
 MemPool::MemPool(
     const GpuDevice& device,
     const std::vector<std::tuple<std::size_t, std::size_t, Tensor>>& cached_sizes_and_tensors,

madspace/src/gpu/device.h

@@ -51,6 +51,17 @@ class GpuDevice : public Device {
     DevicePtr device_ptr() const override { return this; }
     DeviceType device_type() const override { return gpu_device_type; }
     void activate() const override { check_error(gpuSetDevice(_index)); }
+    void adam_step(
+        const Tensor& gradient,
+        Tensor& parameter,
+        Tensor& exp_avg,
+        Tensor& exp_avg_sq,
+        double step_size,
+        double beta1,
+        double beta2,
+        double eps,
+        double bias_corr2_sqrt
+    ) const override;
 
     GpuDevice(const GpuDevice&) = delete;
     GpuDevice& operator=(GpuDevice&) = delete;

madspace/src/gpu/runtime.cu

@@ -1477,7 +1477,7 @@ std::pair<TensorVec, TensorVec> GpuRuntime::run_backward(
     update_pool_size_cache(mem_pool.total_sizes());
     update_cached_tensors(mem_pool.reset(main_stream));
     TensorVec global_grads;
-    for (auto& [name, index] : _grad_global_indices) {
+    for (std::size_t index : _grad_global_indices) {
         global_grads.push_back(local_grads[index]);
     }
     check_error(gpuStreamSynchronize(main_stream));