[AtomicExpandPass] Widen partword atomicrmw or/xor/and before tryExpandAtomicRMW

[asb]

Summary

This patch performs a widening transformation of bitwise atomicrmw or, xor, and and applies it prior to tryExpandAtomicRMW. This operates similarly to convertCmpXchgToIntegerType. For these operations, the i8/i16 atomicrmw can be implemented in terms of the 32-bit atomicrmw by appropriately manipulating the operands. There is no functional change for the handling of partword or/xor, but the transformation for partword and is new.

The advantage of performing this transformation early is that the same code-path can be used regardless of the approach used to expand the atomicrmw (AtomicExpansionKind). i.e. the same logic is used for AtomicExpansionKind::CmpXchg and can also be used by the intrinsic-based expansion in D47882.

[asb]

I've updated this patch so it is independent of the RISC-V atomics lowering. I replace the previous logic in performMaskedAtomicOp with a new helper function widenPartwordAtomicRMW, which is called early on. See the updated summary for more details.

llvm/trunk/lib/CodeGen/AtomicExpandPass.cpp

insertRMWLLSCLoop(IRBuilder<> &Builder, Type *ResultTy, Value *Addr,
                       AtomicOrdering MemOpOrder,
                       function_ref<Value *(IRBuilder<> &, Value *)> PerformOp);
     void expandAtomicOpToLLSC(
         Instruction *I, Type *ResultTy, Value *Addr, AtomicOrdering MemOpOrder,
         function_ref<Value *(IRBuilder<> &, Value *)> PerformOp);
     void expandPartwordAtomicRMW(
         AtomicRMWInst *I,
         TargetLoweringBase::AtomicExpansionKind ExpansionKind);
+    AtomicRMWInst *widenPartwordAtomicRMW(AtomicRMWInst *AI);
     void expandPartwordCmpXchg(AtomicCmpXchgInst *I);
 
     AtomicCmpXchgInst *convertCmpXchgToIntegerType(AtomicCmpXchgInst *CI);
     static Value *insertRMWCmpXchgLoop(
         IRBuilder<> &Builder, Type *ResultType, Value *Addr,
         AtomicOrdering MemOpOrder,
         function_ref<Value *(IRBuilder<> &, Value *)> PerformOp,
         CreateCmpXchgInstFun CreateCmpXchg);
