LLVM freeze instruction between mask and div 5/5

Author: alexbaden

third_party/intel/lib/LLVMIR/CMakeLists.txt

@@ -3,4 +3,4 @@ add_triton_library(TritonIntelLLVMIR
 
         DEPENDS
         LLVMIRIncGen
-        )
+        )

third_party/intel/lib/LLVMIR/LLVMIRFreezeMaskedDivRem.cpp

@@ -1,80 +1,57 @@
 #include "LLVMPasses.h"
-#include "llvm/IR/Instructions.h"
-#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Dominators.h"
+#include "llvm/IR/Instructions.h"
 
 using namespace llvm;
 
-static bool processPhiNode(PHINode *phiNode, BasicBlock& BB) {
-      llvm::errs() << "YOLO: " << *phiNode << "\n";
-
-  const auto phiHasNullValue = any_of(phiNode->incoming_values(), [](Use& U) {
-    if (Constant *C = dyn_cast<Constant>(&U)) {
-      return C->isNullValue();
-    }
-    return false; 
-  });
+static bool processPhiNode(PHINode *PhiNode, BasicBlock &BB) {
+  if (!any_of(PhiNode->incoming_values(), [](Use &U) {
+        if (Constant *C = dyn_cast<Constant>(&U)) {
+          return C->isNullValue();
+        }
+        return false;
+      })) {
+    return false;
+  }
 
   bool Changed = false;
-  if (phiHasNullValue) {
-    for (Instruction &I : BB) {
-      if (I.getOpcode() == Instruction::SDiv || I.getOpcode() == Instruction::SRem) {
-        const size_t OpIdx = 1; // I.getOpcode() == Instruction::SRem ? 0 : 1;
-        if (I.getOperand(OpIdx) == phiNode) {
-          auto *freezePhi = new FreezeInst(phiNode, phiNode->getName() + ".frozen", I.getIterator());
-          I.setOperand(OpIdx, freezePhi);
-          Changed = true;
-        }
+  for (Instruction &I : BB) {
+    if (I.getOpcode() == Instruction::SDiv ||
+        I.getOpcode() == Instruction::SRem) {
+      const size_t OpIdx = 1;
+      if (I.getOperand(OpIdx) == PhiNode) {
+        auto *freezePhi = new FreezeInst(
+            PhiNode, PhiNode->getName() + ".frozen", I.getIterator());
+        I.setOperand(OpIdx, freezePhi);
+        Changed = true;
       }
     }
-#if 0
-        auto FindUse = llvm::find_if(phiNode->users(), [](auto *U) {
-          auto *Use = cast<Instruction>(U);
-          llvm::errs() << "User: " << *Use << "\n";
-          return (Use->getOpcode() == Instruction::SDiv || Use->getOpcode() == Instruction::SRem);
-        });
-        if (FindUse == phiNode->user_end()) {
-          llvm::errs() << "no div :(\n";
-          return false; 
-        }
-        auto *Use = cast<Instruction>(*FindUse);
-        assert(Use->isIntDivRem());
-        const size_t OpIdx = Use->getOpcode() == Instruction::SRem ? 0 : 1;
-        if (Use->getOperand(OpIdx) == phiNode) {
-          llvm::errs() << "Got our user! " << *Use << "\n";
-          llvm::errs() << "Operand 1: " << *Use->getOperand(1) << "\n";
-          auto *freezePhi = new FreezeInst(phiNode, phiNode->getName() + ".frozen", Use->getIterator());
-          Use->setOperand(OpIdx, freezePhi);
-          Changed = true;
-        }
-#endif
   }
-  return Changed; 
+  return Changed;
 }
 
-static bool runOnFunction(Function& F, const TargetTransformInfo &TTI,
-                           const DominatorTree &DT) {
-    bool Changed = false;
+static bool runOnFunction(Function &F) {
+  bool Changed = false;
 
-    SmallVector<PHINode *> PhiNodes;
-    for (BasicBlock &BB : F) {
-    for (Instruction &inst : BB) {
-      if (PHINode *phiNode = dyn_cast<PHINode>(&inst)) {
-        Changed |= processPhiNode(phiNode, BB);
+  SmallVector<PHINode *> PhiNodes;
+  for (BasicBlock &BB : F) {
+    for (Instruction &I : BB) {
+      if (PHINode *PhiNode = dyn_cast<PHINode>(&I)) {
+        Changed |= processPhiNode(PhiNode, BB);
         continue;
       }
       break;
     }
   }
 
-    return Changed;
+  return Changed;
 }
 
- PreservedAnalyses FreezeMaskedDivRemPass::run(Function &F, FunctionAnalysisManager &FAM) {
-    TargetTransformInfo &TTI = FAM.getResult<TargetIRAnalysis>(F);
-    DominatorTree &DT = FAM.getResult<DominatorTreeAnalysis>(F);
-    const auto b = runOnFunction(F, TTI, DT);
+PreservedAnalyses FreezeMaskedDivRemPass::run(Function &F,
+                                              FunctionAnalysisManager &FAM) {
+  const auto b = runOnFunction(F);
 
-    return b ? PreservedAnalyses::none() : PreservedAnalyses::all();
-}
+  return b ? PreservedAnalyses::none() : PreservedAnalyses::all();
+}

third_party/intel/lib/LLVMIR/LLVMPasses.h

@@ -8,4 +8,4 @@ struct FreezeMaskedDivRemPass : PassInfoMixin<FreezeMaskedDivRemPass> {
   static StringRef name() { return "FreezeMaskedDivRemPass"; }
 };
 
-}
+} // namespace llvm

third_party/intel/triton_xpu.cc

@@ -208,22 +208,21 @@ void init_triton_intel(py::module &&m) {
           fpm.addPass(BreakStructPhiNodesPass());
           fpm.addPass(InstCombinePass());
         });
-#if 1
     pb.registerPeepholeEPCallback(
         [&](llvm::FunctionPassManager &fpm, llvm::OptimizationLevel level) {
           // The Triton masked load pattern can generate instances where the
-          // mask false path appears to cause undefined behavior during
-          // computation. Even though the result of that behavior will never be
-          // used, LLVM can choose to optimize away the false path resulting in
-          // an incorrect result for the kernel. Adding `DivRemPairsPass`
-          // introduces freeze instructions which prevent UB from leaking into
-          // div/rem instructions.
-          // fpm.addPass(DivRemPairsPass());
+          // mask value causes undefined behavior in sdiv/srem instructions. The
+          // language allows this UB as the result of those arithmetic
+          // instructions is never used, and control flow to avoid computation
+          // of these instructions would negatively affect performance. But,
+          // LLVM SimplifyCFG aggressively marks code paths with undefined
+          // behavior as dead. This can result in removal of the mask path and
+          // incorrect results from legal Triton kernels due to masked elements
+          // being used in computation. Run a pass to add a freeze instruction
+          // between masked loads and sdiv/srem to signal to LLVM we consider
+          // the sdiv/srem operands to be well defined.
           fpm.addPass(FreezeMaskedDivRemPass());
         });
-#else
-    mpm.addPass(createModuleToFunctionPassAdaptor(FreezeMaskedDivRemPass()));
-#endif
     mpm.addPass(pb.buildPerModuleDefaultPipeline(opt));
     mpm.run(*mod, mam);
   });