Skip to content

Commit

Permalink
Support TSL compilation on Windows with CUDA support.
Browse files Browse the repository at this point in the history
This is following up on #15444. There's still one issue blocking Windows support that I haven't resolved. I've described it here. Any suggestions/advice on how to proceed for that one would be helpful. cc @hawkinsp @ddunl and also fyi @metab0t.

This closes #15499.

PiperOrigin-RevId: 662267938
  • Loading branch information
eaplatanios authored and copybara-github committed Aug 14, 2024
1 parent aa23400 commit 62cfe68
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 5 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,9 @@ def InvokeNvcc(argv, log=False):
nvccopts += ['--keep', '--keep-dir', tempdir]
# Force C++17 dialect (note, everything in just one string!)
nvccopts += ['--std c++17']
# This is so that nvcc does not complain about MSVC or CLANG.
nvccopts += ['-allow-unsupported-compiler']
nvccopts += ['--expt-extended-lambda', '--expt-relaxed-constexpr']
if log:
Log([NVCC_PATH] + nvccopts)

Expand Down
13 changes: 10 additions & 3 deletions third_party/tsl/third_party/gpus/cuda/build_defs.bzl.tpl
Original file line number Diff line number Diff line change
Expand Up @@ -104,9 +104,16 @@ def if_cuda_newer_than(wanted_ver, if_true, if_false = []):
wanted_major = int(wanted_ver.split('_')[0])
wanted_minor = int(wanted_ver.split('_')[1])

configured_version = "%{cuda_version}"
configured_major = int(configured_version.split('.')[0])
configured_minor = int(configured_version.split('.')[1])
# Strip "64_" which appears in the CUDA version on Windows.
configured_version = "%{cuda_version}".rsplit("_", 1)[-1]
configured_version_parts = configured_version.split('.')

# On Windows, the major and minor versions are concatenated without a period and the minor only contains one digit.
if len(configured_version_parts) == 1:
configured_version_parts = [configured_version[0:-1], configured_version[-1:]]

configured_major = int(configured_version_parts[0])
configured_minor = int(configured_version_parts[1])

if %{cuda_is_configured} and (wanted_major, wanted_minor) <= (configured_major, configured_minor):
return select({"//conditions:default": if_true})
Expand Down
5 changes: 3 additions & 2 deletions third_party/tsl/tsl/profiler/lib/nvtx_utils.cc
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ namespace {
// nvtxNameOsThreadA:
// https://nvidia.github.io/NVTX/doxygen/group___r_e_s_o_u_r_c_e___n_a_m_i_n_g.html
// This convention may not match the one in tsl::Env::GetCurrentThreadId().
std::optional<uint32_t> GetCurrentThreadId() {
std::optional<uint32_t> MaybeGetCurrentThreadId() {
#ifdef __linux__
return syscall(SYS_gettid);
#else
Expand All @@ -57,7 +57,8 @@ ProfilerDomainHandle DefaultProfilerDomain() {
}

void NameCurrentThread(const std::string& thread_name) {
if (std::optional<uint32_t> tid = GetCurrentThreadId(); tid.has_value()) {
if (std::optional<uint32_t> tid = MaybeGetCurrentThreadId();
tid.has_value()) {
nvtxNameOsThreadA(*tid, thread_name.c_str());
}
}
Expand Down

0 comments on commit 62cfe68

Please sign in to comment.