@@ -85,6 +85,12 @@ static cl::opt<unsigned> FMAContractLevelOpt(
85
85
" 1: do it 2: do it aggressively" ),
86
86
cl::init(2 ));
87
87
88
+ static cl::opt<bool > DisableFOpTreeReduce (
89
+ " nvptx-disable-fop-tree-reduce" , cl::Hidden,
90
+ cl::desc (" NVPTX Specific: don't emit tree reduction for floating-point "
91
+ " reduction operations" ),
92
+ cl::init(false ));
93
+
88
94
static cl::opt<int > UsePrecDivF32 (
89
95
" nvptx-prec-divf32" , cl::Hidden,
90
96
cl::desc (" NVPTX Specifies: 0 use div.approx, 1 use div.full, 2 use"
@@ -834,6 +840,15 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
834
840
if (STI.allowFP16Math () || STI.hasBF16Math ())
835
841
setTargetDAGCombine (ISD::SETCC);
836
842
843
+ // Vector reduction operations. These are transformed into a tree evaluation
844
+ // of nodes which may or may not be legal.
845
+ for (MVT VT : MVT::fixedlen_vector_valuetypes ()) {
846
+ setOperationAction ({ISD::VECREDUCE_FADD, ISD::VECREDUCE_FMUL,
847
+ ISD::VECREDUCE_FMAX, ISD::VECREDUCE_FMIN,
848
+ ISD::VECREDUCE_FMAXIMUM, ISD::VECREDUCE_FMINIMUM},
849
+ VT, Custom);
850
+ }
851
+
837
852
// Promote fp16 arithmetic if fp16 hardware isn't available or the
838
853
// user passed --nvptx-no-fp16-math. The flag is useful because,
839
854
// although sm_53+ GPUs have some sort of FP16 support in
@@ -1087,6 +1102,10 @@ const char *NVPTXTargetLowering::getTargetNodeName(unsigned Opcode) const {
1087
1102
MAKE_CASE (NVPTXISD::BFI)
1088
1103
MAKE_CASE (NVPTXISD::PRMT)
1089
1104
MAKE_CASE (NVPTXISD::FCOPYSIGN)
1105
+ MAKE_CASE (NVPTXISD::FMAXNUM3)
1106
+ MAKE_CASE (NVPTXISD::FMINNUM3)
1107
+ MAKE_CASE (NVPTXISD::FMAXIMUM3)
1108
+ MAKE_CASE (NVPTXISD::FMINIMUM3)
1090
1109
MAKE_CASE (NVPTXISD::DYNAMIC_STACKALLOC)
1091
1110
MAKE_CASE (NVPTXISD::STACKRESTORE)
1092
1111
MAKE_CASE (NVPTXISD::STACKSAVE)
@@ -2147,6 +2166,108 @@ NVPTXTargetLowering::LowerCONCAT_VECTORS(SDValue Op, SelectionDAG &DAG) const {
2147
2166
return DAG.getBuildVector (Node->getValueType (0 ), dl, Ops);
2148
2167
}
2149
2168
2169
+ // / A generic routine for constructing a tree reduction for a vector operand.
2170
+ // / This method differs from iterative splitting in DAGTypeLegalizer by
2171
+ // / first scalarizing the vector and then progressively grouping elements
2172
+ // / bottom-up. This allows easily building the optimal (minimum) number of nodes
2173
+ // / with different numbers of operands (eg. max3 vs max2).
2174
+ static SDValue BuildTreeReduction (
2175
+ const SDValue &VectorOp,
2176
+ ArrayRef<std::pair<unsigned /* NodeType*/ , unsigned /* NumInputs*/ >> Ops,
2177
+ const SDLoc &DL, const SDNodeFlags Flags, SelectionDAG &DAG) {
2178
+ EVT VectorTy = VectorOp.getValueType ();
2179
+ EVT EltTy = VectorTy.getVectorElementType ();
2180
+ const unsigned NumElts = VectorTy.getVectorNumElements ();
2181
+
2182
+ // scalarize vector
2183
+ SmallVector<SDValue> Elements (NumElts);
2184
+ for (unsigned I = 0 , E = NumElts; I != E; ++I) {
2185
+ Elements[I] = DAG.getNode (ISD::EXTRACT_VECTOR_ELT, DL, EltTy, VectorOp,
2186
+ DAG.getConstant (I, DL, MVT::i64 ));
2187
+ }
2188
+
2189
+ // now build the computation graph in place at each level
2190
+ SmallVector<SDValue> Level = Elements;
2191
+ for (unsigned OpIdx = 0 ; Level.size () > 1 && OpIdx < Ops.size ();) {
2192
+ const auto [DefaultScalarOp, DefaultGroupSize] = Ops[OpIdx];
2193
+
2194
+ // partially reduce all elements in level
2195
+ SmallVector<SDValue> ReducedLevel;
2196
+ unsigned I = 0 , E = Level.size ();
2197
+ for (; I + DefaultGroupSize <= E; I += DefaultGroupSize) {
2198
+ // Reduce elements in groups of [DefaultGroupSize], as much as possible.
2199
+ ReducedLevel.push_back (DAG.getNode (
2200
+ DefaultScalarOp, DL, EltTy,
2201
+ ArrayRef<SDValue>(Level).slice (I, DefaultGroupSize), Flags));
2202
+ }
2203
+
2204
+ if (I < E) {
2205
+ if (ReducedLevel.empty ()) {
2206
+ // The current operator requires more inputs than there are operands at
2207
+ // this level. Pick a smaller operator and retry.
2208
+ ++OpIdx;
2209
+ assert (OpIdx < Ops.size () && " no smaller operators for reduction" );
2210
+ continue ;
2211
+ }
2212
+
2213
+ // Otherwise, we just have a remainder, which we push to the next level.
2214
+ for (; I < E; ++I)
2215
+ ReducedLevel.push_back (Level[I]);
2216
+ }
2217
+ Level = ReducedLevel;
2218
+ }
2219
+
2220
+ return *Level.begin ();
2221
+ }
2222
+
2223
+ // / Lower fadd/fmul vector reductions. Builds a computation graph (tree) and
2224
+ // / serializes it.
2225
+ SDValue NVPTXTargetLowering::LowerVECREDUCE (SDValue Op,
2226
+ SelectionDAG &DAG) const {
2227
+ // If we can't reorder sub-operations, let DAGTypeLegalizer lower this op.
2228
+ if (DisableFOpTreeReduce || !Op->getFlags ().hasAllowReassociation ())
2229
+ return SDValue ();
2230
+
2231
+ EVT EltTy = Op.getOperand (0 ).getValueType ().getVectorElementType ();
2232
+ const bool CanUseMinMax3 = EltTy == MVT::f32 && STI.getSmVersion () >= 100 &&
2233
+ STI.getPTXVersion () >= 88 ;
2234
+ SDLoc DL (Op);
2235
+ SmallVector<std::pair<unsigned /* Op*/ , unsigned /* NumIn*/ >, 2 > Operators;
2236
+ switch (Op->getOpcode ()) {
2237
+ case ISD::VECREDUCE_FADD:
2238
+ Operators = {{ISD::FADD, 2 }};
2239
+ break ;
2240
+ case ISD::VECREDUCE_FMUL:
2241
+ Operators = {{ISD::FMUL, 2 }};
2242
+ break ;
2243
+ case ISD::VECREDUCE_FMAX:
2244
+ if (CanUseMinMax3)
2245
+ Operators.push_back ({NVPTXISD::FMAXNUM3, 3 });
2246
+ Operators.push_back ({ISD::FMAXNUM, 2 });
2247
+ break ;
2248
+ case ISD::VECREDUCE_FMIN:
2249
+ if (CanUseMinMax3)
2250
+ Operators.push_back ({NVPTXISD::FMINNUM3, 3 });
2251
+ Operators.push_back ({ISD::FMINNUM, 2 });
2252
+ break ;
2253
+ case ISD::VECREDUCE_FMAXIMUM:
2254
+ if (CanUseMinMax3)
2255
+ Operators.push_back ({NVPTXISD::FMAXIMUM3, 3 });
2256
+ Operators.push_back ({ISD::FMAXIMUM, 2 });
2257
+ break ;
2258
+ case ISD::VECREDUCE_FMINIMUM:
2259
+ if (CanUseMinMax3)
2260
+ Operators.push_back ({NVPTXISD::FMINIMUM3, 3 });
2261
+ Operators.push_back ({ISD::FMINIMUM, 2 });
2262
+ break ;
2263
+ default :
2264
+ llvm_unreachable (" unhandled vecreduce operation" );
2265
+ }
2266
+
2267
+ return BuildTreeReduction (Op.getOperand (0 ), Operators, DL, Op->getFlags (),
2268
+ DAG);
2269
+ }
2270
+
2150
2271
SDValue NVPTXTargetLowering::LowerBITCAST (SDValue Op, SelectionDAG &DAG) const {
2151
2272
// Handle bitcasting from v2i8 without hitting the default promotion
2152
2273
// strategy which goes through stack memory.
@@ -2935,6 +3056,13 @@ NVPTXTargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
2935
3056
return LowerVECTOR_SHUFFLE (Op, DAG);
2936
3057
case ISD::CONCAT_VECTORS:
2937
3058
return LowerCONCAT_VECTORS (Op, DAG);
3059
+ case ISD::VECREDUCE_FADD:
3060
+ case ISD::VECREDUCE_FMUL:
3061
+ case ISD::VECREDUCE_FMAX:
3062
+ case ISD::VECREDUCE_FMIN:
3063
+ case ISD::VECREDUCE_FMAXIMUM:
3064
+ case ISD::VECREDUCE_FMINIMUM:
3065
+ return LowerVECREDUCE (Op, DAG);
2938
3066
case ISD::STORE:
2939
3067
return LowerSTORE (Op, DAG);
2940
3068
case ISD::LOAD:
0 commit comments