riscv64: Improve ctz/clz/cls codegen (#5854)

* cranelift: Add extra runtests for `clz`/`ctz`

* riscv64: Restrict lowering rules for `ctz`/`clz`

* cranelift: Add `u64` isle helpers

* riscv64: Improve `ctz` codegen

* riscv64: Improve `clz` codegen

* riscv64: Improve `cls` codegen

* riscv64: Improve `clz.i128` codegen

Instead of checking if we have 64 zeros in the top half. Check
if it *is* 0, that way we avoid loading the `64` constant.

* riscv64: Improve `ctz.i128` codegen

Instead of checking if we have 64 zeros in the bottom half. Check
if it *is* 0, that way we avoid loading the `64` constant.

* riscv64: Use extended value in `lower_cls`

* riscv64: Use pattern matches on `bseti`

cranelift/filetests/filetests/isa/riscv64/bitops.clif

@@ -385,28 +385,25 @@ block0(v0: i128):
 
 ; VCode:
 ; block0:
-;   mv t0,a1
-;   clz a2,t0##ty=i64 tmp=a3 step=a1
+;   clz a2,a1##ty=i64 tmp=a4 step=a3
 ;   clz a6,a0##ty=i64 tmp=a4 step=a5
-;   li t3,64
-;   select_reg t0,a6,zero##condition=(t3 eq a2)
-;   add a0,a2,t0
+;   select_reg t3,a6,zero##condition=(a1 eq zero)
+;   add a0,a2,t3
 ;   li a1,0
 ;   ret
 ; 
 ; Disassembled:
 ; block0: ; offset 0x0
-;   ori t0, a1, 0
 ;   ori a2, zero, 0
-;   addi a1, zero, 0x40
-;   addi a3, zero, 1
-;   slli a3, a3, 0x3f
-;   blez a1, 0x1c
-;   and t5, a3, t0
+;   addi a3, zero, 0x40
+;   addi a4, zero, 1
+;   slli a4, a4, 0x3f
+;   blez a3, 0x1c
+;   and t5, a4, a1
 ;   bne zero, t5, 0x14
 ;   addi a2, a2, 1
-;   addi a1, a1, -1
-;   srli a3, a3, 1
+;   addi a3, a3, -1
+;   srli a4, a4, 1
 ;   j -0x18
 ;   ori a6, zero, 0
 ;   addi a5, zero, 0x40
@@ -419,12 +416,11 @@ block0(v0: i128):
 ;   addi a5, a5, -1
 ;   srli a4, a4, 1
 ;   j -0x18
-;   addi t3, zero, 0x40
-;   beq t3, a2, 0xc
-;   ori t0, zero, 0
+;   beqz a1, 0xc
+;   ori t3, zero, 0
 ;   j 8
-;   ori t0, a6, 0
-;   add a0, a2, t0
+;   ori t3, a6, 0
+;   add a0, a2, t3
 ;   mv a1, zero
 ;   ret
 
@@ -438,8 +434,8 @@ block0(v0: i8):
 ; block0:
 ;   slli t2,a0,56
 ;   srai a1,t2,56
-;   not a3,a0
-;   select_reg a5,a3,a0##condition=(a1 slt zero)
+;   not a3,a1
+;   select_reg a5,a3,a1##condition=(a1 slt zero)
 ;   clz t4,a5##ty=i8 tmp=a7 step=t3
 ;   addi a0,t4,-1
 ;   ret
@@ -448,9 +444,9 @@ block0(v0: i8):
 ; block0: ; offset 0x0
 ;   slli t2, a0, 0x38
 ;   srai a1, t2, 0x38
-;   not a3, a0
+;   not a3, a1
 ;   bltz a1, 0xc
-;   ori a5, a0, 0
+;   ori a5, a1, 0
 ;   j 8
 ;   ori a5, a3, 0
 ;   ori t4, zero, 0
@@ -477,8 +473,8 @@ block0(v0: i16):
 ; block0:
 ;   slli t2,a0,48
 ;   srai a1,t2,48
-;   not a3,a0
-;   select_reg a5,a3,a0##condition=(a1 slt zero)
+;   not a3,a1
+;   select_reg a5,a3,a1##condition=(a1 slt zero)
 ;   clz t4,a5##ty=i16 tmp=a7 step=t3
 ;   addi a0,t4,-1
 ;   ret
@@ -487,9 +483,9 @@ block0(v0: i16):
 ; block0: ; offset 0x0
 ;   slli t2, a0, 0x30
 ;   srai a1, t2, 0x30
-;   not a3, a0
+;   not a3, a1
 ;   bltz a1, 0xc
-;   ori a5, a0, 0
+;   ori a5, a1, 0
 ;   j 8
 ;   ori a5, a3, 0
 ;   ori t4, zero, 0
@@ -515,8 +511,8 @@ block0(v0: i32):
 ; VCode:
 ; block0:
 ;   sext.w t2,a0
-;   not a1,a0
-;   select_reg a3,a1,a0##condition=(t2 slt zero)
+;   not a1,t2
+;   select_reg a3,a1,t2##condition=(t2 slt zero)
 ;   clz a7,a3##ty=i32 tmp=a5 step=a6
 ;   addi a0,a7,-1
 ;   ret
@@ -524,9 +520,9 @@ block0(v0: i32):
 ; Disassembled:
 ; block0: ; offset 0x0
 ;   sext.w t2, a0
-;   not a1, a0
+;   not a1, t2
 ;   bltz t2, 0xc
-;   ori a3, a0, 0
+;   ori a3, t2, 0
 ;   j 8
 ;   ori a3, a1, 0
 ;   ori a7, zero, 0
@@ -592,11 +588,10 @@ block0(v0: i128):
 ;   select_reg a6,a4,a1##condition=(a1 slt zero)
 ;   clz t0,a6##ty=i64 tmp=t3 step=t4
 ;   clz a1,a2##ty=i64 tmp=t2 step=a0
-;   li a3,64
-;   select_reg a5,a1,zero##condition=(a3 eq t0)
-;   add a7,t0,a5
-;   li t4,0
-;   addi a0,a7,-1
+;   select_reg a3,a1,zero##condition=(a6 eq zero)
+;   add a5,t0,a3
+;   li a7,0
+;   addi a0,a5,-1
 ;   li a1,0
 ;   ret
 ; 
@@ -632,14 +627,13 @@ block0(v0: i128):
 ;   addi a0, a0, -1
 ;   srli t2, t2, 1
 ;   j -0x18
-;   addi a3, zero, 0x40
-;   beq a3, t0, 0xc
-;   ori a5, zero, 0
+;   beqz a6, 0xc
+;   ori a3, zero, 0
 ;   j 8
-;   ori a5, a1, 0
-;   add a7, t0, a5
-;   mv t4, zero
-;   addi a0, a7, -1
+;   ori a3, a1, 0
+;   add a5, t0, a3
+;   mv a7, zero
+;   addi a0, a5, -1
 ;   mv a1, zero
 ;   ret
 
@@ -759,44 +753,42 @@ block0(v0: i128):
 
 ; VCode:
 ; block0:
-;   mv t0,a0
-;   ctz a2,t0##ty=i64 tmp=a0 step=a3
-;   ctz a6,a1##ty=i64 tmp=a4 step=a5
-;   li t3,64
-;   select_reg t0,a6,zero##condition=(t3 eq a2)
-;   add a0,a2,t0
+;   mv t4,a1
+;   ctz a2,t4##ty=i64 tmp=a3 step=a1
+;   ctz a6,a0##ty=i64 tmp=a4 step=a5
+;   select_reg t3,a2,zero##condition=(a0 eq zero)
+;   add a0,a6,t3
 ;   li a1,0
 ;   ret
 ; 
 ; Disassembled:
 ; block0: ; offset 0x0
-;   ori t0, a0, 0
+;   ori t4, a1, 0
 ;   ori a2, zero, 0
-;   addi a3, zero, 0x40
-;   addi a0, zero, 1
-;   blez a3, 0x1c
-;   and t5, a0, t0
+;   addi a1, zero, 0x40
+;   addi a3, zero, 1
+;   blez a1, 0x1c
+;   and t5, a3, t4
 ;   bne zero, t5, 0x14
 ;   addi a2, a2, 1
-;   addi a3, a3, -1
-;   slli a0, a0, 1
+;   addi a1, a1, -1
+;   slli a3, a3, 1
 ;   j -0x18
 ;   ori a6, zero, 0
 ;   addi a5, zero, 0x40
 ;   addi a4, zero, 1
 ;   blez a5, 0x1c
-;   and t5, a4, a1
+;   and t5, a4, a0
 ;   bne zero, t5, 0x14
 ;   addi a6, a6, 1
 ;   addi a5, a5, -1
 ;   slli a4, a4, 1
 ;   j -0x18
-;   addi t3, zero, 0x40
-;   beq t3, a2, 0xc
-;   ori t0, zero, 0
+;   beqz a0, 0xc
+;   ori t3, zero, 0
 ;   j 8
-;   ori t0, a6, 0
-;   add a0, a2, t0
+;   ori t3, a2, 0
+;   add a0, a6, t3
 ;   mv a1, zero
 ;   ret
 

cranelift/filetests/filetests/isa/riscv64/cls-zbb.clif

@@ -0,0 +1,162 @@
+test compile precise-output
+set unwind_info=false
+target riscv64 has_zbb
+
+
+function %cls_i8(i8) -> i8 {
+block0(v0: i8):
+    v1 = cls v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   sext.b t2,a0
+;   not a1,t2
+;   select_reg a3,a1,t2##condition=(t2 slt zero)
+;   andi a5,a3,255
+;   clz a7,a5
+;   addi t4,a7,-56
+;   addi a0,t4,-1
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0x93, 0x13, 0x45, 0x60
+;   not a1, t2
+;   bltz t2, 0xc
+;   ori a3, t2, 0
+;   j 8
+;   ori a3, a1, 0
+;   andi a5, a3, 0xff
+;   .byte 0x93, 0x98, 0x07, 0x60
+;   addi t4, a7, -0x38
+;   addi a0, t4, -1
+;   ret
+
+function %cls_i16(i16) -> i16 {
+block0(v0: i16):
+    v1 = cls v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   sext.h t2,a0
+;   not a1,t2
+;   select_reg a3,a1,t2##condition=(t2 slt zero)
+;   zext.h a5,a3
+;   clz a7,a5
+;   addi t4,a7,-48
+;   addi a0,t4,-1
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0x93, 0x13, 0x55, 0x60
+;   not a1, t2
+;   bltz t2, 0xc
+;   ori a3, t2, 0
+;   j 8
+;   ori a3, a1, 0
+;   .byte 0xbb, 0xc7, 0x06, 0x08
+;   .byte 0x93, 0x98, 0x07, 0x60
+;   addi t4, a7, -0x30
+;   addi a0, t4, -1
+;   ret
+
+function %cls_i32(i32) -> i32 {
+block0(v0: i32):
+    v1 = cls v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   sext.w t2,a0
+;   not a1,t2
+;   select_reg a3,a1,t2##condition=(t2 slt zero)
+;   clzw a5,a3
+;   addi a0,a5,-1
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   sext.w t2, a0
+;   not a1, t2
+;   bltz t2, 0xc
+;   ori a3, t2, 0
+;   j 8
+;   ori a3, a1, 0
+;   .byte 0x9b, 0x97, 0x06, 0x60
+;   addi a0, a5, -1
+;   ret
+
+function %cls_i64(i64) -> i64 {
+block0(v0: i64):
+    v1 = cls v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   not t2,a0
+;   select_reg a1,t2,a0##condition=(a0 slt zero)
+;   clz a3,a1
+;   addi a0,a3,-1
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   not t2, a0
+;   bltz a0, 0xc
+;   ori a1, a0, 0
+;   j 8
+;   ori a1, t2, 0
+;   .byte 0x93, 0x96, 0x05, 0x60
+;   addi a0, a3, -1
+;   ret
+
+function %cls_i128(i128) -> i128 {
+block0(v0: i128):
+    v1 = cls v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   not a2,a0
+;   select_reg a2,a2,a0##condition=(a1 slt zero)
+;   not a4,a1
+;   select_reg a6,a4,a1##condition=(a1 slt zero)
+;   clz t3,a6
+;   clz t0,a2
+;   select_reg t2,t0,zero##condition=(a6 eq zero)
+;   add a1,t3,t2
+;   li a3,0
+;   addi a0,a1,-1
+;   li a1,0
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   not a2, a0
+;   bltz a1, 8
+;   ori a2, a0, 0
+;   not a4, a1
+;   bltz a1, 0xc
+;   ori a6, a1, 0
+;   j 8
+;   ori a6, a4, 0
+;   .byte 0x13, 0x1e, 0x08, 0x60
+;   .byte 0x93, 0x12, 0x06, 0x60
+;   beqz a6, 0xc
+;   ori t2, zero, 0
+;   j 8
+;   ori t2, t0, 0
+;   add a1, t3, t2
+;   mv a3, zero
+;   addi a0, a1, -1
+;   mv a1, zero
+;   ret
+

cranelift/filetests/filetests/isa/riscv64/clz-zbb.clif

@@ -0,0 +1,103 @@
+test compile precise-output
+set unwind_info=false
+target riscv64 has_zbb
+
+function %clz_i8(i8) -> i8 {
+block0(v0: i8):
+    v1 = clz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   andi t2,a0,255
+;   clz a1,t2
+;   addi a0,a1,-56
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   andi t2, a0, 0xff
+;   .byte 0x93, 0x95, 0x03, 0x60
+;   addi a0, a1, -0x38
+;   ret
+
+function %clz_i16(i16) -> i16 {
+block0(v0: i16):
+    v1 = clz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   zext.h t2,a0
+;   clz a1,t2
+;   addi a0,a1,-48
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0xbb, 0x43, 0x05, 0x08
+;   .byte 0x93, 0x95, 0x03, 0x60
+;   addi a0, a1, -0x30
+;   ret
+
+function %clz_i32(i32) -> i32 {
+block0(v0: i32):
+    v1 = clz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   clzw a0,a0
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0x1b, 0x15, 0x05, 0x60
+;   ret
+
+function %clz_i64(i64) -> i64 {
+block0(v0: i64):
+    v1 = clz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   clz a0,a0
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0x13, 0x15, 0x05, 0x60
+;   ret
+
+function %clz_i128(i128) -> i128 {
+block0(v0: i128):
+    v1 = clz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   clz a2,a1
+;   clz a3,a0
+;   select_reg a4,a3,zero##condition=(a1 eq zero)
+;   add a0,a2,a4
+;   li a1,0
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0x13, 0x96, 0x05, 0x60
+;   .byte 0x93, 0x16, 0x05, 0x60
+;   beqz a1, 0xc
+;   ori a4, zero, 0
+;   j 8
+;   ori a4, a3, 0
+;   add a0, a2, a4
+;   mv a1, zero
+;   ret
+

cranelift/filetests/filetests/isa/riscv64/ctz-zbb-zbs.clif

@@ -0,0 +1,71 @@
+test compile precise-output
+set unwind_info=false
+target riscv64 has_zbb has_zbs
+
+function %ctz_i8(i8) -> i8 {
+block0(v0: i8):
+    v1 = ctz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   bseti t2,a0,8
+;   ctzw a0,t2
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0x93, 0x13, 0x85, 0x28
+;   .byte 0x1b, 0x95, 0x13, 0x60
+;   ret
+
+function %ctz_i16(i16) -> i16 {
+block0(v0: i16):
+    v1 = ctz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   bseti t2,a0,16
+;   ctzw a0,t2
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0x93, 0x13, 0x05, 0x29
+;   .byte 0x1b, 0x95, 0x13, 0x60
+;   ret
+
+function %ctz_i32(i32) -> i32 {
+block0(v0: i32):
+    v1 = ctz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   ctzw a0,a0
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0x1b, 0x15, 0x15, 0x60
+;   ret
+
+function %ctz_i64(i64) -> i64 {
+block0(v0: i64):
+    v1 = ctz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   ctz a0,a0
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0x13, 0x15, 0x15, 0x60
+;   ret

cranelift/filetests/filetests/isa/riscv64/ctz-zbb.clif

@@ -0,0 +1,102 @@
+test compile precise-output
+set unwind_info=false
+target riscv64 has_zbb
+
+
+function %ctz_i8(i8) -> i8 {
+block0(v0: i8):
+    v1 = ctz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   ori t2,a0,256
+;   ctzw a0,t2
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   ori t2, a0, 0x100
+;   .byte 0x1b, 0x95, 0x13, 0x60
+;   ret
+
+function %ctz_i16(i16) -> i16 {
+block0(v0: i16):
+    v1 = ctz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   lui t2,16
+;   or a1,a0,t2
+;   ctzw a0,a1
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   lui t2, 0x10
+;   or a1, a0, t2
+;   .byte 0x1b, 0x95, 0x15, 0x60
+;   ret
+
+function %ctz_i32(i32) -> i32 {
+block0(v0: i32):
+    v1 = ctz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   ctzw a0,a0
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0x1b, 0x15, 0x15, 0x60
+;   ret
+
+function %ctz_i64(i64) -> i64 {
+block0(v0: i64):
+    v1 = ctz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   ctz a0,a0
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0x13, 0x15, 0x15, 0x60
+;   ret
+
+function %ctz_i128(i128) -> i128 {
+block0(v0: i128):
+    v1 = ctz v0
+    return v1
+}
+
+; VCode:
+; block0:
+;   ctz a1,a1
+;   ctz a2,a0
+;   select_reg a4,a1,zero##condition=(a0 eq zero)
+;   add a0,a2,a4
+;   li a1,0
+;   ret
+; 
+; Disassembled:
+; block0: ; offset 0x0
+;   .byte 0x93, 0x95, 0x15, 0x60
+;   .byte 0x13, 0x16, 0x15, 0x60
+;   beqz a0, 0xc
+;   ori a4, zero, 0
+;   j 8
+;   ori a4, a1, 0
+;   add a0, a2, a4
+;   mv a1, zero
+;   ret
+

cranelift/filetests/filetests/runtests/cls.clif

@@ -2,6 +2,7 @@ test interpret
 test run
 target aarch64
 target riscv64
+target riscv64 has_zbb
 target s390x
 ; not implemented on `x86_64`
 

cranelift/filetests/filetests/runtests/clz.clif

@@ -5,12 +5,14 @@ target s390x
 target x86_64
 target x86_64 has_lzcnt
 target riscv64
+target riscv64 has_zbb
 
 function %clz_i8(i8) -> i8 {
 block0(v0: i8):
     v1 = clz v0
     return v1
 }
+; run: %clz_i8(0) == 8
 ; run: %clz_i8(1) == 7
 ; run: %clz_i8(0x40) == 1
 ; run: %clz_i8(-1) == 0
@@ -20,6 +22,7 @@ block0(v0: i16):
     v1 = clz v0
     return v1
 }
+; run: %clz_i16(0) == 16
 ; run: %clz_i16(1) == 15
 ; run: %clz_i16(0x4000) == 1
 ; run: %clz_i16(-1) == 0
@@ -29,6 +32,7 @@ block0(v0: i32):
     v1 = clz v0
     return v1
 }
+; run: %clz_i32(0) == 32
 ; run: %clz_i32(1) == 31
 ; run: %clz_i32(0x40000000) == 1
 ; run: %clz_i32(-1) == 0
@@ -38,6 +42,7 @@ block0(v0: i64):
     v1 = clz v0
     return v1
 }
+; run: %clz_i64(0) == 64
 ; run: %clz_i64(1) == 63
 ; run: %clz_i64(0x4000000000000000) == 1
 ; run: %clz_i64(-1) == 0

cranelift/filetests/filetests/runtests/ctz.clif

@@ -3,14 +3,17 @@ test run
 target aarch64
 target s390x
 target x86_64
-target riscv64
 target x86_64 has_bmi1
+target riscv64
+target riscv64 has_zbb
+target riscv64 has_zbb has_zbs
 
 function %ctz_i8(i8) -> i8 {
 block0(v0: i8):
     v1 = ctz v0
     return v1
 }
+; run: %ctz_i8(0) == 8
 ; run: %ctz_i8(1) == 0
 ; run: %ctz_i8(0x40) == 6
 ; run: %ctz_i8(-1) == 0
@@ -20,6 +23,7 @@ block0(v0: i16):
     v1 = ctz v0
     return v1
 }
+; run: %ctz_i16(0) == 16
 ; run: %ctz_i16(1) == 0
 ; run: %ctz_i16(0x4000) == 14
 ; run: %ctz_i16(-1) == 0
@@ -29,6 +33,7 @@ block0(v0: i32):
     v1 = ctz v0
     return v1
 }
+; run: %ctz_i32(0) == 32
 ; run: %ctz_i32(1) == 0
 ; run: %ctz_i32(0x40000000) == 30
 ; run: %ctz_i32(-1) == 0
@@ -38,6 +43,7 @@ block0(v0: i64):
     v1 = ctz v0
     return v1
 }
+; run: %ctz_i64(0) == 64
 ; run: %ctz_i64(1) == 0
 ; run: %ctz_i64(0x4000000000000000) == 62
 ; run: %ctz_i64(-1) == 0

cranelift/filetests/filetests/runtests/i128-bitops-count.clif

@@ -4,6 +4,8 @@ target aarch64
 target s390x
 target x86_64
 target riscv64
+target riscv64 has_zbb
+target riscv64 has_zbb has_zbs
 
 function %ctz_i128(i128) -> i128 {
 block0(v0: i128):

cranelift/filetests/filetests/runtests/i128-cls.clif

@@ -1,6 +1,7 @@
 test run
 target aarch64
-target riscv64 
+target riscv64
+target riscv64 has_zbb
 target s390x
 
 function %cls_i128(i128) -> i128 {