[llvm] [AArch64] Combine and and lsl into ubfiz (PR #118974)

[Cullen Rhodes via llvm-commits]

c-rhodes wrote:

> > Not sure I follow, what should handle other mask sizes? In the second link you sent inst-combine does canonicalize to the same IR: https://godbolt.org/z/qqqosjP98
> > but this musn't be happening in compilation pipeline? Also, your example is and/shl, whereas the original example is shl(zext). I noticed in SDAG there's no `ZERO_EXTEND_INREG` node as an AND with constant is used instead, are you using and/shl as that's the canonical form or something?
> 
> Yeah sorry. My point is that this adds patterns for `shl(and(x, 0xff))` and `shl(and(x, 0xffff))`, but that does not handle all the other masks that could be used, but would be valid to transform to UBFM. Maybe they don't come up in practice because of the canonicalization that the mid-end does? If DAG isn't canonicalizing `shl(and)` to `and(shl)` then I can imagine they could, maybe we should be either trying to canonicalize it in DAG (to capture this case, might cause other problems), or if it is possible to generalize the pattern or use DAG2DAG to handle the `shl(and)` patterns with any and mask. If not this seems OK.

Thanks I see your point now. Ok so yes as you've mentioned we do have patterns for `and(shl)`, for example this IR is equivalent to the original C example on the issue and is lowered to `ubfiz`: https://godbolt.org/z/qc7xosPxW

```
; uint64_t u8(uint8_t x) { return x << 1; }
define i64 @f(i64 %x) {
  %and = shl i64 %x, 1
  %shl = and i64 %and, u0x1fe
  ret i64 %shl
}
```

and the mid-end can canonicalize the inverse `shl(and) -> and(shl)`: https://godbolt.org/z/7fM53Ydov
```
define i64 @f(i64 %x) {
  %and = and i64 %x, u0xff
  %shl = shl i64 %and, 1
  ret i64 %shl
}

; opt -passes=instcombine

define i64 @f(i64 %x) {
  %and = shl i64 %x, 1
  %shl = and i64 %and, 510 ; u0x1fe
  ret i64 %shl
}
```

however the IR pre-isel for the original example is:
```
define i64 @u8(i8 %0) {
  %2 = zext i8 %0 to i64
  %3 = shl nuw nsw i64 %2, 1
  ret i64 %3
}
```
the `and` is introduced during isel so mid-end canonicalizer is no help there of course. 
```
Initial selection DAG: %bb.0 'lsl_zext_i8_i64:'
SelectionDAG has 10 nodes:
  t0: ch,glue = EntryToken
          t2: i32,ch = CopyFromReg t0, Register:i32 %0
        t3: i8 = truncate t2
      t4: i64 = zero_extend t3
    t6: i64 = shl nuw nsw t4, Constant:i64<1>
  t8: ch,glue = CopyToReg t0, Register:i64 $x0, t6
  t9: ch = AArch64ISD::RET_GLUE t8, Register:i64 $x0, t8:1

...

Combining: t4: i64 = zero_extend t3
Creating new node: t10: i64 = any_extend t2
Creating constant: t11: i64 = Constant<255>
Creating new node: t12: i64 = and t10, Constant:i64<255>
 ... into: t12: i64 = and t10, Constant:i64<2

...

Optimized lowered selection DAG: %bb.0 'lsl_zext_i8_i64:'
SelectionDAG has 11 nodes:
  t0: ch,glue = EntryToken
          t2: i32,ch = CopyFromReg t0, Register:i32 %0
        t10: i64 = any_extend t2
      t12: i64 = and t10, Constant:i64<255>
    t6: i64 = shl nuw nsw t12, Constant:i64<1>
  t8: ch,glue = CopyToReg t0, Register:i64 $x0, t6
  t9: ch = AArch64ISD::RET_GLUE t8, Register:i64 $x0, t8:1
```

So adding a `shl(and) -> and(shl)` combine in isel seems like a sensible approach. I sat down with @david-arm on Tuesday as well and he had the same insight, so I'll give this a try. It seems preferable to adding more patterns in DAG2DAG.

https://github.com/llvm/llvm-project/pull/118974