[mlir][GPU] Add known_block_size and known_grid_size to gpu.func
authorKrzysztof Drewniak <Krzysztof.Drewniak@amd.com>
Fri, 2 Dec 2022 20:38:39 +0000 (20:38 +0000)
committerKrzysztof Drewniak <Krzysztof.Drewniak@amd.com>
Thu, 22 Dec 2022 21:41:46 +0000 (21:41 +0000)
commit85e38d7cd670371206f6067772dc822049d2cbd8
treea22a5ea7a1837b8c28e5815d6a314442096cef28
parent520a63893e2188bec6dbd6c3c6992661c71f32e1
[mlir][GPU] Add known_block_size and known_grid_size to gpu.func

In many cases, the the number of workgroups (the grid size) and the
number of workitems within each group (the block size) that a GPU
kernel will be launched with are known. For example, if gpu.launch is
called with constant block and grid sizes, we know that those are the
only possible sizes that will be used to launch that kernel. In other
cases, a custom code-generation pipeline that eventually produces GPU
kernels may know the launch dimensions of those kernels, or at least
may be able to provide an upper bound on them.

Other GPU programming systems, such as OpenCL, allow capturing such
information to enable compiler optimizations - see
reqd_work_group_size, but MLIR currently has no mechanism for doing so.

This set of attributes is the first step in enabling optimizations
based on the known launch dimensions of kernels. It extends the kernel
outline pass to set these bounds on kernels with constant launch
dimensions and extends integer range inference for GPU index
operations to account for the bounds when they are known.

Subsequent revisions will use this data when lowering GPU operations
to the ROCDL dialect.

Reviewed By: antiagainst

Differential Revision: https://reviews.llvm.org/D139865
mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
mlir/lib/Dialect/GPU/IR/InferIntRangeInterfaceImpls.cpp
mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
mlir/test/Dialect/GPU/int-range-interface.mlir
mlir/test/Dialect/GPU/invalid.mlir
mlir/test/Dialect/GPU/outlining.mlir