support dyanmic case

Author: zhczhong

lib/gc/Transforms/DecomposeTensorOperation.cpp

@@ -37,10 +37,37 @@ namespace {
 ///      %extracted = tensor.extract %0[%in, %17] : tensor<7x128xf16>
 ///      linalg.yield %extracted : f16
 ///    } -> tensor<1x7x128xf16>
-
 struct DecomposeGatherOp : public OpRewritePattern<tensor::GatherOp> {
   using OpRewritePattern<tensor::GatherOp>::OpRewritePattern;
 
+  SmallVector<OpFoldResult> getDstMixedSizes(PatternRewriter &rewriter,
+                                             Location loc,
+                                             tensor::GatherOp gatherOp) const {
+    SmallVector<OpFoldResult> dstSize =
+        tensor::getMixedSizes(rewriter, loc, gatherOp.getResult());
+    SmallVector<OpFoldResult> indexSize =
+        tensor::getMixedSizes(rewriter, loc, gatherOp.getIndices());
+    SmallVector<OpFoldResult> srcSize =
+        tensor::getMixedSizes(rewriter, loc, gatherOp.getSource());
+    SmallVector<int64_t> gatherDims(gatherOp.getGatherDims());
+    bool isShrinkDst = (indexSize.size() - 1) + srcSize.size() ==
+                       dstSize.size() + gatherDims.size();
+    for (size_t i = 0; i < indexSize.size() - 1; i++) {
+      dstSize[i] = indexSize[i];
+    }
+    auto cnt = 0;
+    for (size_t i = indexSize.size() - 1; i < dstSize.size(); i++) {
+      while (isShrinkDst && llvm::find(gatherDims, cnt) != gatherDims.end()) {
+        cnt++;
+      }
+      dstSize[i] = llvm::find(gatherDims, cnt) == gatherDims.end()
+                       ? srcSize[cnt]
+                       : getAsIndexOpFoldResult(rewriter.getContext(), 1);
+      cnt++;
+    }
+    return dstSize;
+  }
+
   LogicalResult matchAndRewrite(tensor::GatherOp gatherOp,
                                 PatternRewriter &rewriter) const override {
     OpBuilder::InsertionGuard g(rewriter);
@@ -51,7 +78,7 @@ struct DecomposeGatherOp : public OpRewritePattern<tensor::GatherOp> {
     // create destination tensor for linalg out
     RankedTensorType dstType = gatherOp.getResultType();
     Value dstTensor = rewriter.create<tensor::EmptyOp>(
-        loc, tensor::getMixedSizes(rewriter, loc, gatherOp.getResult()),
+        loc, getDstMixedSizes(rewriter, loc, gatherOp),
         dstType.getElementType());
 
     // split index tensor to create the linalg input
@@ -113,8 +140,8 @@ struct DecomposeGatherOp : public OpRewritePattern<tensor::GatherOp> {
                              dstRank + gatherDims.size();
           int cnt = 0;
           for (auto i = indexTensorSize.size() - 1; i < dstRank; i++) {
-            while (llvm::find(gatherDims, cnt) != gatherDims.end() &&
-                   isShrinkDst) {
+            while (isShrinkDst &&
+                   llvm::find(gatherDims, cnt) != gatherDims.end()) {
               cnt++;
             }
             indexValues[cnt] = b.create<linalg::IndexOp>(loc, i);

test/mlir/test/gc/Transforms/DecomposeTensorOperation.mlir

@@ -41,3 +41,32 @@ func.func @gather_multiple_gather_dim_no_shrink(%arg0: tensor<2x2x2x2xf32>, %arg
   %1 = tensor.gather %arg0[%arg1] gather_dims([1, 2]) : (tensor<2x2x2x2xf32>, tensor<2x3x2xindex>) -> tensor<2x3x2x1x1x2xf32>
   return %1 : tensor<2x3x2x1x1x2xf32>
 }
+
+// -----
+
+/// CHECK-LABEL: @gather_single_gather_dim_dynamic
+func.func @gather_single_gather_dim_dynamic(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<2x3x1xindex>) -> tensor<2x3x?x?x?xf32> {
+  /// CHECK: %[[DIM1:.*]] = tensor.dim
+  /// CHECK: %[[DIM2:.*]] = tensor.dim
+  /// CHECK: %[[DIM3:.*]] = tensor.dim
+  /// CHECK: %[[EMPTY:.*]] = tensor.empty(%[[DIM1:.*]], %[[DIM2:.*]], %[[DIM3:.*]]) : tensor<2x3x?x?x?xf32>
+  /// CHECK: linalg.generic {{.*}} iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%[[ARG0:.*]] : tensor<2x3x1xindex>) outs(%[[EMPTY:.*]] : tensor<2x3x?x?x?xf32>)
+  %1 = tensor.gather %arg0[%arg1] gather_dims([1]) : (tensor<?x?x?x?xf32>, tensor<2x3x1xindex>) -> tensor<2x3x?x?x?xf32>
+  return %1 : tensor<2x3x?x?x?xf32>
+}
+
+// -----
+
+/// CHECK-LABEL: @gather_multiple_gather_dim_no_shrink_dynamic
+func.func @gather_multiple_gather_dim_no_shrink_dynamic(%arg0: tensor<2x2x2x2xf32>, %arg1: tensor<?x?x2xindex>) -> tensor<?x?x2x1x1x2xf32> {
+  /// CHECK: %[[DIM1:.*]] = tensor.dim
+  /// CHECK: %[[DIM2:.*]] = tensor.dim
+  /// CHECK: %[[EMPTY:.*]] = tensor.empty(%[[DIM1:.*]], %[[DIM2:.*]]) : tensor<?x?x2x1x1x2xf32>
+  /// CHECK: %[[DIM3:.*]] = tensor.dim
+  /// CHECK: %[[DIM4:.*]] = tensor.dim
+  /// CHECK: %[[EXTRACTSLICE1:.*]] = tensor.extract_slice %[[ARG1:.*]][0, 0, 0] [%[[DIM3:.*]], %[[DIM4:.*]], 1] [1, 1, 1] : tensor<?x?x2xindex> to tensor<?x?x1xindex>
+  /// CHECK: %[[EXTRACTSLICE2:.*]] = tensor.extract_slice %[[ARG1:.*]][0, 0, 1] [%[[DIM3:.*]], %[[DIM4:.*]], 1] [1, 1, 1] : tensor<?x?x2xindex> to tensor<?x?x1xindex>
+  /// CHECK: linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%[[EXTRACTSLICE1:.*]], %[[EXTRACTSLICE2:.*]] : tensor<?x?x1xindex>, tensor<?x?x1xindex>) outs(%[[EMPTY:.*]] : tensor<?x?x2x1x1x2xf32>)
+  %1 = tensor.gather %arg0[%arg1] gather_dims([1, 2]) : (tensor<2x2x2x2xf32>, tensor<?x?x2xindex>) -> tensor<?x?x2x1x1x2xf32>
+  return %1 : tensor<?x?x2x1x1x2xf32>
+}