1 files changed, 33 insertions, 19 deletions

diff --git a/arm_compute/runtime/GLES_COMPUTE/GCTensorAllocator.h b/arm_compute/runtime/GLES_COMPUTE/GCTensorAllocator.h
index ce52cbbbdc..fc14f04ac2 100644
--- a/arm_compute/runtime/GLES_COMPUTE/GCTensorAllocator.h
+++ b/arm_compute/runtime/GLES_COMPUTE/GCTensorAllocator.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017 ARM Limited.
+ * Copyright (c) 2017-2018 ARM Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -27,17 +27,37 @@
 
 #include "arm_compute/core/GLES_COMPUTE/OpenGLES.h"
 #include "arm_compute/runtime/ITensorAllocator.h"
+#include "arm_compute/runtime/MemoryGroupBase.h"
 
 #include <memory>
 
 namespace arm_compute
 {
+class GCTensor;
+template <typename>
+class MemoryGroupBase;
+using GCMemoryGroup = MemoryGroupBase<GCTensor>;
+
+class GLBufferWrapper
+{
+public:
+    GLBufferWrapper()
+        : _ssbo_name(0)
+    {
+        ARM_COMPUTE_GL_CHECK(glGenBuffers(1, &_ssbo_name));
+    }
+    ~GLBufferWrapper()
+    {
+        ARM_COMPUTE_GL_CHECK(glDeleteBuffers(1, &_ssbo_name));
+    }
+    GLuint _ssbo_name;
+};
 /** Basic implementation of a GLES memory tensor allocator. */
 class GCTensorAllocator : public ITensorAllocator
 {
 public:
     /** Default constructor. */
-    GCTensorAllocator();
+    GCTensorAllocator(GCTensor *owner = nullptr);
 
     /** Prevent instances of this class from being copied (As this class contains pointers). */
     GCTensorAllocator(const GCTensorAllocator &) = delete;
@@ -52,7 +72,7 @@ public:
     GCTensorAllocator &operator=(GCTensorAllocator &&) = default;
 
     /** Default destructor */
-    ~GCTensorAllocator() = default;
+    ~GCTensorAllocator();
 
     /** Interface to be implemented by the child class to return the pointer to the mapped data. */
     uint8_t *data();
@@ -95,6 +115,12 @@ public:
      */
     void free() override;
 
+    /** Associates the tensor with a memory group
+     *
+     * @param[in] associated_memory_group Memory group to associate the tensor with
+     */
+    void set_associated_memory_group(GCMemoryGroup *associated_memory_group);
+
 protected:
     /** Call map() on the SSBO.
      *
@@ -106,22 +132,10 @@ protected:
     void unlock() override;
 
 private:
-    class GLBufferWrapper
-    {
-    public:
-        GLBufferWrapper()
-            : _ssbo_name(0)
-        {
-            ARM_COMPUTE_GL_CHECK(glGenBuffers(1, &_ssbo_name));
-        }
-        ~GLBufferWrapper()
-        {
-            ARM_COMPUTE_GL_CHECK(glDeleteBuffers(1, &_ssbo_name));
-        }
-        GLuint _ssbo_name;
-    };
-    std::unique_ptr<GLBufferWrapper> _gl_buffer;
-    uint8_t                         *_mapping;
+    GCMemoryGroup                   *_associated_memory_group; /**< Registered memory group */
+    std::unique_ptr<GLBufferWrapper> _gl_buffer;               /**< OpenGL ES object containing the tensor data. */
+    uint8_t                         *_mapping;                 /**< Pointer to the CPU mapping of the OpenGL ES buffer. */
+    GCTensor                        *_owner;                   /**< Owner of the allocator */
 };
 }