[llvm] [InstrProf] Change step from 64-bit to pointer-sized int (PR #83239)

[Jovan Dmitrović via llvm-commits]

https://github.com/jdmitrovic-syrmia updated https://github.com/llvm/llvm-project/pull/83239

>From 530c1a66b0f531ed3f5d024cb8b8723bb5e34650 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jovan=20Dmitrovi=C4=87?= <jovan.dmitrovic at syrmia.com>
Date: Wed, 28 Feb 2024 09:50:20 +0100
Subject: [PATCH] [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.
---
 llvm/include/llvm/IR/IntrinsicInst.h            |  1 +
 llvm/lib/IR/IntrinsicInst.cpp                   | 17 +++++++++++++++++
 .../Instrumentation/InstrProfiling.cpp          |  2 +-
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/IR/IntrinsicInst.h b/llvm/include/llvm/IR/IntrinsicInst.h
index c07b83a81a63e1..227a726aa9e5e1 100644
--- a/llvm/include/llvm/IR/IntrinsicInst.h
+++ b/llvm/include/llvm/IR/IntrinsicInst.h
@@ -1506,6 +1506,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.
diff --git a/llvm/lib/IR/IntrinsicInst.cpp b/llvm/lib/IR/IntrinsicInst.cpp
index 89403e1d7fcb4d..7fe7f89bc4b4ef 100644
--- a/llvm/lib/IR/IntrinsicInst.cpp
+++ b/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;
diff --git a/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp b/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
index c42c53edd51190..ee55e1b1c32d5d 100644
--- a/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
+++ b/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
@@ -958,7 +958,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();