Skip to content

Add BLAS thread control and update docstring#36

Merged
AsymmetryChou merged 4 commits into
deepmodeling:mainfrom
AsymmetryChou:rgf_acc
Jul 12, 2026
Merged

Add BLAS thread control and update docstring#36
AsymmetryChou merged 4 commits into
deepmodeling:mainfrom
AsymmetryChou:rgf_acc

Conversation

@AsymmetryChou

Copy link
Copy Markdown
Collaborator

This pull request adds an autotuning mechanism for BLAS threads in the self-energy calculation pipeline, improving parallel performance and resource utilization. It introduces a new function to select the optimal number of BLAS threads per worker by benchmarking the real code path, updates the worker and parallel execution logic to use this value, and enhances the configuration options for users. Additionally, the code is refactored for clarity and maintainability.

Performance and Parallelism Improvements:

  • Added _autotune_blas_threads, which benchmarks the self-energy computation at several BLAS thread counts and selects the fastest, ensuring efficient use of CPU resources without oversubscription. This is run in the parent process before launching parallel workers.
  • Updated compute_all_self_energy to accept a blas_threads parameter and to use the autotuned value per worker, ensuring each process is optimally configured for BLAS threading. [1] [2] [3]
  • Replaced _self_energy_worker_blas1 with _self_energy_worker_blas, which sets the BLAS thread count per worker according to the autotuned or user-specified value.

API and Configuration Enhancements:

  • Extended compute_all_self_energy to support new parameters: ek_batch_size (renamed from batch_size for clarity) and blas_threads (for user override or autotune of BLAS threads per worker).
  • Improved docstrings and parameter lists to reflect the new options and clarify the behavior of memory and thread management. [1] [2]

Refactoring and Maintenance:

  • Refactored code for clarity, including splitting lines for readability and adding helper functions for principal-layer dimension extraction. [1] [2]
  • Added missing import for time to support benchmarking in autotuning.

These changes collectively make parallel self-energy calculations more robust, efficient, and configurable.

@AsymmetryChou AsymmetryChou self-assigned this Jul 11, 2026
@AsymmetryChou

Copy link
Copy Markdown
Collaborator Author

I add an autotune_blas_threads mechanism in lead_property.py, to choosing the best thread automatically. How do you like it? @Lonya0

@AsymmetryChou
AsymmetryChou merged commit 5009bf7 into deepmodeling:main Jul 12, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant