Add a few missing unroll directives

ROCm · Mar 25, 2024 · a7164ca · a7164ca
1 parent 20f8bd1
commit a7164ca
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/csrc/layernorm_kernels.cu b/csrc/layernorm_kernels.cu
@@ -59,6 +59,8 @@ struct _halfVec {
 
   __device__ inline _halfVec& operator+=(const _halfVec<width>& other) {
     if constexpr (width % 2 == 0) {
+      // Hint to compiler to use packed operations
+      #pragma unroll
       for (int i = 0; i < width; i += 2) {
         __half2 z = __half2{data[i], data[i+1]};
         z += __half2{other.data[i], other.data[i+1]};
@@ -75,6 +77,8 @@ struct _halfVec {
 
   __device__ inline _halfVec& operator*=(const _halfVec<width>& other) {
     if constexpr (width % 2 == 0) {
+      // Hint to compiler to use packed operations
+      #pragma unroll
       for (int i = 0; i < width; i += 2) {
         __half2 z = __half2{data[i], data[i+1]};
         z *= __half2{other.data[i], other.data[i+1]};
@@ -91,6 +95,7 @@ struct _halfVec {
 
   __device__ inline _halfVec& operator*=(const float scale) {
     if constexpr (width % 2 == 0) {
+      // Hint to compiler to use packed operations
       #pragma unroll
       for (int i = 0; i < width; i += 2) {
         float2 zf = __half22float2(__half2{data[i], data[i+1]});
@@ -109,6 +114,7 @@ struct _halfVec {
   __device__ inline float sum_squares() const {
     float result = 0.0f;
     if constexpr (width % 2 == 0) {
+      // Hint to compiler to use packed operations
       #pragma unroll
       for (int i = 0; i < width; i += 2) {
         float2 z = __half22float2(__half2{data[i], data[i+1]});