[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

559–560

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

560

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

[rkayaith]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

559–560

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

 static LogicalResult
 elementwiseMatchAndRewriteHelper(Operation *operation,
                                  PatternRewriter &rewriter) {
   auto loc = operation->getLoc();
 
   assert(operation->getNumResults() == 1 &&
          "All TOSA elementwise ops should only return a single result.");
 
-  auto results = operation->getResults();
-  auto resultTy = dyn_cast<ShapedType>(operation->getResult(0).getType());
+  auto result = operation->getResult(0);
+  auto resultTy = dyn_cast<RankedTensorType>(result.getType());
 
   if (!resultTy)
-    return rewriter.notifyMatchFailure(operation,
-                                       "All results must be a shaped type");
+    return rewriter.notifyMatchFailure(
+        operation, "All results must be a ranked tensor type");
 
   unsigned rank = resultTy.getRank();
 
   // Construct the indexing maps needed for linalg.generic ops.
   SmallVector<Type> bodyArgTypes;
 
   for (Value in : operation->getOperands())
     bodyArgTypes.emplace_back(getElementTypeOrSelf(in.getType()));
 
   SmallVector<Type> opResultTypes;
   SmallVector<Value> emptyTensors;
 
   SmallVector<Value> dynDims;
-  dynDims.resize(cast<ShapedType>(results.front().getType()).getRank());
+  dynDims.resize(rank);
 
   for (auto arg : operation->getOperands()) {
     auto operandTy = cast<ShapedType>(arg.getType());
     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));
-    opResultTypes.push_back(result.getType());
-  }
+  emptyTensors.push_back(
+      rewriter.create<tensor::EmptyOp>(loc, resultTy, filteredDims));
 
   auto bodyResultTypes = llvm::to_vector<4>(llvm::map_range(
       emptyTensors, [](Value v) { return getElementTypeOrSelf(v); }));
 
   SmallVector<Value, 2> operands;
   SmallVector<AffineMap, 2> indexingMaps;
   indexingMaps.reserve(operation->getNumOperands() + bodyResultTypes.size());
 

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

-// RUN: mlir-opt --split-input-file -pass-pipeline="builtin.module(func.func(tosa-to-linalg-named))" %s -verify-diagnostics
+// RUN: mlir-opt --split-input-file -pass-pipeline="builtin.module(func.func(tosa-to-linalg))" %s -verify-diagnostics
 
 // CHECK-LABEL: @avg_pool2d_with_unsupported_quant_type
 func.func @avg_pool2d_with_unsupported_quant_type(%arg0: tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>>) -> tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>> {
   // expected-error@+1 {{failed to legalize operation 'tosa.avg_pool2d'}}
   %0 = "tosa.avg_pool2d"(%arg0) {acc_type = i32, kernel = array<i64: 2, 2>, pad = array<i64: 0, 1, 0, 1>, stride = array<i64: 1, 1>} : (tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>>) -> tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>>
   return %0 : tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>>
 }
+
+// -----
+
+// CHECK-LABEL: @tensor_with_unknown_rank
+func.func @tensor_with_unknown_rank(%arg0: tensor<*xi8>) -> tensor<*xi8> {
+  // expected-error@+1 {{failed to legalize operation 'tosa.abs'}}
+  %0 = "tosa.abs"(%arg0) : (tensor<*xi8>) -> tensor<*xi8>
+  return %0 : tensor<*xi8>
+}

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(%arg0: tensor<2xi8, #SparseVector>) -> tensor<2xi8, #SparseVector> {
+  // CHECK: linalg.generic
+  // CHECK: sparse_tensor
+  %0 = "tosa.abs"(%arg0) : (tensor<2xi8, #SparseVector>) -> tensor<2xi8, #SparseVector>
+  return %0 : tensor<2xi8, #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> {