PR #28782: [XLA:GPU] Annotate cuBLAS/cuDNN outputs to avoid initcheck failures

Imported from GitHub PR openxla/xla#28782

Upgrades NVTX to v3.2.1 and marks the outputs of cuBLAS/cuDNN as initialized (as compute-sanitizer may emit false positives for kernels using TMA).
Copybara import of the project:

--
55977057d4c3bc3008649cdedc7ddb7923780958 by Sergey Kozub <[email protected]>:

[XLA:GPU] Annotate cuBLAS/cuDNN outputs to avoid initcheck failures

Merging this change closes #28782

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#28782 from openxla:skozub/nvtx_init_annotation 55977057d4c3bc3008649cdedc7ddb7923780958
PiperOrigin-RevId: 788806680

Author: sergey-kozub

tsl/profiler/lib/nvtx_utils.cc

@@ -119,4 +119,21 @@ StringHandle RegisterString(ProfilerDomainHandle domain,
   buffer.append(suffix);
   return impl(buffer.c_str());
 }
+
+void MarkMemoryInitialized(void const* address, size_t size,
+                           StreamHandle stream) {
+  auto domain = DefaultProfilerDomain();
+  nvtxMemVirtualRangeDesc_t range_desc{size, address};
+  nvtxMemMarkInitializedBatch_t regions_desc{
+      NVTX_EXT_COMPATID_MEM,
+      sizeof(nvtxMemMarkInitializedBatch_t),
+      NVTX_MEM_TYPE_VIRTUAL_ADDRESS,
+      /*regionDescCount=*/1,
+      sizeof(nvtxMemVirtualRangeDesc_t),
+      &range_desc};
+  nvtxMemCudaMarkInitialized(reinterpret_cast<nvtxDomainHandle_t>(domain),
+                             reinterpret_cast<cudaStream_t>(stream),
+                             /*isPerThreadStream=*/false, &regions_desc);
+}
+
 }  // namespace tsl::profiler

tsl/profiler/lib/nvtx_utils.h

@@ -79,4 +79,9 @@ void RangePush(ProfilerDomainHandle domain, StringHandle title,
 // version of RangePush
 uint64_t RegisterSchema(ProfilerDomainHandle domain, const void* schemaAttr);
 }  // namespace tsl::profiler
+
+// Mark a memory region as initialized.
+// This mitigates false positives from the compute sanitizer (initcheck).
+void MarkMemoryInitialized(void const* address, size_t size,
+                           tsl::profiler::StreamHandle stream);
 #endif  // TENSORFLOW_TSL_PROFILER_LIB_NVTX_UTILS_H_