<table border="1" cellspacing="0" cellpadding="8">
    <tr>
        <th>Issue</th>
        <td>
            <a href=https://github.com/llvm/llvm-project/issues/59789>59789</a>
        </td>
    </tr>

    <tr>
        <th>Summary</th>
        <td>
            Inefficient mask generation for llvm.masked.gather on avx2 targets
        </td>
    </tr>

    <tr>
      <th>Labels</th>
      <td>
            new issue
      </td>
    </tr>

    <tr>
      <th>Assignees</th>
      <td>
      </td>
    </tr>

    <tr>
      <th>Reporter</th>
      <td>
          jhorstmann
      </td>
    </tr>
</table>

<pre>
    In some cases llvm knows how to efficiently create a wide vector mask (`<4 x i32>`) from a bit mask (`<4x i1>`). For example in

```llvm
define void @mask_from_bits(i8 %mask, ptr noalias nocapture noundef writeonly dereferenceable(16) %out) {
  %1 = trunc i8 %mask to i4
  %2 = bitcast i4 %1 to <4 x i1>
  %3 = sext <4 x i1> %2 to <4 x i64>
  store <4 x i64> %3, ptr %out, align 4
  ret void
}
```

The assembly with `-mcpu=skylake` broadcasts the mask and compares it against a predefined vector of lane indices:

```asm
.LCPI3_0:
  .quad 1 # 0x1
  .quad 2 # 0x2
 .quad 4 # 0x4
  .quad 8 # 0x8
mask_from_bits: # @mask_from_bits
 vmovd xmm0, edi
  vpbroadcastd ymm0, xmm0
  vmovdqa ymm1, ymmword ptr [rip + .LCPI3_0] # ymm1 = [1,2,4,8]
  vpand ymm0, ymm0, ymm1
  vpcmpeqq ymm0, ymm0, ymm1
  vmovdqu ymmword ptr [rsi], ymm0
  vzeroupper
 ret
```

This also works in combination with the avx2 specific `llvm.x86.avx2.gather` instructions:

```llvm
define void @gather_masked_avx(ptr noalias nocapture noundef readonly align 4 dereferenceable(16) %indices, i8 %mask, ptr noalias noundef nonnull readonly align 4 %data, ptr noalias nocapture noundef writeonly align 4 dereferenceable(16) %output) {
  %1 = trunc i8 %mask to i4
  %2 = bitcast i4 %1 to <4 x i1>

  %3 = load <4 x i32>, ptr %indices, align 4
  %4 = zext <4 x i32> %3 to <4 x i64>

  %5 = getelementptr i32, ptr %data, <4 x i64> %4

  %6 = sext <4 x i1> %2 to <4 x i32>
  %7 = tail call <4 x i32> @llvm.x86.avx2.gather.q.d.256(<4 x i32> zeroinitializer, ptr %data, <4 x i64> %4, <4 x i32> %6, i8 4)
  store <4 x i32> %7, ptr %output
  ret void
}

declare <4 x i32> @llvm.x86.avx2.gather.q.d.256(<4 x i32>, ptr, <4 x i64>, <4 x i32>, i8) nounwind readonly
```

However, when using the more portable `llvm.masked.gather` instructions, the mask generation uses a scalar sequence of bit tests and vector inserts:

```llvm
define void @gather_masked_portable(ptr noalias nocapture noundef readonly align 4 dereferenceable(16) %indices, i8 %mask, ptr noalias noundef nonnull readonly align 4 %data, ptr noalias nocapture noundef writeonly align 4 dereferenceable(16) %output) {
  %1 = trunc i8 %mask to i4
  %2 = bitcast i4 %1 to <4 x i1>

  %3 = load <4 x i32>, ptr %indices, align 4
  %4 = zext <4 x i32> %3 to <4 x i64>

  %5 = getelementptr i32, ptr %data, <4 x i64> %4

  %6 = tail call <4 x i32> @llvm.masked.gather.v4i32.v4p0(<4 x ptr> %5, i32 4, <4 x i1> %2, <4 x i32> zeroinitializer)
  store <4 x i32> %6, ptr %output
  ret void
}

declare <4 x i32> @llvm.masked.gather.v4i32.v4p0(<4 x ptr>, i32 immarg, <4 x i1>, <4 x i32>)
```

```asm
gather_masked_portable: # @gather_masked_portable
  mov eax, esi
  and eax, 1
  neg eax
 vmovd xmm0, eax
  mov eax, esi
  shr al
  movzx eax, al
  and eax, 1
  neg eax
  vpinsrd xmm0, xmm0, eax, 1
  mov eax, esi
 shr al, 2
  movzx eax, al
  and eax, 1
  neg eax
  vpinsrd xmm0, xmm0, eax, 2
  and sil, 8
  shr sil, 3
  movzx eax, sil
  neg eax
  vpinsrd xmm0, xmm0, eax, 3
  vpmovzxdq ymm1, xmmword ptr [rdi] # ymm1 = mem[0],zero,mem[1],zero,mem[2],zero,mem[3],zero
  vpxor xmm2, xmm2, xmm2
  vpgatherqd xmm2, xmmword ptr [rdx + 4*ymm1], xmm0
  vmovdqa xmmword ptr [rcx], xmm2
  vzeroupper
  ret
```

