Repository: bytecodealliance/wasmtime
Pull request #7710 (2023-12-20 01:54:01)

Add `sextend_maybe` and `uextend_maybe` to opt ISLE
Working on something else, I noticed that the `(x < y) != 0` opt was written in a way that only works for C

https://github.com/bytecodealliance/wasmtime/blob/f8c9f6711f6316001801480d3883888734835faf/cranelift/codegen/src/opts/icmp.isle#L16-L25

whereas in Rust, which doesn't have the "usual arithmetic conversions", the result of the comparison stays as `I8` and thus without the extension it didn't match that pattern.

This picks up a suggestion that alex had made in a previous PR (https://github.com/bytecodealliance/wasmtime/pull/7636#discussion_r1416208525) to make something like lowering ISLE's `maybe_uextend` extractor for optimization ISLE.

They're then used to improve the mentioned icmp-of-icmp pattern

https://github.com/bytecodealliance/wasmtime/commit/8ef0b8181ca23b9133a464da6ae10fa32ffdbf15#diff-fa420f878e321dc6579ac911f9ba8820a9381e4c49615584f6e6909f0e329dfcL124-R125

as well as simplify some of the duplication in the `<=>` patterns.



Patch:
diff --git a/RELEASES.md b/RELEASES.md
index 03f2b857f..6ced868ba 100644
--- a/RELEASES.md
+++ b/RELEASES.md
@@ -12,7 +12,7 @@ Unreleased.
 
 ## 16.0.0
 
-Unreleased.
+Released 2023-12-20
 
 ### Added
 
diff --git a/cranelift/codegen/src/isa/aarch64/lower/isle.rs b/cranelift/codegen/src/isa/aarch64/lower/isle.rs
index 91ea3a735..788213824 100644
--- a/cranelift/codegen/src/isa/aarch64/lower/isle.rs
+++ b/cranelift/codegen/src/isa/aarch64/lower/isle.rs
@@ -307,8 +307,8 @@ impl Context for IsleContext<'_, '_, MInst, AArch64Backend> {
                     self.lower_ctx.add_range_fact(
                         rd.to_reg(),
                         64,
-                        u64::from(lower_halfword),
-                        u64::from(lower_halfword),
+                        u64::from(value),
+                        u64::from(value),
                     );
                 }
             } else {
diff --git a/cranelift/codegen/src/opts.rs b/cranelift/codegen/src/opts.rs
index 8796b62e1..2d05057fc 100644
--- a/cranelift/codegen/src/opts.rs
+++ b/cranelift/codegen/src/opts.rs
@@ -5,6 +5,7 @@ use crate::ir::condcodes;
 pub use crate::ir::condcodes::{FloatCC, IntCC};
 use crate::ir::dfg::ValueDef;
 pub use crate::ir::immediates::{Ieee32, Ieee64, Imm64, Offset32, Uimm8, V128Imm};
+use crate::ir::instructions::InstructionFormat;
 pub use crate::ir::types::*;
 pub use crate::ir::{
     dynamic_to_fixed, AtomicRmwOp, BlockCall, Constant, DynamicStackSlot, FuncRef, GlobalValue,
@@ -121,6 +122,65 @@ where
     }
 }
 
