[InstrProf] Change step from 64-bit to data layout dependent size

Fixed 64-bit step can lead to creating atomic instructions
unsupported by the target architecture (see rust-lang/rust#112313).

When using atomics, type of the step is a pointer-sized integer.
Otherwise, type of the step is of a largest legal integer type.

Author: jdmitrovic-syrmia

llvm/include/llvm/IR/IntrinsicInst.h

@@ -1472,6 +1472,7 @@ class InstrProfIncrementInst : public InstrProfCntrInstBase {
     return isa<IntrinsicInst>(V) && classof(cast<IntrinsicInst>(V));
   }
   Value *getStep() const;
+  Value *getAtomicStep() const;
 };
 
 /// This represents the llvm.instrprof.increment.step intrinsic.

llvm/lib/IR/IntrinsicInst.cpp

@@ -288,9 +288,26 @@ Value *InstrProfIncrementInst::getStep() const {
   }
   const Module *M = getModule();
   LLVMContext &Context = M->getContext();
+  const auto &DL = M->getDataLayout();
+  Type* LargestLegalIntTy = DL.getLargestLegalIntType(Context);
+
+  if (LargestLegalIntTy) {
+    return ConstantInt::get(LargestLegalIntTy, 1);
+  }
+
   return ConstantInt::get(Type::getInt64Ty(Context), 1);
 }
 
+Value *InstrProfIncrementInst::getAtomicStep() const {
+  if (InstrProfIncrementInstStep::classof(this)) {
+    return const_cast<Value *>(getArgOperand(4));
+  }
+  const Module *M = getModule();
+  LLVMContext &Context = M->getContext();
+  const auto &DL = M->getDataLayout();
+  return ConstantInt::get(DL.getIntPtrType(Context), 1);
+}
+
 std::optional<RoundingMode> ConstrainedFPIntrinsic::getRoundingMode() const {
   unsigned NumOperands = arg_size();
   Metadata *MD = nullptr;

llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp

@@ -935,7 +935,7 @@ void InstrLowerer::lowerIncrement(InstrProfIncrementInst *Inc) {
   IRBuilder<> Builder(Inc);
   if (Options.Atomic || AtomicCounterUpdateAll ||
       (Inc->getIndex()->isZeroValue() && AtomicFirstCounter)) {
-    Builder.CreateAtomicRMW(AtomicRMWInst::Add, Addr, Inc->getStep(),
+    Builder.CreateAtomicRMW(AtomicRMWInst::Add, Addr, Inc->getAtomicStep(),
                             MaybeAlign(), AtomicOrdering::Monotonic);
   } else {
     Value *IncStep = Inc->getStep();