[Mlir-commits] [clang] [llvm] [mlir] [clang][NVPTX] Add support for pzo in f32 to f16/bf16 conversions (PR #214667)

[Srinivasa Ravi]

================
@@ -2280,68 +2280,81 @@ def : Pat<(int_nvvm_ui2f_rz i32:$a), (CVT_f32_u32 $a, CvtRZ)>;
 def : Pat<(int_nvvm_ui2f_rm i32:$a), (CVT_f32_u32 $a, CvtRM)>;
 def : Pat<(int_nvvm_ui2f_rp i32:$a), (CVT_f32_u32 $a, CvtRP)>;
 
-def : Pat<(int_nvvm_ff2bf16x2_rn f32:$a, f32:$b),      (CVT_bf16x2_f32 $a, $b, CvtRN)>;
-def : Pat<(int_nvvm_ff2bf16x2_rn_relu f32:$a, f32:$b), (CVT_bf16x2_f32 $a, $b, CvtRN_RELU)>;
-def : Pat<(int_nvvm_ff2bf16x2_rz f32:$a, f32:$b),      (CVT_bf16x2_f32 $a, $b, CvtRZ)>;
-def : Pat<(int_nvvm_ff2bf16x2_rz_relu f32:$a, f32:$b), (CVT_bf16x2_f32 $a, $b, CvtRZ_RELU)>;
-let Predicates = [PTX81, SM80] in {
-  def : Pat<(int_nvvm_ff2bf16x2_rn_satfinite f32:$a, f32:$b), (CVT_bf16x2_f32_sf $a, $b, CvtRN)>;
-  def : Pat<(int_nvvm_ff2bf16x2_rn_relu_satfinite f32:$a, f32:$b), (CVT_bf16x2_f32_sf $a, $b, CvtRN_RELU)>;
-  def : Pat<(int_nvvm_ff2bf16x2_rz_satfinite f32:$a, f32:$b), (CVT_bf16x2_f32_sf $a, $b, CvtRZ)>;
-  def : Pat<(int_nvvm_ff2bf16x2_rz_relu_satfinite f32:$a, f32:$b), (CVT_bf16x2_f32_sf $a, $b, CvtRZ_RELU)>;
-}
-let Predicates = [hasConvertWithStochasticRounding] in {
-def : Pat<(int_nvvm_ff2bf16x2_rs f32:$a, f32:$b, i32:$c),
-          (CVT_bf16x2_f32_rs $a, $b, $c, CvtRS)>;
-def : Pat<(int_nvvm_ff2bf16x2_rs_relu f32:$a, f32:$b, i32:$c),
-          (CVT_bf16x2_f32_rs $a, $b, $c, CvtRS_RELU)>;
-def : Pat<(int_nvvm_ff2bf16x2_rs_satfinite f32:$a, f32:$b, i32:$c), 
-          (CVT_bf16x2_f32_rs_sf $a, $b, $c, CvtRS)>;
-def : Pat<(int_nvvm_ff2bf16x2_rs_relu_satfinite f32:$a, f32:$b, i32:$c),  
-          (CVT_bf16x2_f32_rs_sf $a, $b, $c, CvtRS_RELU)>;
-}
-
-def : Pat<(int_nvvm_ff2f16x2_rn f32:$a, f32:$b),      (CVT_f16x2_f32 $a, $b, CvtRN)>;
-def : Pat<(int_nvvm_ff2f16x2_rn_relu f32:$a, f32:$b), (CVT_f16x2_f32 $a, $b, CvtRN_RELU)>;
-def : Pat<(int_nvvm_ff2f16x2_rz f32:$a, f32:$b),      (CVT_f16x2_f32 $a, $b, CvtRZ)>;
-def : Pat<(int_nvvm_ff2f16x2_rz_relu f32:$a, f32:$b), (CVT_f16x2_f32 $a, $b, CvtRZ_RELU)>;
-let Predicates = [PTX81, SM80] in {
-  def : Pat<(int_nvvm_ff2f16x2_rn_satfinite f32:$a, f32:$b), (CVT_f16x2_f32_sf $a, $b, CvtRN)>;
-  def : Pat<(int_nvvm_ff2f16x2_rn_relu_satfinite f32:$a, f32:$b), (CVT_f16x2_f32_sf $a, $b, CvtRN_RELU)>;
-  def : Pat<(int_nvvm_ff2f16x2_rz_satfinite f32:$a, f32:$b), (CVT_f16x2_f32_sf $a, $b, CvtRZ)>;
-  def : Pat<(int_nvvm_ff2f16x2_rz_relu_satfinite f32:$a, f32:$b), (CVT_f16x2_f32_sf $a, $b, CvtRZ_RELU)>;
+foreach rnd = ["rn", "rz"] in {
+  foreach relu = ["", "_relu"] in {
+    defvar BaseMode = !cast<PatLeaf>("Cvt" # !toupper(rnd # relu));
+
+    foreach pzo = [0, 1] in {
+      // An i1 immediate of one is matched as the sign-extended value -1.
+      defvar PZO = !if(pzo, -1, 0);
+      defvar Mode = CvtModeWithPZO<BaseMode, pzo>.Value;
+      defvar PZOPreds =
+          !if(pzo, [hasConvertWithPZOSupport]<Predicate>, []<Predicate>);
+
+      defvar FF2BF16X2 = !cast<Intrinsic>("int_nvvm_ff2bf16x2_" # rnd # relu);
+      defvar FF2F16X2 = !cast<Intrinsic>("int_nvvm_ff2f16x2_" # rnd # relu);
+      defvar F2BF16 = !cast<Intrinsic>("int_nvvm_f2bf16_" # rnd # relu);
+      defvar F2F16 = !cast<Intrinsic>("int_nvvm_f2f16_" # rnd # relu);
+
+      def : Pat<(FF2BF16X2 f32:$a, f32:$b, PZO),
+                (CVT_bf16x2_f32 $a, $b, Mode)>, Requires<PZOPreds>;
+      def : Pat<(FF2F16X2 f32:$a, f32:$b, PZO),
+                (CVT_f16x2_f32 $a, $b, Mode)>, Requires<PZOPreds>;
+      def : Pat<(F2BF16 f32:$a, PZO), (CVT_bf16_f32 $a, Mode)>,
+            Requires<PZOPreds>;
+      def : Pat<(F2F16 f32:$a, PZO), (CVT_f16_f32 $a, Mode)>,
+            Requires<PZOPreds>;
----------------
Wolfram70 wrote:

It seems we usually don't name these pattern matching records and keep them anonymous, so I'm not sure what a good name for these could be, since they also don't seem to be first-class entities.

https://github.com/llvm/llvm-project/pull/214667