Fix grid_dim_x calculations

[CC-Yeh]

Proposed changes

Fix grid_dim_x calculations which only shows error when launching 2D grids.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[awni]

Looks good but is there a test case for this? Or do we only launch 2D grids for very large inputs (which we can't really test in CI)?

[CC-Yeh]

Looks good but is there a test case for this? Or do we only launch 2D grids for very large inputs (which we can't really test in CI)?

Yeah, I think it would be too big.

bool large = w.size() > UINT_MAX;
------
std::tuple<dim3, uint> get_launch_args(
    size_t size,
    const Shape& shape,
    const Strides& strides,
    bool large,
    int work_per_thread /* = 1 */,
    uint max_block_dim /* = 1024 */) {
  size_t nthreads = cuda::ceil_div(size, work_per_thread);
  uint block_dim = max_block_dim < nthreads ? max_block_dim : nthreads;
  dim3 num_blocks;
  if (large) {
    num_blocks = get_2d_grid_dims(shape, strides, work_per_thread);
    num_blocks.x = cuda::ceil_div(num_blocks.x, block_dim);
  } else {
    num_blocks.x = cuda::ceil_div(nthreads, block_dim);
  }
  return std::make_tuple(num_blocks, block_dim);
}

[awni]

@nastya236 do you mind double checking this? I thought you changed it recently to fix a bug I had written but maybe the change was also not exact?

[nastya236]

Thanks @CC-Yeh! @awni you remember correctly, but I fixed only in fp_quanitze / fp_dequantize. Also @CC-Yeh qqmm_utils.cu would be fixed by this PR.

[awni]

Ty!

[awni]

Do you mind rebasing to resolve the conflict, then we can run tests + merge?

[CC-Yeh]

Do you mind rebasing to resolve the conflict, then we can run tests + merge?

Done!