if (LI) {
       // There are two different ways of expanding RMW instructions:
       // - into a load if it is idempotent
       // - into a Cmpxchg/LL-SC loop otherwise
       // we try them in that order.
 
       if (isIdempotentRMW(RMWI) && simplifyIdempotentRMW(RMWI)) {
         MadeChange = true;
       } else {
+        unsigned MinCASSize = TLI->getMinCmpXchgSizeInBits() / 8;
+        unsigned ValueSize = getAtomicOpSize(RMWI);
+        AtomicRMWInst::BinOp Op = RMWI->getOperation();
+        if (ValueSize < MinCASSize &&
+            (Op == AtomicRMWInst::Or || Op == AtomicRMWInst::Xor ||
+             Op == AtomicRMWInst::And)) {
+          RMWI = widenPartwordAtomicRMW(RMWI);
+          MadeChange = true;
+        }
+
         MadeChange |= tryExpandAtomicRMW(RMWI);
       }
     } else if (CASI) {
       // TODO: when we're ready to make the change at the IR level, we can
       // extend convertCmpXchgToInteger for floating point too.
       assert(!CASI->getCompareOperand()->getType()->isFloatingPointTy() &&
              "unimplemented - floating point not legal at IR level");
       if (CASI->getCompareOperand()->getType()->isPointerTy() ) {
static Value *performMaskedAtomicOp(AtomicRMWInst::BinOp Op,
   switch (Op) {
   case AtomicRMWInst::Xchg: {
     Value *Loaded_MaskOut = Builder.CreateAnd(Loaded, PMV.Inv_Mask);
     Value *FinalVal = Builder.CreateOr(Loaded_MaskOut, Shifted_Inc);
     return FinalVal;
   }
   case AtomicRMWInst::Or:
   case AtomicRMWInst::Xor:
-    // Or/Xor won't affect any other bits, so can just be done
-    // directly.
-    return performAtomicOp(Op, Builder, Loaded, Shifted_Inc);
+  case AtomicRMWInst::And:
+    llvm_unreachable("Or/Xor/And handled by widenPartwordAtomicRMW");
   case AtomicRMWInst::Add:
   case AtomicRMWInst::Sub:
-  case AtomicRMWInst::And:
   case AtomicRMWInst::Nand: {
     // The other arithmetic ops need to be masked into place.
     Value *NewVal = performAtomicOp(Op, Builder, Loaded, Shifted_Inc);
     Value *NewVal_Masked = Builder.CreateAnd(NewVal, PMV.Mask);
     Value *Loaded_MaskOut = Builder.CreateAnd(Loaded, PMV.Inv_Mask);
     Value *FinalVal = Builder.CreateOr(Loaded_MaskOut, NewVal_Masked);
     return FinalVal;
   }
Value *OldResult =
       insertRMWCmpXchgLoop(Builder, PMV.WordType, PMV.AlignedAddr, MemOpOrder,
                            PerformPartwordOp, createCmpXchgInstFun);
   Value *FinalOldResult = Builder.CreateTrunc(
       Builder.CreateLShr(OldResult, PMV.ShiftAmt), PMV.ValueType);
   AI->replaceAllUsesWith(FinalOldResult);
   AI->eraseFromParent();
 }
 
+// Widen the bitwise atomicrmw (or/xor/and) to the minimum supported width.
+AtomicRMWInst *AtomicExpand::widenPartwordAtomicRMW(AtomicRMWInst *AI) {
+  IRBuilder<> Builder(AI);
+  AtomicRMWInst::BinOp Op = AI->getOperation();
+
+  assert((Op == AtomicRMWInst::Or || Op == AtomicRMWInst::Xor ||
+          Op == AtomicRMWInst::And) &&
+         "Unable to widen operation");
+
+  PartwordMaskValues PMV =
+      createMaskInstrs(Builder, AI, AI->getType(), AI->getPointerOperand(),
+                       TLI->getMinCmpXchgSizeInBits() / 8);
+
+  Value *ValOperand_Shifted =
+      Builder.CreateShl(Builder.CreateZExt(AI->getValOperand(), PMV.WordType),
+                        PMV.ShiftAmt, "ValOperand_Shifted");
+
+  Value *NewOperand;
+
+  if (Op == AtomicRMWInst::And)
+    NewOperand =
+        Builder.CreateOr(PMV.Inv_Mask, ValOperand_Shifted, "AndOperand");
+  else
+    NewOperand = ValOperand_Shifted;
+
+  AtomicRMWInst *NewAI = Builder.CreateAtomicRMW(Op, PMV.AlignedAddr,
+                                                 NewOperand, AI->getOrdering());
+
+  Value *FinalOldResult = Builder.CreateTrunc(
+      Builder.CreateLShr(NewAI, PMV.ShiftAmt), PMV.ValueType);
+  AI->replaceAllUsesWith(FinalOldResult);
+  AI->eraseFromParent();
+  return NewAI;
+}
+
 void AtomicExpand::expandPartwordCmpXchg(AtomicCmpXchgInst *CI) {
   // The basic idea here is that we're expanding a cmpxchg of a
   // smaller memory size up to a word-sized cmpxchg. To do this, we
   // need to add a retry-loop for strong cmpxchg, so that
   // modifications to other parts of the word don't cause a spurious
   // failure.
 
   // This generates code like the following:

llvm/trunk/test/Transforms/AtomicExpand/SPARC/partword.ll

 ; CHECK:  %6 = cmpxchg i32* %AlignedAddr, i32 %loaded, i32 %new monotonic monotonic
 ; CHECK:atomicrmw.end:
 define i16 @test_xor_i16(i16* %arg, i16 %val) {
 entry:
   %ret = atomicrmw xor i16* %arg, i16 %val seq_cst
   ret i16 %ret
 }
 
+; CHECK-LABEL: @test_or_i16(
+; (I'm going to just assert on the bits that differ from add, above.)
+; CHECK:atomicrmw.start:
+; CHECK:  %new = or i32 %loaded, %ValOperand_Shifted
+; CHECK:  %6 = cmpxchg i32* %AlignedAddr, i32 %loaded, i32 %new monotonic monotonic
+; CHECK:atomicrmw.end:
+define i16 @test_or_i16(i16* %arg, i16 %val) {
+entry:
+  %ret = atomicrmw or i16* %arg, i16 %val seq_cst
+  ret i16 %ret
+}
+
+; CHECK-LABEL: @test_and_i16(
+; (I'm going to just assert on the bits that differ from add, above.)
+; CHECK:  %AndOperand = or i32 %Inv_Mask, %ValOperand_Shifted
+; CHECK:atomicrmw.start:
+; CHECK:  %new = and i32 %loaded, %AndOperand
+; CHECK:  %6 = cmpxchg i32* %AlignedAddr, i32 %loaded, i32 %new monotonic monotonic
+; CHECK:atomicrmw.end:
+define i16 @test_and_i16(i16* %arg, i16 %val) {
+entry:
+  %ret = atomicrmw and i16* %arg, i16 %val seq_cst
+  ret i16 %ret
+}
+
 ; CHECK-LABEL: @test_min_i16(
 ; CHECK:atomicrmw.start:
 ; CHECK:  %6 = lshr i32 %loaded, %ShiftAmt
 ; CHECK:  %7 = trunc i32 %6 to i16
 ; CHECK:  %8 = icmp sle i16 %7, %val
 ; CHECK:  %new = select i1 %8, i16 %7, i16 %val
 ; CHECK:  %9 = zext i16 %new to i32
 ; CHECK:  %10 = shl i32 %9, %ShiftAmt