[DAGCombiner][ARM][RISCV] Teach ShrinkLoadReplaceStoreWithStore to use truncstore.

[craig.topper]

Summary

The VT we want to shrink to may not be legal especially after type
legalization.

Fixes PR56110.

[RKSimon]

LGTM with one minor

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

17372–17373

Might be worth extending this comment a little for the new truncstore case

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

ShrinkLoadReplaceStoreWithStore(const std::pair<unsigned, unsigned> &MaskInfo,
   // Check to see if IVal is all zeros in the part being masked in by the 'or'
   // that uses this.  If not, this is not a replacement.
   APInt Mask = ~APInt::getBitsSet(IVal.getValueSizeInBits(),
                                   ByteShift*8, (ByteShift+NumBytes)*8);
   if (!DAG.MaskedValueIsZero(IVal, Mask)) return SDValue();
 
   // Check that it is legal on the target to do this.  It is legal if the new
   // VT we're shrinking to (i8/i16/i32) is legal or we're still before type
-  // legalization (and the target doesn't explicitly think this is a bad idea).
+  // legalization. If the source type is legal, but the store type isn't, see
+  // if we can use a truncating store.
   MVT VT = MVT::getIntegerVT(NumBytes * 8);
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
-  if (!DC->isTypeLegal(VT))
+  bool UseTruncStore;
+  if (DC->isTypeLegal(VT))
+    UseTruncStore = false;
+  else if (TLI.isTypeLegal(IVal.getValueType()) &&
+           TLI.isTruncStoreLegal(IVal.getValueType(), VT))
+    UseTruncStore = true;
+  else
     return SDValue();
+  // Check that the target doesn't think this is a bad idea.
   if (St->getMemOperand() &&
       !TLI.allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), VT,
                               *St->getMemOperand()))
     return SDValue();
 
   // Okay, we can do this!  Replace the 'St' store with a store of IVal that is
   // shifted by ByteShift and truncated down to NumBytes.
   if (ByteShift) {
else
     StOffset = IVal.getValueType().getStoreSize() - ByteShift - NumBytes;
 
   SDValue Ptr = St->getBasePtr();
   if (StOffset) {
     SDLoc DL(IVal);
     Ptr = DAG.getMemBasePlusOffset(Ptr, TypeSize::Fixed(StOffset), DL);
   }
 
+  ++OpsNarrowed;
+  if (UseTruncStore)
+    return DAG.getTruncStore(St->getChain(), SDLoc(St), IVal, Ptr,
+                             St->getPointerInfo().getWithOffset(StOffset),
+                             VT, St->getOriginalAlign());
+
   // Truncate down to the new size.
   IVal = DAG.getNode(ISD::TRUNCATE, SDLoc(IVal), VT, IVal);
 
-  ++OpsNarrowed;
   return DAG
       .getStore(St->getChain(), SDLoc(St), IVal, Ptr,
                 St->getPointerInfo().getWithOffset(StOffset),
                 St->getOriginalAlign());
 }
 
 /// Look for sequence of load / op / store where op is one of 'or', 'xor', and
 /// 'and' of immediates. If 'op' is only touching some of the loaded bits, try

llvm/test/CodeGen/ARM/illegal-bitfield-loadstore.ll

 ; LE:       @ %bb.0:
 ; LE-NEXT:    ldr r1, [r0]
 ; LE-NEXT:    orr r1, r1, #384
 ; LE-NEXT:    str r1, [r0]
 ; LE-NEXT:    mov pc, lr
 ;
 ; BE-LABEL: i56_or:
 ; BE:       @ %bb.0:
