Skip to content

Allow user to limit the number of threads that OpenMP spawns. #1362

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 15 commits into from
Apr 11, 2024
Merged

Allow user to limit the number of threads that OpenMP spawns. #1362

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
4612b4f
Refactor the initialization of GUC parameters.
higuoxing Mar 5, 2024
58e2aec
Allow user to limit the number of threads that OpenMP spawns.
higuoxing Mar 8, 2024
99bebe2
Merge branch 'master' into set-omp-num-threads
higuoxing Mar 15, 2024
58f8a9a
Merge branch 'master' into set-omp-num-threads
higuoxing Mar 23, 2024
2dc1964
Merge branch 'master' into set-omp-num-threads
higuoxing Apr 1, 2024
1361bce
Merge branch 'master' into set-omp-num-threads
higuoxing Apr 4, 2024
4e84e08
Address review comments.
higuoxing Apr 5, 2024
2a1853c
Address review comments
higuoxing Apr 5, 2024
3768179
Simplify global variables.
higuoxing Apr 6, 2024
6e3fe80
Fixup!!
higuoxing Apr 6, 2024
2e1891c
Merge branch 'master' into set-omp-num-threads
higuoxing Apr 10, 2024
d05644a
Update pgml-extension/src/bindings/transformers/whitelist.rs
higuoxing Apr 10, 2024
663ed25
Update pgml-extension/src/bindings/transformers/whitelist.rs
higuoxing Apr 10, 2024
c81c42e
Address review comments.
higuoxing Apr 10, 2024
ede3d92
Merge branch 'master' into set-omp-num-threads
kczimm Apr 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Address review comments. 
- Set the minimum value of OMP_NUM_THREADS to 1.
- Move PGML_OMP_NUM_THREADS.get() into initialize_server_params().
  • Loading branch information
@higuoxing
higuoxing committed Apr 10, 2024
commit c81c42e2a64588dd5496a27d2e36b65222584dd9
13 changes: 11 additions & 2 deletions pgml-extension/src/config.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,11 @@ pub static PGML_HF_WHITELIST: GucSetting<Option<&'static CStr>> = GucSetting::<O
pub static PGML_HF_TRUST_REMOTE_CODE: GucSetting<bool> = GucSetting::<bool>::new(false);
pub static PGML_HF_TRUST_REMOTE_CODE_WHITELIST: GucSetting<Option<&'static CStr>> =
GucSetting::<Option<&'static CStr>>::new(None);
pub static PGML_OMP_NUM_THREADS: GucSetting<i32> = GucSetting::<i32>::new(0);
pub static PGML_OMP_NUM_THREADS: GucSetting<i32> = GucSetting::<i32>::new(1);

extern "C" {
fn omp_set_num_threads(num_threads: i32);
}

pub fn initialize_server_params() {
GucRegistry::define_string_guc(
Expand Down Expand Up @@ -53,11 +57,16 @@ pub fn initialize_server_params() {
"Specifies the number of threads used by default of underlying OpenMP library. Only positive integers are valid",
"",
&PGML_OMP_NUM_THREADS,
0,
1,
i32::max_value(),
GucContext::Backend,
GucFlags::default(),
);

let omp_num_threads = PGML_OMP_NUM_THREADS.get();
unsafe {
omp_set_num_threads(omp_num_threads);
}
}

#[cfg(any(test, feature = "pg_test"))]
Expand Down
10 changes: 0 additions & 10 deletions pgml-extension/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,20 +21,10 @@ pg_module_magic!();

extension_sql_file!("../sql/schema.sql", name = "schema");

extern "C" {
fn omp_set_num_threads(num_threads: i32);
}

#[cfg(not(feature = "use_as_lib"))]
#[pg_guard]
pub extern "C" fn _PG_init() {
config::initialize_server_params();
let omp_num_threads = config::PGML_OMP_NUM_THREADS.get();
if omp_num_threads > 0 {
unsafe {
omp_set_num_threads(omp_num_threads);
}
}
bindings::python::activate().expect("Error setting python venv");
orm::project::init();
}
Expand Down