[GlobalISel] Combine and (lshr x, cst), mask -> ubfx x, cst, width

[paquette]

Summary

Only do this post-legalization for now

This is generic now.

[arsenm]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

3914 ↗

(On Diff #333050)

The amounts could be a different type from the value, so I guess I'll need to fix that in the instruction definition

3939 ↗

(On Diff #333050)

getScalarSizeInBits

[paquette]

Rebase + move to AArch64-only to match the G_SBFX combine. Maybe this can be cannibalized into something more generic.

[jroelofs]

llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerCombiner.cpp

289

What is the range of AndImm at this point? I'm worried that we'll have UB if the addition overflows.

[paquette]

• Fix UB by doing all the computations on unsigned ints (which is what SDAG does anyway)
• Simplify the logic a little bit more
• Add boundary testcases which use signed + unsigned max integer values

(Maybe there should be a m_UICst(...) matcher or something...?)

[arsenm]

llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerCombiner.cpp

282–285

This is the only target specific thing here which can be replaced with a legality check on UBFX

308

buildUbfx?

llvm/test/CodeGen/AArch64/GlobalISel/form-bitfield-extract-from-and.mir

4

Why does this need asserts?

[paquette]

llvm/test/CodeGen/AArch64/GlobalISel/form-bitfield-extract-from-and.mir

4

The only-enable-rule stuff only works with asserts.

[paquette]

llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerCombiner.cpp

282–285

The annoying part (for AArch64) is that the legality checks don't work with custom legalization.

I'm not sure what the best way to handle custom actions is.

[arsenm]

llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerCombiner.cpp

282–285

That's what we usually end up adding TLI hooks for, although they're terrible

[paquette]

Make this a generic combine + add a target hook to define when it should be enabled.

[jroelofs]

LGTM

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

4002 ↗

(On Diff #348599)

Depending on subtarget-specific latencies, this may be worthwhile even when the shift has more than one use if lat(ubfx) < lat(and) + lat(lshr).

llvm/lib/Target/AArch64/AArch64Combine.td

 def build_vector_lowering : GICombineGroup<[build_vector_to_dup]>;
 
 def bitfield_extract_from_sext_inreg : GICombineRule<
   (defs root:$root, build_fn_matchinfo:$info),
   (match (wip_match_opcode G_SEXT_INREG):$root,
     [{ return matchBitfieldExtractFromSExtInReg(*${root}, MRI, ${info}); }]),
   (apply [{ return Helper.applyBuildFn(*${root}, ${info}); }])>;
 
-def form_bitfield_extract : GICombineGroup<[bitfield_extract_from_sext_inreg]>;
+def bitfield_extract_from_and : GICombineRule<
+  (defs root:$root, build_fn_matchinfo:$info),
+  (match (wip_match_opcode G_AND):$root,
+    [{ return matchBitfieldExtractFromAnd(*${root}, MRI, ${info}); }]),
+  (apply [{ return Helper.applyBuildFn(*${root}, ${info}); }])>;
+
+def form_bitfield_extract : GICombineGroup<[bitfield_extract_from_sext_inreg,
+                                            bitfield_extract_from_and]>;
 
 def lower_vector_fcmp : GICombineRule<
   (defs root:$root),
   (match (wip_match_opcode G_FCMP):$root,
     [{ return lowerVectorFCMP(*${root}, MRI, B); }]),
   (apply [{}])>;
 
 def form_truncstore_matchdata : GIDefMatchData<"Register">;

llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerCombiner.cpp

static bool matchBitfieldExtractFromSExtInReg(
   MatchInfo = [=](MachineIRBuilder &B) {
     auto Cst1 = B.buildConstant(Ty, ShiftImm);
     auto Cst2 = B.buildConstant(Ty, Width);
     B.buildInstr(TargetOpcode::G_SBFX, {Dst}, {ShiftSrc, Cst1, Cst2});
   };
   return true;
 }
 
+/// Form a G_UBFX from a G_AND fed by a right shift.
+///
+/// e.g.
+///
+/// and (lshr x, cst), mask -> ubfx x, cst, width
+static bool matchBitfieldExtractFromAnd(
+    MachineInstr &MI, MachineRegisterInfo &MRI,
+    std::function<void(MachineIRBuilder &)> &MatchInfo) {
+  assert(MI.getOpcode() == TargetOpcode::G_AND);
+  Register LHS = MI.getOperand(1).getReg();
+  LLT Ty = MRI.getType(LHS);
+
+  // Only legal for s32 and s64.
+  if (Ty != LLT::scalar(32) && Ty != LLT::scalar(64))
+    return false;
+
+  // The mask is a mask of the low bits iff imm & (imm+1) == 0.
+  int64_t AndImm;
+  if (!mi_match(MI.getOperand(2).getReg(), MRI, m_ICst(AndImm)) ||
+      (AndImm & (AndImm + 1)))
+    return false;
+
+  // Check if we have a LSB that fits within the size of the register.
+  const unsigned Size = Ty.getScalarSizeInBits();
+  Register ShiftSrc;
+  int64_t LSBImm;
+  if (!mi_match(LHS, MRI,
+                m_OneNonDBGUse(m_GLShr(m_Reg(ShiftSrc), m_ICst(LSBImm)))) ||
+      static_cast<uint64_t>(LSBImm) >= Size)
+    return false;
+
+  Register Dst = MI.getOperand(0).getReg();
+  int64_t Width = APInt(Size, AndImm).countTrailingOnes();
+  MatchInfo = [=](MachineIRBuilder &B) {
+    auto WidthCst = B.buildConstant(Ty, Width);
+    auto LSBCst = B.buildConstant(Ty, LSBImm);
+    B.buildInstr(TargetOpcode::G_UBFX, {Dst}, {ShiftSrc, LSBCst, WidthCst});
+  };
+  return true;
+}
+
 #define AARCH64POSTLEGALIZERCOMBINERHELPER_GENCOMBINERHELPER_DEPS
 #include "AArch64GenPostLegalizeGICombiner.inc"
 #undef AARCH64POSTLEGALIZERCOMBINERHELPER_GENCOMBINERHELPER_DEPS
 
 namespace {
 #define AARCH64POSTLEGALIZERCOMBINERHELPER_GENCOMBINERHELPER_H
 #include "AArch64GenPostLegalizeGICombiner.inc"
 #undef AARCH64POSTLEGALIZERCOMBINERHELPER_GENCOMBINERHELPER_H

llvm/test/CodeGen/AArch64/GlobalISel/form-bitfield-extract-from-and.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple aarch64 -run-pass=aarch64-postlegalizer-combiner -verify-machineinstrs %s -o - | FileCheck %s
+
+# Check that we can combine
+#
+# and (lshr x, cst), mask -> ubfx x, cst, width
+
+...
+---
+name:            ubfx_s32
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; LSB = 5
+    ; Width = LSB + trailing_ones(255) - 1 =
+    ;         5 + 8 - 1 = 12
+
+    ; CHECK-LABEL: name: ubfx_s32
+    ; CHECK: liveins: $w0
+    ; CHECK: %x:_(s32) = COPY $w0
+    ; CHECK: %lsb:_(s32) = G_CONSTANT i32 5
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 8
+    ; CHECK: %and:_(s32) = G_UBFX %x, %lsb(s32), [[C]]
+    ; CHECK: $w0 = COPY %and(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %lsb:_(s32) = G_CONSTANT i32 5
+    %mask:_(s32) = G_CONSTANT i32 255
+    %shift:_(s32) = G_LSHR %x, %lsb
+    %and:_(s32) = G_AND %shift, %mask
+    $w0 = COPY %and
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            ubfx_s64
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $x0
+
+    ; LSB = 5
+    ; Width = LSB + trailing_ones(1) - 1 =
+    ;         5 + 1 - 1 = 5
+
+    ; CHECK-LABEL: name: ubfx_s64
+    ; CHECK: liveins: $x0
+    ; CHECK: %x:_(s64) = COPY $x0
+    ; CHECK: %lsb:_(s64) = G_CONSTANT i64 5
+    ; CHECK: %mask:_(s64) = G_CONSTANT i64 1
+    ; CHECK: %and:_(s64) = G_UBFX %x, %lsb(s64), %mask
+    ; CHECK: $x0 = COPY %and(s64)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %x:_(s64) = COPY $x0
+    %lsb:_(s64) = G_CONSTANT i64 5
+    %mask:_(s64) = G_CONSTANT i64 1
+    %shift:_(s64) = G_LSHR %x, %lsb
+    %and:_(s64) = G_AND %shift, %mask
+    $x0 = COPY %and
+    RET_ReallyLR implicit $x0
+
+...
+---
+name:            dont_combine_no_and_cst
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $w0, $w1
+
+    ; UBFX needs to be selected to UBFMWri/UBFMXri, so we need constants.
+
+    ; CHECK-LABEL: name: dont_combine_no_and_cst
+    ; CHECK: liveins: $w0, $w1
+    ; CHECK: %x:_(s32) = COPY $w0
+    ; CHECK: %y:_(s32) = COPY $w1
+    ; CHECK: %lsb:_(s32) = G_CONSTANT i32 5
+    ; CHECK: %shift:_(s32) = G_LSHR %x, %lsb(s32)
+    ; CHECK: %and:_(s32) = G_AND %shift, %y
+    ; CHECK: $w0 = COPY %and(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %y:_(s32) = COPY $w1
+    %lsb:_(s32) = G_CONSTANT i32 5
+    %shift:_(s32) = G_LSHR %x, %lsb
+    %and:_(s32) = G_AND %shift, %y
+    $w0 = COPY %and
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            dont_combine_and_cst_not_mask
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $w0
+    ; CHECK-LABEL: name: dont_combine_and_cst_not_mask
+    ; CHECK: liveins: $w0
+    ; CHECK: %x:_(s32) = COPY $w0
+    ; CHECK: %lsb:_(s32) = G_CONSTANT i32 5
+    ; CHECK: %not_a_mask:_(s32) = G_CONSTANT i32 2
+    ; CHECK: %shift:_(s32) = G_LSHR %x, %lsb(s32)
+    ; CHECK: %and:_(s32) = G_AND %shift, %not_a_mask
+    ; CHECK: $w0 = COPY %and(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %lsb:_(s32) = G_CONSTANT i32 5
+    %not_a_mask:_(s32) = G_CONSTANT i32 2
+    %shift:_(s32) = G_LSHR %x, %lsb
+    %and:_(s32) = G_AND %shift, %not_a_mask
+    $w0 = COPY %and
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            dont_combine_shift_more_than_one_use
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $x0
+    ; CHECK-LABEL: name: dont_combine_shift_more_than_one_use
+    ; CHECK: liveins: $x0
+    ; CHECK: %x:_(s64) = COPY $x0
+    ; CHECK: %lsb:_(s64) = G_CONSTANT i64 5
+    ; CHECK: %mask:_(s64) = G_CONSTANT i64 1
+    ; CHECK: %shift:_(s64) = G_LSHR %x, %lsb(s64)
+    ; CHECK: %and:_(s64) = G_AND %shift, %mask
+    ; CHECK: %sub:_(s64) = G_SUB %and, %shift
+    ; CHECK: $x0 = COPY %sub(s64)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %x:_(s64) = COPY $x0
+    %lsb:_(s64) = G_CONSTANT i64 5
+    %mask:_(s64) = G_CONSTANT i64 1
+    %shift:_(s64) = G_LSHR %x, %lsb
+    %and:_(s64) = G_AND %shift, %mask
+    %sub:_(s64) = G_SUB %and, %shift
+    $x0 = COPY %sub
+    RET_ReallyLR implicit $x0
+
+...
+---
+name:            dont_combine_negative_lsb
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; LSB must be in [0, reg_size)
+
+    ; CHECK-LABEL: name: dont_combine_negative_lsb
+    ; CHECK: liveins: $w0
+    ; CHECK: %x:_(s32) = COPY $w0
+    ; CHECK: %negative:_(s32) = G_CONSTANT i32 -1
+    ; CHECK: %mask:_(s32) = G_CONSTANT i32 255
+    ; CHECK: %shift:_(s32) = G_LSHR %x, %negative(s32)
+    ; CHECK: %and:_(s32) = G_AND %shift, %mask
+    ; CHECK: $w0 = COPY %and(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %negative:_(s32) = G_CONSTANT i32 -1
+    %mask:_(s32) = G_CONSTANT i32 255
+    %shift:_(s32) = G_LSHR %x, %negative
+    %and:_(s32) = G_AND %shift, %mask
+    $w0 = COPY %and
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            dont_combine_lsb_too_large
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $w0
+
+    ; LSB must be in [0, reg_size)
+
+    ; CHECK-LABEL: name: dont_combine_lsb_too_large
+    ; CHECK: liveins: $w0
+    ; CHECK: %x:_(s32) = COPY $w0
+    ; CHECK: %too_large:_(s32) = G_CONSTANT i32 32
+    ; CHECK: %mask:_(s32) = G_CONSTANT i32 255
+    ; CHECK: %shift:_(s32) = G_LSHR %x, %too_large(s32)
+    ; CHECK: %and:_(s32) = G_AND %shift, %mask
+    ; CHECK: $w0 = COPY %and(s32)
+    ; CHECK: RET_ReallyLR implicit $w0
+    %x:_(s32) = COPY $w0
+    %too_large:_(s32) = G_CONSTANT i32 32
+    %mask:_(s32) = G_CONSTANT i32 255
+    %shift:_(s32) = G_LSHR %x, %too_large
+    %and:_(s32) = G_AND %shift, %mask
+    $w0 = COPY %and
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            dont_combine_vector
+tracksRegLiveness: true
+legalized: true
+body:             |
+  bb.0:
+    liveins: $d0
+    ; CHECK-LABEL: name: dont_combine_vector
+    ; CHECK: liveins: $d0
+    ; CHECK: %x:_(<2 x s32>) = COPY $d0
+    ; CHECK: %lsb_cst:_(s32) = G_CONSTANT i32 5
+    ; CHECK: %lsb:_(<2 x s32>) = G_BUILD_VECTOR %lsb_cst(s32), %lsb_cst(s32)
+    ; CHECK: %mask_cst:_(s32) = G_CONSTANT i32 255
+    ; CHECK: %mask:_(<2 x s32>) = G_BUILD_VECTOR %mask_cst(s32), %mask_cst(s32)
+    ; CHECK: %shift:_(<2 x s32>) = G_LSHR %x, %lsb(<2 x s32>)
+    ; CHECK: %and:_(<2 x s32>) = G_AND %shift, %mask
+    ; CHECK: $d0 = COPY %and(<2 x s32>)
+    ; CHECK: RET_ReallyLR implicit $d0
+    %x:_(<2 x s32>) = COPY $d0
+    %lsb_cst:_(s32) = G_CONSTANT i32 5
+    %lsb:_(<2 x s32>) = G_BUILD_VECTOR %lsb_cst, %lsb_cst
+    %mask_cst:_(s32) = G_CONSTANT i32 255
+    %mask:_(<2 x s32>) = G_BUILD_VECTOR %mask_cst, %mask_cst
+    %shift:_(<2 x s32>) = G_LSHR %x, %lsb
+    %and:_(<2 x s32>) = G_AND %shift, %mask
+    $d0 = COPY %and
+    RET_ReallyLR implicit $d0