[PATCH] Transform (+ base (* index {2, 4, 8})) to use LEA instruction

Hi,

I'd like to submit a patch that adds IR1 transforms to recognize scaled
addition patterns and convert them to %LEA, generating a single LEA
instruction instead of separate multiply/shift and add operations.

Patterns matched:

  (+ base (* index scale))  where scale in {2, 4, 8}
  (+ base (ash index n))    where n in {1, 2, 3}

Both argument orderings are handled. The ash pattern is needed because
(* x 2) gets transformed to (ash x 1) by the power-of-2 multiply
optimization.

Before:

  (lambda (base index)
    (declare (fixnum base index))
    (+ base (* index 2)))

generated LEA + ADD (two instructions).

After:

  LEA RDX, [RAX+RDI*2]

(single instruction with scaled index addressing).

This was inspired by Matt Godbolt's "Addressing the Adding Situation"
blog post: https://xania.org/202512/02-adding-integers

Tests are included in x86-64-codegen.impure.lisp.

------------------------

Add IR1 transforms that recognize scaled addition patterns and convert
them to %LEA, generating a single LEA instruction instead of separate
multiply/shift and add operations.

Patterns matched:
  (+ base (* index scale))  where scale in {2, 4, 8}
  (+ base (ash index n))    where n in {1, 2, 3}

Both argument orderings are handled. The ash pattern is needed because
(* x 2) gets transformed to (ash x 1) by the power-of-2 multiply
optimization.

Inspired by Matt Godbolt's "Addressing the Adding Situation" blog post:
https://xania.org/202512/02-adding-integers
---
 src/compiler/x86-64/arith.lisp   | 71 ++++++++++++++++++++++++++++++++
 tests/x86-64-codegen.impure.lisp | 70 +++++++++++++++++++++++++++++++
 2 files changed, 141 insertions(+)

