Merge pull request #392 from 9prady9/fix_Program_getIR

Fix Program::GetIR to handle programs with multiple devices

2 files changed, 28 insertions, 4 deletions

diff --git a/CHANGELOG b/CHANGELOG
index 9cedb779..9bdf6fa2 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -4,6 +4,7 @@ Development version (next version)
 - Added batched routines to pyclblast
 - Added CLBLAST_VERSION_MAJOR/MINOR/PATCH defines in headers to store version numbering
 - Various minor fixes and enhancements
+- Fixed a bug in the caching when using a context with multiple devices
 
 Version 1.5.1
 - Implemented single-kernel version of convolution as GEMM
diff --git a/src/clpp11.hpp b/src/clpp11.hpp
index 4ed157ea..2a25606c 100644
--- a/src/clpp11.hpp
+++ b/src/clpp11.hpp
@@ -509,12 +509,35 @@ class Program {
 
   // Retrieves a binary or an intermediate representation of the compiled program
   std::string GetIR() const {
-    auto bytes = size_t{0};
-    CheckError(clGetProgramInfo(program_, CL_PROGRAM_BINARY_SIZES, sizeof(size_t), &bytes, nullptr));
+    cl_uint num_devices = 0;
+    CheckError(clGetProgramInfo(program_, CL_PROGRAM_NUM_DEVICES,
+                sizeof(cl_uint), &num_devices, nullptr));
+
+    std::vector<size_t> binSizesInBytes(num_devices, 0);
+    CheckError(clGetProgramInfo(program_, CL_PROGRAM_BINARY_SIZES,
+                num_devices * sizeof(size_t), binSizesInBytes.data(), nullptr));
+
+    auto bytes       = size_t{0};
+    auto binSizeIter = size_t{0};
+    // Loop over the program binary sizes to find a binary whose size is > 0.
+    // The current logic assumes that there ever is only one valid program binary
+    // in a given cl_program. This should be the case unless the cl_program
+    // is built for all or a subset of devices associated to a given cl_program
+    for (; binSizeIter < binSizesInBytes.size(); ++binSizeIter) {
+        if (binSizesInBytes[binSizeIter] > 0) {
+            bytes = binSizesInBytes[binSizeIter];
+            break;
+        }
+    }
     auto result = std::string{};
     result.resize(bytes);
-    auto result_ptr = result.data();
-    CheckError(clGetProgramInfo(program_, CL_PROGRAM_BINARIES, sizeof(char*), &result_ptr, nullptr));
+
+    std::vector<char*> out(num_devices, nullptr);
+    out[binSizeIter] = const_cast<char*>(result.data());
+
+    CheckError(clGetProgramInfo(program_, CL_PROGRAM_BINARIES,
+                                num_devices * sizeof(char*),
+                                out.data(), nullptr));
     return result;
   }