Address review feedback

openxla · May 2, 2023 · 273cef0 · 273cef0
1 parent 42309e1
commit 273cef0
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 1 deletion.
diff --git a/docs/spec.md b/docs/spec.md
@@ -250,7 +250,8 @@ values of type `tensor<T>`).
   * `f8E4M3FNUZ` and `f8E5M2FNUZ` types corresponding to the `E4M3` and `E5M2`
     encodings of the FP8 formats described in
     [8-bit Numerical Formats for Deep Neural Networks](https://arxiv.org/abs/2206.02915).
-  * `f8E4M3B11FNUZ` type corresponding to the `E4M3` encoding of the FP8 formats described in
+  * `f8E4M3B11FNUZ` type corresponding to the `E4M3` encoding of the FP8 formats
+    described in
     [Hybrid 8-bit Floating Point (HFP8) Training and Inference for Deep Neural Networks](https://proceedings.neurips.cc/paper_files/paper/2019/file/65fc9fb4897a89789352e211ca2d398f-Paper.pdf).
   * `bf16` type corresponding to the `bfloat16` format described in
     [BFloat16: The secret to high performance on Cloud TPUs](https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus).

diff --git a/stablehlo/tests/stablehlo_legalize_to_vhlo.mlir b/stablehlo/tests/stablehlo_legalize_to_vhlo.mlir
@@ -2159,6 +2159,13 @@ func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>)
 }
 // CHECK-LABEL: "type_f8E4M3FNUZ"
 
+func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
+  // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
+  func.return %0 : tensor<f8E4M3B11FNUZ>
+}
+// CHECK-LABEL: "type_f8E4M3B11FNUZ"
+
 func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
   // CHECK: "vhlo.add_v1"(%arg0, %arg1) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>

diff --git a/stablehlo/tests/vhlo_to_version_downgrade_invalid.0_10_0.mlir b/stablehlo/tests/vhlo_to_version_downgrade_invalid.0_10_0.mlir
@@ -0,0 +1,7 @@
+// RUN: stablehlo-opt --stablehlo-legalize-to-vhlo --vhlo-to-version='target=0.10.0' --verify-diagnostics --split-input-file %s
+
+// expected-error @+1 {{failed to legalize operation 'vhlo.func_v1' that was explicitly marked illegal}}
+func.func @type_fp8_E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
+  %0 = stablehlo.add %arg0, %arg0 : tensor<f8E4M3B11FNUZ>
+  func.return %0 : tensor<f8E4M3B11FNUZ>
+}