diff --git a/src/compiler/x86-64/arith.lisp b/src/compiler/x86-64/arith.lisp
index 2ebe22648..90127177d 100644
--- a/src/compiler/x86-64/arith.lisp
+++ b/src/compiler/x86-64/arith.lisp
@@ -3640,6 +3640,77 @@
   "recode as leas, shifts and adds"
   (*-transformer (lvar-value y) node 'sb-vm::%lea-modfx))

+;;; Transform (+ base (* index scale)) -> (%lea base index scale 0)
+;;; and       (+ base (ash index n))   -> (%lea base index 2^n 0)
+;;; where scale is 2, 4, or 8 (n is 1, 2, or 3).
+;;; This generates a single LEA instruction instead of SHL/IMUL + ADD.
+;;;
+;;; We match both * and ash because:
+;;; - (ash x n) is typically transformed to (* x 2^n) for n in {2,3}
+;;; - But (* x 2) gets transformed BACK to (ash x 1) by the generic
power-of-2 transform
+;;; See Matt Godbolt's "Addressing the Adding Situation" blog post:
+;;; https://xania.org/202512/02-adding-integers
+(defun try-scaled-add-lea-transform (scaled-arg scaled-is-second-arg node
lea-fun)
+  "Try to transform addition with scaled operand to LEA.
+   Returns the transform form or NIL if not applicable.
+   SCALED-IS-SECOND-ARG indicates position of the scaled arg in the +
call."
+  (let ((use (lvar-uses scaled-arg)))
+    (when (combination-p use)
+      (let ((fun-name (lvar-fun-name (combination-fun use))))
+        (cond
+          ;; Match (* index scale) where scale in {2, 4, 8}
+          ((eq fun-name '*)
+           (let* ((args (combination-args use))
+                  (scale-arg (and (= (length args) 2) (second args))))
+             (when (and scale-arg (constant-lvar-p scale-arg))
+               (let ((scale (lvar-value scale-arg)))
+                 (when (member scale '(2 4 8))
+                   (delay-ir1-transform node :constraint)
+                   (splice-fun-args scaled-arg '* 2)
+                   (if scaled-is-second-arg
+                       `(lambda (base index scale)
+                          (declare (ignore scale) (fixnum base index))
+                          (truly-the fixnum (,lea-fun base index ,scale
0)))
+                       `(lambda (index scale base)
+                          (declare (ignore scale) (fixnum base index))
+                          (truly-the fixnum (,lea-fun base index ,scale
0)))))))))
+          ;; Match (ash index n) where n in {1, 2, 3} -> scale in {2, 4, 8}
+          ;; This catches the (* x 2) -> (ash x 1) case
+          ((eq fun-name 'ash)
+           (let* ((args (combination-args use))
+                  (shift-arg (and (= (length args) 2) (second args))))
+             (when (and shift-arg (constant-lvar-p shift-arg))
+               (let ((shift (lvar-value shift-arg)))
+                 (when (member shift '(1 2 3))
+                   (let ((scale (ash 1 shift)))
+                     (delay-ir1-transform node :constraint)
+                     (splice-fun-args scaled-arg 'ash 2)
+                     (if scaled-is-second-arg
+                         `(lambda (base index shift)
+                            (declare (ignore shift) (fixnum base index))
+                            (truly-the fixnum (,lea-fun base index ,scale
0)))
+                         `(lambda (index shift base)
+                            (declare (ignore shift) (fixnum base index))
+                            (truly-the fixnum (,lea-fun base index ,scale
0)))))))))))))))
+
+;; Match fixnum + integer since (* fixnum const) produces wider integer
type
+(deftransform + ((x y) (fixnum integer) * :node node)
+  "use LEA for addition with scaled operand"
+  (or (try-scaled-add-lea-transform y t node '%lea)
+      (give-up-ir1-transform)))
+
+(deftransform + ((x y) (integer fixnum) * :node node)
+  "use LEA for addition with scaled operand"
+  (or (try-scaled-add-lea-transform x nil node '%lea)
+      (give-up-ir1-transform)))
+
+;; Also match fixnum + fixnum for ash case (which doesn't widen the type)
+(deftransform + ((x y) (fixnum fixnum) * :node node)
+  "use LEA for addition with shifted operand"
+  (or (try-scaled-add-lea-transform y t node '%lea)
+      (try-scaled-add-lea-transform x nil node '%lea)
+      (give-up-ir1-transform)))
+
 (defun exactly-one-read-p (results)
   (when results
     (let ((refs (tn-reads (tn-ref-tn results))))
diff --git a/tests/x86-64-codegen.impure.lisp
b/tests/x86-64-codegen.impure.lisp
index dbc2df515..7d2267ac9 100644
--- a/tests/x86-64-codegen.impure.lisp
+++ b/tests/x86-64-codegen.impure.lisp
@@ -1383,3 +1383,73 @@

 (with-test (:name :signed-vops) (test-signed))
 (with-test (:name :unsigned-vops) (test-unsigned))
+
+;;; Test that (+ base (* index {2,4,8})) and (+ base (ash index {1,2,3}))
+;;; compile to LEA instructions with scaled index addressing mode.
+;;; See Matt Godbolt's "Addressing the Adding Situation" blog post:
+;;; https://xania.org/202512/02-adding-integers
+(with-test (:name :scaled-add-uses-lea)
+  (flet ((has-scaled-lea-p (lambda-expr scale)
+           ;; Look for LEA with scaled index, e.g. "LEA RDX, [RAX+RDI*2]"
+           ;; We search for the scale pattern like "*2]" or "*4]" on a LEA
line
+           (let ((scale-pattern (format nil "*~d]" scale)))
+             (loop for line in (disassembly-lines (compile nil
lambda-expr))
+                   thereis (and (search "LEA" line)
+                                (search scale-pattern line)))))
+         (no-imul-p (lambda-expr)
+           (not (loop for line in (disassembly-lines (compile nil
lambda-expr))
+                      thereis (search "IMUL" line)))))
+    ;; Test (* index scale) forms with scales 2, 4, 8
+    (dolist (scale '(2 4 8))
+      (let ((expr `(lambda (base index)
+                     (declare (fixnum base index))
+                     (+ base (* index ,scale)))))
+        (assert (has-scaled-lea-p expr scale)
+                () "Expected LEA with scale ~d for (+ base (* index ~d))"
scale scale)
+        (assert (no-imul-p expr)
+                () "Unexpected IMUL for (+ base (* index ~d))" scale))
+      ;; Also test with arguments reversed: (+ (* index scale) base)
+      (let ((expr `(lambda (base index)
+                     (declare (fixnum base index))
+                     (+ (* index ,scale) base))))
+        (assert (has-scaled-lea-p expr scale)
+                () "Expected LEA with scale ~d for (+ (* index ~d) base)"
scale scale)))
+    ;; Test (ash index shift) forms with shifts 1, 2, 3 -> scales 2, 4, 8
+    (dolist (shift '(1 2 3))
+      (let ((scale (ash 1 shift))
+            (expr `(lambda (base index)
+                     (declare (fixnum base index))
+                     (+ base (ash index ,shift)))))
+        (assert (has-scaled-lea-p expr scale)
+                () "Expected LEA with scale ~d for (+ base (ash index
~d))" scale shift))
+      ;; Also test with arguments reversed
+      (let ((scale (ash 1 shift))
+            (expr `(lambda (base index)
+                     (declare (fixnum base index))
+                     (+ (ash index ,shift) base))))
+        (assert (has-scaled-lea-p expr scale)
+                () "Expected LEA with scale ~d for (+ (ash index ~d)
base)" scale shift)))))
+
+;;; Test correctness of scaled-add LEA optimization
+(with-test (:name :scaled-add-lea-correctness)
+  (let ((test-values '(-1000 -100 -10 -1 0 1 10 100 1000)))
+    ;; Test (* index scale) forms
+    (dolist (scale '(2 4 8))
+      (let ((f (compile nil `(lambda (base index)
+                               (declare (fixnum base index))
+                               (the fixnum (+ base (* index ,scale)))))))
+        (dolist (base test-values)
+          (dolist (index test-values)
+            (assert (= (funcall f base index)
+                       (+ base (* index scale)))
+                    () "Mismatch for (+ ~d (* ~d ~d))" base index
scale)))))
+    ;; Test (ash index shift) forms
+    (dolist (shift '(1 2 3))
+      (let ((f (compile nil `(lambda (base index)
+                               (declare (fixnum base index))
+                               (the fixnum (+ base (ash index ,shift)))))))
+        (dolist (base test-values)
+          (dolist (index test-values)
+            (assert (= (funcall f base index)
+                       (+ base (ash index shift)))
+                    () "Mismatch for (+ ~d (ash ~d ~d))" base index
shift)))))))
--
2.51.2

_______________________________________________
Sbcl-devel mailing list
[email protected]
https://lists.sourceforge.net/lists/listinfo/sbcl-devel