Skip to content

Commit a00289e

Browse files
Prince781Prince781
committed
[NVPTX] handle more cases for loads and stores
Split unaligned stores and loads of v2f32. Add DAGCombiner rules for: - target-independent stores that store a v2f32 BUILD_VECTOR. We scalarize the value and rewrite the store Fix test cases.
1 parent 7dd1bec commit a00289e

File tree

5 files changed

+52
-25
lines changed

5 files changed

+52
-25
lines changed

‎llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp‎

Lines changed: 42 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -828,7 +828,7 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
828828
setTargetDAGCombine({ISD::ADD, ISD::AND, ISD::EXTRACT_VECTOR_ELT, ISD::FADD,
829829
ISD::MUL, ISD::SHL, ISD::SREM, ISD::UREM, ISD::VSELECT,
830830
ISD::BUILD_VECTOR, ISD::ADDRSPACECAST, ISD::FP_ROUND,
831-
ISD::TRUNCATE, ISD::LOAD, ISD::BITCAST});
831+
ISD::TRUNCATE, ISD::LOAD, ISD::STORE, ISD::BITCAST});
832832

833833
// setcc for f16x2 and bf16x2 needs special handling to prevent
834834
// legalizer's attempt to scalarize it due to v2i1 not being legal.
@@ -2992,10 +2992,10 @@ SDValue NVPTXTargetLowering::LowerLOAD(SDValue Op, SelectionDAG &DAG) const {
29922992
if (Op.getValueType() == MVT::i1)
29932993
return LowerLOADi1(Op, DAG);
29942994

2995-
// v2f16/v2bf16/v2i16/v4i8 are legal, so we can't rely on legalizer to handle
2996-
// unaligned loads and have to handle it here.
2995+
// v2f16/v2bf16/v2i16/v4i8/v2f32 are legal, so we can't rely on legalizer to
2996+
// handle unaligned loads and have to handle it here.
29972997
EVT VT = Op.getValueType();
2998-
if (Isv2x16VT(VT) || VT == MVT::v4i8) {
2998+
if (Isv2x16VT(VT) || VT == MVT::v4i8 || VT == MVT::v2f32) {
29992999
LoadSDNode *Load = cast<LoadSDNode>(Op);
30003000
EVT MemVT = Load->getMemoryVT();
30013001
if (!allowsMemoryAccessForAlignment(*DAG.getContext(), DAG.getDataLayout(),
@@ -3039,15 +3039,15 @@ SDValue NVPTXTargetLowering::LowerSTORE(SDValue Op, SelectionDAG &DAG) const {
30393039
if (VT == MVT::i1)
30403040
return LowerSTOREi1(Op, DAG);
30413041

3042-
// v2f16 is legal, so we can't rely on legalizer to handle unaligned
3043-
// stores and have to handle it here.
3044-
if ((Isv2x16VT(VT) || VT == MVT::v4i8) &&
3042+
// v2f16/v2bf16/v2i16/v4i8/v2f32 are legal, so we can't rely on legalizer to
3043+
// handle unaligned stores and have to handle it here.
3044+
if ((Isv2x16VT(VT) || VT == MVT::v4i8 || VT == MVT::v2f32) &&
30453045
!allowsMemoryAccessForAlignment(*DAG.getContext(), DAG.getDataLayout(),
30463046
VT, *Store->getMemOperand()))
30473047
return expandUnalignedStore(Store, DAG);
30483048

3049-
// v2f16, v2bf16 and v2i16 don't need special handling.
3050-
if (Isv2x16VT(VT) || VT == MVT::v4i8)
3049+
// v2f16/v2bf16/v2i16/v4i8/v2f32 don't need special handling.
3050+
if (Isv2x16VT(VT) || VT == MVT::v4i8 || VT == MVT::v2f32)
30513051
return SDValue();
30523052

30533053
if (VT.isVector())
@@ -3056,8 +3056,8 @@ SDValue NVPTXTargetLowering::LowerSTORE(SDValue Op, SelectionDAG &DAG) const {
30563056
return SDValue();
30573057
}
30583058

3059-
SDValue
3060-
NVPTXTargetLowering::LowerSTOREVector(SDValue Op, SelectionDAG &DAG) const {
3059+
static SDValue convertVectorStore(SDValue Op, SelectionDAG &DAG,
3060+
const SmallVectorImpl<SDValue> &Elements) {
30613061
SDNode *N = Op.getNode();
30623062
SDValue Val = N->getOperand(1);
30633063
SDLoc DL(N);
@@ -3124,6 +3124,8 @@ NVPTXTargetLowering::LowerSTOREVector(SDValue Op, SelectionDAG &DAG) const {
31243124
SDValue SubVector = DAG.getBuildVector(EltVT, DL, SubVectorElts);
31253125
Ops.push_back(SubVector);
31263126
}
3127+
} else if (!Elements.empty()) {
3128+
Ops.insert(Ops.end(), Elements.begin(), Elements.end());
31273129
} else {
31283130
for (unsigned i = 0; i < NumElts; ++i) {
31293131
SDValue ExtVal = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, EltVT, Val,
@@ -3141,10 +3143,19 @@ NVPTXTargetLowering::LowerSTOREVector(SDValue Op, SelectionDAG &DAG) const {
31413143
DAG.getMemIntrinsicNode(Opcode, DL, DAG.getVTList(MVT::Other), Ops,
31423144
MemSD->getMemoryVT(), MemSD->getMemOperand());
31433145

3144-
// return DCI.CombineTo(N, NewSt, true);
31453146
return NewSt;
31463147
}
31473148

3149+
// Default variant where we don't pass in elements.
3150+
static SDValue convertVectorStore(SDValue Op, SelectionDAG &DAG) {
3151+
return convertVectorStore(Op, DAG, SmallVector<SDValue>{});
3152+
}
3153+
3154+
SDValue NVPTXTargetLowering::LowerSTOREVector(SDValue Op,
3155+
SelectionDAG &DAG) const {
3156+
return convertVectorStore(Op, DAG);
3157+
}
3158+
31483159
// st i1 v, addr
31493160
// =>
31503161
// v1 = zxt v to i16
@@ -5289,6 +5300,9 @@ static SDValue PerformStoreCombineHelper(SDNode *N,
52895300
// -->
52905301
// StoreRetvalV2 {a, b}
52915302
// likewise for V2 -> V4 case
5303+
//
5304+
// We also handle target-independent stores, which require us to first
5305+
// convert to StoreV2.
52925306

52935307
std::optional<NVPTXISD::NodeType> NewOpcode;
52945308
switch (N->getOpcode()) {
@@ -5314,8 +5328,8 @@ static SDValue PerformStoreCombineHelper(SDNode *N,
53145328
SDValue CurrentOp = N->getOperand(I);
53155329
if (CurrentOp->getOpcode() == ISD::BUILD_VECTOR) {
53165330
assert(CurrentOp.getValueType() == MVT::v2f32);
5317-
NewOps.push_back(CurrentOp.getNode()->getOperand(0));
5318-
NewOps.push_back(CurrentOp.getNode()->getOperand(1));
5331+
NewOps.push_back(CurrentOp.getOperand(0));
5332+
NewOps.push_back(CurrentOp.getOperand(1));
53195333
} else {
53205334
NewOps.clear();
53215335
break;
@@ -6086,6 +6100,18 @@ static SDValue PerformBITCASTCombine(SDNode *N,
60866100
return SDValue();
60876101
}
60886102

6103+
static SDValue PerformStoreCombine(SDNode *N,
6104+
TargetLowering::DAGCombinerInfo &DCI) {
6105+
// check if the store'd value can be scalarized
6106+
SDValue StoredVal = N->getOperand(1);
6107+
if (StoredVal.getValueType() == MVT::v2f32 &&
6108+
StoredVal.getOpcode() == ISD::BUILD_VECTOR) {
6109+
SmallVector<SDValue> Elements(StoredVal->op_values());
6110+
return convertVectorStore(SDValue(N, 0), DCI.DAG, Elements);
6111+
}
6112+
return SDValue();
6113+
}
6114+
60896115
SDValue NVPTXTargetLowering::PerformDAGCombine(SDNode *N,
60906116
DAGCombinerInfo &DCI) const {
60916117
CodeGenOptLevel OptLevel = getTargetMachine().getOptLevel();
@@ -6115,6 +6141,8 @@ SDValue NVPTXTargetLowering::PerformDAGCombine(SDNode *N,
61156141
case NVPTXISD::LoadParam:
61166142
case NVPTXISD::LoadParamV2:
61176143
return PerformLoadCombine(N, DCI);
6144+
case ISD::STORE:
6145+
return PerformStoreCombine(N, DCI);
61186146
case NVPTXISD::StoreParam:
61196147
case NVPTXISD::StoreParamV2:
61206148
case NVPTXISD::StoreParamV4:

‎llvm/test/CodeGen/NVPTX/aggregate-return.ll‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,9 @@ define void @test_v2f32(<2 x float> %input, ptr %output) {
1010
; CHECK-LABEL: @test_v2f32
1111
%call = tail call <2 x float> @barv(<2 x float> %input)
1212
; CHECK: .param .align 8 .b8 retval0[8];
13-
; CHECK: ld.param.v2.f32 {[[E0:%f[0-9]+]], [[E1:%f[0-9]+]]}, [retval0];
13+
; CHECK: ld.param.b64 [[E0_1:%rd[0-9]+]], [retval0];
1414
store <2 x float> %call, ptr %output, align 8
15-
; CHECK: st.v2.f32 [{{%rd[0-9]+}}], {[[E0]], [[E1]]}
15+
; CHECK: st.b64 [{{%rd[0-9]+}}], [[E0_1]]
1616
ret void
1717
}
1818

‎llvm/test/CodeGen/NVPTX/f32x2-instructions.ll‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -512,14 +512,13 @@ define <2 x float> @test_frem_ftz(<2 x float> %a, <2 x float> %b) #2 {
512512
define void @test_ldst_v2f32(ptr %a, ptr %b) #0 {
513513
; CHECK-LABEL: test_ldst_v2f32(
514514
; CHECK: {
515-
; CHECK-NEXT: .reg .f32 %f<3>;
516-
; CHECK-NEXT: .reg .b64 %rd<3>;
515+
; CHECK-NEXT: .reg .b64 %rd<4>;
517516
; CHECK-EMPTY:
518517
; CHECK-NEXT: // %bb.0:
519518
; CHECK-NEXT: ld.param.u64 %rd2, [test_ldst_v2f32_param_1];
520519
; CHECK-NEXT: ld.param.u64 %rd1, [test_ldst_v2f32_param_0];
521-
; CHECK-NEXT: ld.v2.f32 {%f1, %f2}, [%rd1];
522-
; CHECK-NEXT: st.v2.f32 [%rd2], {%f1, %f2};
520+
; CHECK-NEXT: ld.b64 %rd3, [%rd1];
521+
; CHECK-NEXT: st.b64 [%rd2], %rd3;
523522
; CHECK-NEXT: ret;
524523
%t1 = load <2 x float>, ptr %a
525524
store <2 x float> %t1, ptr %b, align 32

‎llvm/test/CodeGen/NVPTX/load-with-non-coherent-cache.ll‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -110,9 +110,9 @@ define ptx_kernel void @foo10(ptr noalias readonly %from, ptr %to) {
110110
}
111111

112112
; SM20-LABEL: .visible .entry foo11(
113-
; SM20: ld.global.v2.f32
113+
; SM20: ld.global.b64
114114
; SM35-LABEL: .visible .entry foo11(
115-
; SM35: ld.global.nc.v2.f32
115+
; SM35: ld.global.nc.b64
116116
define ptx_kernel void @foo11(ptr noalias readonly %from, ptr %to) {
117117
%1 = load <2 x float>, ptr %from
118118
store <2 x float> %1, ptr %to

‎llvm/test/CodeGen/NVPTX/misaligned-vector-ldst.ll‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,15 +18,15 @@ define <4 x float> @t1(ptr %p1) {
1818
define <4 x float> @t2(ptr %p1) {
1919
; CHECK-NOT: ld.v4
2020
; CHECK-NOT: ld.v2
21-
; CHECK: ld.f32
21+
; CHECK: ld.u32
2222
%r = load <4 x float>, ptr %p1, align 4
2323
ret <4 x float> %r
2424
}
2525

2626
; CHECK-LABEL: t3
2727
define <4 x float> @t3(ptr %p1) {
2828
; CHECK-NOT: ld.v4
29-
; CHECK: ld.v2
29+
; CHECK: ld.b64
3030
%r = load <4 x float>, ptr %p1, align 8
3131
ret <4 x float> %r
3232
}
@@ -111,7 +111,7 @@ define void @s1(ptr %p1, <4 x float> %v) {
111111
define void @s2(ptr %p1, <4 x float> %v) {
112112
; CHECK-NOT: st.v4
113113
; CHECK-NOT: st.v2
114-
; CHECK: st.f32
114+
; CHECK: st.u32
115115
store <4 x float> %v, ptr %p1, align 4
116116
ret void
117117
}

0 commit comments

Comments
 (0)