From 5d2468d3fb91e5bccef3206e407922bec6918565 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 2 Apr 2024 20:48:07 +0000 Subject: [PATCH] [Debug][Disco] Check if a PackedFunc exists before calling it Prior to this commit, attempting to execute the result of `sess.get_global_func` for a non-existing function name would result in a segfault. While the equivalent `tvm.get_global_func` can throw an exception when looking up the function, Disco returns a `DFunction` immediately. This `DFunction` may resolve to a null pointer, and should be checked in the worker process before calling it. --- src/runtime/disco/disco_worker.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/runtime/disco/disco_worker.cc b/src/runtime/disco/disco_worker.cc index e8ba351e791f..b281a3aca7da 100644 --- a/src/runtime/disco/disco_worker.cc +++ b/src/runtime/disco/disco_worker.cc @@ -77,7 +77,9 @@ struct DiscoWorker::Impl { } case DiscoAction::kCallPacked: { int func_reg_id = args[2]; + CHECK_LT(func_reg_id, self->register_file.size()); PackedFunc func = GetReg(self, func_reg_id); + CHECK(func.defined()); CallPacked(self, reg_id, func, TVMArgs(args.values + 3, args.type_codes + 3, args.num_args - 3)); break;