Skip to content

Commit 98b7c87

Browse files
Prince781Prince781
committed
[NVPTX] fold movs into loads and stores
1 parent d9f7979 commit 98b7c87

23 files changed

+2969
-2315
lines changed

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

Lines changed: 250 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 = DAG.getNode(LoadVT.isVector() ? ISD::EXTRACT_SUBVECTOR
3498+
: ISD::EXTRACT_VECTOR_ELT,
3499+
dl, LoadVT, P,
3500+
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,236 @@ 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; // non-glue, non-chain outputs
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+
NewVTs.append(LD->value_begin() + OldNumValues, LD->value_end());
5126+
5127+
// Create the new load
5128+
SDValue NewLoad =
5129+
DCI.DAG.getMemIntrinsicNode(Opcode, DL, DCI.DAG.getVTList(NewVTs),
5130+
Operands, MemVT, LD->getMemOperand());
5131+
5132+
// Now we use a combination of BUILD_VECTORs and a MERGE_VALUES node to keep
5133+
// the outputs the same. These nodes will be optimized away in later
5134+
// DAGCombiner iterations.
5135+
SmallVector<SDValue> Results;
5136+
unsigned I;
5137+
for (I = 0; I < OldNumValues * 2; I += 2) {
5138+
Results.push_back(DCI.DAG.getBuildVector(
5139+
ElemVT, DL, {NewLoad.getValue(I), NewLoad.getValue(I + 1)}));
5140+
}
5141+
// Add remaining chain and glue nodes
5142+
Results.push_back(NewLoad.getValue(I++));
5143+
if (I < NewLoad->getNumValues())
5144+
Results.push_back(NewLoad.getValue(I++));
5145+
5146+
return DCI.DAG.getMergeValues(Results, DL);
5147+
}
5148+
5149+
/// Fold a packing mov into a store. This may help lower register pressure.
5150+
///
5151+
/// ex:
5152+
/// v: v2f16 = build_vector a:f16, b:f16
5153+
/// StoreRetval v
5154+
///
5155+
/// ...is turned into...
5156+
///
5157+
/// StoreRetvalV2 a:f16, b:f16
5158+
static SDValue combinePackingMovIntoStore(SDNode *N,
5159+
TargetLowering::DAGCombinerInfo &DCI,
5160+
unsigned Front, unsigned Back) {
5161+
// We want to run this as late as possible since other optimizations may
5162+
// eliminate the BUILD_VECTORs.
5163+
if (!DCI.isAfterLegalizeDAG())
5164+
return SDValue();
5165+
5166+
// Get the type of the operands being stored.
5167+
EVT ElementVT = N->getOperand(Front).getValueType();
5168+
5169+
if (!Isv2x16VT(ElementVT))
5170+
return SDValue();
5171+
5172+
auto *ST = cast<MemSDNode>(N);
5173+
EVT MemVT = ElementVT.getVectorElementType();
5174+
5175+
// The new opcode after we double the number of operands.
5176+
NVPTXISD::NodeType Opcode;
5177+
switch (N->getOpcode()) {
5178+
case ISD::STORE:
5179+
// Any packed type is legal, so the legalizer will not have lowered
5180+
// ISD::STORE -> NVPTXISD::Store (unless it's under-aligned). We have to do
5181+
// it here.
5182+
MemVT = ST->getMemoryVT();
5183+
Opcode = NVPTXISD::StoreV2;
5184+
break;
5185+
case NVPTXISD::StoreParam:
5186+
Opcode = NVPTXISD::StoreParamV2;
5187+
break;
5188+
case NVPTXISD::StoreParamV2:
5189+
Opcode = NVPTXISD::StoreParamV4;
5190+
break;
5191+
case NVPTXISD::StoreRetval:
5192+
Opcode = NVPTXISD::StoreRetvalV2;
5193+
break;
5194+
case NVPTXISD::StoreRetvalV2:
5195+
Opcode = NVPTXISD::StoreRetvalV4;
5196+
break;
5197+
case NVPTXISD::StoreV2:
5198+
MemVT = ST->getMemoryVT();
5199+
Opcode = NVPTXISD::StoreV4;
5200+
break;
5201+
case NVPTXISD::StoreV4:
5202+
// PTX doesn't support v8 for 16-bit values
5203+
case NVPTXISD::StoreParamV4:
5204+
case NVPTXISD::StoreRetvalV4:
5205+
case NVPTXISD::StoreV8:
5206+
// PTX doesn't support the next doubling of operands for these opcodes.
5207+
return SDValue();
5208+
default:
5209+
llvm_unreachable("Unhandled store opcode");
5210+
}
5211+
5212+
// Scan the operands and if they're all BUILD_VECTORs, we'll have gathered
5213+
// their elements.
5214+
SmallVector<SDValue, 4> Operands(N->ops().take_front(Front));
5215+
for (SDValue BV : N->ops().drop_front(Front).drop_back(Back)) {
5216+
if (BV.getOpcode() != ISD::BUILD_VECTOR)
5217+
return SDValue();
5218+
5219+
// If the operand has multiple uses, this optimization can increase register
5220+
// pressure.
5221+
if (!BV.hasOneUse())
5222+
return SDValue();
5223+
5224+
// DAGCombiner visits nodes bottom-up. Check the BUILD_VECTOR operands for
5225+
// any signs they may be folded by some other pattern or rule.
5226+
for (SDValue Op : BV->ops()) {
5227+
// Peek through bitcasts
5228+
if (Op.getOpcode() == ISD::BITCAST)
5229+
Op = Op.getOperand(0);
5230+
5231+
// This may be folded into a PRMT.
5232+
if (Op.getValueType() == MVT::i16 && Op.getOpcode() == ISD::TRUNCATE &&
5233+
Op->getOperand(0).getValueType() == MVT::i32)
5234+
return SDValue();
5235+
5236+
// This may be folded into cvt.bf16x2
5237+
if (Op.getOpcode() == ISD::FP_ROUND)
5238+
return SDValue();
5239+
}
5240+
Operands.append({BV.getOperand(0), BV.getOperand(1)});
5241+
}
5242+
Operands.append(N->op_end() - Back, N->op_end());
5243+
5244+
// Now we replace the store
5245+
return DCI.DAG.getMemIntrinsicNode(Opcode, SDLoc(N), N->getVTList(), Operands,
5246+
MemVT, ST->getMemOperand());
5247+
}
5248+
5249+
static SDValue PerformStoreCombineHelper(SDNode *N,
5250+
TargetLowering::DAGCombinerInfo &DCI,
5251+
unsigned Front, unsigned Back) {
50525252
if (all_of(N->ops().drop_front(Front).drop_back(Back),
50535253
[](const SDUse &U) { return U.get()->isUndef(); }))
50545254
// Operand 0 is the previous value in the chain. Cannot return EntryToken
50555255
// as the previous value will become unused and eliminated later.
50565256
return N->getOperand(0);
50575257

5058-
return SDValue();
5258+
return combinePackingMovIntoStore(N, DCI, Front, Back);
5259+
}
5260+
5261+
static SDValue PerformStoreCombine(SDNode *N,
5262+
TargetLowering::DAGCombinerInfo &DCI) {
5263+
return combinePackingMovIntoStore(N, DCI, 1, 2);
50595264
}
50605265