+#[derive(Default)]
+pub(crate) struct MaybeUnaryEtorIter<'a, 'b, 'c> {
+    opcode: Option<Opcode>,
+    inner: InstDataEtorIter<'a, 'b, 'c>,
+    fallback: Option<Value>,
+}
+
+impl MaybeUnaryEtorIter<'_, '_, '_> {
+    fn new(opcode: Opcode, value: Value) -> Self {
+        debug_assert_eq!(opcode.format(), InstructionFormat::Unary);
+        Self {
+            opcode: Some(opcode),
+            inner: InstDataEtorIter::new(value),
+            fallback: Some(value),
+        }
+    }
+}
+
+impl<'a, 'b, 'c> ContextIter for MaybeUnaryEtorIter<'a, 'b, 'c>
+where
+    'b: 'a,
+    'c: 'b,
+{
+    type Context = IsleContext<'a, 'b, 'c>;
+    type Output = (Type, Value);
+
+    fn next(&mut self, ctx: &mut IsleContext<'a, 'b, 'c>) -> Option<Self::Output> {
+        debug_assert_ne!(self.opcode, None);
+        while let Some((ty, inst_def)) = self.inner.next(ctx) {
+            let InstructionData::Unary { opcode, arg } = inst_def else {
+                continue;
+            };
+            if Some(opcode) == self.opcode {
+                self.fallback = None;
+                return Some((ty, arg));
+            }
+        }
+
+        self.fallback.take().map(|value| {
+            let ty = generated_code::Context::value_type(ctx, value);
+            (ty, value)
+        })
+    }
+}
+
+impl<'a, 'b, 'c> IntoContextIter for MaybeUnaryEtorIter<'a, 'b, 'c>
+where
+    'b: 'a,
+    'c: 'b,
+{
+    type Context = IsleContext<'a, 'b, 'c>;
+    type Output = (Type, Value);
+    type IntoIter = Self;
+
+    fn into_context_iter(self) -> Self {
+        self
+    }
+}
+
 impl<'a, 'b, 'c> generated_code::Context for IsleContext<'a, 'b, 'c> {
     isle_common_prelude_methods!();
 
@@ -193,4 +253,14 @@ impl<'a, 'b, 'c> generated_code::Context for IsleContext<'a, 'b, 'c> {
         let imm = V128Imm(val.to_le_bytes());
         self.ctx.func.dfg.constants.insert(imm.into())
     }
+
+    type sextend_maybe_etor_returns = MaybeUnaryEtorIter<'a, 'b, 'c>;
+    fn sextend_maybe_etor(&mut self, value: Value, returns: &mut Self::sextend_maybe_etor_returns) {
+        *returns = MaybeUnaryEtorIter::new(Opcode::Sextend, value);
+    }
+
+    type uextend_maybe_etor_returns = MaybeUnaryEtorIter<'a, 'b, 'c>;
+    fn uextend_maybe_etor(&mut self, value: Value, returns: &mut Self::uextend_maybe_etor_returns) {
+        *returns = MaybeUnaryEtorIter::new(Opcode::Uextend, value);
+    }
 }
diff --git a/cranelift/codegen/src/opts/icmp.isle b/cranelift/codegen/src/opts/icmp.isle
index 65429870e..fc1d0a157 100644
--- a/cranelift/codegen/src/opts/icmp.isle
+++ b/cranelift/codegen/src/opts/icmp.isle
@@ -15,12 +15,12 @@
 
 ;; Optimize icmp-of-icmp.
 (rule (simplify (ne ty
-                      (uextend _ inner @ (icmp ty _ _ _))
+                      (uextend_maybe _ inner @ (icmp ty _ _ _))
                       (iconst_u _ 0)))
       (subsume inner))
 
 (rule (simplify (eq ty
-                      (uextend _ (icmp ty cc x y))
+                      (uextend_maybe _ (icmp ty cc x y))
                       (iconst_u _ 0)))
       (subsume (icmp ty (intcc_complement cc) x y)))
 
