Skip to content

[mlir][GPU] Add RecursiveMemoryEffects to gpu.launch #75315

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Dec 20, 2023
Merged

[mlir][GPU] Add RecursiveMemoryEffects to gpu.launch #75315

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -672,7 +672,8 @@ def GPU_LaunchFuncOp :GPU_Op<"launch_func", [

def GPU_LaunchOp : GPU_Op<"launch", [
AutomaticAllocationScope, AttrSizedOperandSegments, GPU_AsyncOpInterface,
DeclareOpInterfaceMethods<InferIntRangeInterface>]>,
DeclareOpInterfaceMethods<InferIntRangeInterface>,
RecursiveMemoryEffects]>,
Arguments<(ins Variadic<GPU_AsyncToken>:$asyncDependencies,
Index:$gridSizeX, Index:$gridSizeY, Index:$gridSizeZ,
Index:$blockSizeX, Index:$blockSizeY, Index:$blockSizeZ,
Expand Down
33 changes: 33 additions & 0 deletions mlir/test/Dialect/GPU/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ func.func @fold_wait_op_test1() {
}
// CHECK-NOT: gpu.wait

// -----

// Erase duplicate barriers.
// CHECK-LABEL: func @erase_barriers
// CHECK-NEXT: gpu.barrier
Expand All @@ -21,6 +23,8 @@ func.func @erase_barriers() {
return
}

// -----

// Replace uses of gpu.wait op with its async dependency.
// CHECK-LABEL: func @fold_wait_op_test2
func.func @fold_wait_op_test2(%arg0: i1) -> (memref<5xf16>, memref<5xf16>) {
Expand All @@ -38,6 +42,8 @@ func.func @fold_wait_op_test2(%arg0: i1) -> (memref<5xf16>, memref<5xf16>) {
// CHECK-NEXT: gpu.alloc async [%[[TOKEN1]]] ()
// CHECK-NEXT: return

// -----

// CHECK-LABEL: func @fold_memcpy_op
func.func @fold_memcpy_op(%arg0: i1) {
%cst = arith.constant 0.000000e+00 : f16
Expand All @@ -60,6 +66,8 @@ func.func @fold_memcpy_op(%arg0: i1) {
}
// CHECK-NOT: gpu.memcpy

// -----

// We cannot fold memcpy here as dest is a block argument.
// CHECK-LABEL: func @do_not_fold_memcpy_op1
func.func @do_not_fold_memcpy_op1(%arg0: i1, %arg1: memref<2xf16>) {
Expand All @@ -75,6 +83,8 @@ func.func @do_not_fold_memcpy_op1(%arg0: i1, %arg1: memref<2xf16>) {
}
// CHECK: gpu.memcpy

// -----

// We cannot fold gpu.memcpy as it is used by an op having read effect on dest.
// CHECK-LABEL: func @do_not_fold_memcpy_op2
func.func @do_not_fold_memcpy_op2(%arg0: i1, %arg1: index) -> f16 {
Expand All @@ -92,6 +102,8 @@ func.func @do_not_fold_memcpy_op2(%arg0: i1, %arg1: index) -> f16 {
}
// CHECK: gpu.memcpy

// -----

// We cannot fold gpu.memcpy, as the defining op if dest is not a alloc like op.
// CHECK-LABEL: func @do_not_fold_memcpy_op3
func.func @do_not_fold_memcpy_op3(%arg0: memref<1xi8>, %arg1: memref<i1>) {
Expand All @@ -102,6 +114,8 @@ func.func @do_not_fold_memcpy_op3(%arg0: memref<1xi8>, %arg1: memref<i1>) {
}
// CHECK: gpu.memcpy

// -----

// CHECK-LABEL: @memcpy_after_cast
func.func @memcpy_after_cast(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
// CHECK-NOT: memref.cast
Expand All @@ -112,6 +126,8 @@ func.func @memcpy_after_cast(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
return
}

// -----

// CHECK-LABEL: @memset_after_cast
func.func @memset_after_cast(%arg0: memref<10xf32>, %arg1: f32) {
// CHECK-NOT: memref.cast
Expand Down Expand Up @@ -227,3 +243,20 @@ func.func @make_subgroup_reduce_uniform() {
}
return
}

// -----

// The GPU kernel does not have any side effecting ops, so the entire
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not very familiar with RecursiveMemoryEffects, what exactly it does?

Is it safe with managed memory or unified memory? In these cases, a pointer address between gpu and cpu remains the same. Once can imagine there is a llvm.ptr and it can used to build memref descriptor. It can be used inside gpu.launch body, or something else, is it safe in these cases?

Copy link
Member Author

@matthias-springer matthias-springer Dec 15, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

RecursiveMemoryEffects indicates that the side effects of this op are defined by the side effects of nested ops.

In essence: gpu.launch by itself does not have any side effects. E.g., it does not read/write to memory etc. But the GPU kernel may have ops that read/write to memory. When we use the C++ API to query the side effects of a gpu.launch, it will automatically check the side effects of the ops inside the kernel and return those.

If the ops in the kernel have no side effects, then the gpu.launch has no side effects. That's why the gpu.launch in the newly added test case folds away. It does not read/write to memory etc. It can just be removed.

Once can imagine there is a llvm.ptr and it can used to build memref descriptor. It can be used inside gpu.launch body, or something else, is it safe in these cases?

When you use that memref value, e.g., with a memref.load, you have an op with a side effect and the entire kernel will be considered to have that side effect.

Copy link
Member

@grypp grypp Dec 17, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for your patience and explanation.

gpu.launch doesn't write to memory directly, it performs the following actions. It'll still need to read kernel configurations even wtih empty body. However, I think it's safe to fold.

  1. Reads host memory for kernel configuration and parameters.
  2. Copies kernel parameters to the device memory.
  3. Waits for stream(s).

Thinking about the 3rd option with IR below. If taskC (that has empty body) is folded , could taskD run before taskB due to not waiting for stream2 anymore? (Currently, MLIR utilizes single stream, but one can imagine implementing support for multiple streams .)

%taskA = gpu.launch async [%stream1] ...
%taskB = gpu.launch async [%stream2] ...
%taskC = gpu.launch async [%stream1, %stream2] ...
%taskD = gpu.launch async [%stream1] ...

task execution graph:

a   b
 \ /
  c
  |
  d

Is it overkill to think about this case?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The kernel configuration and parameters (e.g., number of blocks, etc.) are passed to gpu.launch as integer operands, so the op does not read anything from memory (memref) at this point, right? (The op may be lowered to something that reads from memory; that op would then have a memory side effect.)

gpu.launch takes async token arguments and produces an async token result. If there is a particular order in which kernels must be launched (and waited for), I would expect this to be modeled with async tokens. IMO, execution order due to "stream is available or not" does not have to be preserved when folding gpu.launch ops.

%taskA = gpu.launch async [] ...
%taskB = gpu.launch async [] ...
%taskC = gpu.launch async [%taskA, %taskB] ...
%taskD = gpu.launch async [%taskC] ...

In the above example, %taskC can canonicalized away if it has no side effects, but then %taskD will depend on %taskA and %taskB.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We discussed offline. The way CUDA streams are expressed in the GPU dialect is via async tokens. It would not make sense for gpu.launch and other ops to take an additional %stream operand in the future. We use async tokens (%task) for that.

Coming back to RecursiveMemoryEffects, this trait should be safe to attach. There are no side effects wrt. streams, because streams are encoded with async tokens.

Empty GPU kernels like the one in the test case will not simply DCE away if they impose a dependency via async tokens. (A folder or canonicalization pattern may find a way to remove empty kernels in a way that preserves the dependencies expressed by the async tokens.)

// gpu.launch op can fold away.

// CHECK-LABEL: func @gpu_launch_without_side_effects
// CHECK-NOT: gpu.launch
func.func @gpu_launch_without_side_effects() {
%0:6 = "test.test1"() : () -> (index, index, index, index, index, index)
gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %0#0, %arg7 = %0#1, %arg8 = %0#2)
threads(%arg3, %arg4, %arg5) in (%arg9 = %0#3, %arg10 = %0#4, %arg11 = %0#5) {
%1 = arith.addi %arg0, %arg1 : index
gpu.terminator
}
return
}