Missed optimization for `*(ptr << 2)`

[WaffleLapkin]

For this function (which is essentially *(ptr << 2)) llvm generates the following assembly:

define noundef i32 @shift_read2(ptr noundef nonnull %packed) unnamed_addr #1 personality ptr @rust_eh_personality {
start:
  %n.i.i = ptrtoint ptr %packed to i64
  %offset.i.i.i = mul i64 %n.i.i, 3
  %0 = getelementptr i8, ptr %packed, i64 %offset.i.i.i
  %1 = icmp ne ptr %0, null
  tail call void @llvm.assume(i1 %1)
  %2 = load i32, ptr %0, align 4 ; , !noundef !10
  ret i32 %2
}

shift_read2:                            # @shift_read2
        lea     rax, [rdi + 2*rdi]
        mov     eax, dword ptr [rdi + rax]
        ret

This looks suboptimal, since a single instruction can be generated instead:

shift_read2:                            # @shift_read2
        mov     eax, dword ptr [4*rdi]
        ret

Moreover if you change the function slightly (replace 3 by 7 which is equivalent to replacing *(ptr << 2) by *(ptr << 3)) the generated assembly is better:

define noundef i32 @shift_read3(ptr noundef nonnull %packed) unnamed_addr #1 personality ptr @rust_eh_personality {
start:
  %n.i.i = ptrtoint ptr %packed to i64
  %offset.i.i.i = mul i64 %n.i.i, 7
  %0 = getelementptr i8, ptr %packed, i64 %offset.i.i.i
  %1 = icmp ne ptr %0, null
  tail call void @llvm.assume(i1 %1)
  %2 = load i32, ptr %0, align 4 ; , !noundef !10
  ret i32 %2
}

shift_read3:                            # @shift_read3
        mov     eax, dword ptr [8*rdi]
        ret

• Compiler explorer
• Original Rust code that reproduces the issue
• llvm-mca showing that there is no benefit in compiling this to two instructions

[llvmbot]

@llvm/issue-subscribers-backend-x86

[RKSimon]

SelectionDAG has 12 nodes:
  t0: ch,glue = EntryToken
  t2: i64,ch = CopyFromReg t0, Register:i64 %0
        t13: i64 = X86ISD::MUL_IMM t2, Constant:i64<3>
      t5: i64 = add t2, t13
    t8: i32,ch = load<(load (s32) from %ir.0)> t0, t5, undef:i64
  t11: ch,glue = CopyToReg t0, Register:i32 $eax, t8
  t12: ch = X86ISD::RET_GLUE t11, TargetConstant:i32<0>, Register:i32 $eax, t11:1

We can probably handle this in X86DAGToDAGISel::matchAddressRecursively

[WaffleLapkin]

I've also tested with other constants:

example::shift_read1:
        mov     eax, dword ptr [rdi + rdi]
        ret

example::shift_read2:
        lea     rax, [rdi + 2*rdi]
        mov     eax, dword ptr [rdi + rax]
        ret

example::shift_read3:
        mov     eax, dword ptr [8*rdi]
        ret

example::shift_read4:
        lea     rax, [rdi + 4*rdi]
        lea     rax, [rax + 2*rax]
        mov     eax, dword ptr [rdi + rax]
        ret

example::shift_read5:
        shl     rdi, 5
        mov     eax, dword ptr [rdi]
        ret

example::shift_read6:
        shl     rdi, 6
        mov     eax, dword ptr [rdi]
        ret

• shift_read1 and shift_read3 look good
• shift_read2 should be mov eax, dword ptr [4*rdi] instead
• shift_read4 looks like it should use shl too, although I'm not entirely sure
• shift_read5 and shift_read6 look as expected (it looks like [x+8*x] is the biggest thing that can be encoded in the instruction, so it makes sense to just use the shift)

[RKSimon]

@nikic has posted a candidate patch: https://reviews.llvm.org/D148341