-; BE-NEXT:    mov r1, r0
-; BE-NEXT:    ldr r12, [r0]
-; BE-NEXT:    ldrh r2, [r1, #4]!
-; BE-NEXT:    ldrb r3, [r1, #2]
+; BE-NEXT:    ldr r1, [r0]
+; BE-NEXT:    strb r1, [r0, #3]
+; BE-NEXT:    ldrh r2, [r0, #4]!
+; BE-NEXT:    ldrb r3, [r0, #2]
 ; BE-NEXT:    orr r2, r3, r2, lsl #8
-; BE-NEXT:    orr r2, r2, r12, lsl #24
-; BE-NEXT:    orr r2, r2, #384
-; BE-NEXT:    strb r2, [r1, #2]
-; BE-NEXT:    lsr r3, r2, #8
-; BE-NEXT:    strh r3, [r1]
-; BE-NEXT:    bic r1, r12, #255
-; BE-NEXT:    orr r1, r1, r2, lsr #24
-; BE-NEXT:    str r1, [r0]
+; BE-NEXT:    orr r1, r2, r1, lsl #24
+; BE-NEXT:    orr r1, r1, #384
+; BE-NEXT:    strb r1, [r0, #2]
+; BE-NEXT:    lsr r1, r1, #8
+; BE-NEXT:    strh r1, [r0]
 ; BE-NEXT:    mov pc, lr
   %aa = load i56, i56* %a
   %b = or i56 %aa, 384
   store i56 %b, i56* %a
   ret void
 }
 
 define void @i56_and_or(i56* %a) {
 ; LE-LABEL: i56_and_or:
 ; LE:       @ %bb.0:
 ; LE-NEXT:    ldr r1, [r0]
 ; LE-NEXT:    orr r1, r1, #384
 ; LE-NEXT:    bic r1, r1, #127
 ; LE-NEXT:    str r1, [r0]
 ; LE-NEXT:    mov pc, lr
 ;
 ; BE-LABEL: i56_and_or:
 ; BE:       @ %bb.0:
-; BE-NEXT:    mov r1, r0
+; BE-NEXT:    ldr r1, [r0]
 ; BE-NEXT:    mov r2, #128
-; BE-NEXT:    ldrh r12, [r1, #4]!
-; BE-NEXT:    ldrb r3, [r1, #2]
-; BE-NEXT:    strb r2, [r1, #2]
+; BE-NEXT:    strb r1, [r0, #3]
+; BE-NEXT:    ldrh r12, [r0, #4]!
+; BE-NEXT:    ldrb r3, [r0, #2]
+; BE-NEXT:    strb r2, [r0, #2]
 ; BE-NEXT:    orr r2, r3, r12, lsl #8
-; BE-NEXT:    ldr r12, [r0]
-; BE-NEXT:    orr r2, r2, r12, lsl #24
-; BE-NEXT:    orr r2, r2, #384
-; BE-NEXT:    lsr r3, r2, #8
-; BE-NEXT:    strh r3, [r1]
-; BE-NEXT:    bic r1, r12, #255
-; BE-NEXT:    orr r1, r1, r2, lsr #24
-; BE-NEXT:    str r1, [r0]
+; BE-NEXT:    orr r1, r2, r1, lsl #24
+; BE-NEXT:    orr r1, r1, #384
+; BE-NEXT:    lsr r1, r1, #8
+; BE-NEXT:    strh r1, [r0]
 ; BE-NEXT:    mov pc, lr
 
   %b = load i56, i56* %a, align 1
   %c = and i56 %b, -128
   %d = or i56 %c, 384
   store i56 %d, i56* %a, align 1
   ret void
 }
 
 define void @i56_insert_bit(i56* %a, i1 zeroext %bit) {
 ; LE-LABEL: i56_insert_bit:
 ; LE:       @ %bb.0:
 ; LE-NEXT:    ldr r2, [r0]
 ; LE-NEXT:    bic r2, r2, #8192
 ; LE-NEXT:    orr r1, r2, r1, lsl #13
 ; LE-NEXT:    str r1, [r0]
 ; LE-NEXT:    mov pc, lr
 ;
 ; BE-LABEL: i56_insert_bit:
 ; BE:       @ %bb.0:
-; BE-NEXT:    .save {r11, lr}
-; BE-NEXT:    push {r11, lr}
-; BE-NEXT:    mov r2, r0
-; BE-NEXT:    ldr lr, [r0]
-; BE-NEXT:    ldrh r12, [r2, #4]!
-; BE-NEXT:    ldrb r3, [r2, #2]
-; BE-NEXT:    orr r12, r3, r12, lsl #8
-; BE-NEXT:    orr r3, r12, lr, lsl #24
-; BE-NEXT:    bic r3, r3, #8192
-; BE-NEXT:    orr r1, r3, r1, lsl #13
-; BE-NEXT:    lsr r3, r1, #8
-; BE-NEXT:    strh r3, [r2]
-; BE-NEXT:    bic r2, lr, #255
-; BE-NEXT:    orr r1, r2, r1, lsr #24
-; BE-NEXT:    str r1, [r0]
-; BE-NEXT:    pop {r11, lr}
+; BE-NEXT:    ldr r2, [r0]
+; BE-NEXT:    strb r2, [r0, #3]
+; BE-NEXT:    ldrh r12, [r0, #4]!
+; BE-NEXT:    ldrb r3, [r0, #2]
+; BE-NEXT:    orr r3, r3, r12, lsl #8
+; BE-NEXT:    orr r2, r3, r2, lsl #24
+; BE-NEXT:    bic r2, r2, #8192
+; BE-NEXT:    orr r1, r2, r1, lsl #13
+; BE-NEXT:    lsr r1, r1, #8
+; BE-NEXT:    strh r1, [r0]
 ; BE-NEXT:    mov pc, lr
   %extbit = zext i1 %bit to i56
   %b = load i56, i56* %a, align 1
   %extbit.shl = shl nuw nsw i56 %extbit, 13
   %c = and i56 %b, -8193
   %d = or i56 %c, %extbit.shl
   store i56 %d, i56* %a, align 1
   ret void
 }
 

llvm/test/CodeGen/RISCV/pr56110.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=riscv32 | FileCheck %s
+; RUN: llc < %s -mtriple=riscv32 -mattr=+unaligned-scalar-mem | FileCheck %s
+
+define void @foo_set(ptr nocapture noundef %a, i32 noundef %v) {
+; CHECK-LABEL: foo_set:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    srli a2, a1, 8
+; CHECK-NEXT:    sb a1, 3(a0)
+; CHECK-NEXT:    sb a2, 4(a0)
+; CHECK-NEXT:    ret
+entry:
+  %bf.load = load i96, ptr %a, align 1
+  %0 = and i32 %v, 65535
+  %bf.value = zext i32 %0 to i96
+  %bf.shl = shl nuw nsw i96 %bf.value, 24
+  %bf.clear = and i96 %bf.load, -1099494850561
+  %bf.set = or i96 %bf.clear, %bf.shl
+  store i96 %bf.set, ptr %a, align 1
+  ret void
+}