[Mlir-commits] [mlir] [mlir][gpu] Skip address space checks for memrefs between launchFuncOp and kernel func (PR #102925)

[Victor Perez]

victor-eds wrote:

After taking a look at the MLIR Vulkan runner (also SPIR-V-based), it faces a similar issue and it's solution is to run a pass to set the expected address spaces both in host and device as part of the conversion pipeline. I haven't dived deep enough into this, but my first perception is this works "by chance" (or good planning in advance!).

The address space they are assuming for default (`memref<f32>`) memory spaces is `StorageBuffer`. So, if they have:
```
gpu.func @kernel(%arg0: memref<f32>) ...
...
gpu.launch @kernel(%arg0: memref<f32>)
```

This would  be changed to:

```
gpu.func @kernel(%arg0: memref<f32, #spirv.StorageClass<StorageBuffer>>) ...
...
gpu.launch @kernel(%arg0: memref<f32, #spirv.StorageClass<StorageBuffer>>)
```

And, when lowered to LLVM dialect (pseudocode):

```
llvm.func @kernel(%arg0: !llvm.ptr) ...
...
"launch_op" @kernel(%arg0: !llvm.ptr)
```

It works fine because the `StorageBuffer` storage class is converted to address space 0 in LLVM. On the host, these programming models (AFAIK) use the default (0) address space both for host-side allocations and device-side (needed for this parameter passing). It happens to be the case the device side uses LLVM address space 0 to represent this particular storage class.

In our case, we'd need to use the `CrossWorkgroup` storage class there, mapping to LLVM address space 1 and, well, we'd be using address space 1 on the host (something these programming models do not do at all). 

IMO being explicit with the memory spaces is better than any other tricks we come up with to avoid this issue. I am not that familiar with other programming models and I don't know how CUDA/HIP pipelines handle this, but, as there are examples of address space mismatches between host and device in LLVM, I would reconsider the approach proposed by @kurapov-peter here.


https://github.com/llvm/llvm-project/pull/102925