[llvm] [NNVMReflect] Force dead branch elimination in NNVMReflect (PR #81189)

Joseph Huber via llvm-commits llvm-commits at lists.llvm.org
Thu Feb 8 13:07:49 PST 2024


https://github.com/jhuber6 created https://github.com/llvm/llvm-project/pull/81189

Summary:
The `__nvvm_reflect` function is used to guard invalid code that varies
between architectures. One problem with this feature is that if it is
used without optimizations, it will leave invalid code in the module
that will then make it to the backend. The `__nvvm_reflect` pass is
already mandatory, so it should do some trivial branch removal to ensure
that constants are handled correctly. This dead branch elimination only
works in the trivial case of a compare on a branch and does not touch
any conditionals that were not realted to the `__nvvm_reflect` call in
order to preserve `O0` semantics as much as possible. This should allow
the following to work on NVPTX targets

```c
int foo() {
  if (__nvvm_reflect__("__CUDA_ARCH") >= 700)
    asm("valid;\n");
}
```


>From 06bf4fa269961910da3b695077e078ef35ff5154 Mon Sep 17 00:00:00 2001
From: Joseph Huber <huberjn at outlook.com>
Date: Thu, 8 Feb 2024 14:59:17 -0600
Subject: [PATCH] [NNVMReflect] Force dead branch elimination in NNVMReflect

Summary:
The `__nvvm_reflect` function is used to guard invalid code that varies
between architectures. One problem with this feature is that if it is
used without optimizations, it will leave invalid code in the module
that will then make it to the backend. The `__nvvm_reflect` pass is
already mandatory, so it should do some trivial branch removal to ensure
that constants are handled correctly. This dead branch elimination only
works in the trivial case of a compare on a branch and does not touch
any conditionals that were not realted to the `__nvvm_reflect` call in
order to preserve `O0` semantics as much as possible. This should allow
the following to work on NVPTX targets

```c
int foo() {
  if (__nvvm_reflect__("__CUDA_ARCH") >= 700)
    asm("valid;\n");
}
```
---
 llvm/lib/Target/NVPTX/NVVMReflect.cpp         |  63 +++++++++++
 .../CodeGen/NVPTX/nvvm-reflect-arch-O0.ll     | 102 ++++++++++++++++++
 llvm/test/CodeGen/NVPTX/nvvm-reflect-arch.ll  |   1 -
 3 files changed, 165 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/NVPTX/nvvm-reflect-arch-O0.ll

diff --git a/llvm/lib/Target/NVPTX/NVVMReflect.cpp b/llvm/lib/Target/NVPTX/NVVMReflect.cpp
index 7d2678ae592748..993d022324cbdb 100644
--- a/llvm/lib/Target/NVPTX/NVVMReflect.cpp
+++ b/llvm/lib/Target/NVPTX/NVVMReflect.cpp
@@ -20,6 +20,7 @@
 
 #include "NVPTX.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
@@ -36,6 +37,8 @@
 #include "llvm/Support/raw_os_ostream.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Local.h"
 #include <sstream>
 #include <string>
 #define NVVM_REFLECT_FUNCTION "__nvvm_reflect"
@@ -87,6 +90,7 @@ static bool runNVVMReflect(Function &F, unsigned SmVersion) {
   }
 
   SmallVector<Instruction *, 4> ToRemove;
+  SmallVector<ICmpInst *, 4> ToSimplify;
 
   // Go through the calls in this function.  Each call to __nvvm_reflect or
   // llvm.nvvm.reflect should be a CallInst with a ConstantArray argument.
@@ -171,6 +175,12 @@ static bool runNVVMReflect(Function &F, unsigned SmVersion) {
     } else if (ReflectArg == "__CUDA_ARCH") {
       ReflectVal = SmVersion * 10;
     }
+
+    // If the immediate user is a simple comparison we want to simplify it.
+    for (User *U : Call->users())
+      if (ICmpInst *I = dyn_cast<ICmpInst>(U))
+        ToSimplify.push_back(I);
+
     Call->replaceAllUsesWith(ConstantInt::get(Call->getType(), ReflectVal));
     ToRemove.push_back(Call);
   }
