Update block size calculation

src/targets/gpu/compile_hip_code_object.cpp

@@ -144,10 +144,9 @@ compute_global_for(context& ctx, std::size_t n, std::size_t over)
 
 std::size_t compute_block_size(std::size_t n, std::size_t max_block_size)
 {
-    size_t block_size = 128;
-    while(block_size <= max_block_size and block_size <= n)
-        block_size *= 2;
-    return block_size / 2;
+    const std::size_t min_block_size = 64;
+    auto block_size = (((n - 1) / min_block_size + 1)) * min_block_size;
+    return std::min(block_size, max_block_size);
 }
 
 operation compile_hip_code_object(const std::string& content, hip_compile_options options)

src/targets/gpu/jit/layernorm.cpp

@@ -62,7 +62,7 @@ struct layernorm_compiler : compiler<layernorm_compiler>
         auto preloads   = preload::broadcasts(axis, inputs);
         auto relements  = inputs[0].lens()[axis] / vec.size;
         auto nelements  = inputs.back().elements() / relements;
-        auto block_size = compute_block_size(relements, 256);
+        auto block_size = compute_block_size(relements, 512);
         hip_compile_options options;
         options.set_launch_params(
             v, compute_global_for(ctx, nelements * block_size, 256), block_size);

src/targets/gpu/jit/softmax.cpp

@@ -58,7 +58,7 @@ struct softmax_compiler : compiler<softmax_compiler>
         }
         auto relements  = inputs[0].lens()[axis] / vec.size;
         auto nelements  = inputs.back().elements() / relements;
-        auto block_size = compute_block_size(relements, 256);
+        auto block_size = compute_block_size(relements, 512);
         hip_compile_options options;
         options.set_launch_params(
             v, compute_global_for(ctx, nelements * block_size, 256), block_size);