Add dynamic buffer support to OCL Backend (#3765)

Summary:
The OpenCL Backend uses a static memory allocation strategy of allocating a single large buffer and then using offsets into it, which is good for the general case, but doesn't allow us to get the most benefit out of Device Resident Tensors (when we'd like to leave an output on the device to be used as the input to another network). This PR adds a more dynamic mapping of device buffers to the OCL backend via OpenCL SubBuffers, which are similar to Glow TensorViews in that they provide access to a region without additional allocations.

There is no behavioural change in this PR, but it provides infrastructure to reference buffers outside of the range of the DeviceBuffer in the future, which we need to get DRT perf wins.

The immediate benefit is that I was able to simplify the OCL kernel code, deleting about 25% of kernels.cl.

Documentation:  NFC
Pull Request resolved: #3765

Test Plan: tests in release and ASAN

Differential Revision: D18465407

Pulled By: nickgg

fbshipit-source-id: 1b5416c4f389885bae4d5e1533a65bef8ab60122

Author: nickgg

.circleci/build.sh

@@ -140,7 +140,7 @@ elif [[ "$CIRCLE_JOB" == "PYTORCH" ]]; then
     cd build
 elif [[ "$CIRCLE_JOB" == "OPENCL" ]]; then
     install_pocl
-    CMAKE_ARGS+=("-DGLOW_WITH_OPENCL=ON")
+    CMAKE_ARGS+=("-DGLOW_WITH_OPENCL=ON" "-DGLOW_OPENCL_ALIGN=128")
 else
     CMAKE_ARGS+=("-DCMAKE_BUILD_TYPE=Debug")
     if [[ "${CIRCLE_JOB}" == "SHARED" ]]; then

lib/Backends/OpenCL/OpenCL.cpp

@@ -145,8 +145,7 @@ class OpenCLFunction final : public CompiledFunction {
   /// Fill the device \p buffer with a given \p value.
   /// \param len number of buffer elements to be filled by the \p value.
   /// Elements are considered to be of the type described by \p elemKind.
-  void fillBuffer(cl_mem buffer, uint64_t start, uint64_t len, float value,
-                  ElemKind elemKind,
+  void fillBuffer(cl_mem buffer, uint64_t len, float value, ElemKind elemKind,
                   runtime::OpenCLDeviceBindings *devBindings);
 
   /// Execution a convolution instruction which uses NCHW format.
@@ -242,10 +241,13 @@ namespace runtime {
 /// device specific information used to run a compiled function on a specific
 /// device.
 struct OpenCLDeviceBindings : DeviceBindings {
-  OpenCLDeviceBindings(cl_mem buffer, cl_command_queue commands,
-                       cl_device_id device, cl_context ctx, cl_program prog)
+  OpenCLDeviceBindings(
+      cl_mem buffer, cl_command_queue commands, cl_device_id device,
+      cl_context ctx, cl_program prog,
+      const std::unordered_map<std::string, cl_mem> &subBuffers)
       : DeviceBindings(OCLBackend::getName()), deviceBuffer{buffer},
-        commandQueue{commands}, deviceId{device}, context{ctx}, program{prog} {}
+        commandQueue{commands}, deviceId{device}, context{ctx}, program{prog},
+        weightBuffers(subBuffers) {}
 
   /// CL memory buffer. Currently this contains both mutable and immutable
   /// weights, the buffer is allocated once when the network is added.
@@ -269,6 +271,12 @@ struct OpenCLDeviceBindings : DeviceBindings {
 
   /// A list of kernels and their associated events.
   std::vector<KernelLaunch> kernelLaunches;
+
+  /// Buffers or subBuffers associated with symbols.
+  std::unordered_map<std::string, cl_mem> weightBuffers;
+
+  /// /returns the subBufffer assciated with a Value.
+  cl_mem getBuffer(glow::Value *v);
 };
 } // namespace runtime
 } // namespace glow

lib/Backends/OpenCL/OpenCL.h

@@ -55,7 +55,29 @@ DeviceManager *createOCLDeviceManager(const DeviceConfig &config) {
   return new OpenCLDeviceManager(config);
 }
 
-OpenCLBuffer::~OpenCLBuffer() { clReleaseMemObject(buffer_); }
+OpenCLBuffer::~OpenCLBuffer() {
+  for (auto buf : subBuffers_) {
+    clReleaseMemObject(buf.second);
+  }
+  subBuffers_.clear();
+
+  clReleaseMemObject(buffer_);
+}
+
+/// Add a mapping from a Symbol name to an offset into buffer_;
+bool OpenCLBuffer::addSubBuffer(std::string name, size_t offset, size_t size) {
+  cl_buffer_region region({offset, size});
+  cl_int err;
+  auto buf = clCreateSubBuffer(buffer_, CL_MEM_READ_WRITE,
+                               CL_BUFFER_CREATE_TYPE_REGION, &region, &err);
+  auto res = subBuffers_.emplace(name, buf);
+  if (!res.second) {
+    llvm::dbgs() << "OpenCLBuffer: failed to add subBuffer for symbol " << name
+                 << "\n";
+    return false;
+  }
+  return true;
+}
 } // namespace runtime
 } // namespace glow
 
