Adapt Stablehlo to recent upstream Tosa changes.

akuegel · TensorFlow MLIR Team · commit 8a47e90ae4cf · 2025-03-21T03:46:09.000-07:00
Changes needed for llvm/llvm-project#130337, llvm/llvm-project#129720, and llvm/llvm-project#130340 PiperOrigin-RevId: 739122578
diff --git a/stablehlo/BUILD b/stablehlo/BUILD
@@ -1542,8 +1542,6 @@ glob_lit_tests(
     driver = "@llvm-project//mlir:run_lit.sh",
     exclude = [
         # TODO: remove the following excludes once #2751 is fixed.
-        "stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir",
-        "stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir",
         "stablehlo/conversions/tosa/tests/rescale_interpreter.mlir",
         "stablehlo/conversions/tosa/tests/binary.mlir",
         "stablehlo/conversions/tosa/tests/nullary.mlir",
@@ -1638,9 +1636,8 @@ cc_library(
     srcs = [
         "stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.cpp",
         "stablehlo/conversions/tosa/transforms/StablehloPrepareForTosa.cpp",
-        # TODO: un-comment the following once #2751 is fixed
-        # "stablehlo/conversions/tosa/transforms/StablehloQuantLegalizeToTosaRescale.cpp",
-        # "stablehlo/conversions/tosa/transforms/TosaRescaleLegalizeToStablehlo.cpp",
+        "stablehlo/conversions/tosa/transforms/StablehloQuantLegalizeToTosaRescale.cpp",
+        "stablehlo/conversions/tosa/transforms/TosaRescaleLegalizeToStablehlo.cpp",
     ],
     hdrs = [
         "stablehlo/conversions/tosa/transforms/Passes.h",
diff --git a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
@@ -1577,9 +1577,8 @@ cc_library(
     srcs = [
         "stablehlo/conversions/tosa/transforms/StablehloLegalizeToTosa.cpp",
         "stablehlo/conversions/tosa/transforms/StablehloPrepareForTosa.cpp",
-        # TODO: un-comment the following once #2751 is fixed
-        # "stablehlo/conversions/tosa/transforms/StablehloQuantLegalizeToTosaRescale.cpp",
-        # "stablehlo/conversions/tosa/transforms/TosaRescaleLegalizeToStablehlo.cpp",
+        "stablehlo/conversions/tosa/transforms/StablehloQuantLegalizeToTosaRescale.cpp",
+        "stablehlo/conversions/tosa/transforms/TosaRescaleLegalizeToStablehlo.cpp",
     ],
     hdrs = [
         "stablehlo/conversions/tosa/transforms/Passes.h",
diff --git a/stablehlo/docs/generated/stablehlo_passes.md b/stablehlo/docs/generated/stablehlo_passes.md
@@ -351,3 +351,17 @@ _Convert between versions of VHLO._
 ```
 -target : The target version. Must be a version of the form #.#.# .
 ```
+
+### `-stablehlo-quant-legalize-to-tosa-rescale`
+
+_Legalize StableHLO Quantized operations to TOSA rescale operations_
+
+This pass rewrites StableHLO quantized operations to integer operations
+by inserting TOSA rescale operations at the inputs and outputs of the
+integer operations.
+
+### `-tosa-rescale-legalize-to-stablehlo`
+
+_Legalize TOSA rescales to StableHlo primitive math operations_
+
+This pass rewrites TOSA rescale operations to StableHLO primitive math operations.
diff --git a/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel b/stablehlo/stablehlo/conversions/tosa/tests/BUILD.bazel
@@ -59,8 +59,6 @@ RUNFILES_DIR = LITE_CFG_PY.parents[4].absolute().as_posix()''',
     for src in glob(
         ["**/*.mlir"],
         exclude = [
-            "legalize_quant_ops_to_tosa_rescale.mlir",
-            "legalize_tosa_rescale_to_stablehlo.mlir",
             "rescale_interpreter.mlir",
             "binary.mlir",
             "nullary.mlir",
diff --git a/stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir b/stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir
diff --git a/stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir b/stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir
@@ -3,8 +3,12 @@
 // -----
 // CHECK-LABEL: @rescale
 func.func @rescale(%arg0 : tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>) -> tensor<2x2xi32> {
-  %0 = tosa.rescale %arg0 {double_round = false, input_zp = -1 : i32, multiplier = array<i32: 1431655765>, input_unsigned = false, output_unsigned = false, output_zp = 0 : i32, per_channel = false, scale32 = true, shift = array<i8: 13>} :
-            (tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>) -> tensor<2x2xi32>
+  %multiplier = "tosa.const"() {values = dense<1431655765> : tensor<1xi32> } : () -> tensor<1xi32>
+  %shift = "tosa.const"() {values = dense<13> : tensor<1xi8> } : () -> tensor<1xi8>
+  %input_zp = "tosa.const"() {values = dense<-1> : tensor<1xi8>} : () -> tensor<1xi8>
+  %output_zp = "tosa.const"() {values = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
+  %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp {rounding_mode = "SINGLE_ROUND", input_unsigned = false, output_unsigned = false, per_channel = false, scale32 = true} :
+            (tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>, tensor<1xi32>, tensor<1xi8>, tensor<1xi8>, tensor<1xi32>) -> tensor<2x2xi32>
 
   // convert input quantized type to storage type
   // CHECK-DAG: %[[arg:.+]] = stablehlo.bitcast_convert %arg0 : (tensor<2x2x!quant.uniform<i8:f32, 2.500000e-02:-1>>) -> tensor<2x2xi8>
diff --git a/stablehlo/stablehlo/conversions/tosa/tests/lit.cfg.py b/stablehlo/stablehlo/conversions/tosa/tests/lit.cfg.py
@@ -26,7 +26,12 @@
 config.suffixes = ['.mlir']
 config.test_source_root = os.path.dirname(__file__)
 # TODO: remove the following once #2751 is fixed.
-config.excludes = ['legalize_quant_ops_to_tosa_rescale.mlir', 'legalize_tosa_rescale_to_stablehlo.mlir', 'rescale_interpreter.mlir', 'binary.mlir', 'nullary.mlir', 'unary.mlir']
+config.excludes = [
+    'rescale_interpreter.mlir',
+    'binary.mlir',
+    'nullary.mlir',
+    'unary.mlir',
+]
 
 # Disallow reusing variables across CHECK-LABEL matches.
 # A variable can eschew this (be made "global") by prefixing its name with $.
diff --git a/stablehlo/stablehlo/conversions/tosa/transforms/CMakeLists.txt b/stablehlo/stablehlo/conversions/tosa/transforms/CMakeLists.txt
@@ -25,11 +25,9 @@ add_mlir_pdll_library(StablehloTOSAPDLLPatternsIncGen
 add_mlir_library(StablehloTOSATransforms
   StablehloLegalizeToTosa.cpp
   StablehloPrepareForTosa.cpp
-  # TODO: un-comment the following once #2751 is fixed.
-  # StablehloQuantLegalizeToTosaRescale.cpp
-  # TosaRescaleLegalizeToStablehlo.cpp
+  StablehloQuantLegalizeToTosaRescale.cpp
+  TosaRescaleLegalizeToStablehlo.cpp
 
-  PARTIAL_SOURCES_INTENDED
   DEPENDS
   StablehloTOSATransformsPassIncGen
   StablehloTOSAPDLLPatternsIncGen
diff --git a/stablehlo/stablehlo/conversions/tosa/transforms/Passes.td b/stablehlo/stablehlo/conversions/tosa/transforms/Passes.td
@@ -33,25 +33,24 @@ def StablehloPrepareForTosaPass : Pass<"stablehlo-prepare-for-tosa", "mlir::func
   let dependentDialects = ["::mlir::tosa::TosaDialect"];
 }
 
-// TODO: un-comment the following once #2751 is fixed.
-// def StablehloQuantLegalizeToTosaRescalePass : Pass<"stablehlo-quant-legalize-to-tosa-rescale", "mlir::func::FuncOp"> {
-//   let summary = "Legalize StableHLO Quantized operations to TOSA rescale operations";
-//   let description = [{
-//     This pass rewrites StableHLO quantized operations to integer operations
-//     by inserting TOSA rescale operations at the inputs and outputs of the
-//     integer operations.
-//   }];
-//   let dependentDialects = [
-//     "::mlir::tosa::TosaDialect",
-//   ];
-// }
-//
-// def TosaRescaleLegalizeToStablehloPass : Pass<"tosa-rescale-legalize-to-stablehlo", "mlir::func::FuncOp"> {
-//   let summary = "Legalize TOSA rescales to StableHlo primitive math operations";
-//   let description = [{
-//     This pass rewrites TOSA rescale operations to StableHLO primitive math operations.
-//   }];
-//   let dependentDialects = [
-//     "::mlir::stablehlo::StablehloDialect"
-//   ];
-// }
+def StablehloQuantLegalizeToTosaRescalePass : Pass<"stablehlo-quant-legalize-to-tosa-rescale", "mlir::func::FuncOp"> {
+  let summary = "Legalize StableHLO Quantized operations to TOSA rescale operations";
+  let description = [{
+    This pass rewrites StableHLO quantized operations to integer operations
+    by inserting TOSA rescale operations at the inputs and outputs of the
+    integer operations.
+  }];
+  let dependentDialects = [
+    "::mlir::tosa::TosaDialect",
+  ];
+}
+
+def TosaRescaleLegalizeToStablehloPass : Pass<"tosa-rescale-legalize-to-stablehlo", "mlir::func::FuncOp"> {
+  let summary = "Legalize TOSA rescales to StableHlo primitive math operations";
+  let description = [{
+    This pass rewrites TOSA rescale operations to StableHLO primitive math operations.
+  }];
+  let dependentDialects = [
+    "::mlir::stablehlo::StablehloDialect"
+  ];
+}
diff --git a/stablehlo/stablehlo/conversions/tosa/transforms/StablehloQuantLegalizeToTosaRescale.cpp b/stablehlo/stablehlo/conversions/tosa/transforms/StablehloQuantLegalizeToTosaRescale.cpp
@@ -45,13 +45,24 @@ Value buildRescale(PatternRewriter &rewriter, Location loc,
                    ShapedType outputType, Value inputVal, int32_t multiplier,
                    int32_t shift, int64_t inputZp, int64_t outputZp,
                    bool doubleRound, bool scale32, bool perChannel) {
+  auto multiplierVal = getConstTensorInt<int32_t>(rewriter, loc, {multiplier});
+  auto shiftVal =
+      getConstTensorInt<int8_t>(rewriter, loc, static_cast<int8_t>(shift));
+  auto inputZpVal =
+      createZeroPointTensor(rewriter, loc, inputVal.getType(), inputZp);
+  if (!inputZpVal) {
+    (void)emitError(loc,
+                    "Failed to create input zero point tensor for RescaleOp");
+  }
+  auto outputZpVal = createZeroPointTensor(rewriter, loc, outputType, outputZp);
+  if (!outputZpVal) {
+    (void)emitError(loc,
+                    "Failed to create output zero point tensor for RescaleOp");
+  }
   auto rescale_op = rewriter.create<RescaleOp>(
-      loc, outputType, inputVal,
-      rewriter.getI32IntegerAttr(static_cast<int32_t>(inputZp)),
-      rewriter.getI32IntegerAttr(static_cast<int32_t>(outputZp)),
-      rewriter.getDenseI32ArrayAttr({multiplier}),
-      rewriter.getDenseI8ArrayAttr({static_cast<int8_t>(shift)}),
-      rewriter.getBoolAttr(scale32), rewriter.getBoolAttr(doubleRound),
+      loc, outputType, inputVal, multiplierVal, shiftVal, inputZpVal.value(),
+      outputZpVal.value(), rewriter.getBoolAttr(scale32),
+      rewriter.getStringAttr(doubleRound ? "DOUBLE_ROUND" : "SINGLE_ROUND"),
       rewriter.getBoolAttr(perChannel),
       /*input_unsigned=*/rewriter.getBoolAttr(false),
       /*output_unsigned=*/rewriter.getBoolAttr(false));
diff --git a/stablehlo/stablehlo/conversions/tosa/transforms/TosaRescaleLegalizeToStablehlo.cpp b/stablehlo/stablehlo/conversions/tosa/transforms/TosaRescaleLegalizeToStablehlo.cpp
@@ -65,7 +65,7 @@ LogicalResult ConvertTosaRescaleToStablehlo::matchAndRewrite(
   }
 
   bool scale32 = op.getScale32();
-  bool doubleRound = op.getDoubleRound();
+  bool doubleRound = op.getRoundingMode() == "DOUBLE_ROUND";
   bool perChannel = op.getPerChannel();
 
   if (perChannel || doubleRound || !scale32) {
@@ -106,18 +106,48 @@ LogicalResult ConvertTosaRescaleToStablehlo::matchAndRewrite(
   auto i32Type = inputType.clone(rewriter.getI32Type());
   auto i64Type = inputType.clone(rewriter.getI64Type());
 
-  // construct multiplier, shift constant values from op attrs
-  // for scale32, multiplier is tensor of i32
+  // construct multiplier, shift constant values for scale32, multiplier and
+  // shift are constant tensors of i32 or i8, respectively.
+  DenseElementsAttr multiplierElems;
+  if (!matchPattern(op.getMultiplier(), m_Constant(&multiplierElems))) {
+    return rewriter.notifyMatchFailure(
+        op, "tosa.rescale requires constant multiplier input values");
+  }
+  llvm::SmallVector<int32_t> multiplierValues =
+      llvm::to_vector(multiplierElems.getValues<int32_t>());
   Value multiplier = getStablehloConstantOp(
-      rewriter, loc, DenseElementsAttr::get(i32Type, op.getMultiplier()));
+      rewriter, loc,
+      DenseElementsAttr::get(
+          i32Type, rewriter.getI32IntegerAttr(multiplierValues.front())));
+  DenseElementsAttr shiftElems;
+  if (!matchPattern(op.getShift(), m_Constant(&shiftElems))) {
+    return rewriter.notifyMatchFailure(
+        op, "tosa.rescale requires constant shift input values");
+  }
+  llvm::SmallVector<int8_t> shiftValues =
+      llvm::to_vector(shiftElems.getValues<int8_t>());
   Value shift = getStablehloConstantOp(
-      rewriter, loc, DenseElementsAttr::get(i8Type, op.getShift()));
+      rewriter, loc,
+      DenseElementsAttr::get(i8Type,
+                             rewriter.getI8IntegerAttr(shiftValues.front())));
 
-  // construct inputZp and outputZp from op attrs
+  // construct inputZp and outputZp
+  FailureOr<int64_t> maybeIZp = op.getInputZeroPoint();
+  if (failed(maybeIZp)) {
+    return rewriter.notifyMatchFailure(
+        op, "input zero point cannot be statically determined");
+  }
   Value inputZpI32 = getStablehloConstantOp(
-      rewriter, loc, DenseElementsAttr::get(i32Type, op.getInputZpAttr()));
+      rewriter, loc,
+      DenseElementsAttr::get(i32Type, rewriter.getI32IntegerAttr(*maybeIZp)));
+  FailureOr<int64_t> maybeOZp = op.getOutputZeroPoint();
+  if (failed(maybeOZp)) {
+    return rewriter.notifyMatchFailure(
+        op, "output zero point cannot be statically determined");
+  }
   Value outputZpI32 = getStablehloConstantOp(
-      rewriter, loc, DenseElementsAttr::get(i32Type, op.getOutputZpAttr()));
+      rewriter, loc,
+      DenseElementsAttr::get(i32Type, rewriter.getI32IntegerAttr(*maybeOZp)));
 
   // construct constant 1, min and max tensors
   Value onesI64 = getStablehloConstantOp(rewriter, loc,
