EnzymeAD / EnzymeAD/Enzyme-JAX

ConvertParallelToGPU puts the block count on gridDim.y (capped at 65535) instead of gridDim.x

Open
#2,789 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
MLIR
Stars
131
Forks
53
Avg merge
1d 10h
Merged PRs (30d)
193

Description

When `convert-parallel-to-gpu1` regenerates a launch for a 1-D CUDA grid, it emits `gridDim = (1, N, 1)`: the recomputed block count `N` goes to `gridDim.y`.
`gridDim.y` is limited to 65535 (2^16-1) on every CUDA architecture, whereas `gridDim.x` allows 2^31-1. Any kernel needing more than 65535 blocks therefore fails to launch with `cudaErrorInvalidConfiguration` (error 9).
The extent is not "moved" to `y` by a swap. It changes *position* during raising: the original `gridDim.x` is fused away into a flat thread extent, the degenerate `gridDim.y` of 1 survives and shifts up into position 0, and the block count can only land in position 1. Position 0 is `x`, position 1 is `y`.

The issue is triggered by RSBench from `wsmoses/Enzyme-GPU-Tests` (branch `mlir`):
`simulation.cu:33` launches `<<<318750, 32>>>`
the default problem size is 300,000 particles x 34 lookups = 10,200,000 lookups, so 10,200,000 / 32 = 318,750 blocks.

**The root cause**:
```
gridDim = (318750, 1, 1) // 318,750 is in position 0
blockDim = ( 32, 1, 1)
```
When raising, `gridDim.x` is fused away:
Two things happen when the launch becomes a parallel loop.
**(a)** `gridDim.x` and `blockDim.x` are multiplied into a single flat extent:

```mlir
%104 = arith.index_cast %101 : i32 to index // original gridDim.x = 318750
%105 = arith.index_cast %103 : i32 to index // original gridDim.y = 1

%106 = "enzymexla.gpu_wrapper"(%104, %105, %c1, %c32, %c1, %c1) ({
%113 = arith.muli %104, %c32 : index // 318750 * 32 = 10,200,000
```

After the multiply, 318,750 no longer exists as a dimension which it has been absorbed into the thread count.

**(b)** The original `gridDim.y` of 1 is *not* fused. It survives as a dimension of its own, and it is emitted first:

```mlir
scf.parallel (%arg4, %arg5) = (%c0, %c0) to (%105, %113) step (%c1, %c1) {
// ^ dim0 = 1 ^ dim1 = 10,200,000
```

| position | bound | origin |
|----------|------------|---------------------------------------|
| 0 | 1 | the original `gridDim.y` (degenerate) |
| 1 | 10,200,000 | `gridDim.x * blockDim.x`, fused |

The degenerate dimension has moved up into position 0, and the dimension carrying all the actual work sits behind it.

**`convert-parallel-to-gpu1`: the block count is recreated**
The pass splits position 1 back into blocks of `blockDim.x`:
```mlir
%114 = arith.subi %106, %c1 : index
%115 = arith.divui %114, %c32 : index
%116 = arith.addi %115, %c1 : index // block count = 318,750
```
This 318,750 is a newly computed value. It happens to equal the original `gridDim.x`, but a multiply and a divide separate them
```
position 0 = 1 (passed through untouched)
position 1 = 318,750 (just split out)
```
`gridBounds` is filled positionally and index i maps to dimension i:

```mlir
gpu.launch blocks(...) in (%arg10 = %105, %arg11 = %116, %arg12 = %c1)
^ x = 1 ^ y = 318,750
threads(...) in (%arg13 = %c32, %arg14 = %c1, %arg15 = %c1)
```

```
position 0 -> x -> gridDim.x = 1
position 1 -> y -> gridDim.y = 318,750
```

It causes the launch fails:

| dimension | limit |
|-------------|----------------|
| `gridDim.x` | 2,147,483,647 |
| `gridDim.y` | **65,535** |
| `gridDim.z` | 65,535 |

Contributor guide

No contributing guide indexed for this repository

Research direction

Start at the convert-parallel-to-gpu1 pass and trace how positional gridBounds become gpu.launch block dimensions. Reproduce the issue with RSBench from Enzyme-GPU-Tests, using simulation.cu:33 and its <<<318750, 32>>> launch; done means the recomputed block count is emitted in gridDim.x so this launch no longer exceeds the y-dimension limit.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.