Add check to only cast opaque handles to cl::BufferDescriptor at runt…

…ime. (apache#8256)
neo-ai · Jun 17, 2021 · 93039ac · 93039ac
1 parent 5ee4091
commit 93039ac
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/src/runtime/opencl/opencl_module.cc b/src/runtime/opencl/opencl_module.cc
@@ -64,8 +64,13 @@ class OpenCLWrappedFunc {
     }
     // setup arguments.
     for (cl_uint i = 0; i < arg_size_.size(); ++i) {
-      auto* arg = static_cast<cl::BufferDescriptor*>(void_args[i]);
-      OPENCL_CALL(clSetKernelArg(kernel, i, arg_size_[i], arg->buffer));
+      void* arg = nullptr;
+      if (args.type_codes[i] == DLDataTypeCode::kDLOpaqueHandle) {
+        arg = static_cast<cl::BufferDescriptor*>(void_args[i])->buffer;
+      } else {
+        arg = void_args[i];
+      }
+      OPENCL_CALL(clSetKernelArg(kernel, i, arg_size_[i], arg));
     }
     cl_command_queue queue = w_->GetQueue(t->device);
     ThreadWorkLoad wl = thread_axis_cfg_.Extract(args);