[mlir][tosa] tosa-to-tensor pass for tosa.slice fails with known output shape but dynamic size attribute

[sahas3]

Consider the input IR:

func.func @slice(%arg0 : tensor<2x60x59x5xf32>) -> tensor<?x60x58x5xf32> {
    %0 = "tosa.cast"(%arg0) : (tensor<2x60x59x5xf32>) -> tensor<?x60x59x5xf32>
    %1 = tosa.const_shape  {values = dense<0> : tensor<4xindex>} : () -> !tosa.shape<4>
    %2 = tosa.const_shape  {values = dense<[-1, 60, 58, 5]> : tensor<4xindex>} : () -> !tosa.shape<4>
    %3 = tosa.slice %0, %1, %2 : (tensor<?x60x59x5xf32>, !tosa.shape<4>, !tosa.shape<4>) -> tensor<?x60x58x5xf32>
    return %3 : tensor<?x60x58x5xf32>
}

running mlir-opt --tosa-to-tensor succeeds:

func.func @slice(%arg0: tensor<2x60x59x5xf32>) -> tensor<?x60x58x5xf32> {
    %0 = tosa.cast %arg0 : (tensor<2x60x59x5xf32>) -> tensor<?x60x59x5xf32>
    %c0 = arith.constant 0 : index
    %dim = tensor.dim %0, %c0 : tensor<?x60x59x5xf32>
    %c0_0 = arith.constant 0 : index
    %1 = arith.subi %dim, %c0_0 : index
    %extracted_slice = tensor.extract_slice %0[0, 0, 0, 0] [%1, 60, 58, 5] [1, 1, 1, 1] : tensor<?x60x59x5xf32> to tensor<?x60x58x5xf32>
    return %extracted_slice : tensor<?x60x58x5xf32>
  }

but running mlir-opt --tosa-infer-shapes --tosa-to-tensor fails:

// -----// IR Dump After TosaInferShapesPass (tosa-infer-shapes) //----- //
func.func @slice(%arg0: tensor<2x60x59x5xf32>) -> tensor<?x60x58x5xf32> {
  %0 = tosa.cast %arg0 : (tensor<2x60x59x5xf32>) -> tensor<2x60x59x5xf32>
  %1 = tosa.const_shape  {values = dense<0> : tensor<4xindex>} : () -> !tosa.shape<4>
  %2 = tosa.const_shape  {values = dense<[-1, 60, 58, 5]> : tensor<4xindex>} : () -> !tosa.shape<4>
  %3 = tosa.slice %0, %1, %2 : (tensor<2x60x59x5xf32>, !tosa.shape<4>, !tosa.shape<4>) -> tensor<2x60x58x5xf32>
  %cast = tensor.cast %3 : tensor<2x60x58x5xf32> to tensor<?x60x58x5xf32>
  return %cast : tensor<?x60x58x5xf32>
}

error: expected type to be 'tensor<?x60x58x5xf32>' or a rank-reduced version. (size mismatch) 
    %3 = tosa.slice %0, %1, %2 : (tensor<?x60x59x5xf32>, !tosa.shape<4>, !tosa.shape<4>) -> tensor<?x60x58x5xf32>
         ^
