Update according to comments.

Signed-off-by: Alan Li <[email protected]>
iree-org · Jul 16, 2024 · 796888a · 796888a
1 parent 6953441
commit 796888a
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 20 deletions.
diff --git a/compiler/plugins/input/StableHLO/Conversion/StableHLOToLinalgExt.cpp b/compiler/plugins/input/StableHLO/Conversion/StableHLOToLinalgExt.cpp
@@ -440,29 +440,22 @@ struct ReverseOpConversion final
     auto resultTy = cast<ShapedType>(op.getType());
     auto dims = op.getDimensions();
     Location loc = op.getLoc();
-
-    SmallVector<Value> dynDims;
-    for (int i = 0; i < inputTy.getRank(); i++) {
-      if (inputTy.isDynamicDim(i)) {
-        dynDims.push_back(rewriter.create<tensor::DimOp>(loc, input, i));
-      }
-    }
+    auto inputTyRank = inputTy.getRank();
 
     // First fill the output buffer with the init value.
-    auto emptyTensor = rewriter
-                           .create<tensor::EmptyOp>(loc, inputTy.getShape(),
-                                                    inputTy.getElementType(),
-                                                    ArrayRef<Value>({dynDims}))
-                           .getResult();
-    SmallVector<AffineMap, 2> affineMaps = {
+    SmallVector<OpFoldResult> inputMixedSizes =
+        tensor::getMixedSizes(rewriter, loc, input);
+    auto emptyTensor = rewriter.create<tensor::EmptyOp>(
+        loc, inputMixedSizes, inputTy.getElementType());
+    SmallVector<AffineMap> affineMaps = {
         rewriter.getMultiDimIdentityMap(resultTy.getRank())};
 
     rewriter.replaceOpWithNewOp<linalg::GenericOp>(
         op, resultTy, ArrayRef<Value>({}), ValueRange{emptyTensor}, affineMaps,
         getNParallelLoopsAttrs(resultTy.getRank()),
         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
           llvm::SmallVector<Value> indices;
-          for (unsigned int i = 0; i < inputTy.getRank(); i++) {
+          for (unsigned int i = 0; i < inputTyRank; i++) {
             Value index =
                 rewriter.create<linalg::IndexOp>(nestedLoc, i).getResult();
             if (std::find(dims.begin(), dims.end(), i) != dims.end()) {

diff --git a/compiler/plugins/input/StableHLO/Conversion/test/stablehlo_to_linalg_ext.mlir b/compiler/plugins/input/StableHLO/Conversion/test/stablehlo_to_linalg_ext.mlir
@@ -540,12 +540,16 @@ func.func @reverse_multi_dim(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
   } : (tensor<?x?xi32>) -> tensor<?x?xi32>
   return %0 : tensor<?x?xi32>
 }
-// CHECK-DAG:    %[[C0:.+]] = arith.constant 0 : index
-// CHECK-DAG:    %[[C1:.+]] = arith.constant 1 : index
-// CHECK-DAG:    %[[D0:.+]] = tensor.dim %[[IN]], %[[C0]]
-// CHECK-DAG:    %[[D1:.+]] = tensor.dim %[[IN]], %[[C1]]
-// CHECK:        %[[INIT:.+]] = tensor.empty(%[[D0]], %[[D1]]) : tensor<?x?xi32>
-// CHECK:        %[[GEN:.+]] = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel"]} outs(%[[INIT]] : tensor<?x?xi32>) {
+// CHECK:    %[[C0:.+]] = arith.constant 0 : index
+// CHECK:    %[[D:.+]] = tensor.dim %[[IN]], %[[C0]] : tensor<?x?xi32>
+// CHECK:    %[[C1:.+]] = arith.constant 1 : index
+// CHECK:    %[[D0:.+]] = tensor.dim %[[IN]], %[[C1]] : tensor<?x?xi32>
+// CHECK:    %[[C0_1:.+]] = arith.constant 0 : index
+// CHECK:    %[[D2:.+]] = tensor.dim %[[IN]], %[[C0_1]]
+// CHECK:    %[[C1_3:.+]] = arith.constant 1 : index
+// CHECK:    %[[D4:.+]] = tensor.dim %[[IN]], %[[C1_3]]
+// CHECK:    %[[INIT:.+]] = tensor.empty(%[[D2]], %[[D4]]) : tensor<?x?xi32>
+// CHECK:    %[[GEN:.+]] = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel"]} outs(%[[INIT]] : tensor<?x?xi32>) {
 
 // First reverse dimension
 // CHECK:       %[[IDX0:.+]] = linalg.index 0 : index