fix groupnorm int32 index overflow (#1845)

tlogn · crcrpar · web-flow · commit 10417aceddd7 · 2025-11-14T11:06:45.000+09:00
* fix groupnorm int32 index overflow

* add groupnorm test case

---------

Co-authored-by: Masaki Kozuki &lt;mkozuki@nvidia.com&gt;
diff --git a/apex/contrib/csrc/group_norm/group_norm_nhwc.h b/apex/contrib/csrc/group_norm/group_norm_nhwc.h
@@ -111,7 +111,7 @@ struct Group_norm_nhwc_fwd_params {
   // The number of instances in the batch.
   int n;
   // The height and width of each activation map. The number of channels.
-  int h, w, c, hw, hwc;
+  int64_t h, w, c, hw, hwc;
   // The number of groups.
   int groups;
   // Do we apply the Swish activation function?
@@ -178,7 +178,7 @@ struct Group_norm_nhwc_bwd_params {
   // The number of instances in the batch.
   int n;
   // The height and width of each activation map. The number of channels.
-  int h, w, c, hw, hwc;
+  int64_t h, w, c, hw, hwc;
   // The number of groups.
   int groups;
   // Do we apply the Swish activation function?
diff --git a/apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_two_pass.cu b/apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_two_pass.cu
@@ -79,7 +79,7 @@ __global__ void group_norm_nhwc_bwd_sum_kernel(Group_norm_nhwc_bwd_params params
   // The first activation loaded by that block.
   int hw_begin = blockIdx.y * params.acts_per_block;
   // The last activation loaded by that block.
-  int hw_end = min(hw_begin + params.acts_per_block, params.hw);
+  int hw_end = min((int64_t) hw_begin + params.acts_per_block, params.hw);
 
   // The gradients for gamma/beta.
   float2 dgamma = make_float2(0.f, 0.f), dbeta = make_float2(0.f, 0.f);
@@ -364,7 +364,7 @@ __global__ void group_norm_nhwc_bwd_scale_kernel(Group_norm_nhwc_bwd_params para
   // The first activation loaded by that block.
   int hw_begin = blockIdx.y * params.acts_per_block;
   // The last activation loaded by that block.
-  int hw_end = min(hw_begin + params.acts_per_block, params.hw);
+  int hw_end = min((int64_t) hw_begin + params.acts_per_block, params.hw);
 
   // Iterate over the activations to compute the sums.
   for( int hwi = hw_begin; hwi < hw_end; ++hwi ) {
diff --git a/apex/contrib/csrc/group_norm/group_norm_nhwc_fwd_two_pass.cu b/apex/contrib/csrc/group_norm/group_norm_nhwc_fwd_two_pass.cu
@@ -43,7 +43,7 @@ __global__ void group_norm_nhwc_fwd_sum_kernel(Group_norm_nhwc_fwd_params params
   // The first activation loaded by that block.
   int hw_begin = blockIdx.y * params.acts_per_block;
   // The last activation loaded by that block.
-  int hw_end = min(hw_begin + params.acts_per_block, params.hw);
+  int hw_end = min((int64_t) hw_begin + params.acts_per_block, params.hw);
 
   // The sums.
   float sum = 0.f, sum_sq = 0.f;
@@ -273,7 +273,7 @@ __global__ void group_norm_nhwc_fwd_scale_kernel(Group_norm_nhwc_fwd_params para
   // The first activation loaded by that block.
   int hw_begin = blockIdx.y * params.acts_per_block;
   // The last activation loaded by that block.
-  int hw_end = min(hw_begin + params.acts_per_block, params.hw);
+  int hw_end = min((int64_t) hw_begin + params.acts_per_block, params.hw);
 
   // Iterate over the activations to compute the sums.
   for( int hwi = hw_begin; hwi < hw_end; ++hwi ) {
diff --git a/apex/contrib/test/group_norm/test_group_norm.py b/apex/contrib/test/group_norm/test_group_norm.py
@@ -208,6 +208,7 @@ def test_16_groups(self):
             [8, 1920, 32, 32],
             [8, 1920, 16, 16],
             [8, 2560, 8, 8],
+            [1, 128, 16128, 1200],
         ]
         for sz in sizes:
             n, c, h, w = sz

Original file line number	Diff line number	Diff line change
`@@ -208,6 +208,7 @@ def test_16_groups(self):`
`208`	`208`	`[8, 1920, 32, 32],`
`209`	`209`	`[8, 1920, 16, 16],`
`210`	`210`	`[8, 2560, 8, 8],`
	`211`	`+ [1, 128, 16128, 1200],`
`211`	`212`	`]`
`212`	`213`	`for sz in sizes:`
`213`	`214`	`n, c, h, w = sz`