note: see current operation: %4 = "tensor.extract_slice"(%arg0, %3) <{operandSegmentSizes = array<i32: 1, 0, 1, 0>, static_offsets = array<i64: 0, 0, 0, 0>, static_sizes = array<i64: -9223372036854775808, 60, 58, 5>, static_strides = array<i64: 1, 1, 1, 1>}> : (tensor<2x60x59x5xf32>, index) -> tensor<2x60x58x5xf32>
// -----// IR Dump After TosaToTensorPass Failed (tosa-to-tensor) //----- //
"builtin.module"() ({
  "func.func"() <{function_type = (tensor<2x60x59x5xf32>) -> tensor<?x60x58x5xf32>, sym_name = "slice"}> ({
  ^bb0(%arg0: tensor<2x60x59x5xf32>):
    %0 = "arith.constant"() <{value = 0 : index}> : () -> index
    %1 = "tensor.dim"(%arg0, %0) : (tensor<2x60x59x5xf32>, index) -> index
    %2 = "arith.constant"() <{value = 0 : index}> : () -> index
    %3 = "arith.subi"(%1, %2) <{overflowFlags = #arith.overflow<none>}> : (index, index) -> index
    %4 = "tensor.extract_slice"(%arg0, %3) <{operandSegmentSizes = array<i32: 1, 0, 1, 0>, static_offsets = array<i64: 0, 0, 0, 0>, static_sizes = array<i64: -9223372036854775808, 60, 58, 5>, static_strides = array<i64: 1, 1, 1, 1>}> : (tensor<2x60x59x5xf32>, index) -> tensor<2x60x58x5xf32>
    %5 = "tensor.cast"(%4) : (tensor<2x60x58x5xf32>) -> tensor<?x60x58x5xf32>
    "func.return"(%5) : (tensor<?x60x58x5xf32>) -> ()
  }) : () -> ()
}) : () -> ()

I think --tosa-infer-shapes does the correct thing by updating the dynamic dim value in the input and output to 2 but the size attribute of tosa.slice op isn't updated which is correct too as tosa-infer-shapes only updates IO shapes as per my understanding.

I think there are a couple of fixes:

1. Enhance tosa-to-tensor to not consider the dynamic size attribute if output shape is known to be static when creating the tensor.extract_slice op.
2. Alternatively, as part of SliceOp canonicalization, update the size attribute to not have -1 when output shape is known.

Any suggestions @sjarus , @eric-k256, @Tai78641 ?

Thanks!

[llvmbot]

@llvm/issue-subscribers-mlir

Author: Sayan Saha (sahas3)

Consider the input IR:

func.func @<!-- -->slice(%arg0 : tensor&lt;2x60x59x5xf32&gt;) -&gt; tensor&lt;?x60x58x5xf32&gt; {
    %0 = "tosa.cast"(%arg0) : (tensor&lt;2x60x59x5xf32&gt;) -&gt; tensor&lt;?x60x59x5xf32&gt;
    %1 = tosa.const_shape  {values = dense&lt;0&gt; : tensor&lt;4xindex&gt;} : () -&gt; !tosa.shape&lt;4&gt;
    %2 = tosa.const_shape  {values = dense&lt;[-1, 60, 58, 5]&gt; : tensor&lt;4xindex&gt;} : () -&gt; !tosa.shape&lt;4&gt;
    %3 = tosa.slice %0, %1, %2 : (tensor&lt;?x60x59x5xf32&gt;, !tosa.shape&lt;4&gt;, !tosa.shape&lt;4&gt;) -&gt; tensor&lt;?x60x58x5xf32&gt;
    return %3 : tensor&lt;?x60x58x5xf32&gt;
}

running mlir-opt --tosa-to-tensor succeeds:

func.func @<!-- -->slice(%arg0: tensor&lt;2x60x59x5xf32&gt;) -&gt; tensor&lt;?x60x58x5xf32&gt; {
    %0 = tosa.cast %arg0 : (tensor&lt;2x60x59x5xf32&gt;) -&gt; tensor&lt;?x60x59x5xf32&gt;
    %c0 = arith.constant 0 : index
    %dim = tensor.dim %0, %c0 : tensor&lt;?x60x59x5xf32&gt;
    %c0_0 = arith.constant 0 : index
    %1 = arith.subi %dim, %c0_0 : index
    %extracted_slice = tensor.extract_slice %0[0, 0, 0, 0] [%1, 60, 58, 5] [1, 1, 1, 1] : tensor&lt;?x60x59x5xf32&gt; to tensor&lt;?x60x58x5xf32&gt;
    return %extracted_slice : tensor&lt;?x60x58x5xf32&gt;
  }

but running mlir-opt --tosa-infer-shapes --tosa-to-tensor fails:

// -----// IR Dump After TosaInferShapesPass (tosa-infer-shapes) //----- //
func.func @<!-- -->slice(%arg0: tensor&lt;2x60x59x5xf32&gt;) -&gt; tensor&lt;?x60x58x5xf32&gt; {
  %0 = tosa.cast %arg0 : (tensor&lt;2x60x59x5xf32&gt;) -&gt; tensor&lt;2x60x59x5xf32&gt;
  %1 = tosa.const_shape  {values = dense&lt;0&gt; : tensor&lt;4xindex&gt;} : () -&gt; !tosa.shape&lt;4&gt;
  %2 = tosa.const_shape  {values = dense&lt;[-1, 60, 58, 5]&gt; : tensor&lt;4xindex&gt;} : () -&gt; !tosa.shape&lt;4&gt;
  %3 = tosa.slice %0, %1, %2 : (tensor&lt;2x60x59x5xf32&gt;, !tosa.shape&lt;4&gt;, !tosa.shape&lt;4&gt;) -&gt; tensor&lt;2x60x58x5xf32&gt;
  %cast = tensor.cast %3 : tensor&lt;2x60x58x5xf32&gt; to tensor&lt;?x60x58x5xf32&gt;
  return %cast : tensor&lt;?x60x58x5xf32&gt;
}

error: expected type to be 'tensor&lt;?x60x58x5xf32&gt;' or a rank-reduced version. (size mismatch) 
    %3 = tosa.slice %0, %1, %2 : (tensor&lt;?x60x59x5xf32&gt;, !tosa.shape&lt;4&gt;, !tosa.shape&lt;4&gt;) -&gt; tensor&lt;?x60x58x5xf32&gt;
         ^
note: see current operation: %4 = "tensor.extract_slice"(%arg0, %3) &lt;{operandSegmentSizes = array&lt;i32: 1, 0, 1, 0&gt;, static_offsets = array&lt;i64: 0, 0, 0, 0&gt;, static_sizes = array&lt;i64: -9223372036854775808, 60, 58, 5&gt;, static_strides = array&lt;i64: 1, 1, 1, 1&gt;}&gt; : (tensor&lt;2x60x59x5xf32&gt;, index) -&gt; tensor&lt;2x60x58x5xf32&gt;
// -----// IR Dump After TosaToTensorPass Failed (tosa-to-tensor) //----- //
"builtin.module"() ({
  "func.func"() &lt;{function_type = (tensor&lt;2x60x59x5xf32&gt;) -&gt; tensor&lt;?x60x58x5xf32&gt;, sym_name = "slice"}&gt; ({
  ^bb0(%arg0: tensor&lt;2x60x59x5xf32&gt;):
    %0 = "arith.constant"() &lt;{value = 0 : index}&gt; : () -&gt; index
    %1 = "tensor.dim"(%arg0, %0) : (tensor&lt;2x60x59x5xf32&gt;, index) -&gt; index
    %2 = "arith.constant"() &lt;{value = 0 : index}&gt; : () -&gt; index
    %3 = "arith.subi"(%1, %2) &lt;{overflowFlags = #arith.overflow&lt;none&gt;}&gt; : (index, index) -&gt; index
    %4 = "tensor.extract_slice"(%arg0, %3) &lt;{operandSegmentSizes = array&lt;i32: 1, 0, 1, 0&gt;, static_offsets = array&lt;i64: 0, 0, 0, 0&gt;, static_sizes = array&lt;i64: -9223372036854775808, 60, 58, 5&gt;, static_strides = array&lt;i64: 1, 1, 1, 1&gt;}&gt; : (tensor&lt;2x60x59x5xf32&gt;, index) -&gt; tensor&lt;2x60x58x5xf32&gt;
    %5 = "tensor.cast"(%4) : (tensor&lt;2x60x58x5xf32&gt;) -&gt; tensor&lt;?x60x58x5xf32&gt;
    "func.return"(%5) : (tensor&lt;?x60x58x5xf32&gt;) -&gt; ()
  }) : () -&gt; ()
}) : () -&gt; ()

I think --tosa-infer-shapes does the correct thing by updating the dynamic dim value in the input and output to 2 but the size attribute of tosa.slice op isn't updated which is correct too as tosa-infer-shapes only updates IO shapes as per my understanding.

I think there are a couple of fixes:

1. Enhance tosa-to-tensor to not consider the dynamic size attribute if output shape is known to be static when creating the tensor.extract_slice op.
2. Alternatively, as part of SliceOp canonicalization, update the size attribute to not have -1 when output shape is known.

Any suggestions @sjarus , @eric-k256, @Tai78641 ?

Thanks!

[Tai78641]

I think #2 is more general. I imagine this would be added as a new SliceOp canonicalization pattern to:

void SliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                          MLIRContext *context) {
  results.add<ConcatSliceOptimization>(context);
}

[sahas3]

I'll start a PR with approach 2, thanks!

[CoTinker]

Addressed by #135560.