[mlir] Refactor the representation of function-like argument/result attributes.
authorRiver Riddle <riddleriver@gmail.com>
Sat, 8 May 2021 02:30:25 +0000 (19:30 -0700)
committerRiver Riddle <riddleriver@gmail.com>
Sat, 8 May 2021 02:32:31 +0000 (19:32 -0700)
commit53b946aa636a31e9243b8c5bf1703a1f6eae798e
tree2bb090f4a31f3fe4e9936b6812a1bc2125ee8dc4
parent44d14d5de6f1f74246bf66dea5538ddc304e6445
[mlir] Refactor the representation of function-like argument/result attributes.

The current design uses a unique entry for each argument/result attribute, with the name of the entry being something like "arg0". This provides for a somewhat sparse design, but ends up being much more expensive (from a runtime perspective) in-practice. The design requires building a string every time we lookup the dictionary for a specific arg/result, and also requires N attribute lookups when collecting all of the arg/result attribute dictionaries.

This revision restructures the design to instead have an ArrayAttr that contains all of the attribute dictionaries for arguments and another for results. This design reduces the number of attribute name lookups to 1, and allows for O(1) lookup for individual element dictionaries. The major downside is that we can end up with larger memory usage, as the ArrayAttr contains an entry for each element even if that element has no attributes. If the memory usage becomes too problematic, we can experiment with a more sparse structure that still provides a lot of the wins in this revision.

This dropped the compilation time of a somewhat large TensorFlow model from ~650 seconds to ~400 seconds.

Differential Revision: https://reviews.llvm.org/D102035
19 files changed:
mlir/include/mlir/Dialect/GPU/GPUOps.td
mlir/include/mlir/IR/BuiltinAttributes.td
mlir/include/mlir/IR/FunctionImplementation.h
mlir/include/mlir/IR/FunctionSupport.h
mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp
mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
mlir/lib/IR/BuiltinDialect.cpp
mlir/lib/IR/FunctionImplementation.cpp
mlir/lib/IR/FunctionSupport.cpp
mlir/lib/Transforms/Utils/DialectConversion.cpp
mlir/test/Dialect/LLVMIR/func.mlir
mlir/test/IR/invalid-func-op.mlir
mlir/test/IR/test-func-set-type.mlir