[mlir][tosa] Improve lowering support for tosa.concat

[sabauma]

Summary

The existing lowering for tosa.concat fails in some instances when the
output shape contains more information the input shapes. The result is
an illegal tensor.empty operation.

This change bases the output shape on the original tosa.concat
operation, while querying the input tensor shapes to build the slicing
operations.

[rsuderman]

mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

374

Why do you make this a one line helper. This is simple enough it should just be inlined.

mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir

244

This reference %dim but no definition of where it came from. I believe this should be the summation of the two correct?

267

We should see CHECK-DAG in most of these tests. Your current implementation enforces order which doesn't work for constants very well.

274

There are no cases that should computing the resulting dynamic length on the axis dimension. Considering most of your changes focus on supporting concatenation across dynamic dims I would expect to see these behavior tested more directly.

291

This is not a good test. Assuming dynamic ? values and static values being safe to concat is potentially unsafe (in the cases they don't match). We allow these mismatches as an intermediate but it produces possibly illegal behavior.

[sabauma]

Address rsuderman's comments

1. Replace uses of CHECK with CHECK-DAG
2. Fix some test cases to properly match on captured patterns
3. Fix an issue in the computation of the output size along the axis dimension

[sabauma]

Lock down computation of the total size of the output tensor

[sabauma]

mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir

244

Fixed this and a few other similar cases.

267

I converted all the tests I modified to use CHECK-DAG rather than CHECK.

274

There are no cases that should computing the resulting dynamic length on the axis dimension.

The lowering does produce a computation of the total size along the concatenation axis. I've updated the test to lock down that computation. I originally omitted it as the actual value of the computation is not used.

Considering most of your changes focus on supporting concatenation across dynamic dims I would expect to see these behavior tested more directly.

The specific issue I am trying to address is when the output dimensions of the tosa.concat are more specified (fewer dynamic dimensions) than the input dimensions. The old lowering would produce a tensor.empty operation like this:

tensor.empty(%d0, %d1) : tensor<5x1x10xf32>

Which fails to verify, as there are more arguments than dynamic dimensions on the output tensor.

I don't mind adding more if you feel that is warranted, but the other tests I updated are also testing concatenation operations with dynamic dimensions, so I opted for testing the specific case I'm aiming to fix.

291

I agree that it is odd to have more information about the result type of the concatentation than the input types.

Some of the other concat tests have similar issues with potential mismatches along dynamic dimensions, and I'm not really sure what I can do to fix that short of generating runtime logic to verify these constraints.

[sabauma]

Convert additional concat test to use CHECK-DAG

[sabauma]

mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

374

I've eliminated the helper function in factor of direct use.

[sabauma]

@rsuderman Could you elaborate more on your last point in the review:

This is not a good test. Assuming dynamic ? values and static values being safe to concat is potentially unsafe (in the cases they don't match). We allow these mismatches as an intermediate but it produces possibly illegal behavior.

I am uncertain how to address this concern in an operator lowering. Some of the other concat test examples potentially have the same issue.

[rsuderman]

mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir

291

You should have two versions of this test. One where the concatenation is along a ? dimension, and another where a non-concat dimension is ?. We want to see that the offsets involving ? are generated by computing using tensor.dim.

[rsuderman]

mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

375

Either initially set this to 0 or compute within the loop. The axisOffsets entries should match 1-1 with eh operands as their expected offset.

407

Keep the arith::AddIOp here or generate axisOffsets so that correspond with each insertion. It is strange to see mutating the variable depending on the next iteration of the loop.

mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir

232

Your last CHECK* statement always needs to be a CHECK not a CHECK-DAG so it knows to check the values as the root.

254

Ditto: last CHECK-DAG should be a CHECK

[sabauma]

Address rsuderman's comments on test cases

[sabauma]

Addressed @rsuderman's feedback.

mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

 //===- TosaToTensor.cpp - Lowering Tosa to Tensor Dialect -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // These rewriters lower from the Tosa to the Tensor dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/TosaToTensor/TosaToTensor.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace tosa;
 
 static bool findIntermediateShape(ArrayRef<int64_t> lhsShape,

 };
 
 struct ConcatConverter : public OpConversionPattern<tosa::ConcatOp> {
   using OpConversionPattern<tosa::ConcatOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(tosa::ConcatOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    auto inputType = cast<ShapedType>(op.getOperand(0).getType());
     auto resultType = dyn_cast<RankedTensorType>(op.getType());
 
     Location loc = op.getLoc();
     int axis = op.getAxis();
     Value axisValue = rewriter.createOrFold<arith::ConstantOp>(
         loc, rewriter.getIndexAttr(axis));
-    int rank = resultType.getRank();
-    SmallVector<Value, 3> offsets, sizes, strides;
-    sizes.reserve(rank);
-    strides.resize(rank, rewriter.create<arith::ConstantIndexOp>(loc, 1));
-    offsets.resize(rank, rewriter.create<arith::ConstantIndexOp>(loc, 0));
+    int64_t rank = resultType.getRank();
 
-    SmallVector<Value> dynDims;
-    for (int i = 0; i < rank; ++i) {
-      sizes.push_back(rewriter.createOrFold<tensor::DimOp>(
-          loc, adaptor.getOperands()[0], i));
-      if (inputType.isDynamicDim(i)) {
-        dynDims.push_back(
-            rewriter.create<tensor::DimOp>(loc, op.getOperand(0), i));
-      }
-    }
+    SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
+    SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
+    SmallVector<OpFoldResult> sizes = tensor::createDimValues(
+        rewriter, op.getLoc(), adaptor.getOperands()[0]);
+
+    // Pre-compute the offsets along the axis dimension.
+    SmallVector<OpFoldResult> axisOffsets;
+    axisOffsets.push_back(sizes[axis]);
 
-    Value resultDimSize = sizes[axis];
     for (auto arg : adaptor.getOperands().drop_front()) {
       auto size = rewriter.createOrFold<tensor::DimOp>(loc, arg, axisValue);
-      resultDimSize =
-          rewriter.createOrFold<arith::AddIOp>(loc, resultDimSize, size);
+      auto currentOffset =
+          getValueOrCreateConstantIndexOp(rewriter, loc, axisOffsets.back());
+      auto total =
+          rewriter.createOrFold<arith::AddIOp>(loc, currentOffset, size);
+      axisOffsets.push_back(getAsOpFoldResult(total));
+    }
+    sizes[axis] = axisOffsets.back();
+
+    // Compute the dynamic sizes of the tensor.empty operation.
+    // This is based off of the specified result type of the tosa.concat
+    // operation, since we don't want to change the result type of the operation
+    // during the conversion.
+    SmallVector<Value> dynDims;
+    for (int64_t i = 0; i < rank; ++i) {
+      if (resultType.isDynamicDim(i)) {
+        dynDims.push_back(
+            getValueOrCreateConstantIndexOp(rewriter, loc, sizes[i]));
+      }
     }
-    sizes[axis] = resultDimSize;
 
     Value emptyTensor = rewriter.create<tensor::EmptyOp>(
         loc, resultType.getShape(), resultType.getElementType(), dynDims);
 
-    auto toOpFoldResult = [](Value v) -> OpFoldResult {
-      auto op = v.getDefiningOp<arith::ConstantIndexOp>();
-      if (!op)
-        return v;
-      return op.getValue();
-    };
     Value result = emptyTensor;
-    for (auto arg : adaptor.getOperands()) {
-      sizes[axis] = rewriter.createOrFold<tensor::DimOp>(loc, arg, axisValue);
+    for (auto [arg, offset] : llvm::zip(adaptor.getOperands(), axisOffsets)) {
+      auto sizes = tensor::createDimValues(rewriter, op.getLoc(), arg);
       result = rewriter.createOrFold<tensor::InsertSliceOp>(
-          loc, arg, result,
-          llvm::to_vector(llvm::map_range(offsets, toOpFoldResult)),
-          llvm::to_vector(llvm::map_range(sizes, toOpFoldResult)),
-          llvm::to_vector(llvm::map_range(strides, toOpFoldResult)));
-      offsets[axis] =
-          rewriter.createOrFold<arith::AddIOp>(loc, offsets[axis], sizes[axis]);
+          loc, arg, result, offsets, sizes, strides);
+      offsets[axis] = offset;
     }
     rewriter.replaceOp(op, result);
     return success();
   }
 };
 
 } // namespace
 

mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir

 }
 
 // -----
 
 // CHECK-LABEL: @concat
 // CHECK-SAME: %[[ARG0:.+]]: tensor<5x1xf32>
 // CHECK-SAME: %[[ARG1:.+]]: tensor<6x1xf32>
 func.func @concat(%arg0: tensor<5x1xf32>, %arg1: tensor<6x1xf32>) -> () {
-  // CHECK: [[AXIS:%.+]] = arith.constant 0
-  // CHECK: [[STRIDE:%.+]]   = arith.constant 1
-  // CHECK: [[OFFSET:%.+]] = arith.constant 0 : index
-  // CHECK: [[IDX0:%.+]] = arith.constant 0 : index
-  // CHECK: [[IDX1:%.+]] = arith.constant 1 : index
   // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<11x1xf32>
   // CHECK: [[INSERT0:%.+]] = tensor.insert_slice %[[ARG0]] into [[INIT]][0, 0] [5, 1] [1, 1]
   // CHECK: [[INSERT1:%.+]] = tensor.insert_slice %[[ARG1]] into [[INSERT0]][5, 0] [6, 1] [1, 1]
   %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i64} : (tensor<5x1xf32>, tensor<6x1xf32>)  -> (tensor<11x1xf32>)
 
