PaddlePaddle
diff --git a/‎paddle/fluid/imperative/amp_auto_cast.h
Copy file name to clipboard
-1Lines changed: 0 additions & 1 deletion b/‎paddle/fluid/imperative/amp_auto_cast.h
Copy file name to clipboard
-1Lines changed: 0 additions & 1 deletion
diff --git a/‎paddle/fluid/pir/dialect/op_generator/api_gen.py
Copy file name to clipboard
+104-2Lines changed: 104 additions & 2 deletions b/‎paddle/fluid/pir/dialect/op_generator/api_gen.py
Copy file name to clipboard
+104-2Lines changed: 104 additions & 2 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/utils/amp_utils.cc
Copy file name to clipboard
+219Lines changed: 219 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/utils/amp_utils.cc
Copy file name to clipboard
+219Lines changed: 219 additions & 0 deletions
@@ -95,7 +95,6 @@ class AMPState {
   std::string GetAmpDtype() const;
   void SetAmpDtype(std::string amp_dtype);
   phi::DataType GetAmpPhiDtype() const;
-  // void Reset();
 
  private:
   static thread_local bool use_promote_;
 
@@ -32,6 +32,7 @@
 
 H_FILE_TEMPLATE = """
 
+// This file is generated by "paddle/fluid/pir/dialect/op_generator/api_gen.py"
 #pragma once
 
 #include <vector>
@@ -49,12 +50,16 @@
 
 CPP_FILE_TEMPLATE = """
 
+// This file is generated by "paddle/fluid/pir/dialect/op_generator/api_gen.py"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
 #include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 #include "paddle/pir/core/builder.h"
 #include "paddle/pir/core/builtin_op.h"
 #include "paddle/fluid/pir/dialect/operator/utils/utils.h"
+#include "paddle/fluid/imperative/amp_auto_cast.h"
+#include "paddle/fluid/pir/dialect/operator/utils/amp_utils.h"
+#include "paddle/fluid/eager/api/utils/global_utils.h"
 
 {body}
 
@@ -81,6 +86,8 @@
 """
 
 API_INNER_CODE_TEMPLATE = """
+    // AMP Logic
+    {amp_logic}
     {check_data_type}
     {handle_optional_inputs}
     {in_combine}
@@ -91,6 +98,27 @@
     {return_result}"""
 
 
+AMP_LOGIC_TEMPLATE = """
+    if (egr::Controller::Instance().GetCurrentAMPState()->GetAmpLevel() != paddle::imperative::AmpLevel::O0){{
+        VLOG(5) << "Check and Prepare For AMP";
+        auto op_name = phi::TransToFluidOpName("{op_name}");
+        std::vector<std::vector<pir::Value>> amp_values_vector = {{ {no_optional_inputs} }};
+        {optional_inputs}
+        auto amp_dst_dtype = paddle::dialect::GetAmpDestDtype("{op_name}", amp_values_vector);
+        {new_inputs}
+        {{
+            paddle::imperative::AutoCastGuard guard(egr::Controller::Instance().GetCurrentAMPState(), paddle::imperative::AmpLevel::O0);
+            return paddle::dialect::{op_name}({args});
+        }}
+    }}
+"""
+
+AMP_OPTIONAL_INPUTS_TEMPLATE = """if ({optional_input}) amp_values_vector.push_back({vec_optional_input});
+"""
+
+AMP_NEW_INPUTS_TEMPLATE = """auto new_{input} = paddle::dialect::PirAmpAutoCast("{input}", {input}, amp_dst_dtype, op_name);
+"""
+
 OP_DISPATCH_TEMPLATE = """
     if ({cond}) {{
         {inner_code}
@@ -569,6 +597,72 @@ def _gen_return_result(self, ret_list):
         elif len(ret_list) == 0:
             return 'return;'
 
+    def _gen_amp_no_optional_inputs(self, op_info):
+        name_list = op_info.input_name_list
+        type_list = op_info.input_type_list
+        no_optional_inputs = []
+        for name, type in zip(name_list, type_list):
+            if self._is_optional_input(op_info, name):
+                continue
+            if VECTOR_TYPE in type:
+                no_optional_inputs.append(name)
+            else:
+                no_optional_inputs.append('{' + name + '}')
+        return ', '.join(no_optional_inputs)
+
+    def _gen_amp_optional_inputs(self, op_info):
+        name_list = op_info.input_name_list
+        type_list = op_info.input_type_list
+        ret = ''
+        for name, type in zip(name_list, type_list):
+            if not self._is_optional_input(op_info, name):
+                continue
+            if VECTOR_TYPE in type:
+                ret += AMP_OPTIONAL_INPUTS_TEMPLATE.format(
+                    optional_input=name, vec_optional_input='*' + name
+                )
+            else:
+                ret += AMP_OPTIONAL_INPUTS_TEMPLATE.format(
+                    optional_input=name, vec_optional_input='{ *' + name + ' }'
+                )
+        return ret
+
+    def _gen_amp_new_inputs(self, op_info, op_name):
+        name_list = op_info.input_name_list
+        ret = ''
+        for name in name_list:
+            ret += AMP_NEW_INPUTS_TEMPLATE.format(input=name, op_name=op_name)
+        return ret
+
+    def _gen_amp_args(self, op_info, is_mutable_attr):
+        if is_mutable_attr:
+            input_list = op_info.input_name_list
+            attr_list = (
+                op_info.mutable_attribute_name_list
+                + op_info.non_mutable_attribute_name_list
+            )
+        else:
+            input_list = op_info.input_name_list
+            attr_list = op_info.attribute_name_list
+        args = ['new_' + input for input in input_list] + attr_list
+        return ', '.join(args)
+
+    def _gen_amp_logic(self, op_info, op_name, is_mutable_attr):
+        input_list = op_info.input_name_list
+        if not input_list:
+            return (
+                f'VLOG(7) << " No AMP for {op_name} because it has no input. ";'
+            )
+        if op_name.endswith(('_grad', '_grad_')):
+            return 'VLOG(7) << " No AMP for grad apis. ";'
+        return AMP_LOGIC_TEMPLATE.format(
+            op_name=op_name,
+            no_optional_inputs=self._gen_amp_no_optional_inputs(op_info),
+            optional_inputs=self._gen_amp_optional_inputs(op_info),
+            new_inputs=self._gen_amp_new_inputs(op_info, op_name),
+            args=self._gen_amp_args(op_info, is_mutable_attr),
+        )
+
     def _gen_check_data_type(self, op_info, op_name):
         mapping_input_name_to_type = dict(
             zip(op_info.input_name_list, op_info.input_type_list)
@@ -732,6 +826,9 @@ def _gen_one_impl(
                 )
 
                 if_inner_code = API_INNER_CODE_TEMPLATE.format(
+                    amp_logic=self._gen_amp_logic(
+                        op_info, op_name, is_mutable_attr
+                    ),
                     check_data_type=self._gen_check_data_type(
                         op_info, kernel_name
                     ),
@@ -790,6 +887,9 @@ def _gen_one_impl(
             if op_name.endswith('_') and not kernel_name.endswith('_'):
                 kernel_name = kernel_name + '_'
             api_inner_code = API_INNER_CODE_TEMPLATE.format(
+                amp_logic=self._gen_amp_logic(
+                    op_info, op_name, is_mutable_attr
+                ),
                 check_data_type=self._gen_check_data_type(op_info, kernel_name),
                 handle_optional_inputs=self._gen_handle_optional_inputs(
                     op_info
@@ -861,8 +961,10 @@ def gen_h_and_cpp_file(
         self._gen_h_file(op_info_items, namespaces, h_file_path)
         self._gen_cpp_file(op_info_items, namespaces, cpp_file_path)
         try:
-            subprocess.run(['clang-format', '-i', h_file_path])
-            subprocess.run(['clang-format', '-i', cpp_file_path])
+            subprocess.run(['clang-format', '-style=Google', '-i', h_file_path])
+            subprocess.run(
+                ['clang-format', '-style=Google', '-i', cpp_file_path]
+            )
         except Exception as e:
             pass
 
 
@@ -0,0 +1,219 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/pir/dialect/operator/utils/amp_utils.h"
+#include "paddle/fluid/eager/api/utils/global_utils.h"
+#include "paddle/fluid/imperative/amp_auto_cast.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
+#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
+
+namespace paddle {
+namespace dialect {
+
+phi::DataType GetPromoteType(
+    const std::string& op_name,
+    const std::vector<std::vector<pir::Value>>& amp_values_vector,
+    const phi::DataType& amp_dtype) {
+  auto dst_type = amp_dtype;
+  // only consider the dtype of input(X).
+  if (op_name == "batch_norm" || op_name == "layer_norm" ||
+      op_name == "sync_batch_norm" ||
+      op_name == "moving_average_abs_max_scale") {
+    if (GetValueDataType(amp_values_vector[0][0]) == phi::DataType::FLOAT32) {
+      dst_type = phi::DataType::FLOAT32;
+    }
+    return dst_type;
+  }
+
+  if (egr::Controller::Instance().GetCurrentAMPState()->GetAmpDtype() ==
+      "float16") {
+    if (op_name == "fused_attention") {
+      for (size_t i = 0; i < amp_values_vector.size(); i++) {
+        if (i != 3 || i != 4 || i != 9 || i != 10) {
+          if (GetValueDataType(amp_values_vector[i][0]) ==
+              phi::DataType::FLOAT32) {
+            dst_type = phi::DataType::FLOAT32;
+            return dst_type;
+          }
+        }
+      }
+    } else if (op_name == "fused_feedforward") {
+      for (size_t i = 0; i < amp_values_vector.size(); i++) {
+        if (i != 7 || i != 8 || i != 9 || i != 10) {
+          if (GetValueDataType(amp_values_vector[i][0]) ==
+              phi::DataType::FLOAT32) {
+            dst_type = phi::DataType::FLOAT32;
+            return dst_type;
+          }
+        }
+      }
+    }
+  }
+
+  for (const auto& values : amp_values_vector) {
+    for (const auto& value : values) {
+      if (GetValueDataType(value) == phi::DataType::FLOAT32) {
+        dst_type = GetValueDataType(value);
+        break;
+      }
+    }
+  }
+
+  return dst_type;
+}
+
+pir::Value Cast(const pir::Value& input, const phi::DataType& dst_dtype) {
+  paddle::imperative::AutoCastGuard guard(
+      egr::Controller::Instance().GetCurrentAMPState(),
+      paddle::imperative::AmpLevel::O0);
+  return paddle::dialect::cast(input, dst_dtype);
+}
+
+bool NeedCast(const pir::Value& value, const phi::DataType& dst_dtype) {
+  auto data_type = GetValueDataType(value);
+  if ((data_type == phi::DataType::FLOAT32 ||
+       data_type == phi::DataType::FLOAT16 ||
+       data_type == phi::DataType::BFLOAT16) &&
+      (data_type != dst_dtype)) {
+    return true;
+  }
+  return false;
+}
+
+pir::Value PirAmpAutoCast(const std::string& input_name,
+                          const pir::Value& input,
+                          const phi::DataType& dst_dtype,
+                          const std::string& op_name) {
+  VLOG(6) << "AMP AmpAutoCasts:"
+          << " input(" << input_name << " to dst_dtype("
+          << phi::DataTypeToString(dst_dtype) << ").";
+  if ((op_name == "batch_norm" || op_name == "layer_norm" ||
+       op_name == "sync_batch_norm" || op_name == "weight_only_linear") &&
+      input_name != "x") {
+    return input;
+  }
+
+  if (dst_dtype == phi::DataType::FLOAT16) {
+    if (op_name == "run_program") {
+      return input;
+    }
+    if ((op_name == "fused_attention" || op_name == "fused_feedforward")) {
+      if (input_name == "LnScale" || input_name == "LnBias" ||
+          input_name == "Ln2Scale" || input_name == "Ln2Bias" ||
+          input_name == "Ln1Scale" || input_name == "Ln1Bias") {
+        return input;
+      }
+    }
+  }
+  if (NeedCast(input, dst_dtype)) {
+    VLOG(6) << "Input : " << input_name << "NeedCast";
+    return Cast(input, dst_dtype);
+  }
+  return input;
+}
+
+paddle::optional<pir::Value> PirAmpAutoCast(
+    const std::string& input_name,
+    const paddle::optional<pir::Value>& input,
+    const phi::DataType& dst_dtype,
+    const std::string& op_name) {
+  if (input) {
+    return PirAmpAutoCast(input_name, *input, dst_dtype, op_name);
+  }
+  return paddle::none;
+}
+
+std::vector<pir::Value> PirAmpAutoCast(const std::string& inputs_name,
+                                       const std::vector<pir::Value>& inputs,
+                                       const phi::DataType& dst_dtype,
+                                       const std::string& op_name) {
+  VLOG(6) << "AMP AmpAutoCasts:"
+          << " inputs(" << inputs_name << ") dst_dtype("
+          << phi::DataTypeToString(dst_dtype) << ").";
+  std::vector<pir::Value> inputs_casted;
+  for (auto& input : inputs) {
+    if (NeedCast(input, dst_dtype)) {
+      inputs_casted.emplace_back(std::move(Cast(input, dst_dtype)));
+    } else {
+      inputs_casted.emplace_back(input);
+    }
+  }
+  return inputs_casted;
+}
+
+paddle::optional<std::vector<pir::Value>> PirAmpAutoCast(
+    const std::string& inputs_name,
+    const paddle::optional<std::vector<pir::Value>>& inputs,
+    const phi::DataType& dst_dtype,
+    const std::string& op_name) {
+  if (inputs) {
+    return PirAmpAutoCast(inputs_name, *inputs, dst_dtype, op_name);
+  }
+  return paddle::optional<std::vector<pir::Value>>();
+}
+
+phi::DataType GetAmpDestDtype(
+    const std::string& op_name,
+    const std::vector<std::vector<pir::Value>>& amp_values_vector) {
+  auto amp_level = egr::Controller::Instance().GetAMPLevel();
+  auto amp_setting_dtype =
+      egr::Controller::Instance().GetCurrentAMPState()->GetAmpPhiDtype();
+  auto dst_type = amp_setting_dtype;
+
+  bool use_promote = true;
+  if (amp_level == paddle::imperative::AmpLevel::O2) {
+    use_promote =
+        egr::Controller::Instance().GetCurrentAMPState()->GetUsePromote();
+  }
+
+  if (use_promote) {
+    if (paddle::imperative::AmpOperators::Instance()
+            .GetMutableAllowOps()
+            ->count(op_name)) {
+      dst_type = amp_setting_dtype;
+    } else if (paddle::imperative::AmpOperators::Instance()
+                   .GetMutableBlockOps()
+                   ->count(op_name)) {
+      dst_type = phi::DataType::FLOAT32;
+    } else {
+      if (amp_level == paddle::imperative::AmpLevel::OD) {
+        dst_type = phi::DataType::FLOAT32;
+      } else {
+        dst_type =
+            GetPromoteType(op_name, amp_values_vector, amp_setting_dtype);
+      }
+    }
+  } else {
+    // use_promote can be set to false only for O2 training.
+    if (paddle::imperative::AmpOperators::Instance()
+            .GetMutableBlockOps()
+            ->count(op_name)) {
+      dst_type = phi::DataType::FLOAT32;
+    }
+  }
+
+  if (dst_type == amp_setting_dtype &&
+      (paddle::imperative::AmpOperators::Instance()
+           .GetMutableUnsupportedOps(amp_setting_dtype)
+           ->count(op_name))) {
+    dst_type = phi::DataType::FLOAT32;
+  }
+
+  VLOG(6) << "AMP GetAmpDestDtype:"
+          << " op(" << op_name << ") amp_dtype(" << dst_type << ") amp_level("
+          << static_cast<int>(amp_level) << ").";
+  return dst_type;
+}
+}  // namespace dialect
+}  // namespace paddle