[AMDGPU] Fix function pointer argument bug in AMDGPU Propagate Attributes pass.

[jweightman]

Summary

This patch fixes a bug in the AMDGPU Propagate Attributes pass where a call
instruction with a function pointer argument is identified as a user of the
passed function, and illegally replaces the called function of the
instruction with the function argument.

For example, given functions f and g with appropriate types, the following
illegal transformation could occur without this fix:
call void @f(void ()* @g)
-->
call void @g(void ()* @g.1)

The solution introduced in this patch is to prevent the cloning and
substitution if the instruction's called function and the function which
might be cloned do not match.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

267

What if the callee was a bitcast of the function? Can you add another testcase with a constexpr function pointer cast?

[jweightman]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

267

I tried adding this to the LIT test:

    %h_cast = bitcast i32* ()* @h to i8* ()*
    %1 = call i8* %h_cast()
...
define private i32* @h() #1 {
    %1 = alloca i32
    ret i32* %1
}

and the pass does not propagate the attributes like one would probably want. However, this seems pretty orthogonal to the problem I was trying to fix (preventing the illegal callee change), because the "user" is an llvm::BitCastInstr so CI is always nullptr in that case.

In order to handle bitcasts, I think we would need new logic to detect which bitcasted functions need to be replaced, a new data structure to track them in (similar to the existing ToReplace), and new logic to do that replacement — all of which would be on a separate logical branch from the existing code for handling call instructions! Unless you see an easier way, I think such a change is beyond the scope of what I was trying to do.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

267

This example isn't a constant expression cast. If you run opt -S -instcombine on this, it should turn it into the constant expr cast form. In that case the user will still be CallInst, it's just CI->getCalledFunction() will be a constantexpr cast of the function. stripPointerCasts should get you the base function

[jweightman]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

267

I see. I misunderstood your initial comment. This seems like a reasonable thing to try to support with this diff.

