| // Source module with hal.executable.source (extern kernel dispatch). |
| // Used by iree-link-executable-source.mlir. |
| |
| module @module_d { |
| // Extern kernel executable (e.g., a pre-compiled GPU kernel object). |
| hal.executable.source private @extern_kernel attributes { |
| objects = #hal.executable.objects<{}> |
| } { |
| hal.executable.export public @entry ordinal(0) |
| layout(#hal.pipeline.layout<constants = 2, bindings = [ |
| #hal.pipeline.binding<storage_buffer, ReadOnly>, |
| #hal.pipeline.binding<storage_buffer> |
| ]>) count(%device: !hal.device) -> (index, index, index) { |
| %c1 = arith.constant 1 : index |
| hal.return %c1, %c1, %c1 : index, index, index |
| } attributes {workgroup_size = [64 : index, 1 : index, 1 : index]} |
| } |
| |
| // Function that dispatches to the extern kernel. |
| util.func public @transform(%arg0: tensor<4xf32>) -> tensor<4xf32> { |
| %c4 = arith.constant 4 : i32 |
| %c1 = arith.constant 1 : index |
| %0 = flow.dispatch @extern_kernel::@entry(%c4, %arg0) |
| : (i32, tensor<4xf32>) -> tensor<4xf32> |
| util.return %0 : tensor<4xf32> |
| } |
| } |