-  // CHECK: [[AXIS:%.+]] = arith.constant 1
-  // CHECK: [[STRIDE:%.+]]   = arith.constant 1
-  // CHECK: [[OFFSET:%.+]] = arith.constant 0 : index
-  // CHECK: [[IDX0:%.+]] = arith.constant 0 : index
-  // CHECK: [[IDX1:%.+]] = arith.constant 1 : index
   // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<5x2xf32>
   // CHECK: [[INSERT0:%.+]] = tensor.insert_slice %[[ARG0]] into [[INIT]][0, 0] [5, 1] [1, 1]
   // CHECK: [[INSERT1:%.+]] = tensor.insert_slice %[[ARG0]] into [[INSERT0]][0, 1] [5, 1] [1, 1]
   %1 = "tosa.concat"(%arg0, %arg0) { axis = 1 : i64} : (tensor<5x1xf32>, tensor<5x1xf32>)  -> (tensor<5x2xf32>)
   return
 }
 
 // -----
 
 // CHECK-LABEL: @concat_non_axis_dyn
 // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
 // CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]
 func.func @concat_non_axis_dyn(%arg0: tensor<5x?xf32>, %arg1: tensor<6x?xf32>) -> () {
-  // CHECK: %[[AXIS:.+]] = arith.constant 0
-  // CHECK: %[[STRIDE:.+]]   = arith.constant 1
-  // CHECK: %[[OFFSET:.+]] = arith.constant 0 : index
-  // CHECK: %[[IDX0:.+]] = arith.constant 0 : index
-  // CHECK: %[[IDX1:.+]] = arith.constant 1 : index
-  // CHECK: %[[SIZE:.+]] = tensor.dim %[[ARG0]], %[[IDX1]]
-  // CHECK: %[[IDX1_2:.+]] = arith.constant 1 : index
-  // CHECK: %[[DYN:.+]] = tensor.dim %[[ARG0]], %[[IDX1_2]]
-  // CHECK: %[[INIT:.+]] = tensor.empty(%[[DYN]]) : tensor<11x?xf32>
-  // CHECK: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [5, %[[SIZE]]] [1, 1]
-  // CHECK: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][5, 0] [6, %[[SIZE]]] [1, 1]
+  // CHECK-DAG: %[[AXIS:.+]] = arith.constant 0
+  // CHECK-DAG: %[[IDX1:.+]] = arith.constant 1
+  // CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[IDX1]]
+  // CHECK-DAG: %[[INIT:.+]] = tensor.empty(%[[DIM0]]) : tensor<11x?xf32>
+  // CHECK-DAG: %[[IDX1_1:.+]] = arith.constant 1 : index
+  // CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG0]], %[[IDX1_1]]
+  // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [5, %[[DIM1]]] [1, 1]
+  // CHECK-DAG: %[[IDX1_2:.+]] = arith.constant 1 : index
+  // CHECK-DAG: %[[DIM2:.+]] = tensor.dim %[[ARG1]], %[[IDX1_2]] : tensor<6x?xf32>
+  // CHECK-DAG: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][5, 0] [6, %[[DIM2]]] [1, 1]
   %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i64} : (tensor<5x?xf32>, tensor<6x?xf32>)  -> (tensor<11x?xf32>)
   return
 }
 
 // -----
 
 // CHECK-LABEL: @concat_axis_dyn
 // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
 // CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]:
 func.func @concat_axis_dyn(%arg0: tensor<?x3xf32>, %arg1: tensor<?x3xf32>) -> () {
-  // CHECK: %[[AXIS:.+]] = arith.constant 0
-  // CHECK: %[[STRIDE:.+]]   = arith.constant 1
-  // CHECK: %[[OFFSET:.+]] = arith.constant 0 : index
-  // CHECK: %[[IDX0:.+]] = arith.constant 0 : index
-  // CHECK: %[[SIZE:.+]] = tensor.dim %[[ARG0]], %[[IDX0]]
-  // CHECK: %[[IDX0_2:.+]] = arith.constant 0 : index
-  // CHECK: %[[DYN:.+]] = tensor.dim %[[ARG0]], %[[IDX0_2]]
-  // CHECK: %[[IDX1:.+]] = arith.constant 1 : index
-  // CHECK: %[[INIT:.+]] = tensor.empty(%[[DYN]]) : tensor<?x3xf32>
-  // CHECK: %[[DYN1:.+]] = tensor.dim %[[ARG0]], %[[AXIS]]
-  // CHECK: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [%[[DYN1]], 3] [1, 1]
-  // CHECK: %[[SUM:.+]]  = arith.addi %[[OFFSET]], %[[DYN1]]
-  // CHECK: %[[DYN2:.+]] = tensor.dim %[[ARG1]], %[[AXIS]]
-  // CHECK: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][%[[SUM]], 0] [%[[DYN2]], 3] [1, 1]
+  // CHECK-DAG: %[[AXIS:.+]] = arith.constant 0 : index
+  // CHECK-DAG: %[[IDX0:.+]] = arith.constant 0 : index
+  // CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[IDX0]] : tensor<?x3xf32>
+  // CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG1]], %[[AXIS]] : tensor<?x3xf32>
+  // CHECK-DAG: %[[SUM:.+]] = arith.addi %[[DIM0]], %[[DIM1]] : index
+  // CHECK-DAG: %[[INIT:.+]] = tensor.empty(%[[SUM]]) : tensor<?x3xf32>
+  // CHECK-DAG: %[[IDX0_1:.+]] = arith.constant 0 : index
+  // CHECK-DAG: %[[DIM2:.+]] = tensor.dim %[[ARG0]], %[[IDX0_1]] : tensor<?x3xf32>
+  // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0] [%[[DIM2]], 3] [1, 1] : tensor<?x3xf32> into tensor<?x3xf32>
+  // CHECK-DAG: %[[IDX0_2:.+]] = arith.constant 0 : index
+  // CHECK-DAG: %[[DIM3:.+]] = tensor.dim %[[ARG1]], %[[IDX0_2]] : tensor<?x3xf32>
+  // CHECK-DAG: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][%[[DIM0]], 0] [%[[DIM3]], 3] [1, 1] : tensor<?x3xf32> into tensor<?x3xf32>
+
   %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i64} : (tensor<?x3xf32>, tensor<?x3xf32>)  -> (tensor<?x3xf32>)
   return
 }
