From 796888ac6a4e97a19f15d92394e189a67d777778 Mon Sep 17 00:00:00 2001
From: Alan Li <me@alanli.org>
Date: Tue, 16 Jul 2024 02:55:57 +0000
Subject: [PATCH] Update according to comments.

Signed-off-by: Alan Li <me@alanli.org>
---
 .../Conversion/StableHLOToLinalgExt.cpp       | 21 +++++++------------
 .../test/stablehlo_to_linalg_ext.mlir         | 16 ++++++++------
 2 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/compiler/plugins/input/StableHLO/Conversion/StableHLOToLinalgExt.cpp b/compiler/plugins/input/StableHLO/Conversion/StableHLOToLinalgExt.cpp
index 0f32f678bbf85..1b2884bd569b5 100644
--- a/compiler/plugins/input/StableHLO/Conversion/StableHLOToLinalgExt.cpp
+++ b/compiler/plugins/input/StableHLO/Conversion/StableHLOToLinalgExt.cpp
@@ -440,21 +440,14 @@ 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>(
@@ -462,7 +455,7 @@ struct ReverseOpConversion final
         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
index 7c2266b7d51b9..59dc41d3fe4a8 100644
--- 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