‘symbol’ Dialect Passes

-encoding-to-symbol

Convert tensor encoding to bind_symbolic_shape.

-erase-symbol

Erase symbols.

Remove symbols from the function, used to clean up symbolic shape annotations when symbol analysis is not needed

-propagate-symbol

Propagate symbols through operations.

Propagate symbols bound for argument by replacing the usage of tensor.dim, to make the relationships between tensor dynamic dims more explicit.

1. bind symbolic_int to func arguments (if not already bound) ensures each dynamic dim is represented by a symbol

2. use reifyResultShapes to bind each op’s tensor result with bind_symbolic_shape, with dynamic output dims represented by tensor.dim, or affine.apply with tensor.dim operands

3. propagate symbol by replacing tensor.dim with symbolic_int. ensures that all dynamic dims are symbolic

Example:

%dim0 = tensor.dim %arg0, %c0 : tensor<?x640xf16>
%out0 = tensor.empty(%dim0) : tensor<?x640xf16>
%add0 = linalg.elemwise_binary ins(%arg0, %arg0) outs(%out0)
%concat = tensor.concat dim(0) %add0, %add0 : tensor<?x640xf16>
%dim1 = tensor.dim %concat, %c0 : tensor<?x640xf16>
%out1 = tensor.empty(%dim1) : tensor<?x640xf16>
%add1 = linalg.elemwise_binary ins(%concat, %concat) outs(%out1)

will be transformed into:

%S0 = symbol.symbolic_int @S0 : index
%S1 = symbol.symbolic_int @S1 [%S0, %S0], affine_map<()[s0, s1] -> (s0 + s1)>
symbol.bind_symbolic_shape %arg0, [%S0]

%out0 = tensor.empty(%S0) : tensor<?x640xf16>
symbol.bind_symbolic_shape %0, [%S0]
%add0 = linalg.elemwise_binary ins(%arg0, %arg0) outs(%out0)
symbol.bind_symbolic_shape %add0, [%S0]

// result dim for concat will be reified to an `affine.apply` op.
// this pass keeps the affine map in `%S1`
%concat = tensor.concat dim(0) %add0, %add0 : tensor<?x640xf16>
symbol.bind_symbolic_shape %concat, [%S1]

%out1 = tensor.empty(%S1) : tensor<?x640xf16>
symbol.bind_symbolic_shape %out1, [%S1]
%add1 = linalg.elemwise_binary ins(%concat, %concat) outs(%out1)
symbol.bind_symbolic_shape %add1, [%S1]

-symbol-to-encoding

Convert bind_symbolic_shape to tensor encoding.

-unfold-symbolic-int

Replace symbol.symbolic_int op by concrete values

Replace all symbol.symbolic_int based on bind_symbolic_shape. All symbol.symbolic_int should be replaceable by the first use in bind_symbolic_shape. This pass doesn’t do any reifying symbols. Check propagate-symbol pass to do the reifying.

Constraints:

1. symbolic_int should be replaceable by the first bind_symbolic_shape that it’s binded to.

2. first bind_symbolic_shape that a symbolic_int’s binded, should have indentity affine map e.g. symbol.bind_symbolic_shape %arg0, [%S0, %S1], affine_map<()[s0, s1] -> (s0, 640, s1)> : tensor<?x640x?xf16>

3. first bind_symbolic_shape that a symbolic_int’s binded, shouldn’t bind to a tensor.empty

Example of invalid bind_symbolic_shape ops:

symbol.bind_symbolic_shape %arg0, [%S0, %S1], affine_map<()[s0, s1] -> (s0/2, s0/s1, s1+1)> : tensor<?x640x?xf16>
%empty = tensor.empty(%S2, %S3) : tensor<?x?xf16>
symbol.bind_symbolic_shape %empty, [%S2, %S3], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf16>

Example before and after the pass:

func.func @test_already_bind_symbol_0(%arg0: tensor<?x640x?xf16>) -> tensor<?x640x?xf16> {
  %S0 = symbol.symbolic_int @S0 {max_val = 9223372036854775807 : i64, min_val = 0 : i64} : index
  %S1 = symbol.symbolic_int @S1 {max_val = 9223372036854775807 : i64, min_val = 0 : i64} : index
  symbol.bind_symbolic_shape %arg0, [%S0, %S1], affine_map<()[s0, s1] -> (s0, 640, s1)> : tensor<?x640x?xf16>
}

will be transformed into:

func.func @test_already_bind_symbol_0(%arg0: tensor<?x640x?xf16>) -> tensor<?x640x?xf16> {
  %c0 = arith.constant 0 : index
  %dim = tensor.dim %arg0, %c0 : tensor<?x640x?xf16>
  %c2 = arith.constant 2 : index
  %dim_0 = tensor.dim %arg0, %c2 : tensor<?x640x?xf16>
}