+
+// -----
+
+// CHECK-LABEL: @concat_axis_dyn_mixed
+// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
+// CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]:
+// CHECK-SAME:  %[[ARG2:[0-9a-zA-Z_]*]]:
+func.func @concat_axis_dyn_mixed(%arg0: tensor<?x?x?xf32>, %arg1: tensor<1x1x10xf32>, %arg2: tensor<?x?x?xf32>) -> () {
+  // CHECK-DAG: %[[AXIS:.+]] = arith.constant 0 : index
+  // CHECK-DAG: %[[IDX0_0:.+]] = arith.constant 0 : index
+  // CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[IDX0_0]] : tensor<?x?x?xf32>
+  // CHECK-DAG: %[[IDX1_0:.+]] = arith.constant 1 : index
+  // CHECK-DAG: %[[IDX2:.+]] = arith.constant 2 : index
+  // CHECK-DAG: %[[IDX1_3:.+]] = arith.constant 1 : index
+  // CHECK-DAG: %[[OFFSET:.+]] = arith.addi %[[DIM0]], %[[IDX1_3]] : index
+  // CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<5x1x10xf32>
+  // CHECK-DAG: %[[IDX0_1:.+]] = arith.constant 0 : index
+  // CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG0]], %[[IDX0_1]] : tensor<?x?x?xf32>
+  // CHECK-DAG: %[[IDX1_1:.+]] = arith.constant 1 : index
+  // CHECK-DAG: %[[DIM2:.+]] = tensor.dim %[[ARG0]], %[[IDX1_1]] : tensor<?x?x?xf32>
+  // CHECK-DAG: %[[IDX2_0:.+]] = arith.constant 2 : index
+  // CHECK-DAG: %[[DIM3:.+]] = tensor.dim %[[ARG0]], %[[IDX2_0]] : tensor<?x?x?xf32>
+  // CHECK-DAG: %[[INSERT0:.+]] = tensor.insert_slice %[[ARG0]] into %[[INIT]][0, 0, 0] [%[[DIM1]], %[[DIM2]], %[[DIM3]]] [1, 1, 1] : tensor<?x?x?xf32> into tensor<5x1x10xf32>
+  // CHECK-DAG: %[[INSERT1:.+]] = tensor.insert_slice %[[ARG1]] into %[[INSERT0]][%[[DIM0]], 0, 0] [1, 1, 10] [1, 1, 1] : tensor<1x1x10xf32> into tensor<5x1x10xf32>
+  // CHECK-DAG: %[[IDX0_2:.+]] = arith.constant 0 : index
+  // CHECK-DAG: %[[DIM2_0:.+]] = tensor.dim %[[ARG2]], %[[IDX0_2]] : tensor<?x?x?xf32>
+  // CHECK-DAG: %[[IDX1_2:.+]] = arith.constant 1 : index
+  // CHECK-DAG: %[[DIM2_1:.+]] = tensor.dim %[[ARG2]], %[[IDX1_2]] : tensor<?x?x?xf32>
+  // CHECK-DAG: %[[IDX2_1:.+]] = arith.constant 2 : index
+  // CHECK-DAG: %[[DIM2_2:.+]] = tensor.dim %[[ARG2]], %[[IDX2_1]] : tensor<?x?x?xf32>
+  // CHECK-DAG: %[[INSERT2:.+]] = tensor.insert_slice %[[ARG2]] into %[[INSERT1]][%[[OFFSET]], 0, 0] [%[[DIM2_0]], %[[DIM2_1]], %[[DIM2_2]]] [1, 1, 1] : tensor<?x?x?xf32> into tensor<5x1x10xf32>
+
+  %0 = "tosa.concat"(%arg0, %arg1, %arg2) <{axis = 0 : i64}> : (tensor<?x?x?xf32>, tensor<1x1x10xf32>, tensor<?x?x?xf32>) -> tensor<5x1x10xf32>
+  return
+}
+