[OpenCL][Texture] Improved texture memmory planning and runtime memory allocation

Motivated form the fact that textures can be allocated over a clBuffer object and
the size of backing clBuffer can be computed based on hardware image pitch alignment.

This optimizes the overall memory allocation on device and helps greately the models with
large memory requirements.

Improvised the graph memory planner to not differentiate buffer and texture storage
tokens and reuse them across. The texture pool in OpenCL runtime is rebranded as memory pool
that handles allocation for both buffer and image objects.

NDArray to DeviceAPI interface is extended with AllocDataSpaceView and FreeDataSpaceView.
These new API's acommodates accessing same physical memory as clBuffer / clImage objects.

* MemoryPool test cases and lint errors.

* test cases and fallback support.

* bug fix and cpp-runtime tests cases for texture views.

* various cl device info organized

* fix graph plan memory bug and correct the testcase.

* device attribute handling

* Some fallback for texture plan on devices w/o cl_khr_image2d_from_buffer

 * Memory Manager

Move the VM memory manager to the runtime level.
Use this memory manager for graph runtime.

 *  Resolve conflicts for VerifyDataType and Buffer

* review comments

Author: srkreddy1238

apps/android_camera/app/src/main/jni/tvm_runtime.h

@@ -61,10 +61,10 @@
 #include "../src/runtime/workspace_pool.cc"
 
 #ifdef TVM_OPENCL_RUNTIME
+#include "../src/runtime/opencl/memory_pool.cc"
 #include "../src/runtime/opencl/opencl_device_api.cc"
 #include "../src/runtime/opencl/opencl_module.cc"
 #include "../src/runtime/opencl/opencl_wrapper/opencl_wrapper.cc"
-#include "../src/runtime/opencl/texture_pool.cc"
 #include "../src/runtime/source_utils.cc"
 #endif
 

apps/android_deploy/app/src/main/jni/tvm_runtime.h

@@ -49,6 +49,5 @@
 #include "../src/runtime/opencl/opencl_device_api.cc"
 #include "../src/runtime/opencl/opencl_module.cc"
 #include "../src/runtime/opencl/opencl_wrapper/opencl_wrapper.cc"
-#include "../src/runtime/opencl/texture_pool.cc"
 #include "../src/runtime/source_utils.cc"
 #endif

apps/android_rpc/app/src/main/jni/tvm_runtime.h

@@ -66,7 +66,6 @@
 #include "../src/runtime/opencl/opencl_device_api.cc"
 #include "../src/runtime/opencl/opencl_module.cc"
 #include "../src/runtime/opencl/opencl_wrapper/opencl_wrapper.cc"
-#include "../src/runtime/opencl/texture_pool.cc"
 #include "../src/runtime/source_utils.cc"
 #endif
 

include/tvm/runtime/device_api.h

@@ -50,6 +50,7 @@ enum DeviceAttrKind : int {
   kApiVersion = 11,
   kDriverVersion = 12,
   kL2CacheSizeBytes = 13,
+  kImagePitchAlignment = 14
 };
 
 #ifdef TVM_KALLOC_ALIGNMENT
