Skip to content

Commit 4e589e6

Browse files
tlivelytlively
committed
[WebAssembly] Fix SIMD shift unrolling to avoid assertion failure
Summary: Using the default DAG.UnrollVectorOp on v16i8 and v8i16 vectors results in i8 or i16 nodes being inserted into the SelectionDAG. Since those are illegal types, this causes a legalization assertion failure for some code patterns, as uncovered by PR45178. This change unrolls shifts manually to avoid this issue by adding and using a new optional EVT argument to DAG.ExtractVectorElements to control the type of the extract_element nodes. Reviewers: aheejin, dschuff Subscribers: sbc100, jgravelle-google, hiraditya, sunfish, zzheng, llvm-commits Tags: #llvm Differential Revision: https://reviews.llvm.org/D76043
1 parent fe74df0 commit 4e589e6

File tree

4 files changed

+157
-22
lines changed

4 files changed

+157
-22
lines changed

‎llvm/include/llvm/CodeGen/SelectionDAG.h

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1740,10 +1740,13 @@ class SelectionDAG {
17401740
/// Widen the vector up to the next power of two using INSERT_SUBVECTOR.
17411741
SDValue WidenVector(const SDValue &N, const SDLoc &DL);
17421742

1743-
/// Append the extracted elements from Start to Count out of the vector Op
1744-
/// in Args. If Count is 0, all of the elements will be extracted.
1743+
/// Append the extracted elements from Start to Count out of the vector Op in
1744+
/// Args. If Count is 0, all of the elements will be extracted. The extracted
1745+
/// elements will have type EVT if it is provided, and otherwise their type
1746+
/// will be Op's element type.
17451747
void ExtractVectorElements(SDValue Op, SmallVectorImpl<SDValue> &Args,
1746-
unsigned Start = 0, unsigned Count = 0);
1748+
unsigned Start = 0, unsigned Count = 0,
1749+
EVT EltVT = EVT());
17471750

17481751
/// Compute the default alignment value for the given type.
17491752
unsigned getEVTAlignment(EVT MemoryVT) const;

