[ROCm]: Add missing hipStreamWaitEvent API call

Author: Rahul Batra

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))
 

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