text: # Copyright (c) 2022 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
  # 
  #     http://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.
  
  import logging
  import yaml
  import re
  import argparse
  import os
  
  ########################
  ### Global Variables ###
  ########################
  ops_to_fill_zero_for_empty_grads = set([
      "split_grad", "rnn_grad", "matmul_double_grad", "matmul_triple_grad",
      "sigmoid_double_grad", "sigmoid_triple_grad", "add_double_grad",
      "add_triple_grad", "multiply_grad", "multiply_double_grad",
      "multiply_triple_grad", "conv2d_grad_grad", "batch_norm_double_grad",
      "tanh_double_grad", "tanh_triple_grad", "subtract_double_grad",
      "divide_double_grad", "log_double_grad", "elu_double_grad",
      "leaky_relu_double_grad", "sqrt_double_grad", "rsqrt_double_grad",
      "square_double_grad", "celu_double_grad", "pad_double_grad",
      "pad3d_double_grad", "squeeze_double_grad", "unsqueeze_double_grad",
      "conv3d_double_grad", "depthwise_conv2d_grad_grad"
  ])
  
  # For API dispatch used at python-level
  # { op_name : [arg_name, ...] }
  core_ops_returns_info = {}
  core_ops_args_info = {}
  core_ops_args_type_info = {}
  
  yaml_types_mapping = {
      'int' : 'int', 'int32_t' : 'int32_t', 'int64_t' : 'int64_t',  'size_t' : 'size_t', \
      'float' : 'float', 'double' : 'double', 'bool' : 'bool', \
      'str' : 'std::string', \
      'Place' : 'paddle::Place', 'DataLayout' : 'paddle::experimental::DataLayout', 'DataType' : 'paddle::experimental::DataType', \
      'int64_t[]' : 'std::vector<int64_t>', 'int[]' : 'std::vector<int>',
      'Tensor' : 'Tensor',
      'Tensor[]' : 'std::vector<Tensor>',
      'Tensor[Tensor[]]' : 'std::vector<std::vector<Tensor>>',
      'Scalar' : 'paddle::experimental::Scalar',
      'Scalar(int)' : 'paddle::experimental::Scalar',
      'Scalar(int64_t)' : 'paddle::experimental::Scalar',
      'Scalar(float)' : 'paddle::experimental::Scalar',
      'Scalar(double)' : 'paddle::experimental::Scalar',
      'IntArray' : 'paddle::experimental::IntArray'
  }
  
  
  #############################
  ###  File Reader Helpers  ###
  #############################
  def AssertMessage(lhs_str, rhs_str):
      return f"lhs: {lhs_str}, rhs: {rhs_str}"
  
  
  def ReadFwdFile(filepath):
      f = open(filepath, 'r')
      # empty file loaded by yaml is None
      contents = yaml.load(f, Loader=yaml.FullLoader)
      f.close()
      return contents if contents is not None else []
  
  
  def ReadBwdFile(filepath):
      f = open(filepath, 'r')
      contents = yaml.load(f, Loader=yaml.FullLoader)
      ret = {}
      if contents is not None:
          for content in contents:
              assert 'backward_api' in content.keys(), AssertMessage(
                  'backward_api', content.keys())
              if 'backward_api' in content.keys():
                  api_name = content['backward_api']
  
              ret[api_name] = content
      f.close()
      return ret
  
  
  ##################################
  ###  Generic Helper Functions  ###
  ##################################
  def FindGradName(string):
      return string + "_grad"
  
  
  def FindForwardName(string):
      if not string.endswith("_grad"):
          return None
      return string[:-5]
  
  
  def IsGradName(string):
      return string.endswith("_grad")
  
  
  def IsPlainTensorType(string):
      plain_tensor_types = ['Tensor&', 'Tensor', 'const Tensor&', 'const Tensor']
      if string in plain_tensor_types:
          return True
      return False
  
  
  def IsVectorTensorType(string):
      vector_tensor_types = [
          'std::vector<std::vector<Tensor>>', 'std::vector<Tensor>'
      ]
      if string in vector_tensor_types:
          return True
      return False
  
  
  def GetSavedName(string):
      return string + "_"
  
  
  def GetConstReference(string):
      ret = string
      if not string.startswith("const "):
          ret = "const " + string
      if not string.endswith("&"):
          ret += "&"
      return ret
  
  
  def RemoveConstAndReference(string):
      ret = string
      if string.startswith("const "):
          ret = ret[6:]
      if string.endswith("&"):
          ret = ret[:-1]
  
      return ret
  
  
  def GetGradNodeName(string):
      return f"GradNode{string}Final"
  
  
  def GetDygraphForwardFunctionName(string):
      return f"{string}_final_state_dygraph_function"
  
  
  def GetIntermediateAPIFunctionName(string):
      return string + "_intermediate"
  
  
  def GetAutoGradMetaName(string):
      return f"{string}_autograd_meta"
  
  
  def GetAutoGradMetaVectorName(string):
      return f"{string}_autograd_meta_vec"
  
  
  def RemoveSpecialSymbolsInName(string):
      # Remove any name after '@'
      ret = string.split("@")[0]
      return ret
  
  
  def RecoverBaseNameOfInplaceFunction(function_name):
      return function_name[:-1]
  
  
  def GetInplacedFunctionName(function_name):
      inplace_func_name = function_name
      if inplace_func_name[-1] != '_':
          inplace_func_name += '_'
      return inplace_func_name
  
  
  def GetForwardFunctionName(string):
      return f"{string}_final_state_dygraph_function"
  
  
  def GetIndent(num):
      tab = "  "
      return "".join([tab for i in range(num)])
  
  
  ######################
  ###  Yaml Parsers  ###
  ######################
  def ParseYamlArgs(string):
      # Example: const Tensor& x, const Tensor& y, bool transpose_x, bool transpose_y
  
      # inputs_list = [ [arg_name, arg_type, orig_position], ...]
      inputs_list = []
      # attrs_list = [ [arg_name, arg_type, default_value, orig_position], ...]
      attrs_list = []
  
      args = [x.strip() for x in string.strip().split(",")]
      atype = r'((const )?\S+) '
      aname = r'(.*)'
      pattern = f'{atype}{aname}'
      for i in range(len(args)):
          arg = args[i]
          m = re.search(pattern, arg)
          arg_type = m.group(1).strip()
          arg_name = m.group(3).split("=")[0].strip()
          default_value = m.group(3).split("=")[1].strip() if len(
              m.group(3).split("=")) > 1 else None
  
          assert arg_type in yaml_types_mapping.keys(
          ), f"The argument type {arg_type} in yaml config is not supported in yaml_types_mapping."
          if arg_type in ["DataType", "DataLayout"] and default_value is not None:
              default_value = f"paddle::experimental::{default_value}"
          arg_type = yaml_types_mapping[arg_type]
  
          arg_name = RemoveSpecialSymbolsInName(arg_name)
          if "Tensor" in arg_type:
              assert default_value is None
              inputs_list.append([arg_name, arg_type, i])
          else:
              attrs_list.append([arg_name, arg_type, default_value, i])
  
      return inputs_list, attrs_list
  
  
  def ParseYamlReturns(string):
      # Example0: Tensor(out), Tensor(out1)
      # Example1: Tensor, Tensor
      # Example2: Tensor[](out), Tensor
  
      # list = [ [ret_name, ret_type, orig_position], ...]
      returns_list = []
  
      returns = [x.strip() for x in string.strip().split(",")]
  
      for i in range(len(returns)):
          ret = returns[i].split("{")[0].strip()
  
          ret_name = ""
          if "(" in ret and ")" in ret:
              # Remove trailing ')'
              ret = ret[:-1]
              ret_type = ret.split("(")[0].strip()
              ret_name = ret.split("(")[1].strip()
          else:
              ret_type = ret.strip()
  
          assert ret_type in yaml_types_mapping.keys(
          ), f"The return type {ret_type} in yaml config is not supported in yaml_types_mapping."
          ret_type = yaml_types_mapping[ret_type]
  
          assert "Tensor" in ret_type, AssertMessage("Tensor", ret_type)
          ret_name = RemoveSpecialSymbolsInName(ret_name)
          returns_list.append([ret_name, ret_type, i])
  
      return returns_list
  
  
  def ParseYamlForwardFromBackward(string):
      # Example: matmul (const Tensor& x, const Tensor& y, bool transpose_x, bool transpose_y) -> Tensor(out)
  
      fname = r'(.*?)'
      wspace = r'\s*'
      fargs = r'(.*?)'
      frets = r'(.*)'
      pattern = f'{fname}{wspace}\({wspace}{fargs}{wspace}\){wspace}->{wspace}{frets}'
  
      m = re.search(pattern, string)
      function_name = m.group(1)
      function_args = m.group(2)
      function_returns = m.group(3)
  
      forward_inputs_list, forward_attrs_list = ParseYamlArgs(function_args)
      forward_returns_list = ParseYamlReturns(function_returns)
  
      return forward_inputs_list, forward_attrs_list, forward_returns_list
  
  
  def ParseYamlForward(args_str, returns_str):
      # args Example: (const Tensor& x, const Tensor& y, bool transpose_x = false, bool transpose_y = false)
      # returns Example: Tensor, Tensor
  
      fargs = r'(.*?)'
      wspace = r'\s*'
      args_pattern = f'^\({fargs}\)$'
      args_str = re.search(args_pattern, args_str.strip()).group(1)
  
      inputs_list, attrs_list = ParseYamlArgs(args_str)
      returns_list = ParseYamlReturns(returns_str)
  
      return inputs_list, attrs_list, returns_list
  
  
  def ParseYamlBackward(args_str, returns_str):
      # args Example: (const Tensor& x, const Tensor& y, const Tensor& out_grad, bool transpose_x=false, bool transpose_y=false)
      # returns Example: Tensor(x_grad), Tensor(y_grad)
  
      fargs = r'(.*?)'
      wspace = r'\s*'
      args_pattern = f'\({fargs}\)'
      args_str = re.search(args_pattern, args_str).group(1)
  
      inputs_list, attrs_list = ParseYamlArgs(args_str)
      returns_list = ParseYamlReturns(returns_str)
  
      return inputs_list, attrs_list, returns_list
  
  
  def ParseYamlInplaceInfo(string):
      # inplace_map_str: "(x -> out0), (y -> out2)"
      inplace_map = {}
      for pair in string.split(","):
          pair = pair.strip()
          if pair.startswith("("):
              pair = pair[1:]
  
          if pair.endswith(")"):
              pair = pair[:-1]
  
          key = pair.split("->")[0].strip()
          val = pair.split("->")[1].strip()
          inplace_map[key] = val
      return inplace_map
  
  
  ########################
  ###  Generator Base  ###
  ########################
  class FunctionGeneratorBase:
      def __init__(self, forward_api_contents, namespace):
          self.forward_api_contents = forward_api_contents
          self.namespace = namespace
  
          self.forward_api_name = ""
  
          self.orig_forward_inputs_list = [
          ]  #[ [arg_name, arg_type, orig_position], ...]
          self.orig_forward_attrs_list = [
          ]  #[ [attr_name, attr_type, default_value, orig_position], ...]
          self.orig_forward_returns_list = [
          ]  #[ [ret_name, ret_type, orig_position], ...]
  
          # Processed Forward Data
          self.forward_inputs_position_map = {
          }  #{ "name" : [type, fwd_position] }
          self.forward_outputs_position_map = {
          }  #{ "name" : [type, fwd_position] }
  
          # Special Op Attributes
          self.optional_inputs = []  #[name, ...]
          self.no_need_buffers = []  #[name, ...]
          self.intermediate_outputs = []  #[name, ...]    
          self.forward_inplace_map = {}  #{name : name, ...}
  
      def ParseForwardInplaceInfo(self):
          forward_api_contents = self.forward_api_contents
          if 'inplace' not in forward_api_contents.keys(): return
  
          inplace_map_str = forward_api_contents['inplace']
          self.forward_inplace_map = ParseYamlInplaceInfo(inplace_map_str)
  
      def ParseNoNeedBuffer(self):
          grad_api_contents = self.grad_api_contents
  
          if 'no_need_buffer' in grad_api_contents.keys():
              no_need_buffer_str = grad_api_contents['no_need_buffer']
              for name in no_need_buffer_str.split(","):
                  name = name.strip()
                  name = RemoveSpecialSymbolsInName(name)
                  self.no_need_buffers.append(name.strip())
  
      def ParseDispensable(self):
          forward_api_contents = self.forward_api_contents
  
          if 'optional' in forward_api_contents.keys():
              optional_inputs_str = forward_api_contents['optional']
              for name in optional_inputs_str.split(","):
                  name = name.strip()
                  name = RemoveSpecialSymbolsInName(name)
                  self.optional_inputs.append(name)
  
      def ParseIntermediate(self):
          forward_api_contents = self.forward_api_contents
  
          if 'intermediate' in forward_api_contents.keys():
              intermediate_str = forward_api_contents['intermediate']
              for name in intermediate_str.split(","):
                  name = name.strip()
                  name = RemoveSpecialSymbolsInName(name)
                  self.intermediate_outputs.append(name)
  
      def CollectOriginalForwardInfo(self):
          forward_api_contents = self.forward_api_contents
  
          self.forward_api_name = forward_api_contents['api']
          forward_args_str = forward_api_contents['args']
          forward_returns_str = forward_api_contents['output']
  
          assert 'api' in forward_api_contents.keys(
          ), "Unable to find \"api\" in forward_api_contents keys"
          assert 'args' in forward_api_contents.keys(
          ), "Unable to find \"args\" in forward_api_contents keys"
          assert 'output' in forward_api_contents.keys(
          ), "Unable to find \"output\" in forward_api_contents keys"
  
          # Collect Original Forward Inputs/Outputs and then perform validation checks
          self.orig_forward_inputs_list, self.orig_forward_attrs_list, self.orig_forward_returns_list = ParseYamlForward(
              forward_args_str, forward_returns_str)
  
      def DetermineForwardPositionMap(self, forward_inputs_list,
                                      forward_returns_list):
          for i in range(len(forward_inputs_list)):
              forward_input = forward_inputs_list[i]
              input_name = forward_input[0]
              input_type = forward_input[1]
              input_pos = forward_input[2]
  
              self.forward_inputs_position_map[
                  input_name] = [input_type, input_pos]
  
          for i in range(len(forward_returns_list)):
              forward_return = forward_returns_list[i]
              return_name = forward_return[0]
              return_type = forward_return[1]
              return_pos = forward_return[2]
  
              self.forward_outputs_position_map[
                  return_name] = [return_type, return_pos]
  
  
  class GeneratorBase:
      def __init__(self, api_yaml_path):
          self.namespace = ""
          self.api_yaml_path = api_yaml_path
  
          self.forward_api_list = []
  
      def ParseForwardYamlContents(self):
          api_yaml_path = self.api_yaml_path
          self.forward_api_list = ReadFwdFile(api_yaml_path)
  
      def InferNameSpace(self):
          api_yaml_path = self.api_yaml_path
          if re.search(r"sparse[a-zA-Z0-9_]*\.yaml", api_yaml_path):
              self.namespace = "sparse::"
          elif re.search(r"strings[a-zA-Z0-9_]*\.yaml", api_yaml_path):
              self.namespace = "strings::"
  