Move Quantization -> Dialect/QuantOps, FxpMathOps -> Dialect/FxpMathOps.

    Adding the additional layer of directory was discussed offline and matches the Target/ tree. The names match the defacto convention we seem to be following where the C++ namespace is ^(.+)Ops/$ matched against the directory name.

    This is in preparation for patching the Quantizer into this tree, which would have been confusing without moving the Quantization dialect to its more proper home. It is left to others to move other dialects if desired.

    Tested:
      ninja check-mlir

--

PiperOrigin-RevId: 248171982

3 files changed, 260 insertions, 0 deletions

diff --git a/mlir/test/Dialect/FxpMathOps/lower-uniform-casts.mlir b/mlir/test/Dialect/FxpMathOps/lower-uniform-casts.mlir
new file mode 100644
index 00000000000..14d98a0babb
--- /dev/null
+++ b/mlir/test/Dialect/FxpMathOps/lower-uniform-casts.mlir
@@ -0,0 +1,64 @@
+// RUN: mlir-opt %s -split-input-file -fxpmath-lower-uniform-casts | FileCheck %s --dump-input=always
+
+// -----
+// CHECK-LABEL: dequantize_per_layer_fixedpoint
+!type_input = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_result = type tensor<4xf32>
+func @dequantize_per_layer_fixedpoint(%arg0 : !type_input) -> !type_result {
+  // CHECK: %cst = constant splat<tensor<4xf32>, 6.250000e-02>
+  // CHECK-NEXT: %0 = "quant.scast"(%arg0) : (tensor<4x!quant.uniform<i8:f32, 6.250000e-02>>) -> tensor<4xi8>
+  // CHECK-NEXT: %1 = "fxpmath.convertis"(%0) : (tensor<4xi8>) -> tensor<4xi32>
+  // CHECK-NEXT: %2 = "fxpmath.convertistof"(%1) : (tensor<4xi32>) -> tensor<4xf32>
+  // CHECK-NEXT: %3 = mulf %2, %cst : tensor<4xf32>
+  // CHECK-NEXT: return %3 : tensor<4xf32>
+  %0 = "quant.dcast"(%arg0) : (!type_input) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// CHECK-LABEL: dequantize_per_layer_affine
+!type_input = type tensor<4x!quant.uniform<i8:f32, 6.25e-2:-36>>
+!type_result = type tensor<4xf32>
+func @dequantize_per_layer_affine(%arg0 : !type_input) -> !type_result {
+  // CHECK: %cst = constant splat<tensor<4xi32>, 36>
+  // CHECK-NEXT: %cst_0 = constant splat<tensor<4xf32>, 6.250000e-02>
+  // CHECK-NEXT: %0 = "quant.scast"(%arg0) : (tensor<4x!quant.uniform<i8:f32, 6.250000e-02:-36>>) -> tensor<4xi8>
+  // CHECK-NEXT: %1 = "fxpmath.convertis"(%0) : (tensor<4xi8>) -> tensor<4xi32>
+  // CHECK-NEXT: %2 = addi %1, %cst : tensor<4xi32>
+  // CHECK-NEXT: %3 = "fxpmath.convertistof"(%2) : (tensor<4xi32>) -> tensor<4xf32>
+  // CHECK-NEXT: %4 = mulf %3, %cst_0 : tensor<4xf32>
+  // CHECK-NEXT: return %4 : tensor<4xf32>
+  %0 = "quant.dcast"(%arg0) : (!type_input) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// CHECK-LABEL: dequantize_per_axis_fixedpoint
+!type_input = type tensor<4x!quant.uniform<i8:f32:0, {6.25e-2,3.26e-2,4.25e-2,1.23e-2}>>
+!type_result = type tensor<4xf32>
+func @dequantize_per_axis_fixedpoint(%arg0 : !type_input) -> !type_result {
+  // expected-warning@+1 {{unimplemented: per-axis uniform dequantization}}
+  %0 = "quant.dcast"(%arg0) : (!type_input) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// CHECK-LABEL: dequantize_per_axis_affine
+!type_input = type tensor<4x!quant.uniform<i8:f32:0, {6.25e-2,3.26e-2,4.25e-2,1.23e-2}>>
+!type_result = type tensor<4xf32>
+func @dequantize_per_axis_affine(%arg0 : !type_input) -> !type_result {
+  // expected-warning@+1 {{unimplemented: per-axis uniform dequantization}}
+  %0 = "quant.dcast"(%arg0) : (!type_input) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// Noop dequantize should be skipped (will be canonicalized away later).
+// CHECK-LABEL: dequantize_noop
+!type_input = type tensor<4x!quant.uniform<i8:f32, 6.25e-2:-36>>
+!type_result = type tensor<4x!quant.uniform<i8:f32, 6.25e-2:-36>>
+func @dequantize_noop(%arg0 : !type_input) -> !type_result {
+  // CHECK: %0 = "quant.dcast"(%arg0)
+  %0 = "quant.dcast"(%arg0) : (!type_input) -> (!type_result)
+  return %0 : !type_result
+}
diff --git a/mlir/test/Dialect/FxpMathOps/lower-uniform-real-math-addew.mlir b/mlir/test/Dialect/FxpMathOps/lower-uniform-real-math-addew.mlir
new file mode 100644
index 00000000000..5f73979bfa1
--- /dev/null
+++ b/mlir/test/Dialect/FxpMathOps/lower-uniform-real-math-addew.mlir
@@ -0,0 +1,102 @@
+// RUN: mlir-opt %s -split-input-file -fxpmath-lower-uniform-real-math -canonicalize | FileCheck %s --dump-input=always
+
+// -----
+// Verify lowering when operands and result have the same fixedpoint scale.
+// CHECK-LABEL: real_addew_fixedpoint_isomorphic
+!type_lhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_rhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_result = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+func @real_addew_fixedpoint_isomorphic(%arg0 : !type_lhs, %arg1: !type_rhs) -> !type_result {
+  // CHECK-NEXT: %0 = "quant.scast"(%arg0) : (tensor<4x!quant.uniform<i8:f32, 6.250000e-02>>) -> tensor<4xi8>
+  // CHECK-NEXT: %1 = "quant.scast"(%arg1) : (tensor<4x!quant.uniform<i8:f32, 6.250000e-02>>) -> tensor<4xi8>
+  // CHECK-NEXT: %2 = "fxpmath.convertis"(%0) : (tensor<4xi8>) -> tensor<4xi16>
+  // CHECK-NEXT: %3 = "fxpmath.convertis"(%1) : (tensor<4xi8>) -> tensor<4xi16>
+  // CHECK-NEXT: %4 = addi %2, %3 : tensor<4xi16>
+  // CHECK-NEXT: %5 = "fxpmath.clampis"(%4) {clamp_max: 127 : i16, clamp_min: -128 : i16} : (tensor<4xi16>) -> tensor<4xi16>
+  // CHECK-NEXT: %6 = "fxpmath.convertis"(%5) : (tensor<4xi16>) -> tensor<4xi8>
+  // CHECK-NEXT: %7 = "quant.scast"(%6) : (tensor<4xi8>) -> tensor<4x!quant.uniform<i8:f32, 6.250000e-02>>
+  // CHECK-NEXT: return %7 : tensor<4x!quant.uniform<i8:f32, 6.250000e-02>>
+  %0 = "fxpmath.real_add_ew"(%arg0, %arg1) : (!type_lhs, !type_rhs) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// Verify lowering when operands and result have the same fixedpoint scale
+// and non-zero zero points.
+// CHECK-LABEL: real_addew_affine_isomorphic
+!type_lhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2:-5>>
+!type_rhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2:-5>>
+!type_result = type tensor<4x!quant.uniform<i8:f32, 6.25e-2:-5>>
+func @real_addew_affine_isomorphic(%arg0 : !type_lhs, %arg1: !type_rhs) -> !type_result {
+  // CHECK-NEXT: %cst = constant splat<tensor<4xi16>, 5>
+  // CHECK-NEXT: %0 = "quant.scast"(%arg0) : (tensor<4x!quant.uniform<i8:f32, 6.250000e-02:-5>>) -> tensor<4xi8>
+  // CHECK-NEXT: %1 = "quant.scast"(%arg1) : (tensor<4x!quant.uniform<i8:f32, 6.250000e-02:-5>>) -> tensor<4xi8>
+  // CHECK-NEXT: %2 = "fxpmath.convertis"(%0) : (tensor<4xi8>) -> tensor<4xi16>
+  // CHECK-NEXT: %3 = "fxpmath.convertis"(%1) : (tensor<4xi8>) -> tensor<4xi16>
+  // CHECK-NEXT: %4 = addi %2, %3 : tensor<4xi16>
+  // CHECK-NEXT: %5 = addi %4, %cst : tensor<4xi16>
+  // CHECK-NEXT: %6 = "fxpmath.clampis"(%5) {clamp_max: 127 : i16, clamp_min: -128 : i16} : (tensor<4xi16>) -> tensor<4xi16>
+  // CHECK-NEXT: %7 = "fxpmath.convertis"(%6) : (tensor<4xi16>) -> tensor<4xi8>
+  // CHECK-NEXT: %8 = "quant.scast"(%7) : (tensor<4xi8>) -> tensor<4x!quant.uniform<i8:f32, 6.250000e-02:-5>>
+  // CHECK-NEXT: return %8 : tensor<4x!quant.uniform<i8:f32, 6.250000e-02:-5>>
+  %0 = "fxpmath.real_add_ew"(%arg0, %arg1) : (!type_lhs, !type_rhs) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// The RHS quant parameters proscribe a range of [-8..8) so an explicit clamp
+// of [-4..4] should result in an integral clamp range of [-64..64].
+// CHECK-LABEL: real_addew_fixedpoint_clamp
+!type_lhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_rhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_result = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+func @real_addew_fixedpoint_clamp(%arg0 : !type_lhs, %arg1: !type_rhs) -> !type_result {
+  // CHECK-NEXT: %0 = "quant.scast"(%arg0) : (tensor<4x!quant.uniform<i8:f32, 6.250000e-02>>) -> tensor<4xi8>
+  // CHECK-NEXT: %1 = "quant.scast"(%arg1) : (tensor<4x!quant.uniform<i8:f32, 6.250000e-02>>) -> tensor<4xi8>
+  // CHECK-NEXT: %2 = "fxpmath.convertis"(%0) : (tensor<4xi8>) -> tensor<4xi16>
+  // CHECK-NEXT: %3 = "fxpmath.convertis"(%1) : (tensor<4xi8>) -> tensor<4xi16>
+  // CHECK-NEXT: %4 = addi %2, %3 : tensor<4xi16>
+  // CHECK-NEXT: %5 = "fxpmath.clampis"(%4) {clamp_max: 64 : i16, clamp_min: -64 : i16} : (tensor<4xi16>) -> tensor<4xi16>
+  // CHECK-NEXT: %6 = "fxpmath.convertis"(%5) : (tensor<4xi16>) -> tensor<4xi8>
+  // CHECK-NEXT: %7 = "quant.scast"(%6) : (tensor<4xi8>) -> tensor<4x!quant.uniform<i8:f32, 6.250000e-02>>
+  // CHECK-NEXT: return %7 : tensor<4x!quant.uniform<i8:f32, 6.250000e-02>>
+  %0 = "fxpmath.real_add_ew"(%arg0, %arg1) { clamp_min:-4.0, clamp_max:4.0 }
+      : (!type_lhs, !type_rhs) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// CHECK-LABEL: real_addew_unquantized_lhs
+// Verifies that leaves as-is for unquantized lhs.
+!type_lhs = type tensor<4xf32>
+!type_rhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_result = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+func @real_addew_unquantized_lhs(%arg0 : !type_lhs, %arg1: !type_rhs) -> !type_result {
+  // CHECK: %0 = "fxpmath.real_add_ew"(%arg0, %arg1)
+  %0 = "fxpmath.real_add_ew"(%arg0, %arg1) : (!type_lhs, !type_rhs) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// CHECK-LABEL: real_addew_unquantized_rhs
+// Verifies that leaves as-is for unquantized rhs.
+!type_lhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_rhs = type tensor<4xf32>
+!type_result = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+func @real_addew_unquantized_rhs(%arg0 : !type_lhs, %arg1: !type_rhs) -> !type_result {
+  // CHECK: %0 = "fxpmath.real_add_ew"(%arg0, %arg1)
+  %0 = "fxpmath.real_add_ew"(%arg0, %arg1) : (!type_lhs, !type_rhs) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// CHECK-LABEL: real_addew_unquantized_result
+// Verifies that leaves as-is for unquantized result.
+!type_lhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_rhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_result = type tensor<4xf32>
+func @real_addew_unquantized_result(%arg0 : !type_lhs, %arg1: !type_rhs) -> !type_result {
+  // CHECK: %0 = "fxpmath.real_add_ew"(%arg0, %arg1)
+  %0 = "fxpmath.real_add_ew"(%arg0, %arg1) : (!type_lhs, !type_rhs) -> (!type_result)
+  return %0 : !type_result
+}
diff --git a/mlir/test/Dialect/FxpMathOps/lower-uniform-real-math-mulew.mlir b/mlir/test/Dialect/FxpMathOps/lower-uniform-real-math-mulew.mlir
new file mode 100644
index 00000000000..edfdc8b6ab7
--- /dev/null
+++ b/mlir/test/Dialect/FxpMathOps/lower-uniform-real-math-mulew.mlir
@@ -0,0 +1,94 @@
+// RUN: mlir-opt %s -split-input-file -fxpmath-lower-uniform-real-math -canonicalize -verify | FileCheck %s --dump-input=always
+
+// -----
+// Verify lowering when operands and result have the same fixedpoint scale.
+// CHECK-LABEL: real_mulew_fixedpoint
+!type_lhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_rhs = type tensor<4x!quant.uniform<i8:f32, 3.875e-2>>
+!type_result = type tensor<4x!quant.uniform<i8:f32, 1.065e-1>>
+func @real_mulew_fixedpoint(%arg0 : !type_lhs, %arg1: !type_rhs) -> !type_result {
+  // CHECK: %0 = "quant.scast"(%arg0) : (tensor<4x!quant.uniform<i8:f32, 6.250000e-02>>) -> tensor<4xi8>
+  // CHECK-NEXT: %1 = "quant.scast"(%arg1) : (tensor<4x!quant.uniform<i8:f32, 3.875000e-02>>) -> tensor<4xi8>
+  // CHECK-NEXT: %2 = "fxpmath.convertis"(%0) : (tensor<4xi8>) -> tensor<4xi32>
+  // CHECK-NEXT: %3 = "fxpmath.convertis"(%1) : (tensor<4xi8>) -> tensor<4xi32>
+  // CHECK-NEXT: %4 = muli %2, %3 : tensor<4xi32>
+  // CHECK-NEXT: %5 = "fxpmath.vs_saturating_rounding_doubling_high_mulis"(%4) {b: 1562722842 : i32} : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK-NEXT: %6 = "fxpmath.rounding_divide_by_potis"(%5) {exponent: 5 : i32} : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK-NEXT: %7 = "fxpmath.clampis"(%6) {clamp_max: 127 : i32, clamp_min: -128 : i32} : (tensor<4xi32>) -> tensor<4xi32>
+  // CHECK-NEXT: %8 = "fxpmath.convertis"(%7) : (tensor<4xi32>) -> tensor<4xi8>
+  // CHECK-NEXT: %9 = "quant.scast"(%8) : (tensor<4xi8>) -> tensor<4x!quant.uniform<i8:f32, 1.065000e-01>>
+  // CHECK-NEXT: return %9 : tensor<4x!quant.uniform<i8:f32, 1.065000e-01>>
+  %0 = "fxpmath.real_mul_ew"(%arg0, %arg1) : (!type_lhs, !type_rhs) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// Verify lowering when operands and result have the same fixedpoint scale
+// and non-zero zero points.
+// CHECK-LABEL: real_mulew_affine_clamp
+!type_lhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2:-3>>
+!type_rhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2:-5>>
+!type_result = type tensor<4x!quant.uniform<i8:f32, 6.25e-2:-9>>
+func @real_mulew_affine_clamp(%arg0 : !type_lhs, %arg1: !type_rhs) -> !type_result {
+  // Just verify that the affine adds/constants and clamps are present.
+  // CHECK: %cst = constant splat<tensor<4xi32>, 3>
+  // CHECK: %cst_0 = constant splat<tensor<4xi32>, 5>
+  // CHECK: %cst_1 = constant splat<tensor<4xi32>, -9>
+  // CHECK: addi %2, %cst : tensor<4xi32>
+  // CHECK: addi %3, %cst_0 : tensor<4xi32>
+  // CHECK: muli %4, %5 : tensor<4xi32>
+  // CHECK: addi %8, %cst_1 : tensor<4xi32>
+  // CHECK: {clamp_max: 55 : i32, clamp_min: -73 : i32}
+  %0 = "fxpmath.real_mul_ew"(%arg0, %arg1) { clamp_min:-4.0, clamp_max:4.0 } : (!type_lhs, !type_rhs) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// CHECK-LABEL: real_mulew_unquantized_lhs
+// Verifies that leaves as-is for unquantized lhs.
+!type_lhs = type tensor<4xf32>
+!type_rhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_result = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+func @real_mulew_unquantized_lhs(%arg0 : !type_lhs, %arg1: !type_rhs) -> !type_result {
+  // CHECK: %0 = "fxpmath.real_mul_ew"(%arg0, %arg1)
+  %0 = "fxpmath.real_mul_ew"(%arg0, %arg1) : (!type_lhs, !type_rhs) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// CHECK-LABEL: real_mulew_unquantized_rhs
+// Verifies that leaves as-is for unquantized rhs.
+!type_lhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_rhs = type tensor<4xf32>
+!type_result = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+func @real_mulew_unquantized_rhs(%arg0 : !type_lhs, %arg1: !type_rhs) -> !type_result {
+  // CHECK: %0 = "fxpmath.real_mul_ew"(%arg0, %arg1)
+  %0 = "fxpmath.real_mul_ew"(%arg0, %arg1) : (!type_lhs, !type_rhs) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// CHECK-LABEL: real_mulew_unquantized_result
+// Verifies that leaves as-is for unquantized result.
+!type_lhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_rhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_result = type tensor<4xf32>
+func @real_mulew_unquantized_result(%arg0 : !type_lhs, %arg1: !type_rhs) -> !type_result {
+  // CHECK: %0 = "fxpmath.real_mul_ew"(%arg0, %arg1)
+  %0 = "fxpmath.real_mul_ew"(%arg0, %arg1) : (!type_lhs, !type_rhs) -> (!type_result)
+  return %0 : !type_result
+}
+
+// -----
+// Verify lowering when operands and result have the same fixedpoint scale.
+// Note that the multiplier = lhs_scale * rhs_scale / result_scale
+//   = 22.740610328638496
+// CHECK-LABEL: real_mulew_multiplier_gt_1
+!type_lhs = type tensor<4x!quant.uniform<i8:f32, 6.25e-2>>
+!type_rhs = type tensor<4x!quant.uniform<i8:f32, 3.875e-2>>
+!type_result = type tensor<4x!quant.uniform<i8:f32, 1.065e-4>>
+func @real_mulew_multiplier_gt_1(%arg0 : !type_lhs, %arg1: !type_rhs) -> !type_result {
+  // expected-warning@+1 {{unimplemented: cannot multiply with multipler > 1.0}}
+  %0 = "fxpmath.real_mul_ew"(%arg0, %arg1) : (!type_lhs, !type_rhs) -> (!type_result)
+  return %0 : !type_result
+}