Skip to content

Use untagged indices for string-like access in flambda#2930

Merged
dvulakh merged 3 commits intomainfrom
dvulakh.use-untagged-indices-for-string-access-in-flambda
Aug 26, 2024
Merged

Use untagged indices for string-like access in flambda#2930
dvulakh merged 3 commits intomainfrom
dvulakh.use-untagged-indices-for-string-access-in-flambda

Conversation

@dvulakh
Copy link
Copy Markdown
Contributor

@dvulakh dvulakh commented Aug 15, 2024

The primitives added in #2906 for string-like access via unboxed indices generate bad code. For example, %caml_string_get64u#_indexed_by_int64# generates:

salq $1, %rbx
sarq $1, %rbx
movq (%rax,%rbx), %rax
ret

This is because the flambda primitives for string-like load and store operate on tagged indices, and the untagging does not happen until conversion to cmm. So an unboxed index is first tagged when the primitive is translated to flambda, and then untagged by cmm. The simplifications in cmm_helpers.ml can't rewrite a left shift followed by a right shift, so the tag-untag sequence is not eliminated.

Flambda's array-access primitives also operate on tagged indices, but these indices are shifted further left when they're used, so the simplifications in cmm_helpers.ml are able to eliminate the right shift for untagging and combine the left shifts. As a result, %array_unsafe_get_indexed_by_int64# already generates the sensible:

movq (%rax,%rbx,8), %rax
ret

This PR changes the flambda primitives for string-like access to use untagged indices, so %caml_string_get64u#_indexed_by_int64# now generates

movq (%rax,%rbx), %rax
ret

For testing, I skimmed the output of

#!/bin/bash

fin=/tmp/a.ml
fout=/tmp/a.s

for index in int32 int64 nativeint
do
  for width in int16 int32 'int32#' float32 'float32#' int64 'int64#'
  do
    for container in string bigstring bytes
    do
      access_sigil=`echo $width | sed -E 's/(f?)(loat|int)([0-9]+)(#?)/\1\3u\4/'`
      access_type=`echo $width | sed 's/16//'`
      printf "open Bigarray\n`
             `type bigstring = (char, int8_unsigned_elt, c_layout) Array1.t\n`
             `external get_unsafe`
             ` : $container`
             ` -> ${index}#`
             ` -> $access_type`
             ` = \"%%caml_${container}_get${access_sigil}_indexed_by_${index}#\"\n`
             `external set_unsafe`
             ` : $container`
             ` -> ${index}#`
             ` -> $access_type`
             ` -> unit`
             ` = \"%%caml_${container}_set${access_sigil}_indexed_by_${index}#\"\n`
             `let[@inline never] get s i = get_unsafe s i\n`
             `let[@inline never] set s i v = set_unsafe s i v" > $fin
      $INSTALL_DIR/bin/ocamlopt.opt -S -c -O3 $fin
      echo "==== access $width in $container by ${index}# ===="
      sed -nE '/^caml.*__get_.*code/,/^\s*ret$/p' $fout
      sed -nE '/^caml.*__set_.*code/,/^\s*ret$/p' $fout
    done
  done
done

and the various combinations generally look right.

this fixes codegen when indexing by an unboxed integer

Signed-off-by: David Vulakh <dvulakh@janestreet.com>
@dvulakh dvulakh added the flambda2 Prerequisite for, or part of, flambda2 label Aug 15, 2024
@dvulakh dvulakh requested a review from TheNumbat August 15, 2024 21:39
Signed-off-by: David Vulakh <dvulakh@janestreet.com>
Signed-off-by: David Vulakh <dvulakh@janestreet.com>
@mshinwell
Copy link
Copy Markdown
Collaborator

An flambda2 dev will need to review this, I'll make sure that happens.

@dvulakh dvulakh removed the request for review from TheNumbat August 16, 2024 14:49
@mshinwell
Copy link
Copy Markdown
Collaborator

I've asked @Gbury if he can review

@mshinwell mshinwell requested a review from Gbury August 16, 2024 15:57
@dvulakh dvulakh merged commit 087a98a into main Aug 26, 2024
@dvulakh dvulakh deleted the dvulakh.use-untagged-indices-for-string-access-in-flambda branch August 26, 2024 14:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

flambda2 Prerequisite for, or part of, flambda2

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants