Browse Source

fix(mge/random): fix delete_rng_handel

GitOrigin-RevId: a836e13c23
release-1.5
Megvii Engine Team huangxinda 3 years ago
parent
commit
eb66681f45
3 changed files with 5 additions and 3 deletions
  1. +3
    -1
      imperative/python/src/ops.cpp
  2. +1
    -2
      imperative/src/impl/interpreter/interpreter_impl.h
  3. +1
    -0
      imperative/src/include/megbrain/imperative/interpreter.h

+ 3
- 1
imperative/python/src/ops.cpp View File

@@ -467,7 +467,9 @@ void init_ops(py::module m) {
m.def("delete_rng_handle", [](size_t handle){
// RNG op might execute after handle released due to async dispatch, so
// we need sync before delete a handle to avoid memory leak or use-after-free
python::interpreter_for_py->sync();
if(python::interpreter_for_py->check_available()){
python::interpreter_for_py->sync();
}
mgb::CompNode::sync_all();
py_task_q.wait_all_task_finish();
rng::delete_handle(handle);


+ 1
- 2
imperative/src/impl/interpreter/interpreter_impl.h View File

@@ -60,6 +60,7 @@ struct ChannelImpl : Interpreter::Channel {

DeviceTensorND get_dev_tensor(Handle) override;

bool check_available() override;
void sync() override;
void close() override;

@@ -112,8 +113,6 @@ private:
const SmallVector<LogicalTensorDesc>& input_descs,
SmallVector<Handle>* outputs);

bool check_available();

void push_scope(std::string, State&);
void pop_scope(std::string, State&);



+ 1
- 0
imperative/src/include/megbrain/imperative/interpreter.h View File

@@ -41,6 +41,7 @@ struct Interpreter {

virtual DeviceTensorND get_dev_tensor(Handle) = 0;

virtual bool check_available() = 0;
virtual void sync() = 0;
virtual void close() = 0;



Loading…
Cancel
Save