@@ -133,12 +134,32 @@ class TVM_DLL DeviceAPI {
    */
   virtual void* AllocDataSpace(Device dev, int ndim, const int64_t* shape, DLDataType dtype,
                                Optional<String> mem_scope = NullOpt);
+
+  /*!
+   * \brief Create a new view with given spec over existing tensor.
+   * \param dev The device device to perform operation.
+   * \param data The source array.
+   * \param shape The shape of allocated tensor.
+   * \param dtype The type of elements.
+   * \param mem_scope The memory scope of allocated tensor.
+   * \return The allocated device pointer.
+   */
+  virtual void* AllocDataSpaceView(Device dev, void* data, ShapeTuple shape, DLDataType dtype,
+                                   Optional<String> mem_scope = NullOpt);
   /*!
    * \brief Free a data space on device.
    * \param dev The device device to perform operation.
    * \param ptr The data space.
    */
   virtual void FreeDataSpace(Device dev, void* ptr) = 0;
+
+  /*!
+   * \brief Free a view data space on device.
+   * \param dev The device device to perform operation.
+   * \param ptr The data space view.
+   */
+  virtual void FreeDataSpaceView(Device dev, void* ptr);
+
   /*!
    * \brief copy data from one place to another
    * \note This API is designed to support special memory with shape dependent layout.

include/tvm/runtime/memory/memory_manager.h

@@ -160,6 +160,13 @@ class Storage : public ObjectRef {
 };
 
 }  // namespace memory
+
+using memory::Allocator;
+using memory::AllocatorType;
+using memory::Buffer;
+using memory::MemoryManager;
+using memory::StorageObj;
+
 }  // namespace runtime
 }  // namespace tvm
 

include/tvm/runtime/ndarray.h

@@ -129,9 +129,11 @@ class NDArray : public ObjectRef {
    * \brief Create a NDArray that shares the data memory with the current one.
    * \param shape The shape of the new array.
    * \param dtype The data type of the new array.
+   * \param mem_scope The memory scope of the array.
    * \note The memory size of new array must be smaller than the current one.
    */
-  TVM_DLL NDArray CreateView(ShapeTuple shape, DLDataType dtype);
+  TVM_DLL NDArray CreateView(ShapeTuple shape, DLDataType dtype,
+                             Optional<String> mem_scope = NullOpt);
   /*!
    * \brief Create a reference view of NDArray that
    *  represents as DLManagedTensor.

src/relay/backend/graph_plan_memory.cc

@@ -229,6 +229,16 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
     VLOG_CONTEXT << "StorageAllocator";
     VLOG(1) << "planning:" << std::endl << PrettyPrint(func);
     prototype_ = StorageAllocaInit(&arena_).GetInitTokenMap(func);
+    // Backup the virtual devices as token reuse might lost the original memory scope
+    std::unordered_map<const ExprNode*, std::vector<VirtualDevice>> virtual_device_map_;
+    for (const auto& kv : prototype_) {
+      std::vector<VirtualDevice> virtual_devices;
+      virtual_devices.reserve(kv.second.size());
+      for (StorageToken* tok : kv.second) {
+        virtual_devices.push_back(tok->virtual_device);
+      }
+      virtual_device_map_.insert({kv.first, virtual_devices});
+    }
     this->Run(func);
 
     // The value of smap contains two integer arrays where the first array
@@ -252,9 +262,13 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
         }
         num_nodes++;
         storage_ids.push_back(tok->storage_id);
-        virtual_devices.push_back(tok->virtual_device);
         sid_sizes_byte.push_back(allocator_.GetMemorySize(tok));
       }
+      ICHECK(kv.second.size() == virtual_device_map_[kv.first].size())
+          << "Mismatch of tokens and virtual devices";
+      for (auto vdev : virtual_device_map_[kv.first]) {
+        virtual_devices.push_back(vdev);
+      }
       auto storage_info = backend::StorageInfo(std::move(storage_ids), std::move(virtual_devices),
                                                std::move(sid_sizes_byte));
       smap.Set(GetRef<Expr>(kv.first), storage_info);
@@ -356,34 +370,27 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
 
   class TokenAllocator {
    public:
-    StorageToken* Alloc(StorageToken* proto) {
-      return Is2DStorage(proto) ? token_2d_.Alloc(proto, storage_ids_++)
-                                : token_1d_.Alloc(proto, storage_ids_++);
-    }
+    StorageToken* Alloc(StorageToken* proto) { return token_mixed_.Alloc(proto, storage_ids_++); }
     StorageToken* Request(StorageToken* proto) {
-      StorageToken* token =
-          Is2DStorage(proto) ? token_2d_.Request(proto) : token_1d_.Request(proto);
+      StorageToken* token = token_mixed_.Request(proto);
       return token ? token : this->Alloc(proto);
     }
-    void CheckForRelease(StorageToken* tok) {
-      return Is2DStorage(tok) ? token_2d_.CheckForRelease(tok) : token_1d_.CheckForRelease(tok);
-    }
+    void CheckForRelease(StorageToken* tok) { return token_mixed_.CheckForRelease(tok); }
 
     size_t GetMemorySize(StorageToken* tok) {
       // TODO(amalyshe): figure out who requries sizes and for what
       // size in case of texture is not enough - we can return any value if it
       // assumed to be used for memory allocatoion or we can return real size
       // if it is just for information
-      return Is2DStorage(tok) ? 0 : token_1d_.GetMemorySize(tok);
+      return token_mixed_.GetMemorySize(tok);
     }
     static bool Is2DStorage(StorageToken* tok) {
       return relay::Is2DStorage(tok->virtual_device->memory_scope);
     }
 
    private:
     int64_t storage_ids_{0};
-    TokenAllocator1D token_1d_;
-    TokenAllocator2D token_2d_;
+    TokenAllocatorMixed token_mixed_;
   };
 
  private: