"vscode:/vscode.git/clone" did not exist on "c413353e8e4b7f7652c877f4ade69f7e6926a430"
utils.cc 428 Bytes
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
/*!
 *  Copyright (c) 2020 by Contributors
 * \file utils.cc
 * \brief DGL util functions
 */

#include <omp.h>

#include <dgl/packed_func_ext.h>
#include "../c_api_common.h"

using namespace dgl::runtime;

namespace dgl {

DGL_REGISTER_GLOBAL("utils.internal._CAPI_DGLSetOMPThreads")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    int num_threads = args[0];
    omp_set_num_threads(num_threads);
  });

}  // namespace dgl