From a3fc99efe0af66da3eae469e30685ed32b3b6877 Mon Sep 17 00:00:00 2001
From: Hanhan Wang <hanchung@google.com>
Date: Mon, 12 Apr 2021 10:34:01 -0700
Subject: [PATCH] Add support for lowering mhlo.dynamic_slice to Linalg ops.

PiperOrigin-RevId: 368033540
---
 .../mhlo/transforms/legalize_to_linalg.cc     | 65 +++++++++++++++++++
 tests/hlo-legalize-to-linalg.mlir             | 29 +++++++++
 2 files changed, 94 insertions(+)

diff --git a/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc b/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
index 850ef00..e482f2b 100644
--- a/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
+++ b/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
@@ -1134,6 +1134,70 @@ class SliceConverter : public OpConversionPattern<OpTy> {
   }
 };
 
+class DynamicSliceConverter : public OpConversionPattern<mhlo::DynamicSliceOp> {
+ public:
+  using OpConversionPattern<mhlo::DynamicSliceOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      mhlo::DynamicSliceOp dynamic_slice_op, ArrayRef<Value> args,
+      ConversionPatternRewriter& rewriter) const final {
+    auto loc = dynamic_slice_op.getLoc();
+    mhlo::DynamicSliceOp::Adaptor adaptor(args);
+    auto arg_type = adaptor.operand().getType().dyn_cast<ShapedType>();
+    if (!arg_type || !arg_type.hasRank()) {
+      return rewriter.notifyMatchFailure(dynamic_slice_op,
+                                         "require known-rank args");
+    }
+
+    auto index_type = rewriter.getIndexType();
+    SmallVector<OpFoldResult, 3> start_indices, sizes;
+    Value zero = rewriter.create<ConstantOp>(
+        loc, rewriter.getZeroAttr(adaptor.start_indices()[0]
+                                      .getType()
+                                      .cast<RankedTensorType>()
+                                      .getElementType()));
+    for (auto en : llvm::enumerate(
+             llvm::zip(adaptor.start_indices(),
+                       dynamic_slice_op.slice_sizes().getValues<int64_t>()))) {
+      int64_t size = std::get<1>(en.value());
+      sizes.push_back(rewriter.getI64IntegerAttr(size));
+
+      // By mhlo.DynamicSlice definition:
+      //   `start_indices[i] = clamp(start_indices[i],
+      //       0, operand.dimension_size[i] - size_indices[i])`
+      Value start_index =
+          rewriter.create<tensor::ExtractOp>(loc, std::get<0>(en.value()));
+      Value ub = rewriter.createOrFold<memref::DimOp>(loc, adaptor.operand(),
+                                                      en.index());
+      // ClampOp lowering does not support index type, so cast it into integer
+      // type.
+      ub = rewriter.createOrFold<IndexCastOp>(loc, start_index.getType(), ub);
+      ub = rewriter.createOrFold<SubIOp>(
+          loc, ub,
+          rewriter.create<ConstantOp>(
+              loc, rewriter.getIntegerAttr(start_index.getType(), size)));
+      // TODO(hanchung): This is a workaround to use the method because only
+      // lmhlo version is defined. The implementation in
+      // map_lmhlo_to_scalar_op.h requires to pass a mhlo op. It will convert it
+      // to an lmhlo op and call the lmhlo implementation.
+      start_index = lmhlo::HloOpToStdScalarOp::map<lmhlo::ClampOp>(
+          loc, start_index.getType(), ArrayRef<Value>{zero, start_index, ub},
+          &rewriter);
+      start_indices.push_back(
+          rewriter.create<IndexCastOp>(loc, index_type, start_index)
+              .getResult());
+    }
+
+    int64_t rank = arg_type.getRank();
+    SmallVector<OpFoldResult, 3> strides(rank, rewriter.getI64IntegerAttr(1));
+
+    rewriter.replaceOpWithNewOp<SubTensorOp>(
+        dynamic_slice_op, dynamic_slice_op.getType().cast<RankedTensorType>(),
+        adaptor.operand(), start_indices, sizes, strides);
+    return success();
+  }
+};
+
 enum class DotOperationType {
   kVectorDot = 0,
   kMatrixVector = 1,
@@ -2090,6 +2154,7 @@ void populateHLOToLinalgConversionPattern(MLIRContext* context,
       ReshapeOpConverter<mhlo::ReshapeOp, false>,
       ReverseConverter<mhlo::ReverseOp, false>,
       SliceConverter<mhlo::SliceOp, false>,
+      DynamicSliceConverter,
       TransposeConverter<mhlo::TransposeOp, false>,
       DotOpOnTensorsConversion<DotOperationType::kMatrixMatrix,
                                linalg::MatmulOp>,
diff --git a/tests/hlo-legalize-to-linalg.mlir b/tests/hlo-legalize-to-linalg.mlir
index a3eb16b..76a1629 100644
--- a/tests/hlo-legalize-to-linalg.mlir
+++ b/tests/hlo-legalize-to-linalg.mlir
@@ -1532,6 +1532,35 @@ func @slice_stride_part(%arg0: tensor<3x4xi32>) -> tensor<1x2xi32> {
 
 // -----
 
+func @dynamic_slice(%arg: tensor<3x4xf32>, %start1: tensor<i64>, %start2: tensor<i64>) -> tensor<1x4xf32> {
+  %0 = "mhlo.dynamic-slice"(%arg, %start1, %start2) {
+    slice_sizes = dense<[1, 4]> : tensor<2xi64>
+  } : (tensor<3x4xf32>, tensor<i64>, tensor<i64>) -> tensor<1x4xf32>
+  return %0 : tensor<1x4xf32>
+}
+// CHECK-LABEL: func @dynamic_slice
+// CHECK-SAME:    %[[ARG0:[a-zA-Z0-9_]*]]
+// CHECK-SAME:    %[[ARG1:[a-zA-Z0-9_]*]]
+// CHECK-SAME:    %[[ARG2:[a-zA-Z0-9_]*]]
+// CHECK:         %[[C0:.*]] = constant 0 : i64
+// CHECK:         %[[SCALAR1:.*]] = tensor.extract %[[ARG1]][] : tensor<i64>
+// CHECK:         %[[UB1:.*]] = constant 2 : i64
+// CHECK:         %[[COND1:.*]] = cmpi slt, %[[SCALAR1]], %[[UB1]] : i64
+// CHECK:         %[[T1:.*]] = select %[[COND1]], %[[SCALAR1]], %[[UB1]] : i64
+// CHECK:         %[[COND2:.*]] = cmpi sgt, %[[T1]], %[[C0]] : i64
+// CHECK:         %[[CLAMPED1:.*]] = select %[[COND2]], %[[T1]], %[[C0]] : i64
+// CHECK:         %[[START1:.*]] = index_cast %[[CLAMPED1]] : i64 to index
+// CHECK:         %[[SCALAR2:.*]] = tensor.extract %[[ARG2]][] : tensor<i64>
+// CHECK:         %[[UB2:.*]] = constant 0 : i64
+// CHECK:         %[[COND3:.*]] = cmpi slt, %[[SCALAR2]], %[[UB2]] : i64
+// CHECK:         %[[T2:.*]] = select %[[COND3]], %[[SCALAR2]], %[[UB2]] : i64
+// CHECK:         %[[COND4:.*]] = cmpi sgt, %[[T2]], %[[C0]] : i64
+// CHECK:         %[[CLAMPED2:.*]] = select %[[COND4]], %[[T2]], %[[C0]] : i64
+// CHECK:         %[[START2:.*]] = index_cast %[[CLAMPED2]] : i64 to index
+// CHECK:           subtensor %[[ARG0]][%[[START1]], %[[START2]]] [1, 4] [1, 1]
+
+// -----
+
 func @pad_cst(%arg0: tensor<12x4xf32>) -> tensor<18x12xf32> {
   %0 = constant dense<0.0> : tensor<f32>
   %1 = "mhlo.pad"(%arg0, %0) {