[ctxprof] Instrument direct calls to other things than Functions

Author: mtrofin

llvm/include/llvm/IR/InstrTypes.h

@@ -1343,7 +1343,8 @@ class CallBase : public Instruction {
   Use &getCalledOperandUse() { return Op<CalledOperandOpEndIdx>(); }
 
   /// Returns the function called, or null if this is an indirect function
-  /// invocation or the function signature does not match the call signature.
+  /// invocation or the function signature does not match the call signature, or
+  /// the call target is an alias.
   Function *getCalledFunction() const {
     if (auto *F = dyn_cast_or_null<Function>(getCalledOperand()))
       if (F->getValueType() == getFunctionType())

llvm/include/llvm/IR/IntrinsicInst.h

@@ -1536,7 +1536,7 @@ class InstrProfCallsite : public InstrProfCntrInstBase {
   static bool canInstrumentCallsite(const CallBase &CB) {
     return !CB.isInlineAsm() &&
            (CB.isIndirectCall() ||
-            (CB.getCalledFunction() && !CB.getCalledFunction()->isIntrinsic()));
+            (CB.getIntrinsicID() == Intrinsic::not_intrinsic));
   }
   LLVM_ABI Value *getCallee() const;
   LLVM_ABI void setCallee(Value *Callee);

llvm/test/Transforms/PGOProfile/ctx-instrumentation-aliases.ll

@@ -0,0 +1,25 @@
+; REQUIRES: x86-registered-target
+; Test that calls to aliases are instrumented, and the assembly references the
+; aliased function.
+;
+; RUN: opt -passes=ctx-instr-gen,assign-guid,ctx-instr-lower -profile-context-root=an_entrypoint \
+; RUN:   -profile-context-root=another_entrypoint_no_callees \
+; RUN:   -S %s -o %t.ll
+; RUN: llc < %t.ll | FileCheck %s
+target triple = "x86_64-unknown-linux-gnu"
+
+@foo_alias = weak_odr unnamed_addr alias void (), ptr @foo
+
+define void @foo(i32) {
+  ret void
+}
+
+define void @call_alias(ptr %a) {
+entry:
+  call void @foo(i32 0, ptr %a)
+  ret void
+}
+
+; CHECK-LABEL:   call_alias:
+; CHECK:         movq    foo@GOTPCREL(%rip), [[REG:%r[a-z0-9]+]]
+; CHECK-NEXT:    movq    [[REG]], %fs:__llvm_ctx_profile_expected_callee@TPOFF{{.*}}

llvm/test/Transforms/PGOProfile/ctx-instrumentation.ll

@@ -8,6 +8,8 @@
 
 declare void @bar()
 
+;.
+; INSTRUMENT: @foo_alias = weak_odr unnamed_addr alias void (i32, ptr), ptr @foo
 ;.
 ; LOWERING: @__llvm_ctx_profile_callsite = external hidden thread_local global ptr
 ; LOWERING: @__llvm_ctx_profile_expected_callee = external hidden thread_local global ptr
@@ -19,6 +21,8 @@ declare void @bar()
 ; LOWERING: @[[GLOB5:[0-9]+]] = internal global { ptr, ptr, ptr, ptr, i8 } zeroinitializer
 ; LOWERING: @[[GLOB6:[0-9]+]] = internal global { ptr, ptr, ptr, ptr, i8 } zeroinitializer
 ; LOWERING: @[[GLOB7:[0-9]+]] = internal global { ptr, ptr, ptr, ptr, i8 } { ptr null, ptr null, ptr inttoptr (i64 1 to ptr), ptr null, i8 0 }
+; LOWERING: @[[GLOB8:[0-9]+]] = internal global { ptr, ptr, ptr, ptr, i8 } zeroinitializer
+; LOWERING: @foo_alias = weak_odr unnamed_addr alias void (i32, ptr), ptr @foo
 ;.
 define void @foo(i32 %a, ptr %fct) {
 ; INSTRUMENT-LABEL: define void @foo(
@@ -335,6 +339,41 @@ define void @unreachable() {
 ;
   unreachable
 }
+
+@foo_alias = weak_odr unnamed_addr alias void (i32, ptr), ptr @foo
+
+define void @call_alias(ptr %a) {
+; INSTRUMENT-LABEL: define void @call_alias(
+; INSTRUMENT-SAME: ptr [[A:%.*]]) {
+; INSTRUMENT-NEXT:  entry:
+; INSTRUMENT-NEXT:    call void @llvm.instrprof.increment(ptr @call_alias, i64 742261418966908927, i32 1, i32 0)
+; INSTRUMENT-NEXT:    call void @llvm.instrprof.callsite(ptr @call_alias, i64 742261418966908927, i32 1, i32 0, ptr @foo_alias)
+; INSTRUMENT-NEXT:    call void @foo_alias(i32 0, ptr [[A]])
+; INSTRUMENT-NEXT:    ret void
+;
+; LOWERING-LABEL: define void @call_alias(
+; LOWERING-SAME: ptr [[A:%.*]]) !guid [[META10:![0-9]+]] {
+; LOWERING-NEXT:  entry:
+; LOWERING-NEXT:    [[TMP0:%.*]] = call ptr @__llvm_ctx_profile_get_context(ptr @[[GLOB8]], ptr @call_alias, i64 2172368043968427688, i32 1, i32 1)
+; LOWERING-NEXT:    [[TMP1:%.*]] = ptrtoint ptr [[TMP0]] to i64
+; LOWERING-NEXT:    [[TMP2:%.*]] = and i64 [[TMP1]], 1
+; LOWERING-NEXT:    [[TMP3:%.*]] = call ptr @llvm.threadlocal.address.p0(ptr @__llvm_ctx_profile_expected_callee)
+; LOWERING-NEXT:    [[TMP4:%.*]] = getelementptr ptr, ptr [[TMP3]], i64 [[TMP2]]
+; LOWERING-NEXT:    [[TMP5:%.*]] = call ptr @llvm.threadlocal.address.p0(ptr @__llvm_ctx_profile_callsite)
+; LOWERING-NEXT:    [[TMP6:%.*]] = getelementptr i32, ptr [[TMP5]], i64 [[TMP2]]
+; LOWERING-NEXT:    [[TMP7:%.*]] = and i64 [[TMP1]], -2
+; LOWERING-NEXT:    [[TMP8:%.*]] = inttoptr i64 [[TMP7]] to ptr
+; LOWERING-NEXT:    store volatile ptr @foo_alias, ptr [[TMP4]], align 8
+; LOWERING-NEXT:    [[TMP9:%.*]] = getelementptr { { i64, ptr, i32, i32 }, [1 x i64], [1 x ptr] }, ptr [[TMP0]], i32 0, i32 2, i32 0
+; LOWERING-NEXT:    store volatile ptr [[TMP9]], ptr [[TMP6]], align 8
+; LOWERING-NEXT:    call void @foo_alias(i32 0, ptr [[A]])
+; LOWERING-NEXT:    call void @__llvm_ctx_profile_release_context(ptr @[[GLOB8]])
+; LOWERING-NEXT:    ret void
+;
+entry:
+  call void @foo_alias(i32 0, ptr %a)
+  ret void
+}
 ;.
 ; LOWERING: attributes #[[ATTR0]] = { noreturn }
 ; LOWERING: attributes #[[ATTR1:[0-9]+]] = { nounwind }
@@ -353,4 +392,5 @@ define void @unreachable() {
 ; LOWERING: [[META7]] = !{i64 -4680624981836544329}
 ; LOWERING: [[META8]] = !{i64 5519225910966780583}
 ; LOWERING: [[META9]] = !{i64 -565652589829076809}
+; LOWERING: [[META10]] = !{i64 2172368043968427688}
 ;.