[NVPTX] Restrict NVVM bf16 intrinsic upgrades to integer storage types

Project: llvm/llvm-project | Language: C++ | Concepts: backwards-compatibility, error-handling, testing-strategy | Difficulty: mid
Link: https://github.com/llvm/llvm-project/pull/197387

The problem

Older NVVM bf16 intrinsics in LLVM used integer types (like i16) as storage for bfloat16 values, but newer versions use the proper bfloat type. The auto-upgrade logic was too permissive: it would try to upgrade any non-bfloat signature, including invalid ones (like those returning half), which could later cause crashes or miscompilations.

Why it's non-trivial

A naive fix would just check for non-bfloat types and upgrade them all, but this would still incorrectly "upgrade" functions with completely invalid or mismatched signatures (e.g., wrong return type, wrong argument types, or wrong argument counts). The upgrade logic needs to be precise: only legacy signatures that used integer types of the correct bit width should be upgraded. All other cases should be left for the verifier to catch as errors.

The approach

1. Add a precise signature check for legacy integer storage types

A new helper function is introduced to check if the function signature matches the pattern of the legacy integer-based bf16 intrinsic (same bit width as the new bfloat signature):

+static bool shouldUpgradeNVPTXBF16IntrinsicSignature(Function *F,
+                                                     Intrinsic::ID IID) {
+  FunctionType *NewFnTy = Intrinsic::getType(F->getContext(), IID);
+  FunctionType *OldFnTy = F->getFunctionType();
+  auto IsOldBF16StorageTy = [](Type *OldTy, Type *NewTy) {
+    return OldTy->getScalarType()->isIntegerTy() &&
+           OldTy->getPrimitiveSizeInBits() == NewTy->getPrimitiveSizeInBits();
+  };
+
+  if (!IsOldBF16StorageTy(OldFnTy->getReturnType(), NewFnTy->getReturnType()))
+    return false;
+
+  if (OldFnTy->getNumParams() != NewFnTy->getNumParams())
+    return false;
+
+  for (unsigned I = 0, E = OldFnTy->getNumParams(); I != E; ++I)
+    if (!IsOldBF16StorageTy(OldFnTy->getParamType(I), NewFnTy->getParamType(I)))
+      return false;
+
+  return true;
+}

This function only returns true if the old function's return type and all parameter types are integer types of the same bit width as the new intrinsic's bfloat types.

2. Use the new check in the upgrade paths

The auto-upgrade logic now uses this helper to restrict upgrades to only valid legacy signatures:

-      if (!F->getReturnType()->getScalarType()->isBFloatTy()) {
-        Intrinsic::ID IID = shouldUpgradeNVPTXBF16Intrinsic(Name);
-        if (IID != Intrinsic::not_intrinsic) {
+      {
+        Intrinsic::ID IID = shouldUpgradeNVPTXBF16Intrinsic(Name);
+        if (IID != Intrinsic::not_intrinsic &&
+            shouldUpgradeNVPTXBF16IntrinsicSignature(F, IID)) {
           NewFn = nullptr;
           return true;
         }

Now, only functions matching the legacy integer signature pattern will be auto-upgraded.

Similarly, in the call upgrade logic:

-    if (IID != Intrinsic::not_intrinsic &&
-        !F->getReturnType()->getScalarType()->isBFloatTy()) {
+    if (IID != Intrinsic::not_intrinsic &&
+        shouldUpgradeNVPTXBF16IntrinsicSignature(F, IID)) {
       rename(F);
       Function *NewFn = Intrinsic::getOrInsertDeclaration(F->getParent(), IID);
       SmallVector<Value *, 2> Args;

This ensures only valid legacy calls are upgraded.

3. Add targeted tests

• Valid upgrade test: Asserts that legacy integer signatures are upgraded to bfloat correctly.
• Invalid upgrade tests: Asserts that functions with invalid signatures (wrong return type, wrong argument type, wrong number of arguments) are not upgraded, and instead cause verifier errors.

Example invalid test:

; CHECK: intrinsic has incorrect return type!
declare half @llvm.nvvm.neg.bf16(bfloat)

define void @test() {
  %t = call half @llvm.nvvm.neg.bf16(bfloat 1.0)
  ret void
}

This checks that a half return type is not silently upgraded and is caught as an error.

From the review thread

[review] arsenm: Check case with wrong argument instead of return too

[review] arsenm: Check wrong number of argument case

[review] durga4github: Can we name this function something like isLegacyNVPTXBF16IntSignature ?

[review] woruyu: Done!

Takeaways

• Upgrade logic should be narrowly targeted: Only perform automatic upgrades when the legacy signature matches exactly the expected pattern; otherwise, let the verifier catch errors.
• Test both valid and invalid cases: Always add tests to confirm that only correct legacy forms are upgraded, and that invalid forms are rejected.
• Helper functions improve clarity and maintainability: Encapsulate signature checks in a helper function to avoid duplicating logic and to make intent explicit.