From 0c9fc220c5f75a409acb4c6f3b93f61586be59eb Mon Sep 17 00:00:00 2001
From: Arda Coskunses <arda.coskunses@amd.com>
Date: Thu, 12 Sep 2024 14:18:47 -0700
Subject: [PATCH] xHip stream fix for non cuda env

---
 .../gen_ai/src/quantize/ck_extensions/fp8_blockwise_gemm.hip  | 4 +++-
 .../gen_ai/src/quantize/ck_extensions/fp8_rowwise_gemm.hip    | 4 +++-
 .../gen_ai/src/quantize/ck_extensions/fp8_tensorwise_gemm.hip | 4 +++-
 .../src/quantize/ck_extensions/kernels/fp8_rowwise_common.h   | 3 +++
 4 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_blockwise_gemm.hip b/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_blockwise_gemm.hip
index 17d4604834..94c9dbe130 100644
--- a/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_blockwise_gemm.hip
+++ b/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_blockwise_gemm.hip
@@ -12,11 +12,13 @@
 #include <numeric>
 
 #include <ATen/ATen.h>
+#if !defined(USE_ROCM)
 #include <c10/cuda/CUDAStream.h>
+#endif
 #include <torch/torch.h>
 
 #if defined(USE_ROCM)
-
+#include <c10/hip/HIPStream.h>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
diff --git a/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise_gemm.hip b/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise_gemm.hip
index afb3064324..b57563df58 100644
--- a/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise_gemm.hip
+++ b/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise_gemm.hip
@@ -15,11 +15,13 @@
 #include <unordered_map>
 
 #include <ATen/ATen.h>
+#if !defined(USE_ROCM)
 #include <c10/cuda/CUDAStream.h>
+#endif
 #include <torch/torch.h>
 
 #if defined(USE_ROCM)
-
+#include <c10/hip/HIPStream.h>
 #include "kernels/fp8_rowwise_kernel_manifest.h"
 
 namespace fbgemm_gpu {
diff --git a/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_tensorwise_gemm.hip b/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_tensorwise_gemm.hip
index 6170675a24..67966a43cf 100644
--- a/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_tensorwise_gemm.hip
+++ b/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_tensorwise_gemm.hip
@@ -12,11 +12,13 @@
 #include <numeric>
 
 #include <ATen/ATen.h>
+#if !defined(USE_ROCM)
 #include <c10/cuda/CUDAStream.h>
+#endif
 #include <torch/torch.h>
 
 #if defined(USE_ROCM)
-
+#include <c10/hip/HIPStream.h>
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
diff --git a/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/kernels/fp8_rowwise_common.h b/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/kernels/fp8_rowwise_common.h
index 7177205802..da40879894 100644
--- a/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/kernels/fp8_rowwise_common.h
+++ b/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/kernels/fp8_rowwise_common.h
@@ -12,10 +12,13 @@
 #include <numeric>
 
 #include <ATen/ATen.h>
+#if !defined(USE_ROCM)
 #include <c10/cuda/CUDAStream.h>
+#endif
 #include <torch/torch.h>
 
 #if defined(USE_ROCM)
+#include <c10/hip/HIPStream.h>
 
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"