[mlir][tosa] Handle rescale case where shift > 63

[rsuderman]

Summary

It is possible for the shift value to exceed the number of bits. In these
cases we can just multiply by zero. This is relatively rare occurence but
should be handled.

[rsuderman]

Fixed test.

[mehdi_amini]

mlir/lib/Dialect/Tosa/Utils/QuantUtils.cpp

85

Right shifting by 32-bits is invalid, should you limit to 31 in the min expression?
(flagged by coverity)

[mehdi_amini]

mlir/lib/Dialect/Tosa/Utils/QuantUtils.cpp

50

(Same here)

[mehdi_amini]

mlir/lib/Dialect/Tosa/Utils/QuantUtils.cpp

85

Ping?

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

LogicalResult matchAndRewrite(tosa::RescaleOp op,
 
     // The shift and multiplier values.
     SmallVector<int32_t> multiplierValues;
     getValuesFromIntArrayAttribute(op.multiplier(), multiplierValues);
 
     SmallVector<int8_t> shiftValues;
     getValuesFromIntArrayAttribute(op.shift(), shiftValues);
 
+    // If we shift by more than the bitwidth, this just sets to 0.
+    for (int i = 0, s = multiplierValues.size(); i < s; i++) {
+      if (shiftValues[i] > 63) {
+        shiftValues[i] = 0;
+        multiplierValues[i] = 0;
+      }
+    }
+
     // Double round only occurs if shift is greater than 31, check that this
     // is ever true.
     bool doubleRound =
         op.double_round() &&
         llvm::any_of(shiftValues, [](int32_t v) { return v > 31; });
 
     SmallVector<AffineMap> indexingMaps = {
         rewriter.getMultiDimIdentityMap(rank)};

mlir/lib/Dialect/Tosa/Utils/QuantUtils.cpp

static void computeMultiplierAndShiftTosaScale16(double scale,
   // TOSA expects right shift to be positive and embed (1 << 15) into right
   // shift bits.
   shift = (-shift) + 15;
 
   assert(shiftedM <= std::numeric_limits<int32_t>::max() &&
          "Shifted mantissa exceeds 32-bit signed output type");
 
   multiplier = static_cast<int32_t>(shiftedM);
+
+  // Shifting tops out at 63 bits. Right shift to make 63 bits the max.
+  if (shift > 63) {
+    // Shifting the multiplier by more than 32-bits is unnecessary.
+    multiplier = multiplier >> std::min<int32_t>(32, shift - 63);
+    shift = 63;
+  }
 }
 
 /// From a scale value, generates multiplier and shift values where
 /// mantissa is in [-1.0,-0.5] or [0.5, 1.0] such that
 /// multiplier = mantissa*2^shift for 32-bit scaling.
 static void computeMultiplierAndShiftTosaScale32(double scale,
                                                  int32_t &multiplier,
                                                  int32_t &shift) {
static void computeMultiplierAndShiftTosaScale32(double scale,
   // TOSA expects right shift to be positive, and embed (1 << 31) into right
   // shift bits.
   shift = (-shift) + 31;
 
   assert(shiftedM <= std::numeric_limits<int32_t>::max() &&
          "Shifted mantissa exceeds 32-bit signed output type");
 
   multiplier = static_cast<int32_t>(shiftedM);
+
+  // Shifting tops out at 63 bits. Right shift to make 63 bits the max.
+  if (shift > 63) {
+    // Shifting the multiplier by more than 32-bits is unnecessary.
+    multiplier = multiplier >> std::min<int32_t>(32, shift - 63);
+    shift = 63;
+  }
 }
 
 /// Generates a quantized multiplier/shift from double.
 void mlir::tosa::computeMultiplierAndShift(double scale, int32_t &multiplier,
                                            int32_t &shift, int32_t scaleWidth) {
 
   switch (scaleWidth) {
   case 16:

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

func @rescale_ui8(%arg0 : tensor<2xui8>) -> () {
   return
 }
 
 // -----
 
 // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL: @rescale_per_channel
-func @rescale_per_channel(%arg0 : tensor<2xi8>) -> (tensor<2xi8>) {
-  // CHECK: [[MULTIPLIERS:%.+]] = arith.constant dense<[42, 43]>
-  // CHECK: [[SHIFTS:%.+]] = arith.constant dense<[14, 15]>
-  // CHECK: [[INIT:%.+]] = linalg.init_tensor [2]
-  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0, [[MULTIPLIERS]], [[SHIFTS]] : tensor<2xi8>, tensor<2xi32>, tensor<2xi8>) outs([[INIT]] : tensor<2xi8>)
+func @rescale_per_channel(%arg0 : tensor<3xi8>) -> (tensor<3xi8>) {
+  // CHECK: [[MULTIPLIERS:%.+]] = arith.constant dense<[42, 43, 0]>
+  // CHECK: [[SHIFTS:%.+]] = arith.constant dense<[14, 15, 0]>
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [3]
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0, [[MULTIPLIERS]], [[SHIFTS]] : tensor<3xi8>, tensor<3xi32>, tensor<3xi8>) outs([[INIT]] : tensor<3xi8>)
   // CHECK: ^bb0([[IN:%.+]]: i8, [[MULTIPLIER:%.+]]: i32, [[SHIFT:%.+]]: i8, [[UNUSED:%.+]]: i8):
   // CHECK: [[C243:%.+]] = arith.constant 243
   // CHECK: [[C252:%.+]] = arith.constant 252
 
   // CHECK-DAG: [[IN32:%.+]] = arith.extsi [[IN]]
   // CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C243]]
   // CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[MULTIPLIER]], [[SHIFT]]) {double_round = false}
   // CHECK-DAG: [[SCALED_ZEROED:%.+]] = arith.addi [[SCALED]], [[C252]]
   // CHECK-DAG: [[CMIN:%.+]] = arith.constant -128
   // CHECK-DAG: [[CMAX:%.+]] = arith.constant 127
   // CHECK-DAG: [[MINLT:%.+]] = arith.cmpi slt, [[SCALED_ZEROED]], [[CMIN]]
   // CHECK-DAG: [[MAXLT:%.+]] = arith.cmpi slt, [[CMAX]], [[SCALED_ZEROED]]
   // CHECK-DAG: [[LOWER:%.+]] = select [[MINLT]], [[CMIN]], [[SCALED_ZEROED]]
   // CHECK-DAG: [[BOUNDED:%.+]] = select [[MAXLT]], [[CMAX]], [[LOWER]]
   // CHECK-DAG: [[TRUNC:%.+]] = arith.trunci [[BOUNDED]]
   // CHECK-DAG: linalg.yield [[TRUNC]]
-  %0 = "tosa.rescale"(%arg0) {input_zp = 243 : i32, output_zp = 252 : i32, multiplier = [42 : i32, 43 : i32], shift = [14 : i32, 15 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<2xi8>)  -> (tensor<2xi8>)
+  %0 = "tosa.rescale"(%arg0) {input_zp = 243 : i32, output_zp = 252 : i32, multiplier = [42 : i32, 43 : i32, 44 : i32], shift = [14 : i32, 15 : i32, 64 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<3xi8>)  -> (tensor<3xi8>)
 
   // CHECK: return [[GENERIC]]
-  return %0 : tensor<2xi8>
+  return %0 : tensor<3xi8>
 }
 
 // -----
 
 // CHECK-LABEL: @rescaleDoubleRound
 func @rescaleDoubleRound(%arg0 : tensor<2xi8>) -> (tensor<2xi8>) {
   // CHECK: linalg.generic
   // CHECK: "tosa.apply_scale"