Skip to content

Commit faedc0b

Browse files
co63occo63oc
authored
optimize phi::DataType in paddle/fluid/imperative/ (PaddlePaddle#77532)
1 parent 377305b commit faedc0b

5 files changed

Lines changed: 59 additions & 64 deletions

File tree

‎paddle/fluid/imperative/amp_auto_cast.cc‎

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -175,14 +175,13 @@ AmpOperators::GetMutableBlockOps() {
175175
}
176176

177177
std::shared_ptr<std::unordered_set<std::string>>
178-
AmpOperators::GetMutableUnsupportedOps(const phi::DataType& data_type) {
178+
AmpOperators::GetMutableUnsupportedOps(const DataType& data_type) {
179179
PADDLE_ENFORCE_EQ(
180-
data_type == phi::DataType::FLOAT16 ||
181-
data_type == phi::DataType::BFLOAT16,
180+
data_type == DataType::FLOAT16 || data_type == DataType::BFLOAT16,
182181
true,
183182
common::errors::InvalidArgument(
184183
"The data_type mismatch. It should be FLOAT16 or BFLOAT16."));
185-
if (data_type == phi::DataType::FLOAT16) { // NOLINT
184+
if (data_type == DataType::FLOAT16) { // NOLINT
186185
return unsupported_fp16_ops_;
187186
} else {
188187
return unsupported_bf16_ops_;
@@ -230,7 +229,7 @@ thread_local bool AmpAttrs::use_promote_ = false;
230229

231230
thread_local AmpLevel AmpAttrs::amp_level_ = AmpLevel::O0;
232231

233-
thread_local phi::DataType AmpAttrs::amp_dtype_ = phi::DataType::FLOAT32;
232+
thread_local DataType AmpAttrs::amp_dtype_ = DataType::FLOAT32;
234233

235234
AmpAttrs::AmpAttrs() {}
236235

@@ -245,9 +244,9 @@ AmpLevel AmpAttrs::GetAmpLevel() const { return amp_level_; }
245244
void AmpAttrs::SetAmpLevel(AmpLevel level) { amp_level_ = level; }
246245

247246
std::string AmpAttrs::GetAmpDtype() const {
248-
if (amp_dtype_ == phi::DataType::FLOAT16) {
247+
if (amp_dtype_ == DataType::FLOAT16) {
249248
return std::string("float16");
250-
} else if (amp_dtype_ == phi::DataType::BFLOAT16) {
249+
} else if (amp_dtype_ == DataType::BFLOAT16) {
251250
return std::string("bfloat16");
252251
} else {
253252
return std::string("float32");
@@ -256,15 +255,15 @@ std::string AmpAttrs::GetAmpDtype() const {
256255

257256
void AmpAttrs::SetAmpDtype(std::string amp_dtype) {
258257
if (amp_dtype == "float16") {
259-
amp_dtype_ = phi::DataType::FLOAT16;
258+
amp_dtype_ = DataType::FLOAT16;
260259
} else if (amp_dtype == "bfloat16") {
261-
amp_dtype_ = phi::DataType::BFLOAT16;
260+
amp_dtype_ = DataType::BFLOAT16;
262261
} else {
263-
amp_dtype_ = phi::DataType::FLOAT32;
262+
amp_dtype_ = DataType::FLOAT32;
264263
}
265264
}
266265

267-
phi::DataType AmpAttrs::GetAmpPhiDtype() const { return amp_dtype_; }
266+
DataType AmpAttrs::GetAmpPhiDtype() const { return amp_dtype_; }
268267

269268
template <typename VarType>
270269
inline std::string GetDtypeStr(const std::shared_ptr<VarType>& var) {

‎paddle/fluid/imperative/amp_auto_cast.h‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ class AmpOperators {
5656
std::shared_ptr<std::unordered_set<std::string>> GetMutableBlockOps();
5757

5858
std::shared_ptr<std::unordered_set<std::string>> GetMutableUnsupportedOps(
59-
const phi::DataType& data_type);
59+
const DataType& data_type);
6060

6161
std::shared_ptr<std::unordered_set<std::string>>
6262
GetMutableUnsupportedFp16Ops();
@@ -94,12 +94,12 @@ class AmpAttrs {
9494
void SetAmpLevel(AmpLevel level);
9595
std::string GetAmpDtype() const;
9696
void SetAmpDtype(std::string amp_dtype);
97-
phi::DataType GetAmpPhiDtype() const;
97+
DataType GetAmpPhiDtype() const;
9898

9999
private:
100100
static thread_local bool use_promote_;
101101
static thread_local AmpLevel amp_level_;
102-
static thread_local phi::DataType amp_dtype_;
102+
static thread_local DataType amp_dtype_;
103103
};
104104

105105
// NOTE(zhiqiu): AutoCastGuard is used for RAII.

‎paddle/fluid/imperative/amp_utils.h‎

Lines changed: 40 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -29,27 +29,27 @@
2929

3030
namespace paddle {
3131
namespace imperative {
32-
static inline phi::DataType GetDataType(const pir::Value& value) {
32+
static inline DataType GetDataType(const pir::Value& value) {
3333
return paddle::dialect::GetValueDataType(value);
3434
}
3535

36-
static inline phi::DataType GetDataType(const paddle::Tensor& tensor) {
36+
static inline DataType GetDataType(const paddle::Tensor& tensor) {
3737
return tensor.dtype();
3838
}
3939

4040
template <class T>
41-
static inline phi::DataType GetPromoteType(
41+
static inline DataType GetPromoteType(
4242
const std::string& op_name,
4343
const paddle::small_vector<std::vector<T>, egr::kSlotSmallVectorSize>&
4444
amp_tensors_vector,
45-
const phi::DataType& amp_dtype) {
45+
const DataType& amp_dtype) {
4646
auto dst_type = amp_dtype;
4747
// only consider the dtype of input(X).
4848
if (op_name == "batch_norm" || op_name == "layer_norm" ||
4949
op_name == "sync_batch_norm" ||
5050
op_name == "moving_average_abs_max_scale") {
51-
if (GetDataType(amp_tensors_vector[0][0]) == phi::DataType::FLOAT32) {
52-
dst_type = phi::DataType::FLOAT32;
51+
if (GetDataType(amp_tensors_vector[0][0]) == DataType::FLOAT32) {
52+
dst_type = DataType::FLOAT32;
5353
}
5454
return dst_type;
5555
}
@@ -59,17 +59,17 @@ static inline phi::DataType GetPromoteType(
5959
if (op_name == "fused_attention") {
6060
for (size_t i = 0; i < amp_tensors_vector.size(); i++) {
6161
if (i < 3 || (i > 4 && i < 9) || i > 10) {
62-
if (GetDataType(amp_tensors_vector[i][0]) == phi::DataType::FLOAT32) {
63-
dst_type = phi::DataType::FLOAT32;
62+
if (GetDataType(amp_tensors_vector[i][0]) == DataType::FLOAT32) {
63+
dst_type = DataType::FLOAT32;
6464
return dst_type;
6565
}
6666
}
6767
}
6868
} else if (op_name == "fused_feedforward") {
6969
for (size_t i = 0; i < amp_tensors_vector.size(); i++) {
7070
if (i < 7 || i > 10) {
71-
if (GetDataType(amp_tensors_vector[i][0]) == phi::DataType::FLOAT32) {
72-
dst_type = phi::DataType::FLOAT32;
71+
if (GetDataType(amp_tensors_vector[i][0]) == DataType::FLOAT32) {
72+
dst_type = DataType::FLOAT32;
7373
return dst_type;
7474
}
7575
}
@@ -79,7 +79,7 @@ static inline phi::DataType GetPromoteType(
7979

8080
for (const auto& tensors : amp_tensors_vector) {
8181
for (const auto& tensor : tensors) {
82-
if (GetDataType(tensor) == phi::DataType::FLOAT32) {
82+
if (GetDataType(tensor) == DataType::FLOAT32) {
8383
dst_type = GetDataType(tensor);
8484
break;
8585
}
@@ -89,12 +89,12 @@ static inline phi::DataType GetPromoteType(
8989
return dst_type;
9090
}
9191

92-
static inline phi::DataType GetDtypeWithPlace(
92+
static inline DataType GetDtypeWithPlace(
9393
const std::string& op_name,
9494
const paddle::small_vector<std::vector<paddle::Tensor>,
9595
egr::kSlotSmallVectorSize>& amp_tensors_vector,
96-
const phi::DataType amp_dtype) {
97-
if (amp_dtype == phi::DataType::FLOAT32) {
96+
const DataType amp_dtype) {
97+
if (amp_dtype == DataType::FLOAT32) {
9898
return amp_dtype;
9999
}
100100
bool is_right_place = false;
@@ -119,22 +119,22 @@ static inline phi::DataType GetDtypeWithPlace(
119119
if (!is_right_place) {
120120
VLOG(6) << "Change " << op_name << "'s AMP type from " << amp_dtype
121121
<< " to FP32";
122-
return phi::DataType::FLOAT32;
122+
return DataType::FLOAT32;
123123
}
124124
return amp_dtype;
125125
}
126126

127-
static inline phi::DataType GetDtypeWithPlace(
127+
static inline DataType GetDtypeWithPlace(
128128
const std::string& op_name UNUSED,
129129
const paddle::small_vector<std::vector<pir::Value>,
130130
egr::kSlotSmallVectorSize>& amp_tensors_vector
131131
UNUSED,
132-
const phi::DataType amp_dtype) {
132+
const DataType amp_dtype) {
133133
return amp_dtype;
134134
}
135135

136136
template <class T>
137-
inline phi::DataType GetAmpDestDtype(
137+
inline DataType GetAmpDestDtype(
138138
const std::string& op_name,
139139
const paddle::small_vector<std::vector<T>, egr::kSlotSmallVectorSize>&
140140
amp_tensors_vector) {
@@ -157,10 +157,10 @@ inline phi::DataType GetAmpDestDtype(
157157
} else if (paddle::imperative::AmpOperators::Instance()
158158
.GetMutableBlockOps()
159159
->count(op_name)) {
160-
dst_type = phi::DataType::FLOAT32;
160+
dst_type = DataType::FLOAT32;
161161
} else {
162162
if (amp_level == paddle::imperative::AmpLevel::OD) {
163-
dst_type = phi::DataType::FLOAT32;
163+
dst_type = DataType::FLOAT32;
164164
} else {
165165
dst_type =
166166
GetPromoteType(op_name, amp_tensors_vector, amp_setting_dtype);
@@ -171,15 +171,15 @@ inline phi::DataType GetAmpDestDtype(
171171
if (paddle::imperative::AmpOperators::Instance()
172172
.GetMutableBlockOps()
173173
->count(op_name)) {
174-
dst_type = phi::DataType::FLOAT32;
174+
dst_type = DataType::FLOAT32;
175175
}
176176
}
177177

178178
if (dst_type == amp_setting_dtype &&
179179
(paddle::imperative::AmpOperators::Instance()
180180
.GetMutableUnsupportedOps(amp_setting_dtype)
181181
->count(op_name))) {
182-
dst_type = phi::DataType::FLOAT32;
182+
dst_type = DataType::FLOAT32;
183183
}
184184

185185
dst_type = GetDtypeWithPlace(op_name, amp_tensors_vector, dst_type);
@@ -190,7 +190,7 @@ inline phi::DataType GetAmpDestDtype(
190190
}
191191

192192
static inline bool NeedCast(const paddle::Tensor& tensor,
193-
const phi::DataType& dst_dtype) {
193+
const DataType& dst_dtype) {
194194
auto place = tensor.place();
195195
auto data_type = tensor.dtype();
196196
// Except CPU judgment, other conditions should be consistent with
@@ -201,9 +201,8 @@ static inline bool NeedCast(const paddle::Tensor& tensor,
201201
// CudaPinnedPlace is added for varbase created by dataloader
202202
// Cpu place is for different place tensor, when input1 is cpu and input2
203203
// is gpu
204-
if ((data_type == phi::DataType::FLOAT32 ||
205-
data_type == phi::DataType::FLOAT16 ||
206-
data_type == phi::DataType::BFLOAT16) &&
204+
if ((data_type == DataType::FLOAT32 || data_type == DataType::FLOAT16 ||
205+
data_type == DataType::BFLOAT16) &&
207206
(data_type != dst_dtype)) {
208207
return true;
209208
}
@@ -212,11 +211,10 @@ static inline bool NeedCast(const paddle::Tensor& tensor,
212211
}
213212

214213
static inline bool NeedCast(const pir::Value& value,
215-
const phi::DataType& dst_dtype) {
214+
const DataType& dst_dtype) {
216215
auto data_type = paddle::dialect::GetValueDataType(value);
217-
if ((data_type == phi::DataType::FLOAT32 ||
218-
data_type == phi::DataType::FLOAT16 ||
219-
data_type == phi::DataType::BFLOAT16) &&
216+
if ((data_type == DataType::FLOAT32 || data_type == DataType::FLOAT16 ||
217+
data_type == DataType::BFLOAT16) &&
220218
(data_type != dst_dtype)) {
221219
return true;
222220
}
@@ -225,14 +223,14 @@ static inline bool NeedCast(const pir::Value& value,
225223

226224
#if !(defined(PADDLE_NO_PYTHON) && defined(PADDLE_ON_INFERENCE))
227225
static inline paddle::Tensor Cast(const paddle::Tensor& input,
228-
const phi::DataType& dst_dtype,
226+
const DataType& dst_dtype,
229227
const bool trace_backward = true) {
230228
if (input.is_sparse_coo_tensor() || input.is_sparse_csr_tensor()) {
231229
if (trace_backward) {
232-
return sparse::cast_ad_func(input, phi::DataType::UNDEFINED, dst_dtype);
230+
return sparse::cast_ad_func(input, DataType::UNDEFINED, dst_dtype);
233231
} else {
234232
return paddle::experimental::sparse::cast(
235-
input, phi::DataType::UNDEFINED, dst_dtype);
233+
input, DataType::UNDEFINED, dst_dtype);
236234
}
237235
} else {
238236
if (trace_backward) {
@@ -245,7 +243,7 @@ static inline paddle::Tensor Cast(const paddle::Tensor& input,
245243
#endif
246244

247245
static inline pir::Value Cast(const pir::Value& input,
248-
const phi::DataType& dst_dtype,
246+
const DataType& dst_dtype,
249247
const bool trace_backward UNUSED = true) {
250248
paddle::imperative::AutoCastGuard guard(
251249
egr::Controller::Instance().GetCurrentAmpAttrs(),
@@ -256,12 +254,12 @@ static inline pir::Value Cast(const pir::Value& input,
256254
template <class T>
257255
inline std::vector<T> AmpAutoCasts(const std::string& inputs_name,
258256
const std::vector<T>& inputs,
259-
const phi::DataType& dst_dtype,
257+
const DataType& dst_dtype,
260258
std::string op_name UNUSED,
261259
bool trace_backward UNUSED = true) {
262260
VLOG(6) << "AMP AmpAutoCasts:"
263261
<< " inputs(" << inputs_name << ") dst_dtype("
264-
<< phi::DataTypeToString(dst_dtype) << ").";
262+
<< DataTypeToString(dst_dtype) << ").";
265263
std::vector<T> inputs_casted;
266264
for (auto& input : inputs) {
267265
if (NeedCast(input, dst_dtype)) {
@@ -276,13 +274,13 @@ inline std::vector<T> AmpAutoCasts(const std::string& inputs_name,
276274
template <class T>
277275
inline T AmpAutoCast(const std::string& input_name,
278276
const T& input,
279-
const phi::DataType& dst_dtype,
277+
const DataType& dst_dtype,
280278
const std::string& op_name,
281279
bool trace_backward = true) {
282280
VLOG(6) << "AMP AmpAutoCasts: op_name(" << op_name << ")input(" << input_name
283-
<< ") dst_dtype(" << phi::DataTypeToString(dst_dtype) << ").";
281+
<< ") dst_dtype(" << DataTypeToString(dst_dtype) << ").";
284282

285-
if (dst_dtype == phi::DataType::FLOAT16) {
283+
if (dst_dtype == DataType::FLOAT16) {
286284
if (op_name == "run_program") {
287285
return input;
288286
}
@@ -304,7 +302,7 @@ inline T AmpAutoCast(const std::string& input_name,
304302
input_name != "x") {
305303
return input;
306304
}
307-
} else if (dst_dtype == phi::DataType::BFLOAT16) {
305+
} else if (dst_dtype == DataType::BFLOAT16) {
308306
if ((op_name == "batch_norm" || op_name == "layer_norm" ||
309307
op_name == "sync_batch_norm" || op_name == "weight_only_linear") &&
310308
input_name != "x") {
@@ -321,7 +319,7 @@ inline T AmpAutoCast(const std::string& input_name,
321319
template <class T>
322320
inline paddle::optional<T> AmpAutoCast(const std::string& input_name,
323321
const paddle::optional<T>& input,
324-
const phi::DataType& dst_dtype,
322+
const DataType& dst_dtype,
325323
const std::string& op_name,
326324
bool trace_backward = true) {
327325
if (input) {
@@ -334,7 +332,7 @@ template <class T>
334332
inline paddle::optional<std::vector<T>> AmpAutoCasts(
335333
const std::string& inputs_name,
336334
const paddle::optional<std::vector<T>>& inputs,
337-
const phi::DataType& dst_dtype,
335+
const DataType& dst_dtype,
338336
std::string op_name,
339337
bool trace_backward = true) {
340338
if (inputs) {

‎paddle/fluid/imperative/tracer.cc‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -285,23 +285,21 @@ void Tracer::TraceOpImpl(const std::string& type,
285285

286286
std::unique_ptr<NameVarMap<VarType>> ins_amp = nullptr;
287287
if (GetCurrentAmpAttrs()->GetAmpLevel() == AmpLevel::O1) {
288-
if (GetCurrentAmpAttrs()->GetAmpPhiDtype() == phi::DataType::FLOAT16) {
288+
if (GetCurrentAmpAttrs()->GetAmpPhiDtype() == DataType::FLOAT16) {
289289
VLOG(5) << "Float16 Auto Mixed Precision O1 run operator: " << type;
290290
ins_amp = std::make_unique<NameVarMap<VarType>>(
291291
AutoCastInputs<VarType>(type, ins));
292-
} else if (GetCurrentAmpAttrs()->GetAmpPhiDtype() ==
293-
phi::DataType::BFLOAT16) {
292+
} else if (GetCurrentAmpAttrs()->GetAmpPhiDtype() == DataType::BFLOAT16) {
294293
VLOG(5) << "BFloat16 Auto Mixed Precision O1 run operator: " << type;
295294
ins_amp = std::make_unique<NameVarMap<VarType>>(
296295
AutoCastBF16Inputs<VarType>(type, ins));
297296
}
298297
} else if (GetCurrentAmpAttrs()->GetAmpLevel() == AmpLevel::O2) {
299-
if (GetCurrentAmpAttrs()->GetAmpPhiDtype() == phi::DataType::FLOAT16) {
298+
if (GetCurrentAmpAttrs()->GetAmpPhiDtype() == DataType::FLOAT16) {
300299
VLOG(5) << "Float16 Auto Mixed Precision O2 run operator: " << type;
301300
ins_amp = std::make_unique<NameVarMap<VarType>>(
302301
CastPureFp16Inputs<VarType>(type, ins));
303-
} else if (GetCurrentAmpAttrs()->GetAmpPhiDtype() ==
304-
phi::DataType::BFLOAT16) {
302+
} else if (GetCurrentAmpAttrs()->GetAmpPhiDtype() == DataType::BFLOAT16) {
305303
VLOG(5) << "BFloat16 Auto Mixed Precision O2 run operator: " << type;
306304
ins_amp = std::make_unique<NameVarMap<VarType>>(
307305
CastPureBf16Inputs<VarType>(type, ins));
@@ -606,7 +604,7 @@ std::string Tracer::GetAmpDtype() const {
606604
return g_current_amp_attrs->GetAmpDtype();
607605
}
608606

609-
phi::DataType Tracer::GetAmpPhiDtype() const {
607+
DataType Tracer::GetAmpPhiDtype() const {
610608
return g_current_amp_attrs->GetAmpPhiDtype();
611609
}
612610

‎paddle/fluid/imperative/tracer.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,7 @@ class Tracer {
159159

160160
PADDLE_API std::string GetAmpDtype() const;
161161

162-
PADDLE_API phi::DataType GetAmpPhiDtype() const;
162+
PADDLE_API DataType GetAmpPhiDtype() const;
163163

164164
TEST_API void DisableLayoutAutoTune();
165165

0 commit comments

Comments
 (0)