5061-
static SDValue PerformStoreParamCombine(SDNode *N) {
5266+
static SDValue PerformStoreParamCombine(SDNode *N,
5267+
TargetLowering::DAGCombinerInfo &DCI) {
50625268
// Operands from the 3rd to the 2nd last one are the values to be stored.
50635269
// {Chain, ArgID, Offset, Val, Glue}
5064-
return PerformStoreCombineHelper(N, 3, 1);
5270+
return PerformStoreCombineHelper(N, DCI, 3, 1);
50655271
}
50665272

5067-
static SDValue PerformStoreRetvalCombine(SDNode *N) {
5273+
static SDValue PerformStoreRetvalCombine(SDNode *N,
5274+
TargetLowering::DAGCombinerInfo &DCI) {
50685275
// Operands from the 2nd to the last one are the values to be stored
5069-
return PerformStoreCombineHelper(N, 2, 0);
5276+
return PerformStoreCombineHelper(N, DCI, 2, 0);
50705277
}
50715278

50725279
/// PerformADDCombine - Target-specific dag combine xforms for ISD::ADD.
@@ -5697,14 +5904,23 @@ SDValue NVPTXTargetLowering::PerformDAGCombine(SDNode *N,
56975904
return PerformREMCombine(N, DCI, OptLevel);
56985905
case ISD::SETCC:
56995906
return PerformSETCCCombine(N, DCI, STI.getSmVersion());
5907+
case ISD::LOAD:
5908+
case NVPTXISD::LoadParamV2:
5909+
case NVPTXISD::LoadV2:
5910+
case NVPTXISD::LoadV4:
5911+
return combineUnpackingMovIntoLoad(N, DCI);
57005912
case NVPTXISD::StoreRetval:
57015913
case NVPTXISD::StoreRetvalV2:
57025914
case NVPTXISD::StoreRetvalV4:
5703-
return PerformStoreRetvalCombine(N);
5915+
return PerformStoreRetvalCombine(N, DCI);
57045916
case NVPTXISD::StoreParam:
57055917
case NVPTXISD::StoreParamV2:
57065918
case NVPTXISD::StoreParamV4:
5707-
return PerformStoreParamCombine(N);
5919+
return PerformStoreParamCombine(N, DCI);
5920+
case ISD::STORE:
5921+
case NVPTXISD::StoreV2:
5922+
case NVPTXISD::StoreV4:
5923+
return PerformStoreCombine(N, DCI);
57085924
case ISD::EXTRACT_VECTOR_ELT:
57095925
return PerformEXTRACTCombine(N, DCI);
57105926
case ISD::VSELECT:

0 commit comments

Comments
 (0)