Merge pull request #20341 from ROCm:rocm_add_hipStreamWaitEvent

jax authors · jax authors · commit 2848cda34cc1 · 2024-03-21T10:41:38.000-07:00
PiperOrigin-RevId: 617893634
diff --git a/jaxlib/gpu/triton_kernels.cc b/jaxlib/gpu/triton_kernels.cc
@@ -31,6 +31,7 @@
 #include "jaxlib/gpu/vendor.h"
 #include "xla/service/custom_call_status.h"
 #include "xla/stream_executor/gpu/asm_compiler.h"
+#include "tsl/platform/env.h"
 
 #define GPU_RETURN_IF_ERROR(expr) JAX_RETURN_IF_ERROR(JAX_AS_STATUS(expr))
 
diff --git a/jaxlib/gpu/vendor.h b/jaxlib/gpu/vendor.h
@@ -496,6 +496,7 @@ typedef hipsparseDnVecDescr_t gpusparseDnVecDescr_t;
 #define gpuMemcpyHostToDevice hipMemcpyHostToDevice
 #define gpuMemcpyDeviceToHost hipMemcpyDeviceToHost
 #define gpuStreamSynchronize hipStreamSynchronize
+#define gpuStreamWaitEvent hipStreamWaitEvent
 #define gpuSuccess hipSuccess
 
 #define gpuCtxGetDevice hipCtxGetDevice