diff --git a/cranelift/codegen/src/opts/selects.isle b/cranelift/codegen/src/opts/selects.isle
index ec38be739..d0e06f855 100644
--- a/cranelift/codegen/src/opts/selects.isle
+++ b/cranelift/codegen/src/opts/selects.isle
@@ -15,7 +15,7 @@
 (rule (simplify (select ty cmp@(icmp _ cc x y)
                         (iconst_u _ 1)
                         (iconst_u _ 0)))
-    (uextend_from_i8 ty cmp))
+    (uextend_maybe ty cmp))
 ;; if icmp(x, y) { -1 } else { 0 } => uextend(icmp(x, y))
 (rule (simplify (select ty cmp@(icmp _ cc x y)
                         (iconst_s _ -1)
@@ -75,9 +75,3 @@
 (rule (simplify (bor (ty_vec128 ty) (band ty (bnot ty c) y) (band ty x c))) (bitselect ty c x y))
 (rule (simplify (bor (ty_vec128 ty) (band ty y (bnot ty c)) (band ty c x))) (bitselect ty c x y))
 (rule (simplify (bor (ty_vec128 ty) (band ty y (bnot ty c)) (band ty x c))) (bitselect ty c x y))
-
-;; extend from i8 to i8 is invalid CLIF, so this allows fixing that in the output
-;; rather than needing to duplicate rules for the different width categories
-(decl uextend_from_i8 (Type Value) Value)
-(rule 0 (uextend_from_i8 ty val) (uextend ty val))
-(rule 1 (uextend_from_i8 $I8 val) val)
diff --git a/cranelift/codegen/src/opts/spaceship.isle b/cranelift/codegen/src/opts/spaceship.isle
index 69dd8f9cb..a834aa49e 100644
--- a/cranelift/codegen/src/opts/spaceship.isle
+++ b/cranelift/codegen/src/opts/spaceship.isle
@@ -12,22 +12,14 @@
 ;; x < y ? -1 : x != y ? +1 : 0
 (rule (simplify (select ty (ult rty x y)
                            (iconst_s ty -1)
-                           (ne rty x y)))
-      (sextend_from_i8 ty (spaceship_u rty x y)))
-(rule (simplify (select ty (ult rty x y)
-                           (iconst_s ty -1)
-                           (uextend ty (ne rty x y))))
-      (sextend_from_i8 ty (spaceship_u rty x y)))
+                           (uextend_maybe ty (ne rty x y))))
+      (sextend_maybe ty (spaceship_u rty x y)))
 ;; x < y ? -1 : x <= y ? 0 : +1
 ;; x < y ? -1 : x > y ? +1 : 0
 (rule (simplify (select ty (ult rty x y)
                            (iconst_s ty -1)
-                           (ugt rty x y)))
-      (sextend_from_i8 ty (spaceship_u rty x y)))
-(rule (simplify (select ty (ult rty x y)
-                           (iconst_s ty -1)
-                           (uextend ty (ugt rty x y))))
-      (sextend_from_i8 ty (spaceship_u rty x y)))
+                           (uextend_maybe ty (ugt rty x y))))
+      (sextend_maybe ty (spaceship_u rty x y)))
 
 ;; x == y ? 0 : x < y ? -1 : +1
 (rule (simplify (select ty (eq rty x y)
@@ -35,49 +27,49 @@
                            (select ty (ult rty x y)
                                       (iconst_s ty -1)
                                       (iconst_s ty 1))))
-      (sextend_from_i8 ty (spaceship_u rty x y)))
+      (sextend_maybe ty (spaceship_u rty x y)))
 ;; x == y ? 0 : x <= y ? -1 : +1
 (rule (simplify (select ty (eq rty x y)
                            (iconst_s ty 0)
                            (select ty (ule rty x y)
                                       (iconst_s ty -1)
                                       (iconst_s ty 1))))
-      (sextend_from_i8 ty (spaceship_u rty x y)))
+      (sextend_maybe ty (spaceship_u rty x y)))
 ;; x == y ? 0 : x > y ? +1 : -1
 (rule (simplify (select ty (eq rty x y)
                            (iconst_s ty 0)
                            (select ty (ugt rty x y)
                                       (iconst_s ty 1)
                                       (iconst_s ty -1))))
-      (sextend_from_i8 ty (spaceship_u rty x y)))
+      (sextend_maybe ty (spaceship_u rty x y)))
 ;; x == y ? 0 : x >= y ? +1 : -1
 (rule (simplify (select ty (eq rty x y)
                            (iconst_s ty 0)
                            (select ty (uge rty x y)
                                       (iconst_s ty 1)
                                       (iconst_s ty -1))))
-      (sextend_from_i8 ty (spaceship_u rty x y)))
+      (sextend_maybe ty (spaceship_u rty x y)))
 
 ;; x > y ? 1 : x < y ? -1 : 0
 ;; x > y ? 1 : x >= y ? 0 : -1
 (rule (simplify (select ty (ugt rty x y)
                            (iconst_s ty 1)
                            (ineg rty (ult rty x y))))
-      (sextend_from_i8 ty (spaceship_u rty x y)))
+      (sextend_maybe ty (spaceship_u rty x y)))
 (rule (simplify (select ty (ugt rty x y)
                            (iconst_s ty 1)
                            (bmask ty (ult rty x y))))
-      (sextend_from_i8 ty (spaceship_u rty x y)))
+      (sextend_maybe ty (spaceship_u rty x y)))
 ;; x > y ? 1 : x != y ? -1 : 0
 ;; x > y ? 1 : x == y ? 0 : -1
 (rule (simplify (select ty (ugt rty x y)
                            (iconst_s ty 1)
                            (ineg rty (ne rty x y))))
-      (sextend_from_i8 ty (spaceship_u rty x y)))
+      (sextend_maybe ty (spaceship_u rty x y)))
 (rule (simplify (select ty (ugt rty x y)
                            (iconst_s ty 1)
                            (bmask ty (ne rty x y))))
-      (sextend_from_i8 ty (spaceship_u rty x y)))
+      (sextend_maybe ty (spaceship_u rty x y)))
 
 ;; Same, but for signed comparisons this time
 
@@ -85,22 +77,14 @@
 ;; x < y ? -1 : x != y ? +1 : 0
 (rule (simplify (select ty (slt rty x y)
                            (iconst_s ty -1)
-                           (ne rty x y)))
-      (spaceship_s rty x y))
-(rule (simplify (select ty (slt rty x y)
-                           (iconst_s ty -1)
-                           (uextend ty (ne rty x y))))
-      (sextend_from_i8 ty (spaceship_s rty x y)))
+                           (uextend_maybe ty (ne rty x y))))
+      (sextend_maybe ty (spaceship_s rty x y)))
 ;; x < y ? -1 : x <= y ? 0 : +1
 ;; x < y ? -1 : x > y ? +1 : 0
 (rule (simplify (select ty (slt rty x y)
                            (iconst_s ty -1)
-                           (sgt rty x y)))
-      (spaceship_s rty x y))
-(rule (simplify (select ty (slt rty x y)
-                           (iconst_s ty -1)
-                           (uextend ty (sgt rty x y))))
-      (sextend_from_i8 ty (spaceship_s rty x y)))
+                           (uextend_maybe ty (sgt rty x y))))
+      (sextend_maybe ty (spaceship_s rty x y)))
 
 ;; x == y ? 0 : x < y ? -1 : +1
 (rule (simplify (select ty (eq rty x y)
@@ -108,49 +92,49 @@
                            (select ty (slt rty x y)
                                       (iconst_s ty -1)
                                       (iconst_s ty 1))))
-      (sextend_from_i8 ty (spaceship_s rty x y)))
+      (sextend_maybe ty (spaceship_s rty x y)))
 ;; x == y ? 0 : x <= y ? -1 : +1
 (rule (simplify (select ty (eq rty x y)
                            (iconst_s ty 0)
                            (select ty (sle rty x y)
                                       (iconst_s ty -1)
                                       (iconst_s ty 1))))
-      (sextend_from_i8 ty (spaceship_s rty x y)))
+      (sextend_maybe ty (spaceship_s rty x y)))
 ;; x == y ? 0 : x > y ? +1 : -1
 (rule (simplify (select ty (eq rty x y)
                            (iconst_s ty 0)
                            (select ty (sgt rty x y)
                                       (iconst_s ty 1)
                                       (iconst_s ty -1))))
-      (sextend_from_i8 ty (spaceship_s rty x y)))
+      (sextend_maybe ty (spaceship_s rty x y)))
 ;; x == y ? 0 : x >= y ? +1 : -1
 (rule (simplify (select ty (eq rty x y)
                            (iconst_s ty 0)
                            (select ty (sge rty x y)
                                       (iconst_s ty 1)
                                       (iconst_s ty -1))))
-      (sextend_from_i8 ty (spaceship_s rty x y)))
+      (sextend_maybe ty (spaceship_s rty x y)))
 
 ;; x > y ? 1 : x < y ? -1 : 0
 ;; x > y ? 1 : x >= y ? 0 : -1
 (rule (simplify (select ty (sgt rty x y)
                            (iconst_s ty 1)
                            (ineg rty (slt rty x y))))
-      (sextend_from_i8 ty (spaceship_s rty x y)))
+      (sextend_maybe ty (spaceship_s rty x y)))
 (rule (simplify (select ty (sgt rty x y)
                            (iconst_s ty 1)
                            (bmask ty (slt rty x y))))
-      (sextend_from_i8 ty (spaceship_s rty x y)))
+      (sextend_maybe ty (spaceship_s rty x y)))
 ;; x > y ? 1 : x != y ? -1 : 0
 ;; x > y ? 1 : x == y ? 0 : -1
 (rule (simplify (select ty (sgt rty x y)
                            (iconst_s ty 1)
                            (ineg rty (ne rty x y))))
-      (sextend_from_i8 ty (spaceship_s rty x y)))
+      (sextend_maybe ty (spaceship_s rty x y)))
 (rule (simplify (select ty (sgt rty x y)
                            (iconst_s ty 1)
                            (bmask ty (ne rty x y))))
-      (sextend_from_i8 ty (spaceship_s rty x y)))
+      (sextend_maybe ty (spaceship_s rty x y)))
 
 ;; Then once we have it normalized, we can apply some basic simplifications.
 ;; For example, a derived `PartialOrd::lt` on a newtype in Rust will essentially
@@ -208,9 +192,3 @@
       (sle ty x y))
 (rule (simplify (ne _ (spaceship_u ty x y) (iconst_s _ 1)))
       (ule ty x y))
-
-;; extend from i8 to i8 is invalid CLIF, so this allows fixing that in the output
-;; rather than needing to duplicate rules for the different width categories
-(decl sextend_from_i8 (Type Value) Value)
-(rule 0 (sextend_from_i8 ty val) (sextend ty val))
-(rule 1 (sextend_from_i8 $I8 val) val)
diff --git a/cranelift/codegen/src/prelude_opt.isle b/cranelift/codegen/src/prelude_opt.isle
index 5011481af..e0cd14860 100644
--- a/cranelift/codegen/src/prelude_opt.isle
+++ b/cranelift/codegen/src/prelude_opt.isle
@@ -1,5 +1,7 @@
 ;; Prelude definitions specific to the mid-end.
 
+;; Any `extern` definitions here are generally implemented in `src/opts.rs`.
+
 ;;;;; eclass and enode access ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 ;; Extract any node(s) for the given eclass ID.
@@ -92,3 +94,29 @@
 	(if-let $true (u64_le c (ty_umax ty)))
     (iconst ty (imm64 c)))
 (rule 1 (iconst_u $I128 c) (uextend $I128 (iconst_u $I64 c)))
+
+;; These take `Value`, rather than going through `inst_data_tupled`, because
+;; most of the time they want to return the original `Value`, and it would be
+;; a waste to need to re-GVN the instruction data in those cases.
+(decl multi sextend_maybe_etor (Type Value) Value)
+(extern extractor infallible sextend_maybe_etor sextend_maybe_etor)
+(decl multi uextend_maybe_etor (Type Value) Value)
+(extern extractor infallible uextend_maybe_etor uextend_maybe_etor)
+
+;; Match or Construct a possibly-`uextend`ed value.
+;; Gives the extended-to type and inner value when matching something that was
+;; extended, or the input value and its type when the value isn't an extension.
+;; Useful to write a single pattern that can match things that may or may not
+;; have undergone C's "usual arithmetic conversions".
+;; When generating values, extending to the same type is invalid CLIF,
+;; so this avoids doing that where there's no extension actually needed.
+(decl uextend_maybe (Type Value) Value)
+(extractor (uextend_maybe ty val) (uextend_maybe_etor ty val))
+(rule 0 (uextend_maybe ty val) (uextend ty val))
+(rule 1 (uextend_maybe ty val@(value_type ty)) val)
+
+;; Same as `uextend_maybe` above, just for `sextend`.
+(decl sextend_maybe (Type Value) Value)
+(extractor (sextend_maybe ty val) (sextend_maybe_etor ty val))
+(rule 0 (sextend_maybe ty val) (sextend ty val))
+(rule 1 (sextend_maybe ty val@(value_type ty)) val)


Test patch:
diff --git a/cranelift/filetests/filetests/egraph/icmp.clif b/cranelift/filetests/filetests/egraph/icmp.clif
index f82d4aaa3..01e6db2aa 100644
--- a/cranelift/filetests/filetests/egraph/icmp.clif
+++ b/cranelift/filetests/filetests/egraph/icmp.clif
@@ -109,3 +109,19 @@ block0(v0: i8):
 ;     v4 = icmp eq v0, v3  ; v3 = 0
 ;     return v4
 ; }
+
+function %byte_icmp_ucmp(i8) -> i8 {
+block0(v0: i8):
+    v1 = iconst.i8 0
+    v2 = icmp sge v0, v1
+    v3 = icmp eq v2, v1
+    return v3
+}
+
+; function %byte_icmp_ucmp(i8) -> i8 fast {
+; block0(v0: i8):
+;     v1 = iconst.i8 0
+;     v4 = icmp slt v0, v1  ; v1 = 0
+;     return v4
+; }
+
diff --git a/cranelift/filetests/filetests/pcc/succeed/const.clif b/cranelift/filetests/filetests/pcc/succeed/const.clif
index d0c381c5c..529c43eaa 100644
--- a/cranelift/filetests/filetests/pcc/succeed/const.clif
+++ b/cranelift/filetests/filetests/pcc/succeed/const.clif
@@ -28,3 +28,9 @@ block0:
     v0 = iconst.i64 0x9_ffff_ffff
     return v0
 }
+
+function %f3() -> i64 {
+block0:
+    v0 = iconst.i64 0xffff_fff7
+    return v0
+}