[CUDA][HIP] Fix template argument deduction

[yaxunl]

Summary

nvcc allows using std::malloc and std::free in device code.
When std::malloc or std::free is passed as a template
function argument with template argument deduction,
there is no diagnostics. e.g.

#include <memory>

__global__ void kern() {
    void *p = std::malloc(1);
    std::free(p);
}
int main()
{

    std::shared_ptr<float> a;
    a = std::shared_ptr<float>(
      (float*)std::malloc(sizeof(float) * 100),
      std::free
    );
    return 0;
}

However, the same code fails to compile with clang
(https://godbolt.org/z/1roGvo6YY). The reason is
that clang does not have logic to choose a function
argument from an overloaded set of candidates
based on host/device attributes for template argument
deduction.

Currently, clang does have a logic to choose a candidate
based on the constraints of the candidates. This patch
extends that logic to account for the CUDA host/device-based
preference.

[tra]

clang/lib/Sema/SemaOverload.cpp

12758–12764

Maybe CheckCUDAPreference should return -1/0/1 or an enum. std::optional does not seem to be very readable here.

E.g. if(MorePreferableByCUDA) sounds like it's going to be satisfied when FD is a better choice than Result, but it's not the case.
I think this would be easier to follow:

if (CheckCUDAPreference(FD, Result) <= 0) // or `!= CP_BETTER`
     continue;

[yaxunl]

clang/lib/Sema/SemaOverload.cpp

12758–12764

will use an integer for that.

[yaxunl]

revised by comments

clang/lib/Sema/SemaOverload.cpp

Sema::ResolveAddressOfOverloadedFunction(Expr *AddressOfExpr,
   }
 
   if (pHadMultipleCandidates)
     *pHadMultipleCandidates = Resolver.hadMultipleCandidates();
   return Fn;
 }
 
 /// Given an expression that refers to an overloaded function, try to
 /// resolve that function to a single function that can have its address taken.
 /// This will modify `Pair` iff it returns non-null.
 ///
 /// This routine can only succeed if from all of the candidates in the overload
 /// set for SrcExpr that can have their addresses taken, there is one candidate
 /// that is more constrained than the rest.
 FunctionDecl *
 Sema::resolveAddressOfSingleOverloadCandidate(Expr *E, DeclAccessPair &Pair) {
   OverloadExpr::FindResult R = OverloadExpr::find(E);
   OverloadExpr *Ovl = R.Expression;
   bool IsResultAmbiguous = false;
   FunctionDecl *Result = nullptr;
   DeclAccessPair DAP;
   SmallVector<FunctionDecl *, 2> AmbiguousDecls;
 
+  // Return positive for better, negative for worse, 0 for equal preference.
+  auto CheckCUDAPreference = [&](FunctionDecl *FD1, FunctionDecl *FD2) {
+    FunctionDecl *Caller = getCurFunctionDecl(/*AllowLambda=*/true);
+    return static_cast<int>(IdentifyCUDAPreference(Caller, FD1)) -
+           static_cast<int>(IdentifyCUDAPreference(Caller, FD2));
+  };
+
   auto CheckMoreConstrained = [&](FunctionDecl *FD1,
                                   FunctionDecl *FD2) -> std::optional<bool> {
     if (FunctionDecl *MF = FD1->getInstantiatedFromMemberFunction())
       FD1 = MF;
     if (FunctionDecl *MF = FD2->getInstantiatedFromMemberFunction())
       FD2 = MF;
     SmallVector<const Expr *, 1> AC1, AC2;
     FD1->getAssociatedConstraints(AC1);
Sema::resolveAddressOfSingleOverloadCandidate(Expr *E, DeclAccessPair &Pair) {
   for (auto I = Ovl->decls_begin(), E = Ovl->decls_end(); I != E; ++I) {
     auto *FD = dyn_cast<FunctionDecl>(I->getUnderlyingDecl());
     if (!FD)
       return nullptr;
 
     if (!checkAddressOfFunctionIsAvailable(FD))
       continue;
 
+    // If we found a better result, update Result.
+    auto FoundBetter = [&]() {
+      IsResultAmbiguous = false;
+      DAP = I.getPair();
+      Result = FD;
+    };
+
     // We have more than one result - see if it is more constrained than the
     // previous one.
     if (Result) {
+      // Check CUDA preference first. If the candidates have differennt CUDA
+      // preference, choose the one with higher CUDA preference. Otherwise,
+      // choose the one with more constraints.
+      if (getLangOpts().CUDA) {
+        int PreferenceByCUDA = CheckCUDAPreference(FD, Result);
+        // FD has different preference than Result.
+        if (PreferenceByCUDA != 0) {
+          // FD is more preferable than Result.
+          if (PreferenceByCUDA > 0)
+            FoundBetter();
+          continue;
+        }
+      }
+      // FD has the same CUDA prefernece than Result. Continue check
+      // constraints.
       std::optional<bool> MoreConstrainedThanPrevious =
           CheckMoreConstrained(FD, Result);
       if (!MoreConstrainedThanPrevious) {
         IsResultAmbiguous = true;
         AmbiguousDecls.push_back(FD);
         continue;
       }
       if (!*MoreConstrainedThanPrevious)
         continue;
       // FD is more constrained - replace Result with it.
     }
-    IsResultAmbiguous = false;
-    DAP = I.getPair();
-    Result = FD;
+    FoundBetter();
   }
 
   if (IsResultAmbiguous)
     return nullptr;
 
   if (Result) {
     SmallVector<const Expr *, 1> ResultAC;
     // We skipped over some ambiguous declarations which might be ambiguous with
     // the selected result.
-    for (FunctionDecl *Skipped : AmbiguousDecls)
+    for (FunctionDecl *Skipped : AmbiguousDecls) {
+      // If skipped candidate has different CUDA preference than the result,
+      // there is no ambiguity. Otherwise check whether they have different
+      // constraints.
+      if (getLangOpts().CUDA && CheckCUDAPreference(Skipped, Result) != 0)
+        continue;
       if (!CheckMoreConstrained(Skipped, Result))
         return nullptr;
+    }
     Pair = DAP;
   }
   return Result;
 }
 
 /// Given an overloaded function, tries to turn it into a non-overloaded
 /// function reference using resolveAddressOfSingleOverloadCandidate. This
 /// will perform access checks, diagnose the use of the resultant decl, and, if

clang/test/SemaCUDA/template-arg-deduction.cu

+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fsyntax-only -verify %s
+// RUN: %clang_cc1 -triple nvptx64-nvidia-cuda -fsyntax-only -fcuda-is-device -verify %s
+
+// expected-no-diagnostics
+
+#include "Inputs/cuda.h"
+
+void foo();
+__device__ void foo();
+
+template<class F>
+void host_temp(F f);
+
+template<class F>
+__device__ void device_temp(F f);
+
+void host_caller() {
+  host_temp(foo);
+}
+
+__global__ void kernel_caller() {
+  device_temp(foo);
+}
+
+__device__ void device_caller() {
+  device_temp(foo);
+}