fix stack OP benchmark (flagos-ai#2716)

* fix stack OP benchmard

* update

* update

Author: huangyiqun

benchmark/test_stack.py

@@ -33,7 +33,6 @@ def _input_fn(shape, dtype, device):
         yield [inp1, inp2, inp3], {"dim": -1},
 
 
-@pytest.mark.skip(reason="CUDA error - illegal memory access: issue #2675")
 @pytest.mark.stack
 def test_stack():
     bench = StackBenchmark(

src/flag_gems/ops/stack.py

@@ -47,11 +47,17 @@ def stack_copy_func_kernel_4(
         dim_offset = dim_offset_d
         total_elements = total_elements_d
 
-    block_start = pid_x * BLOCK_X
-    offsets = tl.arange(0, BLOCK_X)
-    mask = block_start + offsets < total_elements
-
+    block_start = pid_x.to(tl.int64) * BLOCK_X
+    offsets = tl.arange(0, BLOCK_X).to(tl.int64)
     idx = block_start + offsets
+    scalar_zero = offsets * 0
+
+    dim_size_out = dim_size_out + scalar_zero
+    dim_prod_post = dim_prod_post + scalar_zero
+    dim_offset = dim_offset + scalar_zero
+    total_elements = total_elements + scalar_zero
+
+    mask = idx < total_elements
 
     pre_idx = idx // dim_prod_post
     post_idx = idx % dim_prod_post