[llvm] 8ef5da7 - [WebAssembly] Fix crash when selecting 64 bit lane extract operand

[Luke Lau via llvm-commits]

Author: Luke Lau
Date: 2022-12-19T10:37:19Z
New Revision: 8ef5da701003649d3d31c3b3aa299fba30e4a4b3

URL: https://github.com/llvm/llvm-project/commit/8ef5da701003649d3d31c3b3aa299fba30e4a4b3
DIFF: https://github.com/llvm/llvm-project/commit/8ef5da701003649d3d31c3b3aa299fba30e4a4b3.diff

LOG: [WebAssembly] Fix crash when selecting 64 bit lane extract operand

The tablegen patterns on vector_extract only match i32 constants, but
on wasm64 these come in as i64 constants. In certain situations this
would cause crashes whenever it couldn't select an extract_vector_elt
instruction.
Rather than add duplicate patterns for every instruction, this just
canonicalizes the constant to be i32 when lowering.
Fixes https://github.com/llvm/llvm-project/issues/57577

Differential Revision: https://reviews.llvm.org/D140205

Added: 
    llvm/test/CodeGen/WebAssembly/simd-extract64.ll

Modified: 
    llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp
index a2c7ee19d56f..eee707bcd66e 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp
@@ -2245,11 +2245,16 @@ WebAssemblyTargetLowering::LowerAccessVectorElement(SDValue Op,
                                                     SelectionDAG &DAG) const {
   // Allow constant lane indices, expand variable lane indices
   SDNode *IdxNode = Op.getOperand(Op.getNumOperands() - 1).getNode();
-  if (isa<ConstantSDNode>(IdxNode) || IdxNode->isUndef())
-    return Op;
-  else
-    // Perform default expansion
-    return SDValue();
+  if (isa<ConstantSDNode>(IdxNode) || IdxNode->isUndef()) {
+    // Ensure the index type is i32 to match the tablegen patterns
+    uint64_t Idx = cast<ConstantSDNode>(IdxNode)->getZExtValue();
+    SmallVector<SDValue, 3> Ops(Op.getNode()->ops());
+    Ops[Op.getNumOperands() - 1] =
+        DAG.getConstant(Idx, SDLoc(IdxNode), MVT::i32);
+    return DAG.getNode(Op.getOpcode(), SDLoc(Op), Op.getValueType(), Ops);
+  }
+  // Perform default expansion
+  return SDValue();
 }
 
 static SDValue unrollVectorShift(SDValue Op, SelectionDAG &DAG) {

diff  --git a/llvm/test/CodeGen/WebAssembly/simd-extract64.ll b/llvm/test/CodeGen/WebAssembly/simd-extract64.ll
new file mode 100644
index 000000000000..10b9c4d54f4f
--- /dev/null
+++ b/llvm/test/CodeGen/WebAssembly/simd-extract64.ll
@@ -0,0 +1,18 @@
+; RUN: llc < %s -mattr=+simd128 -verify-machineinstrs -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals
+
+; Regression test for a crash on wasm64 when trying to lower extract_vector_elt
+; with a 64 bit constant:
+;
+; t19: i64 = extract_vector_elt t18, Constant:i64<0>
+
+target triple = "wasm64-unknown-unknown"
+
+define void @foo() {
+  store <4 x i32> zeroinitializer, ptr poison, align 16
+  %1 = load <4 x i32>, ptr poison, align 16
+  %2 = extractelement <4 x i32> %1, i32 0
+  %3 = insertelement <2 x i32> undef, i32 %2, i32 0
+  %4 = insertelement <2 x i32> %3, i32 poison, i32 1
+  store <2 x i32> %4, ptr poison, align 8
+  unreachable
+}