Improve codegen in a few memory-inbounds.wat cases

alexcrichton · Feb 6, 2025 · 63e2c5a · 63e2c5a
1 parent 725d03e
commit 63e2c5a
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 31 deletions.
diff --git a/cranelift/codegen/src/isa/pulley_shared/inst.isle b/cranelift/codegen/src/isa/pulley_shared/inst.isle
@@ -361,11 +361,16 @@
 
 ;; Helper used in `wasm_g32` above to extract `(iadd addr N)` where `N` is a
 ;; constant. If there is no constant then the constant 0 is returned instead.
+;;
+;; Note that this also requires `addr` itself to be an `iadd` internally to
+;; represent the host address plus the guest offset. If `addr` isn't internally
+;; an `iadd` then the `N` here is instead probably the static address in the
+;; guest.
 (decl pure host_offset (Value) HostOffset)
 (type HostOffset (enum (All (a Value) (b i32))))
 (rule 0 (host_offset val) (HostOffset.All val 0))
-(rule 1 (host_offset (iadd val (i32_from_iconst n))) (HostOffset.All val n))
-(rule 2 (host_offset (iadd (i32_from_iconst n) val)) (HostOffset.All val n))
+(rule 1 (host_offset (iadd val @ (iadd _ _) (i32_from_iconst n))) (HostOffset.All val n))
+(rule 2 (host_offset (iadd (i32_from_iconst n) val @ (iadd _ _))) (HostOffset.All val n))
 
 ;; Helper to test whether the first argument `oob` contains a condition for
 ;; matching whether second argument `addr` is out-of-bounds. Searches for a
@@ -388,15 +393,23 @@
 ;; a `uextend` from a 32-bit value.
 (rule 0 (wasm_oob_select oob (iadd base wasm_addr_ext @ (value_type $I64)))
   (if-let (OobCond.All bound n) (wasm_oob_cond wasm_addr_ext oob))
-  (if-let (uextend wasm_addr @ (value_type $I32)) wasm_addr_ext)
+  (if-let wasm_addr (wasm32_addr_for_host64 wasm_addr_ext))
   (if-let (PointerWidth.PointerWidth64) (pointer_width))
   (OobSelect.All base bound wasm_addr n))
 (rule 1 (wasm_oob_select oob (iadd wasm_addr_ext @ (value_type $I64) base))
   (if-let (OobCond.All bound n) (wasm_oob_cond wasm_addr_ext oob))
-  (if-let (uextend wasm_addr @ (value_type $I32)) wasm_addr_ext)
+  (if-let wasm_addr (wasm32_addr_for_host64 wasm_addr_ext))
   (if-let (PointerWidth.PointerWidth64) (pointer_width))
   (OobSelect.All base bound wasm_addr n))
 
+;; Helper to extract a 32-bit `Value` from a 64-bit input. The returned value,
+;; if one is returned, can losslessly represent the original value when zero
+;; extended from 32-bits. That means that this only matches `(uextend ...)` or
+;; an `iconst` which already fits in 32 bits.
+(decl pure partial wasm32_addr_for_host64 (Value) Value)
+(rule (wasm32_addr_for_host64 (uextend addr @ (value_type $I32))) addr)
+(rule (wasm32_addr_for_host64 val @ (u32_from_iconst _)) val)
+
 ;; Helper to search for the first argument, `wasm_addr`, within the oob
 ;; condition second argument, `oob`. It should appear as an integer comparison
 ;; of the address against a particular bound.

diff --git a/tests/disas/pulley/memory-inbounds.wat b/tests/disas/pulley/memory-inbounds.wat
@@ -48,15 +48,10 @@
 ;;
 ;; wasm[0]::function[4]::offset_just_bad:
 ;;       push_frame
-;;       xload64le_o32 x10, x0, 88
-;;       xsub64_u8 x10, x10, 4
-;;       xzero x11
-;;       xload64le_o32 x12, x0, 80
-;;       xadd64_u32 x12, x12, 65533
-;;       xconst32 x8, 65533
-;;       xult64 x10, x10, x8
-;;       xselect64 x12, x10, x11, x12
-;;       xload32le_z x0, x12, 0
+;;       xload64le_o32 x5, x0, 88
+;;       xconst32 x6, 65533
+;;       xload64le_o32 x7, x0, 80
+;;       xload32le_g32 x0, x7, x5, x6, 0
 ;;       pop_frame
 ;;       ret
 ;;
@@ -83,15 +78,10 @@
 ;;
 ;; wasm[0]::function[7]::maybe_inbounds:
 ;;       push_frame
-;;       xload64le_o32 x10, x0, 88
-;;       xsub64_u8 x10, x10, 4
-;;       xzero x11
-;;       xload64le_o32 x12, x0, 80
-;;       xadd64_u32 x12, x12, 131068
-;;       xconst32 x8, 131068
-;;       xult64 x10, x10, x8
-;;       xselect64 x12, x10, x11, x12
-;;       xload32le_z x0, x12, 0
+;;       xload64le_o32 x5, x0, 88
+;;       xconst32 x6, 131068
+;;       xload64le_o32 x7, x0, 80
+;;       xload32le_g32 x0, x7, x5, x6, 0
 ;;       pop_frame
 ;;       ret
 ;;
@@ -111,15 +101,10 @@
 ;;
 ;; wasm[0]::function[9]::never_inbounds:
 ;;       push_frame
-;;       xload64le_o32 x10, x0, 88
-;;       xsub64_u8 x10, x10, 4
-;;       xzero x11
-;;       xload64le_o32 x12, x0, 80
-;;       xadd64_u32 x12, x12, 131069
-;;       xconst32 x8, 131069
-;;       xult64 x10, x10, x8
-;;       xselect64 x12, x10, x11, x12
-;;       xload32le_z x0, x12, 0
+;;       xload64le_o32 x5, x0, 88
+;;       xconst32 x6, 131069
+;;       xload64le_o32 x7, x0, 80
+;;       xload32le_g32 x0, x7, x5, x6, 0
 ;;       pop_frame
 ;;       ret
 ;;