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

That doesn't seem to handle overflow.

On Wed, Dec 3, 2025 at 2:49 AM Anthony Green <[email protected]> wrote:
>
> 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


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