Skip to content

Commit 8412432

Browse files
Prince781Prince781
committed
[NVPTX] fold movs into loads and stores
1 parent 0bd614a commit 8412432

23 files changed

+2978
-2315
lines changed

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

Lines changed: 259 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -238,18 +238,11 @@ getVectorLoweringShape(EVT VectorEVT, bool CanLowerTo256Bit) {
238238
return std::nullopt;
239239
LLVM_FALLTHROUGH;
240240
case MVT::v2i8:
241-
case MVT::v2i16:
242241
case MVT::v2i32:
243242
case MVT::v2i64:
244-
case MVT::v2f16:
245-
case MVT::v2bf16:
246243
case MVT::v2f32:
247244
case MVT::v2f64:
248-
case MVT::v4i8:
249-
case MVT::v4i16:
250245
case MVT::v4i32:
251-
case MVT::v4f16:
252-
case MVT::v4bf16:
253246
case MVT::v4f32:
254247
// This is a "native" vector type
255248
return std::pair(NumElts, EltVT);
@@ -262,6 +255,13 @@ getVectorLoweringShape(EVT VectorEVT, bool CanLowerTo256Bit) {
262255
if (!CanLowerTo256Bit)
263256
return std::nullopt;
264257
LLVM_FALLTHROUGH;
258+
case MVT::v2i16: // <1 x i16x2>
259+
case MVT::v2f16: // <1 x f16x2>
260+
case MVT::v2bf16: // <1 x bf16x2>
261+
case MVT::v4i8: // <1 x i8x4>
262+
case MVT::v4i16: // <2 x i16x2>
263+
case MVT::v4f16: // <2 x f16x2>
264+
case MVT::v4bf16: // <2 x bf16x2>
265265
case MVT::v8i8: // <2 x i8x4>
266266
case MVT::v8f16: // <4 x f16x2>
267267
case MVT::v8bf16: // <4 x bf16x2>
@@ -845,7 +845,8 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
845845
// We have some custom DAG combine patterns for these nodes
846846
setTargetDAGCombine({ISD::ADD, ISD::AND, ISD::EXTRACT_VECTOR_ELT, ISD::FADD,
847847
ISD::MUL, ISD::SHL, ISD::SREM, ISD::UREM, ISD::VSELECT,
848-
ISD::BUILD_VECTOR, ISD::ADDRSPACECAST});
848+
ISD::BUILD_VECTOR, ISD::ADDRSPACECAST, ISD::LOAD,
849+
ISD::STORE});
849850

