Support Dynamic Reshape in StableHLO Quantizer

PiperOrigin-RevId: 609759860
tensorflow · Feb 23, 2024 · 9622017 · 9622017
1 parent 1acfa49
commit 9622017
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/mhlo/transforms/mhlo_quant_legalize_to_int/mhlo_quant_legalize_to_int.cc b/mhlo/transforms/mhlo_quant_legalize_to_int/mhlo_quant_legalize_to_int.cc
@@ -1278,9 +1278,10 @@ class ConvertGenericOp : public ConversionPattern {
       Operation *op, ArrayRef<Value> operands,
       ConversionPatternRewriter &rewriter) const override {
     // This pattern only handle selected ops.
-    if (!isa<mhlo::BroadcastInDimOp, mhlo::ConcatenateOp, mhlo::ConstantOp,
-             mhlo::BitcastConvertOp, mhlo::GatherOp, mhlo::MaxOp, mhlo::MinOp,
-             mhlo::PadOp, mhlo::ReduceWindowOp, mhlo::ReshapeOp, mhlo::ReturnOp,
+    if (!isa<mhlo::BitcastConvertOp, mhlo::BroadcastInDimOp,
+             mhlo::ConcatenateOp, mhlo::ConstantOp, mhlo::DynamicReshapeOp,
+             mhlo::GatherOp, mhlo::MaxOp, mhlo::MinOp, mhlo::PadOp,
+             mhlo::ReduceWindowOp, mhlo::ReshapeOp, mhlo::ReturnOp,
              mhlo::SelectOp, mhlo::SliceOp, mhlo::TransposeOp,
              mhlo::GetDimensionSizeOp, mhlo::DynamicBroadcastInDimOp>(op)) {
       return failure();

diff --git a/...alect/mhlo/convert-mhlo-quant-to-int.mlir → ...lect/mhlo/mhlo-quant-legalize-to-int.mlir b/...alect/mhlo/convert-mhlo-quant-to-int.mlir → ...lect/mhlo/mhlo-quant-legalize-to-int.mlir
@@ -2026,6 +2026,21 @@ func.func @reshape(
 
 // -----
 
+// CHECK-LABEL: func @dynamic_reshape
+func.func @dynamic_reshape(
+    %arg0: tensor<?x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>,
+    %arg1: tensor<2xi32>
+  ) -> tensor<?x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>> {
+  // CHECK: mhlo.dynamic_reshape
+  // CHECK-SAME: (tensor<?x3xi8>, tensor<2xi32>) -> tensor<?x1xi8>
+  %0 = "mhlo.dynamic_reshape"(%arg0, %arg1) : (
+    tensor<?x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>, tensor<2xi32>
+  ) -> tensor<?x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
+  return %0 : tensor<?x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
+}
+
+// -----
+
 // CHECK-LABEL: func @select
 func.func @select(
     %arg0: tensor<1x3xi1>,