Skip to content

Commit

Permalink
Enable prefetching before runpy (#16427)
Browse files Browse the repository at this point in the history
This PR enables prefetching before we execute the `runpy` module and
script code.
  • Loading branch information
galipremsagar authored Jul 30, 2024
1 parent 5feeaf3 commit 0f07b0b
Showing 1 changed file with 10 additions and 11 deletions.
21 changes: 10 additions & 11 deletions python/cudf/cudf/pandas/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,16 @@ def main():
args = parser.parse_args()

rmm_mode = install()
if "managed" in rmm_mode:
for key in {
"column_view::get_data",
"mutable_column_view::get_data",
"gather",
"hash_join",
}:
from cudf._lib import pylibcudf

pylibcudf.experimental.enable_prefetching(key)
with profile(args.profile, args.line_profile, args.args[0]) as fn:
args.args[0] = fn
if args.module:
Expand All @@ -86,17 +96,6 @@ def main():
sys.argv[:] = args.args
runpy.run_path(args.args[0], run_name="__main__")

if "managed" in rmm_mode:
for key in {
"column_view::get_data",
"mutable_column_view::get_data",
"gather",
"hash_join",
}:
from cudf._lib import pylibcudf

pylibcudf.experimental.enable_prefetching(key)


if __name__ == "__main__":
main()

0 comments on commit 0f07b0b

Please sign in to comment.