MLIR Sparse Compiler Progress

MLIR
19

The Google Sparse compiler team is in the process of improving the surface syntax of the Sparse Tensor Encoding Attribute (STEA), following an elegant design by Wren Romano (@wrengr). For example, currently, CSC looks as follows (which always had the ambiguity whether lvlTypes applied before or after the permutation (it is after):

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

But with the new syntax, that ambiguity is gone.

#CSC = #sparse_tensor.encoding<{
  (d0, d1) -> (d1 : dense, d0 : compressed)
}>

In addition, the new surface syntax also provides a very elegant way to define higher order mappings, needed to define, for example, block sparsity. Here, the syntax allows for just defining one mapping from dimensions to levels (implicit) or a form that also specifies the mapping from levels back to dimensions (explicit).

// Block sparsity for 2x3 blocks

#BCSR_implicit = #sparse_tensor.encoding<{
  ( i, j ) ->
  ( i floordiv 2 : compressed,
    j floordiv 3 : compressed,
    i mod 2      : dense,
    j mod 3      : dense
  )
}>

#BCSR_explicit = #sparse_tensor.encoding<{
  {il, jl, ii, jj}
  ( i = il * 2 + ii,
    j = jl * 3 + jj
  ) ->
  ( il = i floordiv 2 : compressed,
    jl = j floordiv 3 : compressed,
    ii = i mod 2      : dense,
    jj = j mod 3      : dense
  )
}>

This new surface syntax will also be proposed for StableHLO (see this StableHLO+Sparsity RFC and contact @burmako if you are interested in seeing this done soon).

7 Likes
Stable HLO Sparsity RFC