[WebAssembly] Fix FastISel of condition in different block (PR51651)

[nikic]

Summary

If the icmp is in a different block, then the register for the icmp operand may not be initialized, as it nominally does not have cross-block uses. Add a check that the icmp is in the same block as the branch, which should be the common case.

This matches what X86 FastISel does: https://github.com/llvm/llvm-project/blob/5b6b090cf2129228f05d7d0f504676b67f7524cf/llvm/lib/Target/X86/X86FastISel.cpp#L1648

Fixes https://bugs.llvm.org/show_bug.cgi?id=51651.

[nikic]

llvm/lib/Target/WebAssembly/WebAssemblyFastISel.cpp

433

This code could principally have the same issue. However, it's also dead code, so I just dropped it entirely.

The reason it's dead is that V must have i1 type here (used either as a branch or select condition), while this checks that the type is i32.

[tlively]

Looks good, thanks!

llvm/lib/Target/WebAssembly/WebAssemblyFastISel.cpp

MVT::SimpleValueType getLegalType(MVT::SimpleValueType VT) {
     }
     return MVT::INVALID_SIMPLE_VALUE_TYPE;
   }
   bool computeAddress(const Value *Obj, Address &Addr);
   void materializeLoadStoreOperands(Address &Addr);
   void addLoadStoreOperands(const Address &Addr, const MachineInstrBuilder &MIB,
                             MachineMemOperand *MMO);
   unsigned maskI1Value(unsigned Reg, const Value *V);
-  unsigned getRegForI1Value(const Value *V, bool &Not);
+  unsigned getRegForI1Value(const Value *V, const BasicBlock *BB, bool &Not);
   unsigned zeroExtendToI32(unsigned Reg, const Value *V,
                            MVT::SimpleValueType From);
   unsigned signExtendToI32(unsigned Reg, const Value *V,
                            MVT::SimpleValueType From);
   unsigned zeroExtend(unsigned Reg, const Value *V, MVT::SimpleValueType From,
                       MVT::SimpleValueType To);
   unsigned signExtend(unsigned Reg, const Value *V, MVT::SimpleValueType From,
                       MVT::SimpleValueType To);
void WebAssemblyFastISel::addLoadStoreOperands(const Address &Addr,
 
   MIB.addMemOperand(MMO);
 }
 
 unsigned WebAssemblyFastISel::maskI1Value(unsigned Reg, const Value *V) {
   return zeroExtendToI32(Reg, V, MVT::i1);
 }
 
-unsigned WebAssemblyFastISel::getRegForI1Value(const Value *V, bool &Not) {
+unsigned WebAssemblyFastISel::getRegForI1Value(const Value *V,
+                                               const BasicBlock *BB,
+                                               bool &Not) {
   if (const auto *ICmp = dyn_cast<ICmpInst>(V))
     if (const ConstantInt *C = dyn_cast<ConstantInt>(ICmp->getOperand(1)))
-      if (ICmp->isEquality() && C->isZero() && C->getType()->isIntegerTy(32)) {
+      if (ICmp->isEquality() && C->isZero() && C->getType()->isIntegerTy(32) &&
+          ICmp->getParent() == BB) {
         Not = ICmp->isTrueWhenEqual();
         return getRegForValue(ICmp->getOperand(0));
       }
 
-  Value *NotV;
-  if (match(V, m_Not(m_Value(NotV))) && V->getType()->isIntegerTy(32)) {
-    Not = true;
-    return getRegForValue(NotV);
-  }
-
   Not = false;
   unsigned Reg = getRegForValue(V);
   if (Reg == 0)
     return 0;
   return maskI1Value(Reg, V);
 }
 
 unsigned WebAssemblyFastISel::zeroExtendToI32(unsigned Reg, const Value *V,
if (!IsVoid)
     updateValueMap(Call, ResultReg);
   return true;
 }
 
 bool WebAssemblyFastISel::selectSelect(const Instruction *I) {
   const auto *Select = cast<SelectInst>(I);
 
   bool Not;
-  unsigned CondReg = getRegForI1Value(Select->getCondition(), Not);
+  unsigned CondReg =
+      getRegForI1Value(Select->getCondition(), I->getParent(), Not);
   if (CondReg == 0)
     return false;
 
   unsigned TrueReg = getRegForValue(Select->getTrueValue());
   if (TrueReg == 0)
     return false;
 
   unsigned FalseReg = getRegForValue(Select->getFalseValue());
if (Br->isUnconditional()) {
     fastEmitBranch(MSucc, Br->getDebugLoc());
     return true;
   }
 
   MachineBasicBlock *TBB = FuncInfo.MBBMap[Br->getSuccessor(0)];
   MachineBasicBlock *FBB = FuncInfo.MBBMap[Br->getSuccessor(1)];
 
   bool Not;
-  unsigned CondReg = getRegForI1Value(Br->getCondition(), Not);
+  unsigned CondReg = getRegForI1Value(Br->getCondition(), Br->getParent(), Not);
   if (CondReg == 0)
     return false;
 
   unsigned Opc = WebAssembly::BR_IF;
   if (Not)
     Opc = WebAssembly::BR_UNLESS;
 
   BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc, TII.get(Opc))

llvm/test/CodeGen/WebAssembly/pr51651.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -O0 -mtriple=wasm32-unknown-unknown -wasm-disable-explicit-locals -wasm-keep-registers < %s | FileCheck %s
+
+define i32 @test(i8* %p, i8* %p2) {
+; CHECK-LABEL: test:
+; CHECK:         .functype test (i32, i32) -> (i32)
+; CHECK-NEXT:  # %bb.0:
+; CHECK-NEXT:    i32.load8_u $3=, 0($0)
+; CHECK-NEXT:    i32.eqz $2=, $3
+; CHECK-NEXT:    i32.store8 0($1), $3
+; CHECK-NEXT:  # %bb.1: # %bb2
+; CHECK-NEXT:    i32.const $4=, 1
+; CHECK-NEXT:    i32.and $5=, $2, $4
+; CHECK-NEXT:    block
+; CHECK-NEXT:    br_if 0, $5 # 0: down to label0
+; CHECK-NEXT:  # %bb.2: # %bb4
+; CHECK-NEXT:    i32.const $6=, 0
+; CHECK-NEXT:    return $6
+; CHECK-NEXT:  .LBB0_3: # %bb3
+; CHECK-NEXT:    end_block # label0:
+; CHECK-NEXT:    i32.const $7=, 1
+; CHECK-NEXT:    return $7
+  %v = load i8, i8* %p
+  %v.ext = zext i8 %v to i32
+  %cond = icmp eq i32 %v.ext, 0
+  ; Cause FastISel abort.
+  %shl = shl i8 %v, 0
+  store i8 %shl, i8* %p2
+  br label %bb2
+
+bb2:
+  br i1 %cond, label %bb3, label %bb4
+
+bb4:
+  ret i32 0
+
+bb3:
+  ret i32 1
+}