fix rocm build

microsoft · Feb 22, 2024 · a4c01a5 · a4c01a5
1 parent b8cf88d
commit a4c01a5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc b/onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc
@@ -6,9 +6,10 @@
 #include "fast_gelu.h"
 #include "core/providers/cuda/tensor/gelu_impl.h"
 #include "contrib_ops/cpu/bert/bias_gelu_helper.h"
-#include "transformer_common.h"
 #ifdef USE_ROCM
 #include "contrib_ops/rocm/bert/elementwise.h"
+#else
+#include "contrib_ops/cuda/bert/transformer_common.h"
 #endif
 
 namespace onnxruntime {
@@ -34,8 +35,10 @@ using namespace ONNX_NAMESPACE;
 
 template <typename T>
 FastGelu<T>::FastGelu(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) {
+#ifdef USE_CUDA
   const TransformerOptions* options = TransformerOptions::GetInstance();
   use_half2_ = !options->DisableHalf2();
+#endif
 }
 
 template <typename T>

diff --git a/onnxruntime/contrib_ops/cuda/bert/fast_gelu.h b/onnxruntime/contrib_ops/cuda/bert/fast_gelu.h
@@ -18,7 +18,7 @@ class FastGelu final : public CudaKernel {
   Status ComputeInternal(OpKernelContext* ctx) const override;
 
  private:
-  bool use_half2_;
+  bool use_half2_;  // Only applicable to CUDA kernel (not ROCM).
 };
 
 }  // namespace cuda