Recommit [AArch64] Improve codegen for shifted mask op

[bcl5980]

Summary

The original change compares APInt to check the constant is the same or not. But shift amount may have different constant types.
So, this patch change to use getZExtValue to compare constant value.

Original comment:
The special case for bit extraction pattern is ((x >> C) & mask) << C.
It can be combined to x & (mask << C) by return true in isDesirableToCommuteWithShift.

Fix: #56427

[Allen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14404

Does it right when (mask << C) overflow ?

[bcl5980]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14404

Yes, this patch only change if we can change the sequence of shift and binop(and, or, add).
Overflow logic should not have any difference.

[bcl5980]

ping

[bcl5980]

Ping

[dmgreen]

llvm/test/CodeGen/AArch64/pr56427.ll

4 ↗

(On Diff #468003)

Can you find a better place for this test?
Do we have a better example that will not be optimized away by opt?

[bcl5980]

update test.

[dmgreen]

Thanks. This LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14404

This comment doesn't parse very well. Perhaps try to reword the sentance.

[nathanchance]

I bisect a crash while compiling the Linux kernel to this change. A simplified C reproducer (sorry, I do not have time at the moment for llvm-reduce):

union drbd_state {
  struct {
    unsigned : 5;
    unsigned : 4;
    unsigned susp_nod : 1;
    unsigned susp_fen : 1;
  };
  int i;
} drbd_req_state_ns;
enum chg_state_flags {
  CS_DC_PDSK
} _drbd_set_state(union drbd_state, enum chg_state_flags, int *);
struct {
  unsigned susp_nod : 1;
  unsigned susp_fen : 1;
} resource;
int apply_mask_val_mask_1, drbd_req_state_done;
enum chg_state_flags drbd_req_state_f;
void drbd_req_state() {
  union drbd_state rv;
  rv.susp_nod = resource.susp_nod;
  rv.susp_fen = resource.susp_fen;
  drbd_req_state_ns.i = rv.i & apply_mask_val_mask_1;
  _drbd_set_state(drbd_req_state_ns, drbd_req_state_f, &drbd_req_state_done);
}

$ clang --target=aarch64-linux-gnu -O2 -c -o /dev/null drbd_state.i
clang: /home/nathan/cbl/src/llvm-project/llvm/include/llvm/ADT/APInt.h:1027: bool llvm::APInt::operator==(const llvm::APInt &) const: Assertion `BitWidth == RHS.BitWidth && "Comparison requires equal bit widths"' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace, preprocessed source, and associated run script.
Stack dump:
0.      Program arguments: clang --target=aarch64-linux-gnu -O2 -c -o /dev/null drbd_state.i
1.      <eof> parser at end of file
2.      Code generation
3.      Running pass 'Function Pass Manager' on module 'drbd_state.i'.
4.      Running pass 'AArch64 Instruction Selection' on function '@drbd_req_state'
 #0 0x000056317d224071 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x408b071)
 #1 0x000056317d221f9e llvm::sys::RunSignalHandlers() (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x4088f9e)
 #2 0x000056317d22344b llvm::sys::CleanupOnSignal(unsigned long) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x408a44b)
 #3 0x000056317d1a6e7e (anonymous namespace)::CrashRecoveryContextImpl::HandleCrash(int, unsigned long) CrashRecoveryContext.cpp:0:0
 #4 0x000056317d1a70d7 CrashRecoverySignalHandler(int) CrashRecoveryContext.cpp:0:0
 #5 0x00007f04d3bc6a00 (/usr/lib/libc.so.6+0x38a00)
 #6 0x00007f04d3c1664c (/usr/lib/libc.so.6+0x8864c)
 #7 0x00007f04d3bc6958 raise (/usr/lib/libc.so.6+0x38958)
 #8 0x00007f04d3bb053d abort (/usr/lib/libc.so.6+0x2253d)
 #9 0x00007f04d3bb045c (/usr/lib/libc.so.6+0x2245c)
#10 0x00007f04d3bbf486 (/usr/lib/libc.so.6+0x31486)
#11 0x000056317ba4ad5e (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x28b1d5e)
#12 0x000056317e3b3d9a (anonymous namespace)::DAGCombiner::visitShiftByConstant(llvm::SDNode*) DAGCombiner.cpp:0:0
#13 0x000056317e35f1f0 (anonymous namespace)::DAGCombiner::visitSHL(llvm::SDNode*) DAGCombiner.cpp:0:0
#14 0x000056317e34b6c8 (anonymous namespace)::DAGCombiner::combine(llvm::SDNode*) DAGCombiner.cpp:0:0
#15 0x000056317e346ad4 llvm::SelectionDAG::Combine(llvm::CombineLevel, llvm::AAResults*, llvm::CodeGenOpt::Level) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x51adad4)
#16 0x000056317e4d1452 llvm::SelectionDAGISel::CodeGenAndEmitDAG() (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x5338452)
#17 0x000056317e4d0082 llvm::SelectionDAGISel::SelectAllBasicBlocks(llvm::Function const&) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x5337082)
#18 0x000056317e4cc6b5 llvm::SelectionDAGISel::runOnMachineFunction(llvm::MachineFunction&) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x53336b5)
#19 0x000056317c60740b llvm::MachineFunctionPass::runOnFunction(llvm::Function&) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x346e40b)
#20 0x000056317caec265 llvm::FPPassManager::runOnFunction(llvm::Function&) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x3953265)
#21 0x000056317caf3e91 llvm::FPPassManager::runOnModule(llvm::Module&) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x395ae91)
#22 0x000056317caecbd8 llvm::legacy::PassManagerImpl::run(llvm::Module&) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x3953bd8)
#23 0x000056317db0846f clang::EmitBackendOutput(clang::DiagnosticsEngine&, clang::HeaderSearchOptions const&, clang::CodeGenOptions const&, clang::TargetOptions const&, clang::LangOptions const&, llvm::StringRef, llvm::Module*, clang::BackendAction, std::unique_ptr<llv
m::raw_pwrite_stream, std::default_delete<llvm::raw_pwrite_stream>>) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x496f46f)
#24 0x000056317e02d1ad clang::BackendConsumer::HandleTranslationUnit(clang::ASTContext&) CodeGenAction.cpp:0:0
#25 0x000056317e8dd789 clang::ParseAST(clang::Sema&, bool, bool) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x5744789)
#26 0x000056317e02b1d4 clang::CodeGenAction::ExecuteAction() (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x4e921d4)
#27 0x000056317df56a7e clang::FrontendAction::Execute() (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x4dbda7e)
#28 0x000056317debbbbf clang::CompilerInstance::ExecuteAction(clang::FrontendAction&) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x4d22bbf)
#29 0x000056317e025d63 clang::ExecuteCompilerInvocation(clang::CompilerInstance*) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x4e8cd63)
#30 0x000056317b9cad97 cc1_main(llvm::ArrayRef<char const*>, char const*, void*) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x2831d97)
#31 0x000056317b9c772b ExecuteCC1Tool(llvm::SmallVectorImpl<char const*>&) driver.cpp:0:0
#32 0x000056317dd20ae2 void llvm::function_ref<void ()>::callback_fn<clang::driver::CC1Command::Execute(llvm::ArrayRef<llvm::Optional<llvm::StringRef>>, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char>>*, bool*) const::$_1>(long) Job.cpp:0:
0
#33 0x000056317d1a6cf9 llvm::CrashRecoveryContext::RunSafely(llvm::function_ref<void ()>) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x400dcf9)
#34 0x000056317dd203de clang::driver::CC1Command::Execute(llvm::ArrayRef<llvm::Optional<llvm::StringRef>>, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char>>*, bool*) const (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x4b8
73de)
#35 0x000056317dce1054 clang::driver::Compilation::ExecuteCommand(clang::driver::Command const&, clang::driver::Command const*&, bool) const (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x4b48054)
#36 0x000056317dce131e clang::driver::Compilation::ExecuteJobs(clang::driver::JobList const&, llvm::SmallVectorImpl<std::pair<int, clang::driver::Command const*>>&, bool) const (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x4b4831e)
#37 0x000056317dcff73f clang::driver::Driver::ExecuteCompilation(clang::driver::Compilation&, llvm::SmallVectorImpl<std::pair<int, clang::driver::Command const*>>&) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x4b6673f)
#38 0x000056317b9c6c64 clang_main(int, char**) (/home/nathan/tmp/tmp.VdDNi49u8X/install/llvm/bad/bin/clang-16+0x282dc64)
#39 0x00007f04d3bb1290 (/usr/lib/libc.so.6+0x23290)
#40 0x00007f04d3bb134a __libc_start_main (/usr/lib/libc.so.6+0x2334a)
#41 0x000056317b9c40a5 _start /build/glibc/src/glibc/csu/../sysdeps/x86_64/start.S:117:0
clang-16: error: clang frontend command failed with exit code 134 (use -v to see invocation)
ClangBuiltLinux clang version 16.0.0 (https://github.com/llvm/llvm-project b4e1466c35d3ca3e04244e8e8b4ffaf0784d6d37)
Target: aarch64-unknown-linux-gnu
Thread model: posix
InstalledDir: /home/nathan/tmp/tmp.VdDNi49u8X/cvise/../install/llvm/bad/bin
clang-16: note: diagnostic msg: Error generating preprocessed source(s) - no preprocessable inputs.

[nathanchance]

Actually, I lied, here is an LLVM IR reproducer that has been passed through llvm-reduce based on whether llc crashed or not.

target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
target triple = "aarch64-unknown-linux-gnu"

define void @drbd_req_state(ptr %drbd_req_state_ns) {
entry:
  %bf.load = load i8, ptr null, align 4
  %0 = and i8 %bf.load, 0
  %bf.lshr = lshr i8 %bf.load, 1
  %1 = and i8 %bf.lshr, 1
  %bf.value7 = zext i8 %1 to i32
  %bf.shl8 = shl nuw i32 %bf.value7, 1
  %2 = load i32, ptr null, align 4
  %and = and i32 %bf.shl8, %2
  store i32 %bf.shl8, ptr %drbd_req_state_ns, align 4
  %coerce.val.ii = zext i32 %and to i64
  ret void
}

$ ../install/llvm/good/bin/llc reduced.ll

$ ../install/llvm/bad/bin/llc reduced.ll
llc: /home/nathan/cbl/src/llvm-project/llvm/include/llvm/ADT/APInt.h:1027: bool llvm::APInt::operator==(const llvm::APInt &) const: Assertion `BitWidth == RHS.BitWidth && "Comparison requires equal bit widths"' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.      Program arguments: ../install/llvm/bad/bin/llc reduced.ll
1.      Running pass 'Function Pass Manager' on module 'reduced.ll'.
2.      Running pass 'AArch64 Instruction Selection' on function '@drbd_req_state'
 #0 0x000055a2e0706641 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (../install/llvm/bad/bin/llc+0x35d4641)
 #1 0x000055a2e07044fe llvm::sys::RunSignalHandlers() (../install/llvm/bad/bin/llc+0x35d24fe)
 #2 0x000055a2e0706ba6 SignalHandler(int) Signals.cpp:0:0
 #3 0x00007f077b3f4a00 (/usr/lib/libc.so.6+0x38a00)
 #4 0x00007f077b44464c (/usr/lib/libc.so.6+0x8864c)
 #5 0x00007f077b3f4958 raise (/usr/lib/libc.so.6+0x38958)
 #6 0x00007f077b3de53d abort (/usr/lib/libc.so.6+0x2253d)
 #7 0x00007f077b3de45c (/usr/lib/libc.so.6+0x2245c)
 #8 0x00007f077b3ed486 (/usr/lib/libc.so.6+0x31486)
 #9 0x000055a2ded20b3e (../install/llvm/bad/bin/llc+0x1beeb3e)
#10 0x000055a2e03c71ba (anonymous namespace)::DAGCombiner::visitShiftByConstant(llvm::SDNode*) DAGCombiner.cpp:0:0
#11 0x000055a2e03723d0 (anonymous namespace)::DAGCombiner::visitSHL(llvm::SDNode*) DAGCombiner.cpp:0:0
#12 0x000055a2e035e8a8 (anonymous namespace)::DAGCombiner::combine(llvm::SDNode*) DAGCombiner.cpp:0:0
#13 0x000055a2e0359cb4 llvm::SelectionDAG::Combine(llvm::CombineLevel, llvm::AAResults*, llvm::CodeGenOpt::Level) (../install/llvm/bad/bin/llc+0x3227cb4)
#14 0x000055a2e051ca22 llvm::SelectionDAGISel::CodeGenAndEmitDAG() (../install/llvm/bad/bin/llc+0x33eaa22)
#15 0x000055a2e051b652 llvm::SelectionDAGISel::SelectAllBasicBlocks(llvm::Function const&) (../install/llvm/bad/bin/llc+0x33e9652)
#16 0x000055a2e0517c85 llvm::SelectionDAGISel::runOnMachineFunction(llvm::MachineFunction&) (../install/llvm/bad/bin/llc+0x33e5c85)
#17 0x000055a2df920ccb llvm::MachineFunctionPass::runOnFunction(llvm::Function&) (../install/llvm/bad/bin/llc+0x27eeccb)
#18 0x000055a2dfe1ae45 llvm::FPPassManager::runOnFunction(llvm::Function&) (../install/llvm/bad/bin/llc+0x2ce8e45)
#19 0x000055a2dfe231a1 llvm::FPPassManager::runOnModule(llvm::Module&) (../install/llvm/bad/bin/llc+0x2cf11a1)
#20 0x000055a2dfe1b848 llvm::legacy::PassManagerImpl::run(llvm::Module&) (../install/llvm/bad/bin/llc+0x2ce9848)
#21 0x000055a2dec1fa84 main (../install/llvm/bad/bin/llc+0x1aeda84)
#22 0x00007f077b3df290 (/usr/lib/libc.so.6+0x23290)
#23 0x00007f077b3df34a __libc_start_main (/usr/lib/libc.so.6+0x2334a)
#24 0x000055a2dec1a675 _start /build/glibc/src/glibc/csu/../sysdeps/x86_64/start.S:117:0
fish: Job 1, '../install/llvm/bad/bin/llc red…' terminated by signal SIGABRT (Abort)

[nathanchance]

Judging from the time zone of the committer, I suspect this will not get dealt with right away so I have reverted this change in 74bace2dfe57d9cf569addf94af4e01a990d2374 due to the above crash.

[bcl5980]

In D136014#3905953, @nathanchance wrote:

Judging from the time zone of the committer, I suspect this will not get dealt with right away so I have reverted this change in 74bace2dfe57d9cf569addf94af4e01a990d2374 due to the above crash.

Thanks for the revert commit and test case. I will trying to fix the issue soon.

[bcl5980]

Fix crash.

[nickdesaulniers]

In D136014#3907080, @bcl5980 wrote:

Fix crash.

Please add a test that would have caught the regression introduced.

[nathanchance]

I had the same thought as Nick that there should be a test added for the regression. Regardless, I just tested your change against the Linux kernel and everything build successfully, thank you for the quick fix! I do not think I am qualified enough to approve though.

[dmgreen]

Yes, please add a test with different bitwidths. It's not obvious what is changing types. Is it during legalization?

If we know that they are always <= 64, it could just use getZExtValue too.

[bcl5980]

In D136014#3909891, @dmgreen wrote:

Yes, please add a test with different bitwidths. It's not obvious what is changing types. Is it during legalization?

The type change comes from ShrinkDemandedOp in truncate store. A normal different type case can't reproduce the issue. I also try to fix it the shrink demanded op. But that's another individual patch. I am not sure which patch need to include the crash test for now.

If we know that they are always <= 64, it could just use getZExtValue too.

Thanks for the mention.Yeah, zext value should enough.

[bcl5980]

update test.

[dmgreen]

Thanks. LGTM

[nickdesaulniers]

LGTM; phab doesn't give me the option to "accept the revision," only "reopen."

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

AArch64TargetLowering::isDesirableToCommuteWithShift(const SDNode *N,
                                                      CombineLevel Level) const {
   assert((N->getOpcode() == ISD::SHL || N->getOpcode() == ISD::SRA ||
           N->getOpcode() == ISD::SRL) &&
          "Expected shift op");
 
   SDValue ShiftLHS = N->getOperand(0);
   EVT VT = N->getValueType(0);
 
-  // If ShiftLHS is unsigned bit extraction: ((x >> C) & mask), then do not combine
-  // it with shift 'N' to let it be lowered to UBFX.
+  // If ShiftLHS is unsigned bit extraction: ((x >> C) & mask), then do not
+  // combine it with shift 'N' to let it be lowered to UBFX.
+  // The special case is ((x >> C) & mask) << C.
+  // It can be combine to x & (mask << C) by return true
   if (ShiftLHS.getOpcode() == ISD::AND && (VT == MVT::i32 || VT == MVT::i64) &&
       isa<ConstantSDNode>(ShiftLHS.getOperand(1))) {
     uint64_t TruncMask = ShiftLHS.getConstantOperandVal(1);
-    if (isMask_64(TruncMask) &&
-        ShiftLHS.getOperand(0).getOpcode() == ISD::SRL &&
-        isa<ConstantSDNode>(ShiftLHS.getOperand(0).getOperand(1)))
+    if (isMask_64(TruncMask)) {
+      SDValue AndLHS = ShiftLHS.getOperand(0);
+      if (AndLHS.getOpcode() == ISD::SRL) {
+        if (auto *SRLC = dyn_cast<ConstantSDNode>(AndLHS.getOperand(1))) {
+          if (N->getOpcode() == ISD::SHL)
+            if (auto *SHLC = dyn_cast<ConstantSDNode>(N->getOperand(1)))
+              return SRLC->getAPIntValue() == SHLC->getAPIntValue();
           return false;
         }
+      }
+    }
+  }
   return true;
 }
 
 bool AArch64TargetLowering::isDesirableToCommuteXorWithShift(
     const SDNode *N) const {
   assert(N->getOpcode() == ISD::XOR &&
          (N->getOperand(0).getOpcode() == ISD::SHL ||
           N->getOperand(0).getOpcode() == ISD::SRL) &&

llvm/test/CodeGen/AArch64/shift-logic.ll

 ; CHECK-NEXT:    str w8, [x2]
 ; CHECK-NEXT:    ret
   %sh0 = lshr i32 %x, 5
   %r = or i32 %sh0, %y
   store i32 %r, i32* %p
   %sh1 = lshr i32 %r, 7
   ret i32 %sh1
 }
+
+define i64 @desirable_to_commute1(i64 %x) {
+; CHECK-LABEL: desirable_to_commute1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    and x0, x0, #0x7fff8
+; CHECK-NEXT:    ret
+  %s1 = lshr i64 %x, 3
+  %a = and i64 %s1, 65535
+  %s2 = shl i64 %a, 3
+  ret i64 %s2
+}
+
+define i64 @desirable_to_commute2(i64* %p, i64 %i) {
+; CHECK-LABEL: desirable_to_commute2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    and x8, x1, #0x1ff8
+; CHECK-NEXT:    ldr x0, [x0, x8]
+; CHECK-NEXT:    ret
+  %lshr = lshr i64 %i, 3
+  %and = and i64 %lshr, 1023
+  %pidx = getelementptr i64, i64* %p, i64 %and
+  %r = load i64, i64* %pidx
+  ret i64 %r
+}