[mlir][tosa] Lowering of tosa.gather operations with dynamic dimensions

[sabauma]

Summary

The existing TOSA->Linalg lowering of tosa.gather only supports gathers
with either a static shape or a single dynamic batch dimension.
This change extends support to arbitrary number of dynamic dimensions on
both the values and indices of the gather operation.

[sabauma]

Code cleanup for output dimensions calculation

[sabauma]

Cleanup

[sabauma]

Rename dimension inference function to be consistent with similar functions

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

 public:
   using OpConversionPattern<tosa::GatherOp>::OpConversionPattern;
   LogicalResult
   matchAndRewrite(tosa::GatherOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
     auto input = adaptor.getOperands()[0];
     auto indices = adaptor.getOperands()[1];
 
+    auto valuesTy = op.getValues().getType();
     auto resultTy = op.getType().cast<ShapedType>();
 
-    auto dynamicDimsOr = checkHasDynamicBatchDims(
-        rewriter, op, {input, indices, op.getOutput()});
-    if (!dynamicDimsOr.has_value())
-      return rewriter.notifyMatchFailure(
-          op, "tosa.gather currently only supports dynamic batch dimensions");
-    SmallVector<Value> dynamicDims = *dynamicDimsOr;
+    if (!valuesTy.dyn_cast<TensorType>().hasRank())
+      return rewriter.notifyMatchFailure(op, "unranked tensors not supported");
+
+    auto dynamicDims = gatherOutputDynamicSizes(
+        rewriter, op.getLoc(), adaptor.getValues(),
+        valuesTy.dyn_cast<RankedTensorType>(), adaptor.getIndices(),
+        op.getIndices().getType());
 
     auto resultElementTy = resultTy.getElementType();
 
     auto loc = op.getLoc();
-
     auto emptyTensor =
         rewriter
             .create<tensor::EmptyOp>(loc, resultTy.getShape(), resultElementTy,
                                      dynamicDims)
             .getResult();
 
     SmallVector<AffineMap, 2> affineMaps = {
         AffineMap::get(
auto genericOp = rewriter.create<linalg::GenericOp>(
           auto index2 = rewriter.create<linalg::IndexOp>(loc, 2);
           Value extract = rewriter.create<tensor::ExtractOp>(
               loc, input, ValueRange{index0, index1, index2});
           rewriter.create<linalg::YieldOp>(loc, extract);
         });
     rewriter.replaceOp(op, genericOp.getResult(0));
     return success();
   }
+
+  static llvm::SmallVector<Value>
+  gatherOutputDynamicSizes(OpBuilder &builder, Location loc, Value values,
+                           RankedTensorType valueType, Value indices,
+                           RankedTensorType indicesType) {
+    auto valueShape = valueType.getShape();
+    auto indicesShape = indicesType.getShape();
+
+    // The shape of the gather op result is n x w x c.
+    auto n = valueShape[0];
+    auto w = indicesShape[1];
+    auto c = valueShape[2];
+
+    llvm::SmallVector<Value> results;
+    if (ShapedType::isDynamic(n)) {
+      results.push_back(builder.create<tensor::DimOp>(loc, values, 0));
+    }
+
+    if (ShapedType::isDynamic(w)) {
+      results.push_back(builder.create<tensor::DimOp>(loc, indices, 1));
+    }
+
+    if (ShapedType::isDynamic(c)) {
+      results.push_back(builder.create<tensor::DimOp>(loc, values, 2));
+    }
+
+    return results;
+  }
 };
 
 // Lowerings the TableOp to a series of gathers and numerica operations. This
 // includes interpolation between the high/low values. For the I8 varient, this
 // simplifies to a single gather operation.
 class TableConverter : public OpRewritePattern<tosa::TableOp> {
 public:
   using OpRewritePattern<tosa::TableOp>::OpRewritePattern;

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

func.func @gather_float_dyn(%arg0: tensor<?x3x2xf32>, %arg1: tensor<?x3xi32>) -> () {
   // CHECK:   %[[EXTRACT:.+]] = tensor.extract %[[ARG0]][%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<?x3x2xf32>
   // CHECK:   linalg.yield %[[EXTRACT]]
   %0 = "tosa.gather"(%arg0, %arg1)  : (tensor<?x3x2xf32>, tensor<?x3xi32>)  -> (tensor<?x3x2xf32>)
   return
 }
 
 // -----
 
+// CHECK-LABEL: @gather_float_all_dynamic
+// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]
+// CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]
+func.func @gather_float_all_dynamic(%arg0: tensor<?x?x?xf32>, %arg1: tensor<?x?xi32>) -> () {
+  // CHECK: %[[C0:.+]] = arith.constant 0
+  // CHECK: %[[BATCH:.+]] = tensor.dim %[[ARG0]], %[[C0]]
+  // CHECK: %[[C1:.+]] = arith.constant 1
+  // CHECK: %[[INDEX:.+]] = tensor.dim %[[ARG1]], %[[C1]]
+  // CHECK: %[[C2:.+]] = arith.constant 2
+  // CHECK: %[[CHANNEL:.+]] = tensor.dim %[[ARG0]], %[[C2]]
+
+  // CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]], %[[INDEX]], %[[CHANNEL]])
+  // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[ARG1]] : tensor<?x?xi32>) outs(%[[INIT]] : tensor<?x?x?xf32>)
+  // CHECK: ^bb0(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: f32)
+  // CHECK:   %[[IDX0:.+]] = linalg.index 0
+  // CHECK:   %[[CAST:.+]] = arith.index_cast %[[BBARG0]]
+  // CHECK:   %[[IDX2:.+]] = linalg.index 2
+  // CHECK:   %[[EXTRACT:.+]] = tensor.extract %[[ARG0]][%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<?x?x?xf32>
+  // CHECK:   linalg.yield %[[EXTRACT]]
+  %0 = "tosa.gather"(%arg0, %arg1)  : (tensor<?x?x?xf32>, tensor<?x?xi32>)  -> (tensor<?x?x?xf32>)
+  return
+}
+
+// -----
+
 // CHECK-LABEL: @gather_int
 // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]
 // CHECK-SAME:  %[[ARG1:[0-9a-zA-Z_]*]]
 func.func @gather_int(%arg0: tensor<2x3x2xi32>, %arg1: tensor<2x3xi32>) -> () {
   // CHECK: %[[INIT:.+]] = tensor.empty()
   // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[ARG1]] : tensor<2x3xi32>) outs(%[[INIT]] : tensor<2x3x2xi32>)
   // CHECK: ^bb0(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: i32)
   // CHECK:   %[[IDX0:.+]] = linalg.index 0