riscv64: Fix replicated_imm5 lowering rule (#12956)

* riscv64: Fix `replicated_imm5` lowering rule

This commit fixes the `replicated_imm5` lowering rule in the riscv64
backend when applied to `vconst` inputs. This is reachable in
WebAssembly via `v128.const`, for example, or via Cranelift
optimizations. The previous iteration of the rule did not take lane
width into account and always reduced the immediate as much as possible
where the reduction needed to halt once the lane size had been reached.

The fix here is to refactor the definition to have a rule-per-lane-width
which handles the input constant. This is accompanied with some minor
refactorings around available ISLE rules and such to handle more bit
widths in more locations, have more faithful signatures, etc.

* Fix winch testing

* Fix test exemption

prtest:full

Author: alexcrichton

cranelift/codegen/src/isa/aarch64/inst.isle

@@ -4112,13 +4112,13 @@
 ;; as it will make it easier to create the immediates in the instructions below.
 (rule 5 (splat_const (u64_replicated_u32 n) (VectorSize.Size64x2))
         (splat_const n (VectorSize.Size32x4)))
-(rule 5 (splat_const (u32_replicated_u16 n) (VectorSize.Size32x4))
+(rule 5 (splat_const (u32_from_u64 (u32_replicated_u16 n)) (VectorSize.Size32x4))
         (splat_const n (VectorSize.Size16x8)))
-(rule 5 (splat_const (u32_replicated_u16 n) (VectorSize.Size32x2))
+(rule 5 (splat_const (u32_from_u64 (u32_replicated_u16 n)) (VectorSize.Size32x2))
         (splat_const n (VectorSize.Size16x4)))
-(rule 5 (splat_const (u16_replicated_u8 n) (VectorSize.Size16x8))
+(rule 5 (splat_const (u16_from_u64 (u16_replicated_u8 n)) (VectorSize.Size16x8))
         (splat_const n (VectorSize.Size8x16)))
-(rule 5 (splat_const (u16_replicated_u8 n) (VectorSize.Size16x4))
+(rule 5 (splat_const (u16_from_u64 (u16_replicated_u8 n)) (VectorSize.Size16x4))
         (splat_const n (VectorSize.Size8x8)))
 
 ;; Special cases for `vec_dup_imm` instructions where the input is either

cranelift/codegen/src/isa/riscv64/inst.isle

@@ -2020,9 +2020,15 @@
 (decl imm5_from_i64 (Imm5) i64)
 (extern extractor imm5_from_i64 imm5_from_i64)
 
-;; Construct a Imm5 from an i8
+;; Construct a Imm5 from an iNN
 (decl pure partial i8_to_imm5 (i8) Imm5)
 (extern constructor i8_to_imm5 i8_to_imm5)
+(decl pure partial i16_to_imm5 (i16) Imm5)
+(rule (i16_to_imm5 (i8_from_i16 n)) (i8_to_imm5 n))
+(decl pure partial i32_to_imm5 (i32) Imm5)
+(rule (i32_to_imm5 (i8_from_i32 n)) (i8_to_imm5 n))
+(decl pure partial i64_to_imm5 (i64) Imm5)
+(rule (i64_to_imm5 (i8_from_i64 n)) (i8_to_imm5 n))
 
 ;; Helper to go directly from a `Value` to an `Imm5`.
 (decl imm5_from_value (Imm5) Value)
@@ -2037,11 +2043,20 @@
 ;; Constructor that matches a `Value` equivalent to a replicated Imm5 on all lanes.
 (decl pure partial replicated_imm5 (Value) Imm5)
 (rule (replicated_imm5 (splat _ (imm5_from_value n))) n)
-(rule (replicated_imm5 (vconst _ (u128_from_constant n128)))
+(rule (replicated_imm5 (vconst (multi_lane 64 _) (u128_from_constant n128)))
   (if-let (u128_replicated_u64 n64) n128)
-  (if-let (u64_replicated_u32 n32) n64)
-  (if-let (u32_replicated_u16 n16) n32)
-  (if-let (u16_replicated_u8 n8) n16)
+  (if-let n (i64_to_imm5 (u64_cast_signed n64)))
+  n)
+(rule (replicated_imm5 (vconst (multi_lane 32 _) (u128_from_constant n128)))
+  (if-let (u128_replicated_u32 n32) n128)
+  (if-let n (i32_to_imm5 (u32_cast_signed n32)))
+  n)
+(rule (replicated_imm5 (vconst (multi_lane 16 _) (u128_from_constant n128)))
+  (if-let (u128_replicated_u16 n16) n128)
+  (if-let n (i16_to_imm5 (u16_cast_signed n16)))
+  n)
+(rule (replicated_imm5 (vconst (multi_lane 8 _) (u128_from_constant n128)))
+  (if-let (u128_replicated_u8 n8) n128)
   (if-let n (i8_to_imm5 (u8_cast_signed n8)))
   n)
 

cranelift/codegen/src/isle_prelude.rs

@@ -792,27 +792,25 @@ macro_rules! isle_common_prelude_methods {
             }
         }
 
-        fn u64_replicated_u32(&mut self, val: u64) -> Option<u64> {
+        fn u64_replicated_u32(&mut self, val: u64) -> Option<u32> {
             let low32 = val as u32 as u64;
             if (low32 | (low32 << 32)) == val {
-                Some(low32)
+                Some(val as u32)
             } else {
                 None
             }
         }
 
-        fn u32_replicated_u16(&mut self, val: u64) -> Option<u64> {
-            let val = val as u32;
+        fn u32_replicated_u16(&mut self, val: u32) -> Option<u16> {
             let low16 = val as u16 as u32;
             if (low16 | (low16 << 16)) == val {
-                Some(low16.into())
+                Some(low16 as u16)
             } else {
                 None
             }
         }
 
-        fn u16_replicated_u8(&mut self, val: u64) -> Option<u8> {
-            let val = val as u16;
+        fn u16_replicated_u8(&mut self, val: u16) -> Option<u8> {
             let low8 = val as u8 as u16;
             if (low8 | (low8 << 8)) == val {
                 Some(low8 as u8)

cranelift/codegen/src/prelude.isle

@@ -115,13 +115,21 @@
 ;; lower half of the input
 (decl u128_replicated_u64 (u64) u128)
 (extern extractor u128_replicated_u64 u128_replicated_u64)
-(decl u64_replicated_u32 (u64) u64)
+(decl u64_replicated_u32 (u32) u64)
 (extern extractor u64_replicated_u32 u64_replicated_u32)
-(decl u32_replicated_u16 (u64) u64)
+(decl u32_replicated_u16 (u16) u32)
 (extern extractor u32_replicated_u16 u32_replicated_u16)
-(decl u16_replicated_u8 (u8) u64)
+(decl u16_replicated_u8 (u8) u16)
 (extern extractor u16_replicated_u8 u16_replicated_u8)
 
+;; Convenience extractors to go directly from u128 to something smaller
+(decl u128_replicated_u32 (u32) u128)
+(extractor (u128_replicated_u32 n) (u128_replicated_u64 (u64_replicated_u32 n)))
+(decl u128_replicated_u16 (u16) u128)
+(extractor (u128_replicated_u16 n) (u128_replicated_u32 (u32_replicated_u16 n)))
+(decl u128_replicated_u8 (u8) u128)
+(extractor (u128_replicated_u8 n) (u128_replicated_u16 (u16_replicated_u8 n)))
+
 ;; Get the low and high bits of a `u128` as `u64`s.
 (decl u128_low_bits (u128) u64)
 (extern constructor u128_low_bits u128_low_bits)