Address review comments

openxla · Jun 8, 2023 · 03f1a16 · 03f1a16
1 parent 844f727
commit 03f1a16
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 37 deletions.
diff --git a/docs/spec.md b/docs/spec.md
@@ -5522,7 +5522,7 @@ Produces a `result` tuple from values `val`.
 #### Semantics
 
 Performs element-wise conversion of uniform quantized tensor `operand` to a
-floating point tensor `result` according to the quantization parameters defined
+floating-point tensor `result` according to the quantization parameters defined
 by the `operand` type.
 
 Formally, `result = (operand - zero_point(operand)) * scale(operand)`.
@@ -5541,8 +5541,8 @@ Formally, `result = (operand - zero_point(operand)) * scale(operand)`.
 
 #### Constraints
 
-* (C1) `expressed_type(operand) = element_type(result)`.
-* (C2) `shape(operand) = shape(result)`.
+* (C1) `shape(operand) = shape(result)`.
+* (C2) `expressed_type(result) = element_type(operand)`.
 
 #### Examples
 
@@ -5562,34 +5562,30 @@ quantization parameters defined by the `result` type.
 
 Formally,
 
-* For `element_type(operand)` a floating-point type,
+* If `is_float(operand)`:
   * `rounded_result = round_nearest_even(operand / scale(result))`.
   * `result = clamp(storage_min(result), rounded_result + zero_point(result), storage_max(result))`.
 
-* For `element_type(operand)` a quantized type,
-  * `float_result = (operand - zero_point(operand)) * scale(operand)`.
-  * `rounded_result = round_nearest_even(float_result / scale(result))`.
-  * `result = clamp(storage_min(result), rounded_result + zero_point(result), storage_max(result))`.
+* If `is_quantized(operand)`: performs `dequantize_op_quantize(lambda operand:
+        operand, operand, type(result))`
 
 #### Inputs
 
-| Label | Name      | Type                                        | Constraints      |
-|-------|-----------|---------------------------------------------|------------------|
-| (I1)  | `operand` | tensor of floating-point or quantized  type | (C1), (C2), (C3) |
+| Label | Name      | Type                                        | Constraints |
+|-------|-----------|---------------------------------------------|-------------|
+| (I1)  | `operand` | tensor of floating-point or quantized  type | (C1), (C2)  |
 
 #### Outputs
 
-| Name     | Type             | Constraints      |
-|----------|------------------|------------------|
-| `result` | quantized tensor | (C1), (C2), (C3) |
+| Name     | Type             | Constraints |
+|----------|------------------|-------------|
+| `result` | quantized tensor | (C1), (C2)  |
 
 #### Constraints
 
-* (C1) If `element_type(operand)` is a floating-point type,
-  * `element_type(operand) = expressed_type(result)`.
-* (C2) If `element_type(operand)` is a quantized type,
-  * `expressed_type(operand) = expressed_type(result)`.
-* (C3) `shape(operand) = shape(result)`.
+* (C1) `shape(operand) = shape(result)`.
+* (C2) `expressed_type(result) = is_float(operand) ? element_type(operand) :
+  expressed_type(operand)`.
 
 #### Examples
 
@@ -5599,8 +5595,8 @@ Formally,
 // %result: [10, 10]
 
 // %operand: [10, 10]
-%result = "stablehlo.uniform_quantize"(%operand) : (tensor<2x!quant.uniform<i8<-128:127>:f32:0, {0.1:-30, 0.5:-20}>>) -> tensor<2x!quant.uniform<i8<-128:127>:f32:0, {0.1:-30, 0.5:-20}>>
-// %result: [10, 10]
+%result = "stablehlo.uniform_quantize"(%operand) : (tensor<2x!quant.uniform<i8<-128:127>:f32:0, {0.1:-30, 0.5:-20}>>) -> tensor<2x!quant.uniform<i8<-128:127>:f32:0, {0.1:-20, 0.2:-30}>>
+// %result: [20, 45]
 ```
 
 ### while

diff --git a/docs/status.md b/docs/status.md
@@ -153,7 +153,7 @@ one of the following tracking labels.
 | triangular_solve         | yes           | revisit      | yes            | no              | no          |
 | tuple                    | yes           | yes          | yes            | yes             | no          |
 | unary_einsum             | no            | revisit      | no             | yes             | no          |
-| uniform_dequantize       | no            | yes\*        | yes\*          | yes             | no          |
-| uniform_quantize         | no            | yes\*        | infeasible     | yes             | no          |
+| uniform_dequantize       | yes           | yes          | yes            | yes             | no          |
+| uniform_quantize         | yes           | no           | infeasible     | yes             | no          |
 | while                    | yes           | revisit      | yes            | revisit         | yes         |
 | xor                      | yes           | yes          | yes            | yes             | yes         |
diff --git a/stablehlo/dialect/StablehloOps.td b/stablehlo/dialect/StablehloOps.td
@@ -3055,16 +3055,13 @@ def StableHLO_UniformQuantizeOp : StableHLO_UnaryElementwiseOp<"uniform_quantize
       HLO_QuantizedIntTensor> {
   let summary = "UniformQuantize operation";
   let description = [{
-    This operation is a work in progress, so it is not yet included in
-    the StableHLO specification: https://github.com/openxla/stablehlo/issues/588.
-
-    Informally, this operation converts floating point tensors or uniform
-    quantized tensors to uniform quantized tensors according to the quantization
-    parameters defined by the result type.
+    Performs element-wise conversion of floating-point tensor or uniform
+    quantized tensor `operand` to a uniform quantized tensor `result`
+    according to the quantization parameters defined by the `result` type.
 
     Example:
     ```mlir
-    %result = stablehlo.uniform_quantize %operand : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>
+    %result = "stablehlo.uniform_quantize"(%operand) : (tensor<2xf32>) -> tensor<2x!quant.uniform<i8<-128:127>:f32:0, {0.1:-30, 0.5:-20}>>
     ```
   }];
 }
@@ -3073,16 +3070,13 @@ def StableHLO_UniformDequantizeOp : StableHLO_UnaryElementwiseOp<"uniform_dequan
       [InferTensorType, Pure], HLO_QuantizedIntTensor, TensorOf<[F32, BF16]>> {
   let summary = "UniformDequantize operation";
   let description = [{
-    This operation is a work in progress, so it is not yet included in
-    the StableHLO specification: https://github.com/openxla/stablehlo/issues/588.
-
-    Informally, this operation converts uniform quantized tensors to floating
-    point tensors according to the quantization parameters defined by the
-    operand type.
+    Performs element-wise conversion of uniform quantized tensor `operand` to a
+    floating-point tensor `result` according to the quantization parameters
+    defined by the `operand` type.
 
     Example:
     ```mlir
-    %result = stablehlo.uniform_dequantize %operand : (tensor<16x16x!quant.uniform<i8:f32, 34.0:16>>) -> tensor<16x16xf32>
+    %result = "stablehlo.uniform_dequantize"(%operand) : (tensor<2x!quant.uniform<i8<-128:127>:f32:0, {0.1:-30, 0.5:-20}>>) -> tensor<2xf32>
     ```
   }];
 }