‎llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9501,12 +9501,13 @@ SDValue SelectionDAG::WidenVector(const SDValue &N, const SDLoc &DL) {
95019501

95029502
void SelectionDAG::ExtractVectorElements(SDValue Op,
95039503
SmallVectorImpl<SDValue> &Args,
9504-
unsigned Start, unsigned Count) {
9504+
unsigned Start, unsigned Count,
9505+
EVT EltVT) {
95059506
EVT VT = Op.getValueType();
95069507
if (Count == 0)
95079508
Count = VT.getVectorNumElements();
9508-
9509-
EVT EltVT = VT.getVectorElementType();
9509+
if (EltVT == EVT())
9510+
EltVT = VT.getVectorElementType();
95109511
SDLoc SL(Op);
95119512
for (unsigned i = Start, e = Start + Count; i != e; ++i) {
95129513
Args.push_back(getNode(ISD::EXTRACT_VECTOR_ELT, SL, EltVT, Op,

‎llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

Lines changed: 19 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1582,22 +1582,25 @@ static SDValue unrollVectorShift(SDValue Op, SelectionDAG &DAG) {
15821582
return DAG.UnrollVectorOp(Op.getNode());
15831583
// Otherwise mask the shift value to get proper semantics from 32-bit shift
15841584
SDLoc DL(Op);
1585-
SDValue ShiftVal = Op.getOperand(1);
1586-
uint64_t MaskVal = LaneT.getSizeInBits() - 1;
1587-
SDValue MaskedShiftVal = DAG.getNode(
1588-
ISD::AND, // mask opcode
1589-
DL, ShiftVal.getValueType(), // masked value type
1590-
ShiftVal, // original shift value operand
1591-
DAG.getConstant(MaskVal, DL, ShiftVal.getValueType()) // mask operand
1592-
);
1593-
1594-
return DAG.UnrollVectorOp(
1595-
DAG.getNode(Op.getOpcode(), // original shift opcode
1596-
DL, Op.getValueType(), // original return type
1597-
Op.getOperand(0), // original vector operand,
1598-
MaskedShiftVal // new masked shift value operand
1599-
)
1600-
.getNode());
1585+
size_t NumLanes = Op.getSimpleValueType().getVectorNumElements();
1586+
SDValue Mask = DAG.getConstant(LaneT.getSizeInBits() - 1, DL, MVT::i32);
1587+
unsigned ShiftOpcode = Op.getOpcode();
1588+
SmallVector<SDValue, 16> ShiftedElements;
1589+
DAG.ExtractVectorElements(Op.getOperand(0), ShiftedElements, 0, 0, MVT::i32);
1590+
SmallVector<SDValue, 16> ShiftElements;
1591+
DAG.ExtractVectorElements(Op.getOperand(1), ShiftElements, 0, 0, MVT::i32);
1592+
SmallVector<SDValue, 16> UnrolledOps;
1593+
for (size_t i = 0; i < NumLanes; ++i) {
1594+
SDValue MaskedShiftValue =
1595+
DAG.getNode(ISD::AND, DL, MVT::i32, ShiftElements[i], Mask);
1596+
SDValue ShiftedValue = ShiftedElements[i];
1597+
if (ShiftOpcode == ISD::SRA)
1598+
ShiftedValue = DAG.getNode(ISD::SIGN_EXTEND_INREG, DL, MVT::i32,
1599+
ShiftedValue, DAG.getValueType(LaneT));
1600+
UnrolledOps.push_back(
1601+
DAG.getNode(ShiftOpcode, DL, MVT::i32, ShiftedValue, MaskedShiftValue));
1602+
}
1603+
return DAG.getBuildVector(Op.getValueType(), DL, UnrolledOps);
16011604
}
16021605

16031606
SDValue WebAssemblyTargetLowering::LowerShift(SDValue Op,

‎llvm/test/CodeGen/WebAssembly/simd-shift-unroll.ll

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
; RUN: llc < %s -asm-verbose=false -verify-machineinstrs -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals -wasm-keep-registers -mattr=+unimplemented-simd128 | FileCheck %s --check-prefixes CHECK,SIMD128,SIMD128-SLOW
2+
3+
;; Test that the custom shift unrolling works correctly in cases that
4+
;; cause assertion failures due to illegal types when using
5+
;; DAG.UnrollVectorOp. Regression test for PR45178.
6+
7+
target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
8+
target triple = "wasm32-unknown-unknown"
9+
10+
; CHECK-LABEL: shl_v16i8:
11+
; CHECK-NEXT: .functype shl_v16i8 (v128) -> (v128)
12+
; CHECK-NEXT: i8x16.extract_lane_u $push0=, $0, 0
13+
; CHECK-NEXT: i32.const $push1=, 3
14+
; CHECK-NEXT: i32.shl $push2=, $pop0, $pop1
15+
; CHECK-NEXT: i8x16.splat $push3=, $pop2
16+
; CHECK-NEXT: i8x16.extract_lane_u $push4=, $0, 1
17+
; CHECK-NEXT: i8x16.replace_lane $push5=, $pop3, 1, $pop4
18+
; ...
19+
; CHECK: i8x16.extract_lane_u $push32=, $0, 15
20+
; CHECK-NEXT: i8x16.replace_lane $push33=, $pop31, 15, $pop32
21+
; CHECK-NEXT: v8x16.shuffle $push34=, $pop33, $0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
22+
; CHECK-NEXT: return $pop34
23+
define <16 x i8> @shl_v16i8(<16 x i8> %in) {
24+
%out = shl <16 x i8> %in,
25+
<i8 3, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0,
26+
i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0>
27+
%ret = shufflevector <16 x i8> %out, <16 x i8> undef, <16 x i32> zeroinitializer
28+
ret <16 x i8> %ret
29+
}
30+
31+
; CHECK-LABEL: shr_s_v16i8:
32+
; CHECK-NEXT: functype shr_s_v16i8 (v128) -> (v128)
33+
; CHECK-NEXT: i8x16.extract_lane_s $push0=, $0, 0
34+
; CHECK-NEXT: i32.const $push1=, 3
35+
; CHECK-NEXT: i32.shr_s $push2=, $pop0, $pop1
36+
; CHECK-NEXT: i8x16.splat $push3=, $pop2
37+
; CHECK-NEXT: i8x16.extract_lane_s $push4=, $0, 1
38+
; CHECK-NEXT: i8x16.replace_lane $push5=, $pop3, 1, $pop4
39+
; ...
40+
; CHECK: i8x16.extract_lane_s $push32=, $0, 15
41+
; CHECK-NEXT: i8x16.replace_lane $push33=, $pop31, 15, $pop32
42+
; CHECK-NEXT: v8x16.shuffle $push34=, $pop33, $0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
43+
; CHECK-NEXT: return $pop34
44+
define <16 x i8> @shr_s_v16i8(<16 x i8> %in) {
45+
%out = ashr <16 x i8> %in,
46+
<i8 3, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0,
47+
i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0>
48+
%ret = shufflevector <16 x i8> %out, <16 x i8> undef, <16 x i32> zeroinitializer
49+
ret <16 x i8> %ret
50+
}
51+
52+
; CHECK-LABEL: shr_u_v16i8:
53+
; CHECK-NEXT: functype shr_u_v16i8 (v128) -> (v128)
54+
; CHECK-NEXT: i8x16.extract_lane_u $push0=, $0, 0
55+
; CHECK-NEXT: i32.const $push1=, 3
56+
; CHECK-NEXT: i32.shr_u $push2=, $pop0, $pop1
57+
; CHECK-NEXT: i8x16.splat $push3=, $pop2
58+
; CHECK-NEXT: i8x16.extract_lane_u $push4=, $0, 1
59+
; CHECK-NEXT: i8x16.replace_lane $push5=, $pop3, 1, $pop4
60+
; ...
61+
; CHECK: i8x16.extract_lane_u $push32=, $0, 15
62+
; CHECK-NEXT: i8x16.replace_lane $push33=, $pop31, 15, $pop32
63+
; CHECK-NEXT: v8x16.shuffle $push34=, $pop33, $0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
64+
; CHECK-NEXT: return $pop34
65+
define <16 x i8> @shr_u_v16i8(<16 x i8> %in) {
66+
%out = lshr <16 x i8> %in,
67+
<i8 3, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0,
68+
i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0, i8 0>
69+
%ret = shufflevector <16 x i8> %out, <16 x i8> undef, <16 x i32> zeroinitializer
70+
ret <16 x i8> %ret
71+
}
72+
73+
; CHECK-LABEL: shl_v8i16:
74+
; CHECK-NEXT: functype shl_v8i16 (v128) -> (v128)
75+
; CHECK-NEXT: i16x8.extract_lane_u $push0=, $0, 0
76+
; CHECK-NEXT: i32.const $push1=, 9
77+
; CHECK-NEXT: i32.shl $push2=, $pop0, $pop1
78+
; CHECK-NEXT: i16x8.splat $push3=, $pop2
79+
; CHECK-NEXT: i16x8.extract_lane_u $push4=, $0, 1
80+
; CHECK-NEXT: i16x8.replace_lane $push5=, $pop3, 1, $pop4
81+
; ...
82+
; CHECK: i16x8.extract_lane_u $push16=, $0, 7
83+
; CHECK-NEXT: i16x8.replace_lane $push17=, $pop15, 7, $pop16
84+
; CHECK-NEXT: v8x16.shuffle $push18=, $pop17, $0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1
85+
; CHECK-NEXT: return $pop18
86+
define <8 x i16> @shl_v8i16(<8 x i16> %in) {
87+
%out = shl <8 x i16> %in, <i16 9, i16 0, i16 0, i16 0, i16 0, i16 0, i16 0, i16 0>
88+
%ret = shufflevector <8 x i16> %out, <8 x i16> undef, <8 x i32> zeroinitializer
89+
ret <8 x i16> %ret
90+
}
91+
92+
; CHECK-LABEL: shr_s_v8i16:
93+
; CHECK-NEXT: functype shr_s_v8i16 (v128) -> (v128)
94+
; CHECK-NEXT: i16x8.extract_lane_s $push0=, $0, 0
95+
; CHECK-NEXT: i32.const $push1=, 9
96+
; CHECK-NEXT: i32.shr_s $push2=, $pop0, $pop1
97+
; CHECK-NEXT: i16x8.splat $push3=, $pop2
98+
; CHECK-NEXT: i16x8.extract_lane_s $push4=, $0, 1
99+
; CHECK-NEXT: i16x8.replace_lane $push5=, $pop3, 1, $pop4
100+
; ...
101+
; CHECK: i16x8.extract_lane_s $push16=, $0, 7
102+
; CHECK-NEXT: i16x8.replace_lane $push17=, $pop15, 7, $pop16
103+
; CHECK-NEXT: v8x16.shuffle $push18=, $pop17, $0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1
104+
; CHECK-NEXT: return $pop18
105+
define <8 x i16> @shr_s_v8i16(<8 x i16> %in) {
106+
%out = ashr <8 x i16> %in, <i16 9, i16 0, i16 0, i16 0, i16 0, i16 0, i16 0, i16 0>
107+
%ret = shufflevector <8 x i16> %out, <8 x i16> undef, <8 x i32> zeroinitializer
108+
ret <8 x i16> %ret
109+
}
110+
111+
; CHECK-LABEL: shr_u_v8i16:
112+
; CHECK-NEXT: functype shr_u_v8i16 (v128) -> (v128)
113+
; CHECK-NEXT: i16x8.extract_lane_u $push0=, $0, 0
114+
; CHECK-NEXT: i32.const $push1=, 9
115+
; CHECK-NEXT: i32.shr_u $push2=, $pop0, $pop1
116+
; CHECK-NEXT: i16x8.splat $push3=, $pop2
117+
; CHECK-NEXT: i16x8.extract_lane_u $push4=, $0, 1
118+
; CHECK-NEXT: i16x8.replace_lane $push5=, $pop3, 1, $pop4
119+
; ...
120+
; CHECK: i16x8.extract_lane_u $push16=, $0, 7
121+
; CHECK-NEXT: i16x8.replace_lane $push17=, $pop15, 7, $pop16
122+
; CHECK-NEXT: v8x16.shuffle $push18=, $pop17, $0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1
123+
; CHECK-NEXT: return $pop18
124+
define <8 x i16> @shr_u_v8i16(<8 x i16> %in) {
125+
%out = lshr <8 x i16> %in, <i16 9, i16 0, i16 0, i16 0, i16 0, i16 0, i16 0, i16 0>
126+
%ret = shufflevector <8 x i16> %out, <8 x i16> undef, <8 x i32> zeroinitializer
127+
ret <8 x i16> %ret
128+
}

0 commit comments

Comments
 (0)