commit 7361340b62ecc37a685a8c6bc4b0fec58ddbf3ba
author Ben Vanik <ben.vanik@gmail.com> Tue Jul 30 09:39:57 2024 -0700
committer GitHub <noreply@github.com> Tue Jul 30 16:39:57 2024 +0000
tree a200ee053ffc253dc968f192554f41d70799e787
parent 370c1994b9d9d27af2f742239e1bb43e7a729825

Adding `hal.device.memoize` op. (#17938)

This is an inline nested region that returns HAL resources for the given
device and affinity pair. The region may only implicitly capture
constants and global values in order for it to be outlined for
initialization, though in the future we may extend this to support
dynamic values by adding an LRU cache.

To start all memoize regions are inlined after conversion. This optional
inlining will also serve as a useful debug feature for isolating
memoization related issues. The initial outlining and proper memoization
pass is present but disabled by default until we can verify all targets
properly support reusable command buffers.

Progress on #17875. Nothing is generating these yet.

Example:
```mlir
util.global private @device = #hal.device.target<"local"> : !hal.device
util.global private @executable : !hal.executable
util.global private @buffer : !hal.buffer
util.func public @memoize_command_buffer() -> !hal.command_buffer {
  %device = util.global.load immutable @device : !hal.device
  %affinity = arith.constant -1 : i64
  %executable = util.global.load immutable @executable : !hal.executable
  %buffer = util.global.load immutable @buffer : !hal.buffer
  %result = hal.device.memoize<%device : !hal.device> affinity(%affinity) -> !hal.command_buffer {
    %cmd = hal.command_buffer.create device(%device : !hal.device) mode(None) categories("Transfer|Dispatch") : !hal.command_buffer
    %dispatch_ordinal = arith.constant 123 : index
    %offset = arith.constant 100 : index
    hal.command_buffer.dispatch.indirect<%cmd : !hal.command_buffer>
        target(%executable : !hal.executable)[%dispatch_ordinal]
        workgroups(%buffer : !hal.buffer)[%offset]
    hal.command_buffer.execution_barrier<%cmd : !hal.command_buffer>
        source(CommandIssue) target(CommandProcess) flags(None)
    hal.command_buffer.finalize<%cmd : !hal.command_buffer>
    hal.return %cmd : !hal.command_buffer
  }
  util.return %result : !hal.command_buffer
}
```
->
```mlir
util.global private @device = #hal.device.target<"local"> : !hal.device
util.global private @executable : !hal.executable
util.global private @global_resource : !hal.buffer
util.func private @__memoize_command_buffer_memoize_apply(%arg0: !hal.device, %arg1: i64, %arg2: !hal.executable, %arg3: !hal.buffer) -> !hal.command_buffer {
  %cmd = hal.command_buffer.create device(%arg0 : !hal.device) mode(OneShot) categories("Transfer|Dispatch") : !hal.command_buffer
  %c123 = arith.constant 123 : index
  %c100 = arith.constant 100 : index
  hal.command_buffer.dispatch.indirect<%cmd : !hal.command_buffer> target(%arg2 : !hal.executable)[%c123] workgroups(%arg3 : !hal.buffer)[%c100]
  hal.command_buffer.execution_barrier<%cmd : !hal.command_buffer> source(CommandIssue) target(CommandProcess) flags("None")
  hal.command_buffer.finalize<%cmd : !hal.command_buffer>
  util.return %cmd : !hal.command_buffer
}
util.global private @__memoize_command_buffer_memoize_result_0_device : !hal.command_buffer
util.initializer {
  %device = util.global.load immutable @device : !hal.device
  %c-1_i64 = arith.constant -1 : i64
  %executable = util.global.load immutable @executable : !hal.executable
  %global_resource = util.global.load immutable @global_resource : !hal.buffer
  %0 = util.call @__memoize_command_buffer_memoize_apply(%device, %c-1_i64, %executable, %global_resource) : (!hal.device, i64, !hal.executable, !hal.buffer) -> !hal.command_buffer
  util.global.store %0, @__memoize_command_buffer_memoize_result_0_device : !hal.command_buffer
  util.return
}
util.func private @__memoize_command_buffer_memoize_lookup(%arg0: !hal.device, %arg1: i64) -> !hal.command_buffer {
  %device = util.global.load @device : !hal.device
  %0 = util.cmp.eq %arg0, %device : !hal.device
  %1 = scf.if %0 -> (!hal.command_buffer) {
    %__memoize_command_buffer_memoize_result_0_device = util.global.load @__memoize_command_buffer_memoize_result_0_device : !hal.command_buffer
    scf.yield %__memoize_command_buffer_memoize_result_0_device : !hal.command_buffer
  } else {
    %2 = util.null : !hal.command_buffer
    scf.yield %2 : !hal.command_buffer
  }
  util.return %1 : !hal.command_buffer
}
util.func public @memoize_command_buffer() -> !hal.command_buffer {
  %device = util.global.load immutable @device : !hal.device
  %c-1_i64 = arith.constant -1 : i64
  %executable = util.global.load immutable @executable : !hal.executable
  %global_resource = util.global.load immutable @global_resource : !hal.buffer
  %0 = util.call @__memoize_command_buffer_memoize_lookup(%device, %c-1_i64) : (!hal.device, i64) -> !hal.command_buffer
  util.return %0 : !hal.command_buffer
}
```