fixed gru memory problem

TODO.md

@@ -26,4 +26,5 @@
 - Trying SonarCube
 - Separate `gather_batch` from the off-policy runner (should be an operation on the replay buffer)
 - Add Github action that test the compilation of the PX4 module [embedded_platforms/px4](embedded_platforms/px4)
-- Remove the memcopy in `containers/tensor/persist.h`
+- Remove the memcopy in `containers/tensor/persist.h`
+- Check all examples with `-fsanitize=address`

include/rl_tools/containers/tensor/operations_generic.h

@@ -13,16 +13,18 @@ namespace rl_tools{
     }
     template <typename DEVICE, typename SPEC>
     void free(DEVICE& device, Tensor<SPEC>& tensor){
-        delete data(tensor);
+        delete[] data(tensor);
     }
 
     template <typename DEVICE, typename SPEC, typename TI, auto DIM=0, auto SIZE=0>
     auto view_range(DEVICE& device, Tensor<SPEC>& tensor, TI index, tensor::ViewSpec<DIM, SIZE> = {}){
         static_assert(SIZE > 0);
-//        using NEW_SHAPE = tensor::Replace<typename SPEC::SHAPE, SIZE, DIM>;
-//        using NEW_STRIDE = typename SPEC::STRIDE;
+        static_assert(get<DIM>(typename SPEC::SHAPE{}) >= SIZE);
         auto offset = index * get<DIM>(typename SPEC::STRIDE{});
-//        using NEW_SPEC = tensor::Specification<typename SPEC::T, typename SPEC::TI, NEW_SHAPE, NEW_STRIDE>;
+#ifdef RL_TOOLS_DEBUG_CONTAINER_CHECK_BOUNDS
+        utils::assert_exit(device, offset < SPEC::SIZE, "Index out of bounds");
+        utils::assert_exit(device, offset + SIZE <= SPEC::SIZE, "Index out of bounds");
+#endif
         Tensor<tensor::spec::view::range::Specification<SPEC, tensor::ViewSpec<DIM, SIZE>>> view;
         data_reference(view) = data(tensor) + offset;
         return view;
@@ -61,7 +63,11 @@ namespace rl_tools{
     template<typename DEVICE, typename SPEC, typename TII>
     typename SPEC::T get(DEVICE& device, Tensor<SPEC>& tensor, TII local_index){
         static_assert(length(typename SPEC::SHAPE{})==1);
-        return *(data(tensor) + index(device, tensor, local_index));
+        auto idx = index(device, tensor, local_index);
+#ifdef RL_TOOLS_DEBUG_CONTAINER_CHECK_BOUNDS
+        utils::assert_exit(device, idx < SPEC::SIZE, "Index out of bounds");
+#endif
+        return *(data(tensor) + idx);
     }
 
     template<typename DEVICE, typename SPEC, typename TII, typename... INDICES>
@@ -78,7 +84,11 @@ namespace rl_tools{
     template<typename DEVICE, typename SPEC, typename TII> //SFINAE actually not required: typename utils::typing::enable_if_t<length(typename SPEC::SHAPE{})==1>* = nullptr>
     void set(DEVICE& device, Tensor<SPEC>& tensor, typename SPEC::T value, TII current_index){
         static_assert(length(typename SPEC::SHAPE{})==1);
-        *(data(tensor) + index(device, tensor, current_index)) = value;
+        auto idx = index(device, tensor, current_index);
+#ifdef RL_TOOLS_DEBUG_CONTAINER_CHECK_BOUNDS
+        utils::assert_exit(device, idx < SPEC::SIZE, "Index out of bounds");
+#endif
+        *(data(tensor) + idx) = value;
     }
 
     template<typename DEVICE, typename SPEC, typename TII, typename... INDICES> //, typename utils::typing::enable_if_t<tensor::RANK_LARGER_THAN<typename SPEC::SHAPE, 1>>* = nullptr>

include/rl_tools/containers/tensor/persist.h

@@ -75,7 +75,7 @@ namespace rl_tools {
         utils::assert_exit(device, tensor::check_dimensions(device, tensor, dims), "Dimension mismatch");
         typename SPEC::T* data_ptr = data(tensor);
         utils::assert_exit(device, data_ptr != nullptr, "Data pointer is null");
-        constexpr bool VIA_VECTOR = false;
+        constexpr bool VIA_VECTOR = true;
         if constexpr(VIA_VECTOR){
             static_assert(!VIA_VECTOR || (length(typename SPEC::SHAPE{}) <= 3));
             if constexpr(length(typename SPEC::SHAPE{}) == 1){
@@ -97,6 +97,7 @@ namespace rl_tools {
             }
         }
         else{
+            utils::assert_exit(device, dataset.getStorageSize() == sizeof(T)*SPEC::SIZE_BYTES, "Storage size mismatch");
             dataset.read(data_ptr);
         }
     }

include/rl_tools/devices/cpu.h

@@ -24,7 +24,7 @@ namespace rl_tools::devices{
         struct Base{
             static constexpr DeviceId DEVICE_ID = DeviceId::CPU;
             using index_t = size_t;
-            static constexpr index_t MAX_INDEX = -1;
+            static constexpr index_t MAX_INDEX = std::numeric_limits<index_t>::max();
         };
     }
     namespace math{

include/rl_tools/nn/layers/gru/layer.h

@@ -33,28 +33,28 @@ namespace rl_tools::nn::layers::gru {
         static constexpr TI INPUT_DIM = SPEC::INPUT_DIM;
         static constexpr TI HIDDEN_DIM = SPEC::HIDDEN_DIM;
         static constexpr TI NUM_WEIGHTS = SPEC::NUM_WEIGHTS;
-        using WEIGHTS_INPUT_CONTAINER_SHAPE = tensor::Shape<TI, HIDDEN_DIM, INPUT_DIM>;
+        using WEIGHTS_INPUT_CONTAINER_SHAPE = tensor::Shape<TI, 3*HIDDEN_DIM, INPUT_DIM>;
         using WEIGHTS_INPUT_CONTAINER_SPEC = tensor::Specification<T, TI, WEIGHTS_INPUT_CONTAINER_SHAPE>;
         using WEIGHTS_INPUT_CONTAINER_TYPE = typename SPEC::CONTAINER_TYPE_TAG::template type<WEIGHTS_INPUT_CONTAINER_SPEC>;
         using WEIGHTS_INPUT_PARAMETER_SPEC = typename SPEC::PARAMETER_TYPE::template spec<WEIGHTS_INPUT_CONTAINER_TYPE, typename SPEC::PARAMETER_GROUP, nn::parameters::categories::Weights>;
         typename SPEC::PARAMETER_TYPE::template instance<WEIGHTS_INPUT_PARAMETER_SPEC> weights_input;
         typename decltype(weights_input.parameters)::template VIEW_RANGE<tensor::ViewSpec<0, HIDDEN_DIM>> W_ir, W_iz, W_in;
 
-        using BIASES_INPUT_CONTAINER_SHAPE = tensor::Shape<TI, HIDDEN_DIM>;
+        using BIASES_INPUT_CONTAINER_SHAPE = tensor::Shape<TI, 3*HIDDEN_DIM>;
         using BIASES_INPUT_CONTAINER_SPEC = tensor::Specification<T, TI, BIASES_INPUT_CONTAINER_SHAPE>;
         using BIASES_INPUT_CONTAINER_TYPE = typename SPEC::CONTAINER_TYPE_TAG::template type<BIASES_INPUT_CONTAINER_SPEC>;
         using BIASES_INPUT_PARAMETER_SPEC = typename SPEC::PARAMETER_TYPE::template spec<BIASES_INPUT_CONTAINER_TYPE, typename SPEC::PARAMETER_GROUP, nn::parameters::categories::Biases>;
         typename SPEC::PARAMETER_TYPE::template instance<BIASES_INPUT_PARAMETER_SPEC> biases_input;
         typename decltype(biases_input.parameters)::template VIEW_RANGE<tensor::ViewSpec<0, HIDDEN_DIM>> b_ir, b_iz, b_in;
 
-        using WEIGHTS_HIDDEN_CONTAINER_SHAPE = tensor::Shape<TI, HIDDEN_DIM, HIDDEN_DIM>;
+        using WEIGHTS_HIDDEN_CONTAINER_SHAPE = tensor::Shape<TI, 3*HIDDEN_DIM, HIDDEN_DIM>;
         using WEIGHTS_HIDDEN_CONTAINER_SPEC = tensor::Specification<T, TI, WEIGHTS_HIDDEN_CONTAINER_SHAPE>;
         using WEIGHTS_HIDDEN_CONTAINER_TYPE = typename SPEC::CONTAINER_TYPE_TAG::template type<WEIGHTS_HIDDEN_CONTAINER_SPEC>;
         using WEIGHTS_HIDDEN_PARAMETER_SPEC = typename SPEC::PARAMETER_TYPE::template spec<WEIGHTS_HIDDEN_CONTAINER_TYPE, typename SPEC::PARAMETER_GROUP, nn::parameters::categories::Weights>;
         typename SPEC::PARAMETER_TYPE::template instance<WEIGHTS_HIDDEN_PARAMETER_SPEC> weights_hidden;
         typename decltype(weights_hidden.parameters)::template VIEW_RANGE<tensor::ViewSpec<0, HIDDEN_DIM>> W_hr, W_hz, W_hn;
 
-        using BIASES_HIDDEN_CONTAINER_SHAPE = tensor::Shape<TI, HIDDEN_DIM>;
+        using BIASES_HIDDEN_CONTAINER_SHAPE = tensor::Shape<TI, 3*HIDDEN_DIM>;
         using BIASES_HIDDEN_CONTAINER_SPEC = tensor::Specification<T, TI, BIASES_HIDDEN_CONTAINER_SHAPE>;
         using BIASES_HIDDEN_CONTAINER_TYPE = typename SPEC::CONTAINER_TYPE_TAG::template type<BIASES_HIDDEN_CONTAINER_SPEC>;
         using BIASES_HIDDEN_PARAMETER_SPEC = typename SPEC::PARAMETER_TYPE::template spec<BIASES_HIDDEN_CONTAINER_TYPE, typename SPEC::PARAMETER_GROUP, nn::parameters::categories::Biases>;

include/rl_tools/nn/layers/gru/operations_generic.h

@@ -28,7 +28,7 @@ namespace rl_tools{
         layer.b_hz = view_range(device, layer.biases_hidden.parameters, 1*SPEC::HIDDEN_DIM, VIEW_SPEC{});
         layer.b_hn = view_range(device, layer.biases_hidden.parameters, 2*SPEC::HIDDEN_DIM, VIEW_SPEC{});
 
-        malloc(device, layer.initial_hidden_state.parameters);
+        malloc(device, layer.initial_hidden_state);
         set_all(device, layer.initial_hidden_state.parameters, 0);
     }
     template<typename DEVICE, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC>
@@ -44,10 +44,11 @@ namespace rl_tools{
     }
     template <typename DEVICE, typename SPEC>
     void free(DEVICE& device, nn::layers::gru::Layer<SPEC>& layer){
-        malloc(device, layer.weights_input);
-        malloc(device, layer.biases_input);
-        malloc(device, layer.weights_hidden);
-        malloc(device, layer.biases_hidden);
+        free(device, layer.weights_input);
+        free(device, layer.biases_input);
+        free(device, layer.weights_hidden);
+        free(device, layer.biases_hidden);
+        free(device, layer.initial_hidden_state);
     }
 }
 RL_TOOLS_NAMESPACE_WRAPPER_END

tests/src/container/tensor.cpp

@@ -104,6 +104,10 @@ TEST(RL_TOOLS_TENSOR_TEST, SHAPE_OPERATIONS){
     using DEVICE = rlt::devices::DefaultCPU;
     using T = double;
     using TI = typename DEVICE::index_t;
+    using TEST_SHAPE = rlt::tensor::Shape<TI, 2, 3, 4>;
+    using TEST_SPEC = rlt::tensor::Specification<T, TI, TEST_SHAPE>;
+    static_assert(TEST_SPEC::SIZE == 2*3*4);
+
     test_shape_operations<rlt::tensor::Shape<TI, 2, 3, 4>>(3);
     test_shape_operations<rlt::tensor::Shape<TI, 2, 3>>(2);
     test_shape_operations<rlt::tensor::Shape<TI, 2>>(1);

tests/src/nn/layers/gru/CMakeLists.txt

@@ -1,14 +1,17 @@
-#add_executable(
-#        test_nn_layers_gru
-#        gru.cpp
-#)
-#
-#target_link_libraries(
-#        test_nn_layers_gru
-#        PRIVATE
-#        RLtools::RLtools
-#        RLtools::Test
-#        GTest::gtest_main
-#)
-#gtest_discover_tests(test_nn_layers_gru)
-#RL_TOOLS_WARNINGS_ARE_ERRORS(test_nn_layers_gru)
+add_executable(
+        test_nn_layers_gru
+        gru.cpp
+)
+
+target_link_libraries(
+        test_nn_layers_gru
+        PRIVATE
+        RLtools::RLtools
+        RLtools::Test
+        GTest::gtest_main
+        asan
+)
+gtest_discover_tests(test_nn_layers_gru)
+RL_TOOLS_WARNINGS_ARE_ERRORS(test_nn_layers_gru)
+target_compile_options(test_nn_layers_gru PRIVATE -fsanitize=address -fno-omit-frame-pointer)
+target_compile_definitions(test_nn_layers_gru PRIVATE -DRL_TOOLS_DEBUG_CONTAINER_CHECK_BOUNDS)

tests/src/nn/layers/gru/gru.cpp

@@ -14,8 +14,8 @@ namespace rlt = rl_tools;
 
 #include "../../../utils/utils.h"
 
+
 TEST(RL_TOOLS_NN_LAYERS_GRU, LOAD_GRU){
-//int main(){
     using DEVICE = rlt::devices::DefaultCPU;
     using T = double;
     using TI = DEVICE::index_t;
@@ -30,14 +30,6 @@ TEST(RL_TOOLS_NN_LAYERS_GRU, LOAD_GRU){
     rlt::Tensor<rlt::tensor::Specification<T, TI, INPUT_SHAPE>> input;
     using OUTPUT_SHAPE = rlt::tensor::Shape<TI, SEQUENCE_LENGTH, BATCH_SIZE, OUTPUT_DIM>;
     rlt::Tensor<rlt::tensor::Specification<T, TI, INPUT_SHAPE>> output;
-    using WI_SHAPE = rlt::tensor::Shape<TI, HIDDEN_DIM*3, INPUT_DIM>;
-    rlt::Tensor<rlt::tensor::Specification<T, TI, WI_SHAPE>> weight_in, weight_in_grad;
-    using BI_SHAPE = rlt::tensor::Shape<TI, HIDDEN_DIM*3>;
-    rlt::Tensor<rlt::tensor::Specification<T, TI, BI_SHAPE>> bias_in, bias_in_grad;
-    using WH_SHAPE = rlt::tensor::Shape<TI, HIDDEN_DIM*3, HIDDEN_DIM>;
-    rlt::Tensor<rlt::tensor::Specification<T, TI, WH_SHAPE>> weight_hidden, weight_hidden_grad;
-    using BH_SHAPE = rlt::tensor::Shape<TI, HIDDEN_DIM*3>;
-    rlt::Tensor<rlt::tensor::Specification<T, TI, BH_SHAPE>> bias_hidden, bias_hidden_grad;
     using WOUT_SHAPE = rlt::tensor::Shape<TI, OUTPUT_DIM, HIDDEN_DIM>;
     rlt::Tensor<rlt::tensor::Specification<T, TI, WOUT_SHAPE>> weight_out, weight_out_grad;
     using BOUT_SHAPE = rlt::tensor::Shape<TI, OUTPUT_DIM>;
@@ -47,41 +39,13 @@ TEST(RL_TOOLS_NN_LAYERS_GRU, LOAD_GRU){
     rlt::nn::layers::gru::Layer<GRU_SPEC> gru;
     rlt::malloc(device, gru);
 
-
     rlt::malloc(device, input);
     rlt::malloc(device, output);
-    rlt::malloc(device, weight_in);
-    rlt::malloc(device, weight_in_grad);
-    rlt::malloc(device, bias_in);
-    rlt::malloc(device, bias_in_grad);
-    rlt::malloc(device, weight_hidden);
-    rlt::malloc(device, weight_hidden_grad);
-    rlt::malloc(device, bias_hidden);
-    rlt::malloc(device, bias_hidden_grad);
     rlt::malloc(device, weight_out);
     rlt::malloc(device, weight_out_grad);
     rlt::malloc(device, bias_out);
     rlt::malloc(device, bias_out_grad);
 
-    auto W_ir = rlt::view_range(device, weight_in, 0*HIDDEN_DIM, rlt::tensor::ViewSpec<0, HIDDEN_DIM>{});
-    auto W_iz = rlt::view_range(device, weight_in, 1*HIDDEN_DIM, rlt::tensor::ViewSpec<0, HIDDEN_DIM>{});
-    auto W_in = rlt::view_range(device, weight_in, 2*HIDDEN_DIM, rlt::tensor::ViewSpec<0, HIDDEN_DIM>{});
-
-    auto b_ir = rlt::view_range(device, bias_in, 0*HIDDEN_DIM, rlt::tensor::ViewSpec<0, HIDDEN_DIM>{});
-    auto b_iz = rlt::view_range(device, bias_in, 1*HIDDEN_DIM, rlt::tensor::ViewSpec<0, HIDDEN_DIM>{});
-    auto b_in = rlt::view_range(device, bias_in, 2*HIDDEN_DIM, rlt::tensor::ViewSpec<0, HIDDEN_DIM>{});
-
-    auto W_hr = rlt::view_range(device, weight_hidden, 0*HIDDEN_DIM, rlt::tensor::ViewSpec<0, HIDDEN_DIM>{});
-    auto W_hz = rlt::view_range(device, weight_hidden, 1*HIDDEN_DIM, rlt::tensor::ViewSpec<0, HIDDEN_DIM>{});
-    auto W_hn = rlt::view_range(device, weight_hidden, 2*HIDDEN_DIM, rlt::tensor::ViewSpec<0, HIDDEN_DIM>{});
-
-    auto b_hr = rlt::view_range(device, bias_hidden, 0*HIDDEN_DIM, rlt::tensor::ViewSpec<0, HIDDEN_DIM>{});
-    auto b_hz = rlt::view_range(device, bias_hidden, 1*HIDDEN_DIM, rlt::tensor::ViewSpec<0, HIDDEN_DIM>{});
-    auto b_hn = rlt::view_range(device, bias_hidden, 2*HIDDEN_DIM, rlt::tensor::ViewSpec<0, HIDDEN_DIM>{});
-
-    auto W_out = rlt::view_range(device, weight_out, 0*OUTPUT_DIM, rlt::tensor::ViewSpec<0, OUTPUT_DIM>{});
-    auto b_out = rlt::view_range(device, bias_out, 0*OUTPUT_DIM, rlt::tensor::ViewSpec<0, OUTPUT_DIM>{});
-
     std::string DATA_FILE_NAME = "gru_training_trace.h5";
     const char *data_path_stub = RL_TOOLS_MACRO_TO_STR(RL_TOOLS_TESTS_DATA_PATH);
     std::string DATA_FILE_PATH = std::string(data_path_stub) + "/" + DATA_FILE_NAME;
@@ -123,9 +87,15 @@ TEST(RL_TOOLS_NN_LAYERS_GRU, LOAD_GRU){
             rlt::load(device, b_hr_ds, gru.b_hr);
             rlt::load(device, b_hz_ds, gru.b_hz);
             rlt::load(device, b_hn_ds, gru.b_hn);
-            rlt::load(device, W_out_ds, W_out);
-            rlt::load(device, b_out_ds, b_out);
-            rlt::print(device, gru.b_hn);
+            rlt::load(device, W_out_ds, weight_out);
+            rlt::load(device, b_out_ds, bias_out);
         }
     }
+    rlt::free(device, input);
+    rlt::free(device, output);
+    rlt::free(device, weight_out);
+    rlt::free(device, weight_out_grad);
+    rlt::free(device, bias_out);
+    rlt::free(device, bias_out_grad);
+    rlt::free(device, gru);
 }