[tvm] branch main updated: Add check to only cast opaque handles to cl::BufferDescriptor at runtime. (#8256)

[ma...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 75d9b78  Add check to only cast opaque handles to cl::BufferDescriptor at runtime. (#8256)
75d9b78 is described below

commit 75d9b78054ca005f95dbbd02dea1395a8c28eac5
Author: Chris Sullivan <cs...@octoml.ai>
AuthorDate: Mon Jun 14 21:33:02 2021 -0700

    Add check to only cast opaque handles to cl::BufferDescriptor at runtime. (#8256)
---
 src/runtime/opencl/opencl_module.cc | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/runtime/opencl/opencl_module.cc b/src/runtime/opencl/opencl_module.cc
index 631d404..4040d82 100644
--- 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);