@@ -356,6 +378,15 @@ void OpenCLDeviceManager::addNetworkImpl(const Module *module,
       clFinish(commands);
     }
     usedMemoryBytes_ += sizeInBytes;
+
+    // Add a sub-buffer for each symbol in the symbol table. OpenCL sub-buffers
+    // are essentially TensorViews in Glow.
+    for (auto &pair : bundle.getSymbolTable()) {
+      bool success = buffer->addSubBuffer(pair.first, pair.second.offset,
+                                          pair.second.size);
+      DCHECK(success);
+    }
+
     // Compile the CL program.
     // Add to the function name lookup map.
     // Add shared pointer to the buffer to buffers. This way the buffer will
@@ -376,6 +407,7 @@ void OpenCLDeviceManager::addNetworkImpl(const Module *module,
     programs_.emplace(func.first, program);
     functions_.emplace(func.first, func.second);
     buffers_.emplace(func.first, buffer);
+
     buffer->incrementUsers();
 
     DCHECK_LE(usedMemoryBytes_, maxMemoryBytes_);
@@ -666,7 +698,7 @@ void OpenCLDeviceManager::runFunctionImpl(
   auto program = programs_[function];
   auto clBindings = glow::make_unique<runtime::OpenCLDeviceBindings>(
       buffers_[function]->getBuffer(), queue.backingQueue, deviceId_, context_,
-      program);
+      program, buffers_[function]->getSubBuffers());
 
   // Copy inputs to the device.
   copyInputsToDevice(func->getRuntimeBundle(), context.get(), clBindings.get());

lib/Backends/OpenCL/OpenCLDeviceManager.cpp

@@ -98,6 +98,9 @@ class OpenCLBuffer {
   /// The OpenCL buffer being stored.
   cl_mem buffer_;
 
+  /// Subbuffers for symbols.
+  std::unordered_map<std::string, cl_mem> subBuffers_;
+
   /// Count of functions using this buffer.
   unsigned int users_{0};
 
@@ -120,6 +123,14 @@ class OpenCLBuffer {
 
   /// Get size of buffer in bytes.
   size_t getSize() { return size_; }
+
+  /// Return the mapping from Symbol name to subBuffer for this Buffer.
+  const std::unordered_map<std::string, cl_mem> &getSubBuffers() {
+    return subBuffers_;
+  }
+
+  /// Add a mapping from a Symbol name to an offset into buffer_;
+  bool addSubBuffer(std::string name, size_t offset, size_t size);
 };
 
 /// A class controlling a single OpenCL device. Many OpenCLFunctions may be