The instcombine pass seems to prefer to cast the argument/return value (I modified the test to try both) rather than the function pointer itself. Going back to the IR Language Reference, I saw the constant expression section (I'm still new at this...), so I tried this IR:

%1 = alloca i8
%2 = call i8* bitcast (i32* ()* @h to i8* ()*)()

However, the Propagate Attributes pass doesn't produce a clone of @h in this situation. I see that the dyn_cast<Instruction>(U) fails because the user of h is just the constexpr bitcast:

> U->dump()
void (i8*)* bitcast (void (i32*)* @g to void (i8*)*)

Sorry for the back and forth, but is there some way to find an enclosing CallBase if it exists?

[jdoerfert]

Drive by only

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

268

The usual way to deal with this is bool CallBase::isCallee(const Use *U). It might make a difference if you have foo(&foo).

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

267

I don't think there's a super easy way to do it. You would have to find the context. The normal way to check that would be to start from the instruction and look into the constantexpr for the global, which starts looking like general indirect call handling.

Ultimately this pass is a grotesque hack that should absolutely not exist, and only does because comgr still does not link the device libraries correctly for opencl. How did you run into this?

[madhur13490]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

267

Ultimately this pass is a grotesque hack that should absolutely not exist, and only does because comgr still does not link the device libraries correctly for opencl. How did you run into this?

I don't think that should be the only reason. In case of indirect calls, we may want to clone functions based on conflicting register budget. So, this should exist but probably with some fixes.

267

User of 'bitcast' operator is CallBase, so you will have to go one level up in user "chain".

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

267

But you don't know how many layers up you would need to go. You could potentially need to do an exhaustive search which touches a lot of the program

[jweightman]

I've attempted to handle constexpr function pointer casts. I had some issues, so I left some TODOs describing the problems I encountered in the code. This version isn't passing the new lit test. Thanks for all the help so far!

[jweightman]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

267

@arsenm to your question of how I ran into this: I ran into this issue while working on the Cray compiler. With the way we have things set up, it is possible to have indirect calls in the IR when this pass runs, though I'm not sure if it's cause the issue with Clang.

267

See the loop that I added above. Hopefully that isn't too invasive, though it's possibly not the most robust way of doing it.

[jweightman]

This should now propagate attributes to functions called through a constexpr bitcast.

[jweightman]

Deleted a stray #include <iostream>

[arsenm]

I think at this point I'd prefer to just not worry about the cast stripping case and go back to just handling the direct callee.

This pass should not be required for correctness, and is only because opencl is still not linking the device libraries correctly such that the linked in functions have the correct subtarget/subtarget features. Long term I would like to move this pass into a form that does conservative merges of attributes, which can safely give up on anything indirect looking

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

259

This doesn't consider if there could be multiple users. It's a whole expression tree. It still feels wrong to be doing this style of search for this.

[jweightman]

I've reverted back to not handling calls to constexpr function pointers. I've also
removed it from the test.

[jweightman]

Cleaned up stray newline

[madhur13490]

In D101847#2755717, @arsenm wrote:

I think at this point I'd prefer to just not worry about the cast stripping case and go back to just handling the direct callee.

This pass should not be required for correctness, and is only because opencl is still not linking the device libraries correctly such that the linked in functions have the correct subtarget/subtarget features. Long term I would like to move this pass into a form that does conservative merges of attributes, which can safely give up on anything indirect looking

I think the new diff can be kept as anyway you have spent effort on it. But I am fine either way.

llvm/test/CodeGen/AMDGPU/propagate-attributes-function-pointer-argument.ll

15

Probably can have CHECK-NOT too to explicitly call out what should not be generated.

[madhur13490]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

267

A comment here, please!

llvm/test/CodeGen/AMDGPU/propagate-attributes-function-pointer-argument.ll

15

CHECK-LABEL followed by CHECK-DAG would be more appropriate here.

[jweightman]

Incorporated code review feedback from @madhur13490

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

269

Add a fixme for the bitcasted callee case

[jweightman]

Added fixme (@arsenm)

[madhur13490]

LGTM

[jweightman]

@arsenm @madhur13490 thank you both for the thorough review! I don't have commit access, so I will need someone to commit it for me.

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

 #include "Utils/AMDGPUBaseInfo.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 
+#include <iostream>
+
 #define DEBUG_TYPE "amdgpu-propagate-attributes"
 
 using namespace llvm;
 
 namespace llvm {
 extern const SubtargetFeatureKV AMDGPUFeatureKV[AMDGPU::NumSubtargetFeatures-1];
 }
 
do {
     Roots.insert(NewRoots.begin(), NewRoots.end());
     NewRoots.clear();
 
     for (auto &F : M.functions()) {
       if (F.isDeclaration())
         continue;
 
       const FnProperties CalleeProps(*TM, F);
-      SmallVector<std::pair<CallBase *, Function *>, 32> ToReplace;
+      SmallVector<std::pair<Use *, Function *>, 32> ToReplace;
       SmallSet<CallBase *, 32> Visited;
 
-      for (User *U : F.users()) {
+      for (Use &OriginalUse : F.uses()) {
+        // If the user is a constant expression, look for an enclosing call.
+        User *U = OriginalUse.getUser();
+        while(isa<ConstantExpr>(U)) {
+          // TODO: I'd like to assert that there's only one user, but it seems
+          // the constexpr cast in propagate-attributes-bitcast-function.ll has
+          // zero users, so this preserves existing behavior. However, this
+          // causes the instruction cast below to fail anyway. Why isn't it
+          // used by the call instruction? How can I clean this up?
+          if(!U->hasOneUser())
+            break;
+          //assert(U->hasOneUser());
+          U = *U->users().begin();
+        }
         Instruction *I = dyn_cast<Instruction>(U);
         if (!I)
           continue;
         CallBase *CI = dyn_cast<CallBase>(I);
-        if (!CI)
+        if (!CI || CI->getCalledOperand()->stripPointerCasts() != &F)
           continue;
         Function *Caller = CI->getCaller();
         if (!Caller || !Visited.insert(CI).second)
           continue;
         if (!Roots.count(Caller) && !NewRoots.count(Caller))
           continue;
 
         const FnProperties CallerProps(*TM, *Caller);
 
         if (CalleeProps == CallerProps) {
for (auto &F : M.functions()) {
             break;
           }
 
           NewF = cloneWithProperties(F, NewProps);
           Clones.push_back(Clone(CallerProps, &F, NewF));
           NewRoots.insert(NewF);
         }
 
-        ToReplace.push_back(std::make_pair(CI, NewF));
+        ToReplace.push_back(std::make_pair(&OriginalUse, NewF));
         Replaced.insert(&F);
 
         Changed = true;
       }
 
       while (!ToReplace.empty()) {
         auto R = ToReplace.pop_back_val();
-        R.first->setCalledFunction(R.second);
+        // TODO: Is there something else I need to do here? Making this change
+        // causes an assert when cleaning up the module:
+        // Assertion failed: (I != Map.end() && "Constant not found in constant table!")
+        R.first->set(R.second);
       }
     }
   } while (!NewRoots.empty());
 
   for (Function *F : Replaced) {
     if (F->use_empty())
       F->eraseFromParent();
   }

llvm/test/CodeGen/AMDGPU/propagate-attributes-function-pointer-argument.ll

+; This is a regression test for a bug in the AMDGPU Propagate Attributes pass
+; where a call instruction's callee could be replaced with a function pointer
+; passed to the original call instruction as an argument.
+;
+; Example:
+; `call void @f(void ()* @g)`
+; could become
+; `call void @g(void ()* @g.1)`
+; which is invalid IR.
+
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -amdgpu-propagate-attributes-late %s | FileCheck %s
+
+; CHECK: define amdgpu_kernel void @thiswasabug() #0
+; CHECK: call void @f(void ()* @g.1)
+; CHECK: call void @g()
+; CHECK: call i8* bitcast (i32* ()* @h to i8* ()*)()
+define amdgpu_kernel void @thiswasabug() #0 {
+    call void @f(void ()* @g)
+    call void @g()
+    call i8* bitcast (i32* ()* @h to i8* ()*)()
+    ret void
+}
+
+define private void @f(void ()* nocapture %0) #0 {
+    ret void
+}
+
+; In order to expose this bug, it is necessary that `g` have one of the
+; propagated attributes, so that a clone and substitution would take place if g
+; were actually the function being called.
+; CHECK-DAG: define private void @g.1() #1
+; CHECK-DAG: define internal void @g() #2
+define private void @g() #1 {
+    ret void
+}
+
+; CHECK-DAG: define internal i32* @h() #2
+define private i32* @h() #1 {
+    %1 = alloca i32
+    ret i32* %1
+}
+
+attributes #0 = { noinline }
+attributes #1 = { noinline "amdgpu-waves-per-eu"="1,10" }