From 9bbe5500327388bd9efb23e0d5456d9aabea83bd Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Fri, 5 Aug 2022 15:31:53 +0800
Subject: [PATCH] fix(opencl/extern_c_opr): fix cl_mem UAF issue when run model
 OpenCL + ExternCOprRunner, for example

graph: part_a(OpenCL) --> part_b(ExternCOprRunner) --> part_c(OpenCL)
GitOrigin-RevId: f754b559a28dd1efedb2e6649d01e35f35ee4439
---
 src/serialization/impl/extern_c_opr.cpp          |  1 -
 .../megbrain/serialization/extern_c_opr_io.h     |  3 +++
 src/serialization/test/extern_c_opr.cpp          | 16 ++++++++++++++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/serialization/impl/extern_c_opr.cpp b/src/serialization/impl/extern_c_opr.cpp
index 2a2da69e6..e9b49597c 100644
--- a/src/serialization/impl/extern_c_opr.cpp
+++ b/src/serialization/impl/extern_c_opr.cpp
@@ -381,7 +381,6 @@ void ExternCOprRunner::check_param() {
 
 void ExternCOprRunner::scn_do_execute() {
     SmallVector<MGBTensor> c_inp(input().size()), c_out(output().size());
-    SmallVector<HostTensorND> cpu_inp, cpu_out;
     check_param();
 
     bool need_copy = false;
diff --git a/src/serialization/include/megbrain/serialization/extern_c_opr_io.h b/src/serialization/include/megbrain/serialization/extern_c_opr_io.h
index 9d4d4e2f5..f95729dd3 100644
--- a/src/serialization/include/megbrain/serialization/extern_c_opr_io.h
+++ b/src/serialization/include/megbrain/serialization/extern_c_opr_io.h
@@ -16,6 +16,9 @@ MGB_DEFINE_OPR_CLASS_WITH_EXPORT(
     //! store dynamic store param
     std::shared_ptr<ExternCOprParam> m_param;
 
+    //!  HostTensorND holder for scn_do_execute
+    SmallVector<HostTensorND> cpu_inp, cpu_out;
+
     void get_output_var_shape(
             const TensorShapeArray& inp_shape,
             TensorShapeArray& out_shape) const override;
diff --git a/src/serialization/test/extern_c_opr.cpp b/src/serialization/test/extern_c_opr.cpp
index a5be95827..191788e00 100644
--- a/src/serialization/test/extern_c_opr.cpp
+++ b/src/serialization/test/extern_c_opr.cpp
@@ -445,6 +445,22 @@ TEST(TestExternCOpr, GPUCompute) {
     run_compute_test(CompNode::load("gpux"), MGB_DTYPE_FLOAT32);
 }
 
+#if MGB_OPENCL
+#include "megcore_opencl.h"
+
+#define REQUIRE_OPENCL()                                                 \
+    do {                                                                 \
+        if (!CompNode::get_device_count(CompNode::DeviceType::OPENCL)) { \
+            return;                                                      \
+        }                                                                \
+    } while (0)
+
+TEST(TestExternCOpr, OPENCLCompute) {
+    REQUIRE_OPENCL();
+    run_compute_test(CompNode::load("openclx"), MGB_DTYPE_FLOAT32);
+}
+#endif
+
 TEST(TestExternCOpr, CPUComputeMultiDtype) {
     run_compute_test(CompNode::load("cpux"), MGB_DTYPE_INT32);
 #if !MEGDNN_DISABLE_FLOAT16
-- 
GitLab