Skip to content

Add a pass that adds shared mem space attribute to allocs in a GPU kernel #393

@dchigarev

Description

@dchigarev

[SUB-TASK OF #360]

We need a pass that adds shared mem space to memref being allocated in a GPU kernel (thus marking those as allocs for SLM memory).

We may also add a check, that the total amount of such allocations does not exceed the total amount of available shared memory for this specific device.

Example of how this pass is expected to work:

// INPUT:
func.func @entry() {
  %0 = memref.alloc() : memref<16x16xf16>
  gpu.launch blocks(%arg5, %arg6, %arg7) in (%arg11 = %0, %arg12 = %1, %arg13 = %c1_0) threads(%arg8, %arg9, %arg10) in (%arg14 = %2, %arg15 = %3, %arg16 = %c1_0) {
    %1 = memref.alloc() : memref<16x16xf16>
    memref.dealloc %1 : memref<16x16xf16>
    gpu.terminator
  } {SCFToGPU_visited}
  return
}

// OUTPUT:
func.func @entry() {
  %0 = memref.alloc() : memref<16x16xf16> // got nothing
  gpu.launch blocks(%arg5, %arg6, %arg7) in (%arg11 = %0, %arg12 = %1, %arg13 = %c1_0) threads(%arg8, %arg9, %arg10) in (%arg14 = %2, %arg15 = %3, %arg16 = %c1_0) {
    %1 = memref.alloc() : memref<16x16xf16, 3> // got 'shared' mem space attribute
    memref.dealloc %1 : memref<16x16xf16, 3>
    gpu.terminator
  } {SCFToGPU_visited}
  return
}

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions