Simple GPU Memory Allocation

MLIR
,
1

Hello. I am having some trouble allocating a simple memory buffer for access on the GPU.

When writing CUDA code, you’d typically do something like the following:

__global__ void kernel(float *mem) { mem[0] = 3.14; }

int main() {
  // Allocate host memory
  float *mem = new float[10];

  // Allocate device memory
  float *d_mem;
  cudaMalloc(&d_mem, 10);

  // Memcpy host -> device
  cudaMemcpy(d_mem, mem, 10, cudaMemcpyHostToDevice);

  // Launch kernel
  kernel<<<1, 1>>>(d_mem);

  // Memcpy device -> host
  cudaMemcpy(mem, d_mem, 10, cudaMemcpyDeviceToHost);

  cudaFree(d_mem);
  delete[] mem;
  return 0;
}

I tried writing a similar program in MLIR and came up with the following:

module {
  func.func @main() {
    %alloc = memref.alloc() : memref<2x2xf32>
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1) threads(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) {
      %f = arith.constant 3.1415 : f32
      memref.store %f, %alloc[%c0, %c1] : memref<2x2xf32>
      gpu.terminator
    }
    return
  }
}

The translation has no problems:

mlir-opt -gpu-lower-to-nvvm-pipeline="cubin-format=fatbin"

module attributes {gpu.container_module} {
  llvm.func @malloc(i64) -> !llvm.ptr
  llvm.func @main() {
    %0 = llvm.mlir.constant(64 : index) : i64
    %1 = llvm.mlir.constant(0 : index) : i64
    %2 = llvm.mlir.constant(1 : index) : i64
    %3 = llvm.mlir.constant(2 : index) : i64
    %4 = llvm.mlir.zero : !llvm.ptr
    %5 = llvm.getelementptr %4[4] : (!llvm.ptr) -> !llvm.ptr, f32
    %6 = llvm.ptrtoint %5 : !llvm.ptr to i64
    %7 = llvm.add %6, %0 : i64
    %8 = llvm.call @malloc(%7) : (i64) -> !llvm.ptr
    %9 = llvm.ptrtoint %8 : !llvm.ptr to i64
    %10 = llvm.sub %0, %2 : i64
    %11 = llvm.add %9, %10 : i64
    %12 = llvm.urem %11, %0 : i64
    %13 = llvm.sub %11, %12 : i64
    %14 = llvm.inttoptr %13 : i64 to !llvm.ptr
    gpu.launch_func  @main_kernel::@main_kernel blocks in (%2, %2, %2) threads in (%2, %2, %2) : i64 args(%8 : !llvm.ptr, %14 : !llvm.ptr, %1 : i64, %3 : i64, %3 : i64, %3 : i64, %2 : i64, %1 : i64, %2 : i64)
    llvm.return
  }
  gpu.binary @main_kernel  [#gpu.object<#nvvm.target, "(omitted)">]
}

However, when I try running this code:

mlir_runner_utils=".../libmlir_cuda_runtime.so.20.0git"
mlir_c_runner_utils=".../libmlir_c_runner_utils.so.20.0git"
mlir_cuda_runtime=".../libcudart.so.11.7.99"
mlir-cpu-runner \
  --shared-libs=$mlir_cuda_runtime \
  --shared-libs=$mlir_runner_utils \
  --shared-libs=$mlir_c_runner_utils \
  --entry-point-result=void

I get the following error:

'cuStreamSynchronize(stream)' failed with 'CUDA_ERROR_ILLEGAL_ADDRESS'
'cuStreamDestroy(stream)' failed with 'CUDA_ERROR_ILLEGAL_ADDRESS'
'cuModuleUnload(module)' failed with 'CUDA_ERROR_ILLEGAL_ADDRESS'

Here’s what I’ve tried:

  • Doing the memref.alloc inside the kernel: works perfectly with no errors, but this is not what I want since the memref is not available to the host
  • Outlining the gpu.launch body into its own kernel function, and using gpu.launch_func (same error)
  • Using the memory-space attribute on the buffer {alignment = 64 : i64, memory_space = #gpu.address_space<global>} (same error)
  • Performing a %a = memref.load instead of a memref.store (same error when I try to gpu.print "%f" %a : f32)
  • Using synchronous gpu.alloc instead. error: (error: unsupported GPU operation: gpu.alloc. error: LLVM Translation failed for operation: gpu.alloc)
  • Reading through some of the GPU dialect tests to see how this is done (didn’t help)

Surely I am missing something here since this is such a simple task. Do you have any debugging tips? I am presented with this error message, but I have no idea where it comes from or what is triggering it.

2

I guess you’re missing gpu.host_register? Alternatively you might be able to do gpu.alloc host_shared.

I would say this deceptive - gpu dialect abstracts over the various runtimes and it’s quite a leaky/thin abstraction I think - it doesn’t automatically handle many of the things superficially you would expect it to from just reading the names/descriptions of the ops (i.e., you have to dive into the impls of the op lowerings to figure out what actual runtime calls you’re gonna get).

3

I also tried using host_register, but I’ve heard it isn’t fully supported.

  func.func @main() {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %alloc = memref.alloc() : memref<4xf32>

    %a0 = memref.cast %alloc : memref<4xf32> to memref<?xf32>
    %a1 = memref.cast %a0    : memref<?xf32> to memref<*xf32>
    gpu.host_register %a1 : memref<*xf32>

    gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1) threads(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) {
      %f = arith.constant 3.1415 : f32
      memref.store %f, %alloc[%c0] : memref<4xf32>

      gpu.terminator
    }
    return
  }

Is this usage correct? I still get the same error…

4

Your updated case should be able to run smoothly, just as the upstream MLIR GPU Integration test.

If you still get the same error, you should check if the GPU supports the host register API.

On the other hand, you could also generate a new gpu.alloc and gpu.memcpy for the original kernel args, and pass the memcpyed value into the kernel, like the SparseGPUCodegen pass does.

5

Note, there is a footgun here where MLIR CPU runner only supports async calls for gpu.alloc and gpu.memcpy. If you look into how SparseGPUCodegen is using the ops, you’ll notice the need for async.

6

See llvm-project/mlir/test/Integration/GPU/CUDA/TensorCore/wmma-matmul-f32-bare-ptr.mlir at main · llvm/llvm-project · GitHub for reference. You can also see its lowered IR at the desired level. gpu.alloc + gpu.memcpy is a standard way to allocate and use/manage GPU global memory from the host side.