[MLIR][Tosa] Pass encoding through `tosa-to-linalg`

[rikhuijzer]

Summary

As pointed out by @Sinclair-Dee in
https://github.com/llvm/llvm-project/issues/62304, the tosa-to-linalg
conversion ignored the encoding attribute.

Also, this patch avoids an assertion error crash on unranked tensors.
Instead, the conversion now throws a "failed to legalize" error.

Fixes #62304 and fixes #63165.

[mehdi_amini]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

561–564

If you dyn_cast that means we expect that this conversion can fail, so it should be checked. Otherwise if we know by construction that this can't fail we should cast.

[rikhuijzer]

Add assertion after dyn_cast

In response to the comment by Mehdi, I added an assertion since I'm
unable to verify that the cast will always succeed. It looks to me
like the only situation in which the cast would fail is in case of an
UnrankedTensorType and I haven't been able to produce a test which
would trigger that. However I'm definitely no MLIR expert, so a kind
request to the reviewers to double check this.

[mehdi_amini]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

562

cast embeds an assertion already: the use of dyn_cast is intended for gracefully handling the result.

[rkayaith]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

561–564

i don't think you need a cast here anyways

[rikhuijzer]

Make assumptions more explicit

Jacques Pienaar pointed out on Discord that the code being ranked was already
assumed at an earlier point in the code. This patch makes that and some other
assumptions more explicit.

[rikhuijzer]

To be clear to the reviewer, note that the line

unsigned rank = resultTy.getRank();

would crash on an unranked tensor anyway, so the introduction of the explicit failure on !rankedType does not change program behavior (before this patch, the program crashed on unranked tensors; after this patch, the program throws an error as is tested in the added test).

[rikhuijzer]

Avoid duplicate cast

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

for (int i = 0; i < operandTy.getRank(); i++) {
       if (operandTy.isDynamicDim(i) && !dynDims[i])
         dynDims[i] = rewriter.create<tensor::DimOp>(loc, arg, i);
     }
   }
 
   SmallVector<Value> filteredDims = condenseValues(dynDims);
 
   for (auto result : results) {
-    auto resultTy = cast<ShapedType>(result.getType());
-    emptyTensors.push_back(rewriter.create<tensor::EmptyOp>(
-        loc, resultTy.getShape(), resultTy.getElementType(), filteredDims));
+    RankedTensorType rankedType = result.getType().dyn_cast<RankedTensorType>();
+    assert(rankedType && "Expected all results to be ranked tensors");
+    emptyTensors.push_back(
+        rewriter.create<tensor::EmptyOp>(loc, rankedType, filteredDims));
     opResultTypes.push_back(result.getType());
   }
 
   auto bodyResultTypes = llvm::to_vector<4>(llvm::map_range(
       emptyTensors, [](Value v) { return getElementTypeOrSelf(v); }));
 
   SmallVector<Value, 2> operands;
   SmallVector<AffineMap, 2> indexingMaps;

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

func.func @test_abs_dyn(%arg0: tensor<2x?xf32>) -> tensor<2x?xf32> {
   // CHECK: %[[C1:.+]] = arith.constant 1
   // CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C1]]
   // CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM]])
   // CHECK: linalg.generic
   // CHECK: math.absf
   %0 = "tosa.abs"(%arg0) : (tensor<2x?xf32>) -> tensor<2x?xf32>
   return %0 : tensor<2x?xf32>
 }
+
 // -----
 
+#SparseVector = #sparse_tensor.encoding<{ lvlTypes = [ "compressed" ] }>
+
+// CHECK-LABEL: @test_encoding_passthrough
+func.func @test_encoding_passthrough() -> tensor<8xi32, #SparseVector> {
+  %0 = tensor.empty() : tensor<8xi32, #SparseVector>
+  // CHECK: linalg.generic
+  // CHECK: sparse_tensor
+  %1 = "tosa.abs"(%0) : (tensor<8xi32, #SparseVector>) -> tensor<8xi32, #SparseVector>
+  return %1 : tensor<8xi32, #SparseVector>
+}
+
+// -----
 
 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> ()>
 // CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL: @test_broadcast
 // CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<1xf32
 // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]: tensor<2xf32>
 func.func @test_broadcast(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {