[WebAssembly] Don't assume that zext/sext result is i32/i64 in fast isel (PR41841)

[nikic]

Summary

Usually this will abort fast-isel at the instruction using the non-legal result, but if the only use is in a different basic block, we'll incorrectly assume that the zext/sext is to i32 (rather than i128 in this case).

[tlively]

Thanks for fixing this! I have a few comments on the test style, but otherwise LGTM.

llvm/test/CodeGen/WebAssembly/PR41841.ll

2

We usually pass -asm-verbose=false in our tests to avoid having to match uninteresting output.

19

Please move the CHECK lines to be right before the top of the function

37

Please standardize the whitespace to make the start of the instructions line up.

[nikic]

Add -asm-verbose=false, move CHECK lines, remove tabs.

[tlively]

Looks good. Feel free to commit yourself or lmk if you want me to commit it for you.

llvm/lib/Target/WebAssembly/WebAssemblyFastISel.cpp

if (To == MVT::i64) {
 
     unsigned Result = createResultReg(&WebAssembly::I64RegClass);
     BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc,
             TII.get(WebAssembly::I64_EXTEND_U_I32), Result)
         .addReg(Reg);
     return Result;
   }
 
+  if (To == MVT::i32)
     return zeroExtendToI32(Reg, V, From);
+
+  return 0;
 }
 
 unsigned WebAssemblyFastISel::signExtend(unsigned Reg, const Value *V,
                                          MVT::SimpleValueType From,
                                          MVT::SimpleValueType To) {
   if (To == MVT::i64) {
     if (From == MVT::i64)
       return copyValue(Reg);
 
     Reg = signExtendToI32(Reg, V, From);
 
     unsigned Result = createResultReg(&WebAssembly::I64RegClass);
     BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc,
             TII.get(WebAssembly::I64_EXTEND_S_I32), Result)
         .addReg(Reg);
     return Result;
   }
 
+  if (To == MVT::i32)
     return signExtendToI32(Reg, V, From);
+
+  return 0;
 }
 
 unsigned WebAssemblyFastISel::getRegForUnsignedValue(const Value *V) {
   MVT::SimpleValueType From = getSimpleType(V->getType());
   MVT::SimpleValueType To = getLegalType(From);
   unsigned VReg = getRegForValue(V);
   if (VReg == 0)
     return 0;

llvm/test/CodeGen/WebAssembly/PR41841.ll

+; RUN: llc < %s -O0 -wasm-disable-explicit-locals -wasm-keep-registers -asm-verbose=false | FileCheck %s
+
+target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+target triple = "wasm32-unknown-unknown"
+
+declare void @foo(i128)
+
+; CHECK-LABEL: test_zext:
+; CHECK-NEXT: .functype test_zext (i32) -> (){{$}}
+; CHECK-NEXT: i64.extend_i32_u $[[TMP3:[0-9]+]]=, $0{{$}}
+; CHECK-NEXT: i64.const $[[TMP4:[0-9]+]]=, 1{{$}}
+; CHECK-NEXT: i64.and $[[TMP1:[0-9]+]]=, $[[TMP3]], $[[TMP4]]{{$}}
+; CHECK-NEXT: i64.const $[[TMP2:[0-9]+]]=, 0{{$}}
+; CHECK-NEXT: call foo, $[[TMP1]], $[[TMP2]]{{$}}
+; CHECK-NEXT: return{{$}}
+define void @test_zext(i1 %b) nounwind {
+  %res = zext i1 %b to i128
+  br label %next
+
+next:                                             ; preds = %start
+  call void @foo(i128 %res)
+  ret void
+}
+
+; CHECK-LABEL: test_sext:
+; CHECK-NEXT:.functype test_sext (i32) -> (){{$}}
+; CHECK-NEXT: i64.extend_i32_u $[[TMP3:[0-9]+]]=, $0{{$}}
+; CHECK-NEXT: i64.const $[[TMP4:[0-9]+]]=, 1{{$}}
+; CHECK-NEXT: i64.and $[[TMP5:[0-9]+]]=, $[[TMP3]], $[[TMP4]]{{$}}
+; CHECK-NEXT: i64.const $[[TMP6:[0-9]+]]=, 0{{$}}
+; CHECK-NEXT: i64.sub $[[TMP1:[0-9]+]]=, $[[TMP6]], $[[TMP5]]{{$}}
+; CHECK-NEXT: local.copy $[[TMP2:[0-9]+]]=, $[[TMP1]]{{$}}
+; CHECK-NEXT: call foo, $[[TMP1]], $[[TMP2]]{{$}}
+; CHECK-NEXT: return{{$}}
+define void @test_sext(i1 %b) nounwind {
+  %res = sext i1 %b to i128
+  br label %next
+
+next:                                             ; preds = %start
+  call void @foo(i128 %res)
+  ret void
+}