[mlir][sparse] Start migration to new surface syntax for STEA
authorAart Bik <ajcbik@google.com>
Wed, 28 Jun 2023 19:54:12 +0000 (12:54 -0700)
committerAart Bik <ajcbik@google.com>
Thu, 29 Jun 2023 18:32:07 +0000 (11:32 -0700)
commit6b88c852b649a4eb3e0eb01efea88f5bac360f1a
tree955859604d821a002d287a2085a78b4619365e35
parente9a9fccc9dde1387b6a8f3d4c83c29d2bd7136bd
[mlir][sparse] Start migration to new surface syntax for STEA

We are in the progress of migrating to a much improved surface syntax for the Sparse Tensor Encoding Attribute (STEA).

You can see a preview of this in the StableHLO RFC at

 https://github.com/openxla/stablehlo/blob/main/rfcs/20230210-sparsity.md

//**This design is courtesy Wren Romano.**//

This initial revision
(1) Introduces the first version of a new parser written by Wren Romano
(2) Introduces a simple "migration plan" using NEW_SYNTAX on the STEA, which will allow us to test the new parser with new examples, as well as migrate existing examples over without the need to rewrite them all

This first "drop" merely provides the entry points to parse the new syntax. The parser is still under active development. For example, we need to address the "lookahead" issue when parsing the lvl spec (viz. do we see l0 = d0 or a direct d0). Another larger task is to actually implement "affine" parsing (since the MLIR affine parser is not accessible in other parts of the tree).

EXAMPLE:

Currently, CSR looks like

  #CSR = #sparse_tensor.encoding<{
    lvlTypes = ["dense","compressed"],
    dimToLvl = affine_map<(i,j) -> (i,j)>
  }>

but you can "force" the new parser with

  #CSR = #sparse_tensor.encoding<{
    NEW_SYNTAX =
    (d0, d1) -> (l0 = d0 : dense, l1 = d1 : compressed)
  }>

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D153997
13 files changed:
mlir/lib/Dialect/SparseTensor/IR/CMakeLists.txt
mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp [new file with mode: 0644]
mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.h [new file with mode: 0644]
mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMapParser.cpp [new file with mode: 0644]
mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMapParser.h [new file with mode: 0644]
mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.cpp [new file with mode: 0644]
mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.h [new file with mode: 0644]
mlir/lib/Dialect/SparseTensor/IR/Detail/TemplateExtras.h [new file with mode: 0644]
mlir/lib/Dialect/SparseTensor/IR/Detail/Var.cpp [new file with mode: 0644]
mlir/lib/Dialect/SparseTensor/IR/Detail/Var.h [new file with mode: 0644]
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel