Add optimizations for leading / trailing zeros

Author: KGrewal1

cranelift/codegen/src/isle_prelude.rs

@@ -27,6 +27,16 @@ macro_rules! isle_common_prelude_methods {
             self.checked_add_with_type(ty, a, b).is_none()
         }
 
+        #[inline]
+        fn imm64_clz(&mut self, ty: Type, a: Imm64) -> Imm64 {
+            let bits = ty.bits();
+            assert!(bits <= 64);
+            let clz_offset = 64 - bits;
+            let a_v: u64 = a.bits().cast_unsigned();
+            let lz = a_v.leading_zeros() - clz_offset;
+            Imm64::new(lz as i64)
+        }
+
         #[inline]
         fn imm64_sdiv(&mut self, ty: Type, x: Imm64, y: Imm64) -> Option<Imm64> {
             // Sign extend `x` and `y`.

cranelift/codegen/src/opts/cprop.isle

@@ -1,5 +1,16 @@
 ;; Constant propagation.
 
+(rule (simplify
+       (clz (fits_in_64 ty)
+             (iconst ty kx)))
+      (subsume (iconst ty (imm64_clz ty kx))))
+
+
+(rule (simplify
+       (ctz (fits_in_64 ty)
+             (iconst ty (u64_from_imm64 kx))))
+      (subsume (iconst ty (imm64_masked ty (u64_trailing_zeros kx)))))
+
 (rule (simplify
        (iadd (fits_in_64 ty)
              (iconst ty (u64_from_imm64 k1))

cranelift/codegen/src/prelude.isle

@@ -96,6 +96,10 @@
 (decl pure imm64_icmp (Type IntCC Imm64 Imm64) Imm64)
 (extern constructor imm64_icmp imm64_icmp)
 
+(decl pure imm64_clz (Type Imm64) Imm64)
+(extern constructor imm64_clz imm64_clz)
+
+
 ;; Each of these extractors tests whether the upper half of the input equals the
 ;; lower half of the input
 (decl u128_replicated_u64 (u64) u128)

cranelift/filetests/filetests/egraph/cprop.clif

@@ -22,6 +22,33 @@ block0:
 ; check: v3 = iconst.i16 -2
 ; nextln: return v3
 
+function %f0() -> i8 {
+block0:
+    v1 = iconst.i8 51
+    v2 = clz.i8 v1
+    return v2
+}
+
+function %f0() -> i16 {
+block0:
+    v1 = iconst.i16 51
+    v2 = clz.i16 v1
+    return v2
+}
+
+; check: v3 = iconst.i16 10
+; nextln: return v3
+
+function %f0() -> i16 {
+block0:
+    v1 = iconst.i16 48
+    v2 = ctz.i16 v1
+    return v2
+}
+
+; check: v3 = iconst.i16 4
+; nextln: return v3
+
 function %ishl() -> i8 {
 block0:
     v0 = iconst.i8 1