[mlir][spirv] Convert tensor.extract for very small tensors
authorLei Zhang <antiagainst@google.com>
Sat, 6 Mar 2021 13:03:27 +0000 (08:03 -0500)
committerLei Zhang <antiagainst@google.com>
Sat, 6 Mar 2021 13:03:36 +0000 (08:03 -0500)
commitbb6f5c8314799a6553829b724c649e825a558caf
tree27c50babc85df913eb4cba97a924279b604bee26
parentcf7cdaff64fb2a97f9af533ad0a5ed9c94b72c8e
[mlir][spirv] Convert tensor.extract for very small tensors

Normally tensors will be stored in buffers before converting to SPIR-V,
given that is how a large amount of data is sent to the GPU. However,
SPIR-V supports converting from tensors directly too. This is for the
cases where the tensor just contains a small amount of elements and it
makes sense to directly inline them as a small data array in the shader.
To handle this, internally the conversion might create new local
variables. SPIR-V consumers in GPU drivers may or may not optimize that
away. So this has implications over register pressure. Therefore, a
threshold is used to control when the patterns should kick in.

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D98052
mlir/include/mlir/Conversion/StandardToSPIRV/StandardToSPIRV.h
mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h
mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp
mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRVPass.cpp
mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
mlir/test/Conversion/StandardToSPIRV/std-ops-to-spirv.mlir