850851
// setcc for f16x2 and bf16x2 needs special handling to prevent
851852
// legalizer's attempt to scalarize it due to v2i1 not being legal.
@@ -3464,19 +3465,16 @@ SDValue NVPTXTargetLowering::LowerFormalArguments(
34643465
unsigned I = 0;
34653466
for (const unsigned NumElts : VectorInfo) {
34663467
const EVT EltVT = VTs[I];
3467-
const EVT LoadVT = [&]() -> EVT {
3468-
// i1 is loaded/stored as i8.
3469-
if (EltVT == MVT::i1)
3470-
return MVT::i8;
3471-
// getLoad needs a vector type, but it can't handle
3472-
// vectors which contain v2f16 or v2bf16 elements. So we must load
3473-
// using i32 here and then bitcast back.
3474-
if (EltVT.isVector())
3475-
return MVT::getIntegerVT(EltVT.getFixedSizeInBits());
3476-
return EltVT;
3477-
}();
3468+
// i1 is loaded/stored as i8
3469+
const EVT LoadVT = EltVT == MVT::i1 ? MVT::i8 : EltVT;
3470+
// If the element is a packed type (ex. v2f16, v4i8, etc) holding
3471+
// multiple elements.
3472+
const unsigned PackingAmt =
3473+
LoadVT.isVector() ? LoadVT.getVectorNumElements() : 1;
3474+
3475+
const EVT VecVT = EVT::getVectorVT(
3476+
F->getContext(), LoadVT.getScalarType(), NumElts * PackingAmt);
34783477

3479-
const EVT VecVT = EVT::getVectorVT(F->getContext(), LoadVT, NumElts);
34803478
SDValue VecAddr = DAG.getObjectPtrOffset(
34813479
dl, ArgSymbol, TypeSize::getFixed(Offsets[I]));
34823480

@@ -3496,8 +3494,10 @@ SDValue NVPTXTargetLowering::LowerFormalArguments(
34963494
if (P.getNode())
34973495
P.getNode()->setIROrder(Arg.getArgNo() + 1);
34983496
for (const unsigned J : llvm::seq(NumElts)) {
3499-
SDValue Elt = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, dl, LoadVT, P,
3500-
DAG.getIntPtrConstant(J, dl));
3497+
SDValue Elt =
3498+
DAG.getNode(LoadVT.isVector() ? ISD::EXTRACT_SUBVECTOR
3499+
: ISD::EXTRACT_VECTOR_ELT,
3500+
dl, LoadVT, P, DAG.getIntPtrConstant(J * PackingAmt, dl));
35013501

35023502
// Extend or truncate the element if necessary (e.g. an i8 is loaded
35033503
// into an i16 register)
@@ -3511,9 +3511,6 @@ SDValue NVPTXTargetLowering::LowerFormalArguments(
35113511
Elt);
35123512
} else if (ExpactedVT.bitsLT(Elt.getValueType())) {
35133513
Elt = DAG.getNode(ISD::TRUNCATE, dl, ExpactedVT, Elt);
3514-
} else {
3515-
// v2f16 was loaded as an i32. Now we must bitcast it back.
3516-
Elt = DAG.getBitcast(EltVT, Elt);
35173514
}
35183515
InVals.push_back(Elt);
35193516
}
@@ -5047,26 +5044,244 @@ PerformFADDCombineWithOperands(SDNode *N, SDValue N0, SDValue N1,
50475044
return SDValue();
50485045
}
50495046

5050-
static SDValue PerformStoreCombineHelper(SDNode *N, std::size_t Front,
5051-
std::size_t Back) {
5047+
/// Combine extractelts into a load by increasing the number of return values.
5048+
static SDValue
5049+
combineUnpackingMovIntoLoad(SDNode *N, TargetLowering::DAGCombinerInfo &DCI) {
5050+
// Don't run this optimization before the legalizer
5051+
if (!DCI.isAfterLegalizeDAG())
5052+
return SDValue();
5053+
5054+
EVT ElemVT = N->getValueType(0);
5055+
if (!Isv2x16VT(ElemVT))
5056+
return SDValue();
5057+
5058+
// Check whether all outputs are either used by an extractelt or are
5059+
// glue/chain nodes
5060+
if (!all_of(N->uses(), [&](SDUse &U) {
5061+
// Skip glue, chain nodes
5062+
if (U.getValueType() == MVT::Glue || U.getValueType() == MVT::Other)
5063+
return true;
5064+
if (U.getUser()->getOpcode() == ISD::EXTRACT_VECTOR_ELT) {
5065+
if (N->getOpcode() != ISD::LOAD)
5066+
return true;
5067+
// Since this is an ISD::LOAD, check all extractelts are used. If
5068+
// any are not used, we don't want to defeat another optimization that
5069+
// will narrow the load.
5070+
//
5071+
// For example:
5072+
//
5073+
// L: v2f16,ch = load <p>
5074+
// e0: f16 = extractelt L:0, 0
5075+
// e1: f16 = extractelt L:0, 1 <-- unused
5076+
// store e0
5077+
//
5078+
// Can be optimized by DAGCombiner to:
5079+
//
5080+
// L: f16,ch = load <p>
5081+
// store L:0
5082+
return !U.getUser()->use_empty();
5083+
}
5084+
5085+
// Otherwise, this use prevents us from splitting a value.
5086+
return false;
5087+
}))
5088+
return SDValue();
5089+
5090+
auto *LD = cast<MemSDNode>(N);
5091+
EVT MemVT = LD->getMemoryVT();
5092+
SDLoc DL(LD);
5093+
5094+
// the new opcode after we double the number of operands
5095+
NVPTXISD::NodeType Opcode;
5096+
SmallVector<SDValue> Operands(LD->ops());
5097+
unsigned OldNumValues;
5098+
switch (LD->getOpcode()) {
5099+
case ISD::LOAD:
5100+
OldNumValues = 1;
5101+
// Any packed type is legal, so the legalizer will not have lowered
5102+
// ISD::LOAD -> NVPTXISD::Load (unless it's under-aligned). We have to do it
5103+
// here.
5104+
Opcode = NVPTXISD::LoadV2;
5105+
Operands.push_back(DCI.DAG.getIntPtrConstant(
5106+
cast<LoadSDNode>(LD)->getExtensionType(), DL));
5107+
break;
5108+
case NVPTXISD::LoadParamV2:
5109+
OldNumValues = 2;
5110+
Opcode = NVPTXISD::LoadParamV4;
5111+
break;
5112+
case NVPTXISD::LoadV2:
5113+
OldNumValues = 2;
5114+
Opcode = NVPTXISD::LoadV4;
5115+
break;
5116+
case NVPTXISD::LoadV4:
5117+
// PTX doesn't support v8 for 16-bit values
5118+
case NVPTXISD::LoadV8:
5119+
// PTX doesn't support the next doubling of outputs
5120+
return SDValue();
5121+
}
5122+
5123+
SmallVector<EVT> NewVTs(OldNumValues * 2, ElemVT.getVectorElementType());
5124+
// add remaining chain and glue values
5125+
for (unsigned I = OldNumValues; I < LD->getNumValues(); ++I)
5126+
NewVTs.push_back(LD->getValueType(I));
5127+
5128+
// Create the new load
5129+
SDValue NewLoad =
5130+
DCI.DAG.getMemIntrinsicNode(Opcode, DL, DCI.DAG.getVTList(NewVTs),
5131+
Operands, MemVT, LD->getMemOperand());
5132+
5133+
// Now we use a combination of BUILD_VECTORs and a MERGE_VALUES node to keep
5134+
// the outputs the same. These nodes will be optimized away in later
5135+
// DAGCombiner iterations.
5136+
SmallVector<SDValue> Results;
5137+
for (unsigned I = 0; I < NewLoad->getNumValues();) {
5138+
if (NewLoad->getValueType(I) == ElemVT.getVectorElementType()) {
5139+
Results.push_back(DCI.DAG.getBuildVector(
5140+
ElemVT, DL, {NewLoad.getValue(I), NewLoad.getValue(I + 1)}));
5141+
I += 2;
5142+
} else {
5143+
Results.push_back(NewLoad.getValue(I));
5144+
I += 1;
5145+
}
5146+
}
5147+
5148+
return DCI.DAG.getMergeValues(Results, DL);
5149+
}
5150+
5151+
/// Fold a packing mov into a store. This may help lower register pressure.
5152+
///
5153+
/// ex:
5154+
/// v: v2f16 = build_vector a:f16, b:f16
5155+
/// StoreRetval v
5156+
///
5157+
/// ...is turned into...
5158+
///
5159+
/// StoreRetvalV2 a:f16, b:f16
5160+
static SDValue combinePackingMovIntoStore(SDNode *N,
5161+
TargetLowering::DAGCombinerInfo &DCI,
5162+
unsigned Front, unsigned Back) {
5163+
// We want to run this as late as possible since other optimizations may
5164+
// eliminate the BUILD_VECTORs.
5165+
if (!DCI.isAfterLegalizeDAG())
5166+
return SDValue();
5167+
5168+
// Get the type of the operands being stored.
5169+
EVT ElementVT = N->getOperand(Front).getValueType();
5170+
5171+
if (!Isv2x16VT(ElementVT))
5172+
return SDValue();
5173+
5174+
auto *ST = cast<MemSDNode>(N);
5175+
EVT MemVT = ElementVT.getVectorElementType();
5176+
5177+
// The new opcode after we double the number of operands.
5178+
NVPTXISD::NodeType Opcode;
5179+
switch (N->getOpcode()) {
5180+
case ISD::STORE:
5181+
// Any packed type is legal, so the legalizer will not have lowered
5182+
// ISD::STORE -> NVPTXISD::Store (unless it's under-aligned). We have to do
5183+
// it here.
5184+
MemVT = ST->getMemoryVT();
5185+
Opcode = NVPTXISD::StoreV2;
5186+
break;
5187+
case NVPTXISD::StoreParam:
5188+
Opcode = NVPTXISD::StoreParamV2;
5189+
break;
5190+
case NVPTXISD::StoreParamV2:
5191+
Opcode = NVPTXISD::StoreParamV4;
5192+
break;
5193+
case NVPTXISD::StoreRetval:
5194+
Opcode = NVPTXISD::StoreRetvalV2;
5195+
break;
5196+
case NVPTXISD::StoreRetvalV2:
5197+
Opcode = NVPTXISD::StoreRetvalV4;
5198+
break;
5199+
case NVPTXISD::StoreV2:
5200+
MemVT = ST->getMemoryVT();
5201+
Opcode = NVPTXISD::StoreV4;
5202+
break;
5203+
case NVPTXISD::StoreV4:
5204+
// PTX doesn't support v8 for 16-bit values
5205+
case NVPTXISD::StoreParamV4:
5206+
case NVPTXISD::StoreRetvalV4:
5207+
case NVPTXISD::StoreV8:
5208+
// PTX doesn't support the next doubling of operands for these opcodes.
5209+
return SDValue();
5210+
default:
5211+
llvm_unreachable("Unhandled store opcode");
5212+
}
5213+
5214+
// Scan the operands and if they're all BUILD_VECTORs, we'll have gathered
5215+
// their elements.
5216+
SmallVector<SDValue, 4> Operands(N->ops().take_front(Front));
5217+
for (SDValue BV : N->ops().drop_front(Front).drop_back(Back)) {
5218+
if (BV.getOpcode() != ISD::BUILD_VECTOR)
5219+
return SDValue();
5220+
5221+
// If the operand has multiple uses, this optimization can increase register
5222+
// pressure.
5223+
if (!BV.hasOneUse())
5224+
return SDValue();
5225+
5226+
// DAGCombiner visits nodes bottom-up. Check the BUILD_VECTOR operands for
5227+
// any signs they may be folded by some other pattern or rule.
5228+
for (SDValue Op : BV->ops()) {
5229+
// Peek through bitcasts
5230+
if (Op.getOpcode() == ISD::BITCAST)
5231+
Op = Op.getOperand(0);
5232+
5233+
// This may be folded into a PRMT.
5234+
if (Op.getValueType() == MVT::i16 && Op.getOpcode() == ISD::TRUNCATE &&
5235+
Op->getOperand(0).getValueType() == MVT::i32)
5236+
return SDValue();
5237+
5238+
// This may be folded into cvt.bf16x2
5239+
if (Op.getOpcode() == ISD::FP_ROUND)
5240+
return SDValue();
5241+
}
5242+
Operands.insert(Operands.end(), {BV.getOperand(0), BV.getOperand(1)});
5243+
}
5244+
for (SDValue Op : N->ops().take_back(Back))
5245+
Operands.push_back(Op);
5246+
5247+
// Now we replace the store
5248+
return DCI.DAG.getMemIntrinsicNode(Opcode, SDLoc(N), N->getVTList(),
5249+
Operands, MemVT, ST->getMemOperand());
5250+
}
5251+
5252+
static SDValue PerformStoreCombineHelper(SDNode *N,
5253+
TargetLowering::DAGCombinerInfo &DCI,
5254+
unsigned Front, unsigned Back) {
50525255
if (all_of(N->ops().drop_front(Front).drop_back(Back),
50535256
[](const SDUse &U) { return U.get()->isUndef(); }))
50545257
// Operand 0 is the previous value in the chain. Cannot return EntryToken
50555258
// as the previous value will become unused and eliminated later.
50565259
return N->getOperand(0);
50575260

5058-
return SDValue();
5261+
return combinePackingMovIntoStore(N, DCI, Front, Back);
50595262
}
50605263

5061-
static SDValue PerformStoreParamCombine(SDNode *N) {
5264+
static SDValue PerformSTORECombine(SDNode *N,
5265+
TargetLowering::DAGCombinerInfo &DCI) {
5266+
return combinePackingMovIntoStore(N, DCI, 1, 2);
5267+
}
5268+
5269+
static SDValue PerformStoreCombine(SDNode *N,
5270+
TargetLowering::DAGCombinerInfo &DCI) {
5271+
return combinePackingMovIntoStore(N, DCI, 1, 2);
5272+
}
5273+
5274+
static SDValue PerformStoreParamCombine(SDNode *N,
5275+
TargetLowering::DAGCombinerInfo &DCI) {
50625276
// Operands from the 3rd to the 2nd last one are the values to be stored.
50635277
// {Chain, ArgID, Offset, Val, Glue}
5064-
return PerformStoreCombineHelper(N, 3, 1);
5278+
return PerformStoreCombineHelper(N, DCI, 3, 1);
50655279
}
50665280

5067-
static SDValue PerformStoreRetvalCombine(SDNode *N) {
5281+
static SDValue PerformStoreRetvalCombine(SDNode *N,
5282+
TargetLowering::DAGCombinerInfo &DCI) {
50685283
// Operands from the 2nd to the last one are the values to be stored
5069-
return PerformStoreCombineHelper(N, 2, 0);
5284+
return PerformStoreCombineHelper(N, DCI, 2, 0);
50705285
}
50715286

50725287
/// PerformADDCombine - Target-specific dag combine xforms for ISD::ADD.
@@ -5697,14 +5912,24 @@ SDValue NVPTXTargetLowering::PerformDAGCombine(SDNode *N,
56975912
return PerformREMCombine(N, DCI, OptLevel);
56985913
case ISD::SETCC:
56995914
return PerformSETCCCombine(N, DCI, STI.getSmVersion());
5915+
case ISD::LOAD:
5916+
case NVPTXISD::LoadParamV2:
5917+
case NVPTXISD::LoadV2:
5918+
case NVPTXISD::LoadV4:
5919+
return combineUnpackingMovIntoLoad(N, DCI);
5920+
case ISD::STORE:
5921+
return PerformSTORECombine(N, DCI);
57005922
case NVPTXISD::StoreRetval:
57015923
case NVPTXISD::StoreRetvalV2:
57025924
case NVPTXISD::StoreRetvalV4:
5703-
return PerformStoreRetvalCombine(N);
5925+
return PerformStoreRetvalCombine(N, DCI);
57045926
case NVPTXISD::StoreParam:
57055927
case NVPTXISD::StoreParamV2:
57065928
case NVPTXISD::StoreParamV4:
5707-
return PerformStoreParamCombine(N);
5929+
return PerformStoreParamCombine(N, DCI);
5930+
case NVPTXISD::StoreV2:
5931+
case NVPTXISD::StoreV4:
5932+
return PerformStoreCombine(N, DCI);
57085933
case ISD::EXTRACT_VECTOR_ELT:
57095934
return PerformEXTRACTCombine(N, DCI);
57105935
case ISD::VSELECT:

0 commit comments

Comments
 (0)