@@ -178,6 +188,59 @@ static bool runNVVMReflect(Function &F, unsigned SmVersion) {
   for (Instruction *I : ToRemove)
     I->eraseFromParent();
 
+  // The code guarded by __nvvm_reflect may be invalid for the target machine.
+  // We need to do some basic dead code elimination to trim invalid code before
+  // it reaches the backend at all optimization levels.
+  SmallVector<BranchInst *> Simplified;
+  for (ICmpInst *Cmp : ToSimplify) {
+    Constant *LHS = dyn_cast<Constant>(Cmp->getOperand(0));
+    Constant *RHS = dyn_cast<Constant>(Cmp->getOperand(1));
+
+    if (!LHS || !RHS)
+      continue;
+
+    // If the comparison is a compile time constat we sipmly propagate it.
+    Constant *C = ConstantFoldCompareInstOperands(
+        Cmp->getPredicate(), LHS, RHS, Cmp->getModule()->getDataLayout());
+
+    if (!C)
+      continue;
+
+    for (User *U : Cmp->users())
+      if (BranchInst *I = dyn_cast<BranchInst>(U))
+        Simplified.push_back(I);
+
+    Cmp->replaceAllUsesWith(C);
+    Cmp->eraseFromParent();
+
+    F.dump();
+  }
+
+  // Each instruction here is a conditional branch off of a constant true or
+  // false value. Simply replace it with an unconditional branch to the
+  // appropriate basic block and delete the rest if it is trivally dead.
+  DenseSet<Instruction *> Removed;
+  for (BranchInst *Branch : Simplified) {
+    if (Removed.contains(Branch))
+      continue;
+
+    ConstantInt *C = dyn_cast<ConstantInt>(Branch->getCondition());
+    if (!C || (!C->isOne() && !C->isZero()))
+      continue;
+
+    BasicBlock *TrueBB =
+        C->isOne() ? Branch->getSuccessor(0) : Branch->getSuccessor(1);
+    BasicBlock *FalseBB =
+        C->isOne() ? Branch->getSuccessor(1) : Branch->getSuccessor(0);
+
+    ReplaceInstWithInst(Branch, BranchInst::Create(TrueBB));
+    if (FalseBB->use_empty() && FalseBB->hasNPredecessors(0) &&
+        FalseBB->getFirstNonPHIOrDbg()) {
+      Removed.insert(FalseBB->getFirstNonPHIOrDbg());
+      changeToUnreachable(FalseBB->getFirstNonPHIOrDbg());
+    }
+  }
+
   return ToRemove.size() > 0;
 }
 
diff --git a/llvm/test/CodeGen/NVPTX/nvvm-reflect-arch-O0.ll b/llvm/test/CodeGen/NVPTX/nvvm-reflect-arch-O0.ll
new file mode 100644
index 00000000000000..3e480aaf89e50f
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/nvvm-reflect-arch-O0.ll
@@ -0,0 +1,102 @@
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_52 -mattr=+ptx64 -O0 | FileCheck %s --check-prefix=SM_52
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_70 -mattr=+ptx64 -O0 | FileCheck %s --check-prefix=SM_70
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_90 -mattr=+ptx72 -O0 | FileCheck %s --check-prefix=SM_90
+
+ at .str = private unnamed_addr constant [12 x i8] c"__CUDA_ARCH\00"
+
+declare i32 @__nvvm_reflect(ptr)
+
+;      SM_52: .visible .func  (.param .b32 func_retval0) foo()
+;      SM_52: mov.b32         %[[REG:.+]], 3;
+; SM_52-NEXT: st.param.b32    [func_retval0+0], %[[REG:.+]];
+; SM_52-NEXT: ret;
+;
+;      SM_70: .visible .func  (.param .b32 func_retval0) foo()
+;      SM_70: mov.b32         %[[REG:.+]], 2;
+; SM_70-NEXT: st.param.b32    [func_retval0+0], %[[REG:.+]];
+; SM_70-NEXT: ret;
+;
+;      SM_90: .visible .func  (.param .b32 func_retval0) foo()
+;      SM_90: mov.b32         %[[REG:.+]], 1;
+; SM_90-NEXT: st.param.b32    [func_retval0+0], %[[REG:.+]];
+; SM_90-NEXT: ret;
+define i32 @foo() {
+entry:
+  %call = call i32 @__nvvm_reflect(ptr @.str)
+  %cmp = icmp uge i32 %call, 900
+  br i1 %cmp, label %if.then, label %if.else
+
+if.then:
+  br label %return
+
+if.else:
+  %call1 = call i32 @__nvvm_reflect(ptr @.str)
+  %cmp2 = icmp uge i32 %call1, 700
+  br i1 %cmp2, label %if.then3, label %if.else4
+
+if.then3:
+  br label %return
+
+if.else4:
+  %call5 = call i32 @__nvvm_reflect(ptr @.str)
+  %cmp6 = icmp uge i32 %call5, 520
+  br i1 %cmp6, label %if.then7, label %if.else8
+
+if.then7:
+  br label %return
+
+if.else8:
+  br label %return
+
+return:
+  %retval.0 = phi i32 [ 1, %if.then ], [ 2, %if.then3 ], [ 3, %if.then7 ], [ 4, %if.else8 ]
+  ret i32 %retval.0
+}
+
+;      SM_52: .visible .func  (.param .b32 func_retval0) bar()
+;      SM_52: mov.b32         %[[REG:.+]], 2;
+; SM_52-NEXT: st.param.b32    [func_retval0+0], %[[REG:.+]];
+; SM_52-NEXT: ret;
+;
+;      SM_70: .visible .func  (.param .b32 func_retval0) bar()
+;      SM_70: mov.b32         %[[REG:.+]], 1;
+; SM_70-NEXT: st.param.b32    [func_retval0+0], %[[REG:.+]];
+; SM_70-NEXT: ret;
+;
+;      SM_90: .visible .func  (.param .b32 func_retval0) bar()
+;      SM_90: mov.b32         %[[REG:.+]], 1;
+; SM_90-NEXT: st.param.b32    [func_retval0+0], %[[REG:.+]];
+; SM_90-NEXT: ret;
+define i32 @bar() {
+entry:
+  %call = call i32 @__nvvm_reflect(ptr @.str)
+  %cmp = icmp uge i32 %call, 700
+  br i1 %cmp, label %if.then, label %if.else
+
+if.then:
+  br label %if.end
+
+if.else:
+  br label %if.end
+
+if.end:
+  %x = phi i32 [ 1, %if.then ], [ 2, %if.else ]
+  ret i32 %x
+}
+
+; SM_52-NOT: valid;
+; SM_70: valid;
+; SM_90: valid;
+define void @baz() {
+entry:
+  %call = call i32 @__nvvm_reflect(ptr @.str)
+  %cmp = icmp uge i32 %call, 700
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  call void asm sideeffect "valid;\0A", ""()
+  br label %if.end
+
+if.end:
+  ret void
+}
diff --git a/llvm/test/CodeGen/NVPTX/nvvm-reflect-arch.ll b/llvm/test/CodeGen/NVPTX/nvvm-reflect-arch.ll
index e8c554c9ed5289..ac5875c6ab1043 100644
--- a/llvm/test/CodeGen/NVPTX/nvvm-reflect-arch.ll
+++ b/llvm/test/CodeGen/NVPTX/nvvm-reflect-arch.ll
@@ -18,4 +18,3 @@ define i32 @foo(float %a, float %b) {
 ; SM35: ret i32 350  
   ret i32 %reflect
 }
-



More information about the llvm-commits mailing list