[Compiler explorer](https://llvm.godbolt.org/z/f56EE6a9a)

This usage of bitmasks is quite common in the [apache arrow in-memory format](https://arrow.apache.org/docs/format/Columnar.html#validity-bitmaps) where an array consists of data buffers and a separate validity bitmap.
</pre>
<img width="1px" height="1px" alt="" src="http://email.email.llvm.org/o/eJzsWEtv4z4O_zTKhahhy48khxym6RQ7wB72sPeCsZlE_8qSI8l59NMvZDl20qSP2cV_TwMUTUvxLfJHKmit2CiiBcsfWf40wdZttVn8tdXGuhqVmqx0dVr8UmB1TVCiJQtS7mt4VfpgYasP4DTQei1KQcrJE5SG0BEgHERFsKfSaQM12ldgfMaKmKXLDI4gUs7Sn_5_Poe10TUgrIR7z3kEkQx8ETxrA3TEupEEQrH4icU_-t9FHH68d4FU0Voogr0WFbAs9ppfvKWXlXCW8ZmYAeO5JzO-hMYZUBqlQAtKl9i41hAo3aqK1nAwwpFW8gQVGVqTIVUSriQxPksKHwPjuW5d99f0MTgAnpgAS5_AmVaVMFr0WRPZBRvv2FbClWgdiCyIOg3nfHVpGPnTjt_S0V1xBFUXYkV2IWedNnR91Ok6x3-OYQkoxUbB4KAh1-Wxz_X06V3SL2_i31sCtJbqlTzBQbgtsCJ-qMumZemTfT1JfCVWxLAyGisfrQW3pXDxqCoodd2gIQvCAW5QKOsAoTEU7rM615Reg0Tl66ASJVmW_rhbD2j7coj-ufzXr_QlHhgBol2LFSTAeArxMbkm857Me3KgZj01u2ae9eRZIL-rtfRHd3xbhEHHvtb7Co51HfvUUyXOuvfNkKQKTv15x3dm8JI79GeJPzvV9UGbKlxm_mhEA4w_whB6_tQ54tm7-mH5o5fjjC8zxpczlj-Ntv1lnI1efCYjR1k3tNt9ztR52N54ZoW3dRY6M7-R0W3TkOkphtynlSYsoLQaDtq8WhDKF89KKHRCq1B7vrRwf-RgGyrFWpTQY0R0nBWRP4k26LZkfEn6WjNt6aU_rKcP8SWoefE3TNUL7o-Mzz4HFUNYdZjSd9tn2HKucr6Ez3ArKFZaqVbKWwOM5xU6_B28-4ZvunXN_wf6bgBQaqzgeqSMWHaRs3d4xniedfJvlwDayQfd9xH0QjzvxDfkSFJNynmTXsFo_ZzoG7jNbpQV3wTzEOAoNw0JRiGhRCnfB5LdLfRoF1URzws_ZK_4ffMJJZxAKd7IfC-UC-qQvqKv0ozx-f3RM7BOr0ePL6MvRk7feqXEW32_F3Bv-iawm5hCOL6-fXschKqGzvoEnf6hD7QPaTxsSUFrhdqEUedT0WjjfCMNeBSA4yM04stxSm5IkQkY1_qFDMGWKNGApV3rG9TPRr9MOfLT1eN4PzOFsmTc_4ptZ9f_ANwfgPsdgPsKqK46INpnIuXRPmvisXF9uwYbeVcoKYcrCBpQ8xaXbtDtK2gq_g5o-m6I5-hEXaPZvA_xDkLNP0Gi2134g34e99QPGPoE1HoPhMduW7XDtupxpqcO65-iTUe7u-gO9A8U2q0BlBc8b8cz10j-2irsG6GsGQ1fOnApdteL3gm-BP63e8IvVVnRWZ1dZqOnpXc98Yf_ldl05OsUVrvhQXF8t7ZX4uYNUVPN8sc4rPO-yxhfBlpyh8bv0NIL2uDJURtvnfdejJ8DR6jRXXV5fO3ssXv9ZIz_6OIJD45776f3kuVxZOYfvk6-eJ6w_HGp60ZIMkDHRmpDplM72zrXdEOYPzP-3AHDRlcrLV2kfbc_vzH-vM6Lnz8LnOPY3OOrp7W4OY9536QWhIVdKxz5J1CtlX8M-YWB5Y_YYOlfQcboAwj1UFOtzQnW2tTo7jnUcUZBrHeo0qX1PgUZ_rzUsq0Vmmjrasl4ukcpKuFOD50_jfWz8rAlQ4DKW8YTlFpZ4fcRvQY_OmDVrtdkwn6CYKlBg47grCqE1kSTapFW83SOE1okxZQXU57Nk8l2MUuzWVKVFKfFfJogn88SzimnHHEeZymfiAWPeRonMU_yJMuTaF0VSZWvqZzGJS_zFctiqlHIqLsCbTYTYW1Li3w-nc0nElckbffVGOeKDtAdMu5LeGIWXuZh1W6sx3Yf2KjFCSdp8UsN34vd7G1rbeB2IIBW4bHq0GzI2Ulr5OL6bjbCbdtVVOq6r5z-46Ex-i8q_d10bvrb6sL4TwAAAP__XajeFA">