wehub-resource-sync
685 lines
21 KiB
Python
685 lines
21 KiB
Python
# Copyright (c) 2023 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 argparse
|
|
import hashlib
|
|
import pathlib
|
|
import sys
|
|
|
|
import jinja2
|
|
import yaml
|
|
|
|
# import from paddle/fluid/operators/generator
|
|
sys.path.append(
|
|
str(pathlib.Path(__file__).resolve().parents[2] / 'operators/generator')
|
|
)
|
|
import filters as op_gen_filters
|
|
import tests_utils as op_gen_tests
|
|
from parse_utils import to_named_dict
|
|
|
|
# import from paddle/fluid/pir/dialect/op_generator/api_gen.py
|
|
sys.path.append(
|
|
str(
|
|
pathlib.Path(__file__).resolve().parents[2] / 'pir/dialect/op_generator'
|
|
)
|
|
)
|
|
|
|
from decomp_interface_gen_op_list import (
|
|
decomp_vjp_interface_implementation_gen_op_list,
|
|
)
|
|
from gen_utils import attr_types_map, to_pascal_case
|
|
from type_mapping import output_type_map
|
|
|
|
# fmt: on
|
|
|
|
|
|
VJPS_BLACK_LIST = [
|
|
'reshape_grad',
|
|
'add_n_grad',
|
|
'fused_attention_grad',
|
|
]
|
|
|
|
BACKENDS_BLACK_LIST = [
|
|
'accuracy_check',
|
|
'copy_to',
|
|
'add_n_grad',
|
|
"allclose",
|
|
"isclose",
|
|
"send_v2",
|
|
"assert",
|
|
"embedding_sparse_grad",
|
|
"embedding_grad",
|
|
"full",
|
|
"partial_send",
|
|
"push_dense",
|
|
"comm_init_all",
|
|
]
|
|
|
|
# which is prim vjp
|
|
# op only has backward decomp_vjp rules
|
|
PRIM_VJP = [
|
|
'abs_grad',
|
|
'add_grad',
|
|
'amax_grad',
|
|
'amin_grad',
|
|
'angle_grad',
|
|
'argsort_grad',
|
|
'assign_grad',
|
|
'atan_grad',
|
|
'atan2_grad',
|
|
'cast_grad',
|
|
'ceil_grad',
|
|
'concat_grad',
|
|
'cos_grad',
|
|
'cumprod_grad',
|
|
'cumsum_grad',
|
|
'divide_grad',
|
|
'dot_grad',
|
|
'elementwise_pow_grad',
|
|
'erf_grad',
|
|
'exp_grad',
|
|
'expm1_grad',
|
|
'expand_grad',
|
|
'floor_grad',
|
|
'fmax_grad',
|
|
'fmin_grad',
|
|
'gather_grad',
|
|
'gather_nd_grad',
|
|
'kron_grad',
|
|
'kthvalue_grad',
|
|
'log_grad',
|
|
'logcumsumexp_grad',
|
|
'logsumexp_grad',
|
|
'masked_select_grad',
|
|
'matmul_grad',
|
|
'linear_v2_grad',
|
|
'max_grad',
|
|
'maximum_grad',
|
|
'minimum_grad',
|
|
'multiply_grad',
|
|
'pad_grad',
|
|
'pow_grad',
|
|
'prod_grad',
|
|
'put_along_axis_grad',
|
|
'reduce_as_grad',
|
|
'reshape_grad',
|
|
'roll_grad',
|
|
'rsqrt_grad',
|
|
'scale_grad',
|
|
"div_scale_grad",
|
|
'scatter_grad',
|
|
'scatter_nd_add_grad',
|
|
'sigmoid_grad',
|
|
'sin_grad',
|
|
'slice_grad',
|
|
'squeeze_grad',
|
|
'split_grad',
|
|
'sqrt_grad',
|
|
'square_grad',
|
|
'subtract_grad',
|
|
'sum_grad',
|
|
'take_along_axis_grad',
|
|
'tanh_grad',
|
|
'tile_grad',
|
|
'topk_grad',
|
|
'transpose_grad',
|
|
'trunc_grad',
|
|
'unsqueeze_grad',
|
|
'where_grad',
|
|
]
|
|
|
|
# which op is custom_vjp?
|
|
# op has forward decomp_rules and backward decomp_vjp rules
|
|
CUSTOM_VJP = [
|
|
'batch_norm_grad',
|
|
'bce_loss_grad',
|
|
'dropout_grad',
|
|
'gelu_grad',
|
|
'group_norm_grad',
|
|
'hardsigmoid_grad',
|
|
'hardswish_grad',
|
|
'instance_norm_grad',
|
|
'layer_norm_grad',
|
|
'leaky_relu_grad',
|
|
'mean_grad',
|
|
'relu_grad',
|
|
'relu6_grad',
|
|
'silu_grad',
|
|
'softmax_grad',
|
|
'softsign_grad',
|
|
'stack_grad',
|
|
'swish_grad',
|
|
'elu_grad',
|
|
'swiglu_grad',
|
|
'p_norm_grad',
|
|
'masked_fill_grad',
|
|
'index_put_grad',
|
|
'index_add_grad',
|
|
"var_grad",
|
|
] # custom vjp list of composite op
|
|
|
|
VJP_COMPS = PRIM_VJP + CUSTOM_VJP
|
|
|
|
|
|
def load(path: pathlib.Path):
|
|
"""Load config from yaml file.
|
|
|
|
Args:
|
|
path (pathlib.Path): The path of yaml config.
|
|
|
|
Returns:
|
|
dict: The config info.
|
|
|
|
"""
|
|
with open(path, 'rt') as f:
|
|
return yaml.safe_load(f)
|
|
|
|
|
|
def render(src_dir: pathlib.Path, dst_dir: pathlib.Path, *args, **kwargs):
|
|
"""Render and save Jinja2 templates to the destination directory.
|
|
|
|
Args:
|
|
src_dir (pathlib.Path): The source directory containing Jinja2 templates.
|
|
dst_dir (pathlib.Path): The destination directory to save rendered files.
|
|
*args: Additional positional arguments passed to the `render` function.
|
|
**kwargs: Additional keyword arguments passed to the `render` function.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
env = jinja2.Environment(
|
|
loader=jinja2.FileSystemLoader(src_dir),
|
|
keep_trailing_newline=True,
|
|
trim_blocks=True,
|
|
lstrip_blocks=True,
|
|
undefined=jinja2.StrictUndefined,
|
|
extensions=['jinja2.ext.do'],
|
|
)
|
|
env.filters.update(
|
|
{
|
|
'to_paddle_attr_type': op_gen_filters.to_paddle_attr_type,
|
|
'to_paddle_input_type': op_gen_filters.to_paddle_input_type,
|
|
'to_paddle_output_type': op_gen_filters.to_paddle_output_type,
|
|
'trip_intermediate': op_gen_filters.filter_intermediate,
|
|
}
|
|
)
|
|
env.tests.update(
|
|
{
|
|
'scalar': op_gen_tests.is_scalar,
|
|
'intarray': op_gen_tests.is_intarray,
|
|
'datatype': op_gen_tests.is_datatype,
|
|
'exist_mutable_attribute': op_gen_tests.exist_mutable_attribute,
|
|
'mutable_attribute': op_gen_tests.is_mutable_attribute,
|
|
'only_composite_op': op_gen_tests.is_only_composite_op,
|
|
}
|
|
)
|
|
for tpl in env.list_templates(
|
|
filter_func=lambda name: ".h" in name or ".cc" in name
|
|
):
|
|
save(
|
|
env.get_template(tpl).render(*args, **kwargs),
|
|
dst_dir / tpl.rstrip('.j2'),
|
|
)
|
|
|
|
|
|
def render_decomp_vjp(
|
|
src_dir: pathlib.Path, dst_dir: pathlib.Path, *args, **kwargs
|
|
):
|
|
"""Render and save Jinja2 templates to the destination directory.
|
|
|
|
Args:
|
|
src_dir (pathlib.Path): The source directory containing Jinja2 templates.
|
|
dst_dir (pathlib.Path): The destination directory to save rendered files.
|
|
*args: Additional positional arguments passed to the `render` function.
|
|
**kwargs: Additional keyword arguments passed to the `render` function.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
env = jinja2.Environment(
|
|
loader=jinja2.FileSystemLoader(src_dir),
|
|
keep_trailing_newline=True,
|
|
trim_blocks=True,
|
|
lstrip_blocks=True,
|
|
undefined=jinja2.StrictUndefined,
|
|
extensions=['jinja2.ext.do'],
|
|
)
|
|
env.filters.update(
|
|
{
|
|
'to_paddle_attr_type': op_gen_filters.to_paddle_attr_type,
|
|
'to_paddle_input_type': op_gen_filters.to_paddle_input_type,
|
|
'to_paddle_output_type': op_gen_filters.to_paddle_output_type,
|
|
'trip_intermediate': op_gen_filters.filter_intermediate,
|
|
}
|
|
)
|
|
env.tests.update(
|
|
{
|
|
'scalar': op_gen_tests.is_scalar,
|
|
'intarray': op_gen_tests.is_intarray,
|
|
'datatype': op_gen_tests.is_datatype,
|
|
'exist_mutable_attribute': op_gen_tests.exist_mutable_attribute,
|
|
'mutable_attribute': op_gen_tests.is_mutable_attribute,
|
|
'only_composite_op': op_gen_tests.is_only_composite_op,
|
|
}
|
|
)
|
|
|
|
decomp_temp = "decomp/generated_decomp_vjp.j2"
|
|
save(
|
|
env.get_template(decomp_temp).render(*args, **kwargs),
|
|
pathlib.Path(dst_dir),
|
|
)
|
|
|
|
|
|
def save(content: str, path: pathlib.Path):
|
|
"""Saves the given string contents to a file in the specified path.
|
|
|
|
Args:
|
|
content (str): The string content that needs to be saved.
|
|
path (pathlib.Path): The path to save the file, a Pathlib path object
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
dst_content = ''
|
|
if path.is_file():
|
|
with open(path, 'r') as f:
|
|
dst_content = f.read()
|
|
|
|
if (
|
|
hashlib.md5(content.encode("UTF-8")).hexdigest()
|
|
!= hashlib.md5(dst_content.encode("UTF-8")).hexdigest()
|
|
):
|
|
with open(path, 'w') as f:
|
|
f.write(content)
|
|
print(f"Generate source file {path}")
|
|
|
|
|
|
def get_inplace_api(apis):
|
|
inplace_apis = []
|
|
for api in apis:
|
|
if (
|
|
'inplace' in api
|
|
and api['inplace'] is not None
|
|
and not api['name'].endswith('_')
|
|
):
|
|
inplace_api = api.copy()
|
|
inplace_api['name'] = api['name'] + '_'
|
|
inplace_apis.append(inplace_api)
|
|
return inplace_apis
|
|
|
|
|
|
def filter_compat_info(items):
|
|
for item in items:
|
|
item['op'] = item['op'].split('(')[0].strip()
|
|
if 'backward' in item:
|
|
item_backwards = item['backward'].split(',')
|
|
for idx, item_backward in enumerate(item_backwards):
|
|
item_backward = item_backward.split('(')[0].strip()
|
|
item_backwards[idx] = item_backward
|
|
item['backward'] = (
|
|
','.join(item_backwards)
|
|
if len(item_backwards) > 0
|
|
else item_backwards[0]
|
|
)
|
|
|
|
|
|
def extend_compat_info(apis, compats):
|
|
for api in apis:
|
|
if api['name'].endswith('sp') or api['name'].endswith('sp_'):
|
|
continue
|
|
attrs = api["attrs"]
|
|
for attr in attrs:
|
|
if op_gen_tests.is_scalar(
|
|
attr['typename']
|
|
) or op_gen_tests.is_intarray(attr['typename']):
|
|
attr["support_tensor"] = False
|
|
apis_dict = to_named_dict(apis)
|
|
for compat_item in compats:
|
|
fwd_op_name = compat_item["op"]
|
|
if fwd_op_name not in apis_dict:
|
|
continue
|
|
fwd_api = apis_dict[fwd_op_name]
|
|
backward_op_names = []
|
|
while fwd_op_name is not None and fwd_op_name in apis_dict:
|
|
backward_op_names.append(apis_dict[fwd_op_name]['backward'])
|
|
fwd_op_name = apis_dict[fwd_op_name]['backward']
|
|
backward_apis = []
|
|
for backward_op_name in backward_op_names:
|
|
if backward_op_name in apis_dict:
|
|
backward_apis.append(apis_dict[backward_op_name])
|
|
support_tensor_attrs_names = []
|
|
compat_attrs_data_type = {}
|
|
if 'scalar' in compat_item and compat_item['op'] != "pow":
|
|
for attr_name, attr_info in compat_item['scalar'].items():
|
|
if (
|
|
'support_tensor' in attr_info
|
|
and attr_info['support_tensor'] is True
|
|
or 'tensor_name' in attr_info
|
|
):
|
|
support_tensor_attrs_names.append(attr_name)
|
|
if 'data_type' in attr_info:
|
|
compat_attrs_data_type.update(
|
|
{attr_name: attr_info['data_type']}
|
|
)
|
|
if 'int_array' in compat_item:
|
|
for attr_name, attr_info in compat_item['int_array'].items():
|
|
if (
|
|
'support_tensor' in attr_info
|
|
and attr_info['support_tensor'] is True
|
|
or 'tensor_name' in attr_info
|
|
or 'tensors_name' in attr_info
|
|
):
|
|
support_tensor_attrs_names.append(attr_name)
|
|
if len(support_tensor_attrs_names) > 0:
|
|
for api in [fwd_api, *backward_apis]:
|
|
attrs = api["attrs"]
|
|
for attr in attrs:
|
|
if attr['name'] in support_tensor_attrs_names:
|
|
attr['support_tensor'] = True
|
|
for api in [fwd_api, *backward_apis]:
|
|
attrs = api["attrs"]
|
|
for attr in attrs:
|
|
if attr['name'] in compat_attrs_data_type:
|
|
attr['data_type'] = compat_attrs_data_type[attr['name']]
|
|
return apis
|
|
|
|
|
|
def process_backward_invoke_info(apis):
|
|
apis_dict = to_named_dict(apis)
|
|
for api in apis:
|
|
if api['is_fwd']:
|
|
continue
|
|
if 'invoke' in api and api['invoke']['func'] in apis_dict:
|
|
args = api['invoke']['args'].split(',')
|
|
args = [arg.strip() for arg in args]
|
|
attrs_dict = to_named_dict(api['attrs'])
|
|
inputs_dict = to_named_dict(api['inputs'])
|
|
arg_inputs = []
|
|
arg_attrs = []
|
|
for arg in args:
|
|
if arg in inputs_dict:
|
|
arg_inputs.append(arg)
|
|
elif arg in attrs_dict and attrs_dict[arg].get(
|
|
"support_tensor", False
|
|
):
|
|
arg_inputs.append(arg + '_')
|
|
else:
|
|
arg_attrs.append(arg)
|
|
args = arg_inputs + arg_attrs
|
|
api['invoke']['args'] = ', '.join(args)
|
|
|
|
|
|
def process_optional_inplace_output_info(apis):
|
|
for api in apis:
|
|
inputs_dict = to_named_dict(api['inputs'])
|
|
for output in api['outputs']:
|
|
if not api['is_fwd']:
|
|
return
|
|
else:
|
|
if (
|
|
api.get("inplace", None)
|
|
and output['name'] in api['inplace']
|
|
and inputs_dict[api['inplace'][output['name']]]['optional']
|
|
):
|
|
output['optional'] = True
|
|
else:
|
|
output['optional'] = False
|
|
|
|
|
|
def update_apis(op_yaml_items, update_yaml_file):
|
|
with open(update_yaml_file, "r") as f:
|
|
update_apis = yaml.safe_load(f)
|
|
for i in range(len(op_yaml_items)):
|
|
for update_api in update_apis:
|
|
if op_yaml_items[i]['name'] == update_api['name']:
|
|
op_yaml_items[i] = update_api
|
|
break
|
|
|
|
|
|
def gen(
|
|
prim_path: pathlib.Path,
|
|
fwd_path: pathlib.Path,
|
|
rev_path: pathlib.Path,
|
|
compat_path: pathlib.Path,
|
|
fwd_pd_op_path: pathlib.Path,
|
|
update_fwd_pd_op_path: pathlib.Path,
|
|
rev_pd_op_path: pathlib.Path,
|
|
fused_op_path: pathlib.Path,
|
|
fused_rev_path: pathlib.Path,
|
|
sparse_op_path: pathlib.Path,
|
|
sparse_rev_op_path: pathlib.Path,
|
|
templates_dir: pathlib.Path,
|
|
destination_dir: pathlib.Path,
|
|
decomp_vjp_destination_dir: pathlib.Path,
|
|
):
|
|
"""The `gen` load jinja2 templates and relative config info, use jinja2
|
|
templating engine to generate c++ code, and save the code into destination.
|
|
|
|
Args:
|
|
prim_path (pathlib.Path): The YAML file path of the primitive API.
|
|
fwd_path (pathlib.Path): The YAML file path of the forward API.
|
|
rev_path (pathlib.Path): The YAML file path of the backward API.
|
|
compat_path: (pathlib.Path): The YAML file path of the ops compat.
|
|
fwd_pd_op_path (pathlib.Path): The YAML file path of the ir forward API.
|
|
update_fwd_pd_op_path (pathlib.Path): The YAML file path of the ir update_ops.
|
|
rev_pd_op_path (pathlib.Path): The YAML file path of the ir backward API.
|
|
fused_op_path (pathlib.Path): The YAML file path of the fused API.
|
|
fused_rev_path (pathlib.Path): The YAML file path of the fused backward API.
|
|
sparse_op_path (pathlib.Path): The YAML file path of the sparse API.
|
|
sparse_rev_op_path (pathlib.Path): The YAML file path of the sparse backward API.
|
|
templates_dir (pathlib.Path): The directory of the templates.
|
|
destination_dir (pathlib.Path): The Directory of the generated file.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
(
|
|
prims,
|
|
fwds,
|
|
revs,
|
|
compats,
|
|
ir_fwds,
|
|
ir_revs,
|
|
ir_update_fwds,
|
|
fused_fwds,
|
|
fused_revs,
|
|
sparse_fwds,
|
|
sparse_revs,
|
|
) = (
|
|
load(prim_path),
|
|
load(fwd_path),
|
|
load(rev_path),
|
|
load(compat_path),
|
|
load(fwd_pd_op_path),
|
|
load(rev_pd_op_path),
|
|
load(update_fwd_pd_op_path),
|
|
load(fused_op_path),
|
|
load(fused_rev_path),
|
|
load(sparse_op_path),
|
|
load(sparse_rev_op_path),
|
|
)
|
|
filter_compat_info(compats)
|
|
for sparse_op in sparse_fwds:
|
|
if sparse_op['name'].endswith("_"):
|
|
sparse_op['name'] += 'sp_'
|
|
if sparse_op['backward'] is not None:
|
|
sparse_op['backward'] += '_sp'
|
|
else:
|
|
sparse_op['name'] += '_sp'
|
|
if sparse_op['backward'] is not None:
|
|
sparse_op['backward'] += '_sp'
|
|
fwd_apis = fwds + ir_fwds + ir_update_fwds + fused_fwds + sparse_fwds
|
|
|
|
for sparse_op in sparse_revs:
|
|
sparse_op['name'] += '_sp'
|
|
if sparse_op['forward']['name'].endswith("_"):
|
|
sparse_op['forward']['name'] += 'sp_'
|
|
if sparse_op.get('invoke') is not None:
|
|
sparse_op['invoke']['func'] += 'sp_'
|
|
else:
|
|
sparse_op['forward']['name'] += '_sp'
|
|
if sparse_op.get('invoke') is not None:
|
|
sparse_op['invoke']['func'] += '_sp'
|
|
apis = [{**api, **{'is_fwd': True}} for api in fwd_apis]
|
|
apis = apis + [
|
|
{**api, **{'is_fwd': False}}
|
|
for api in revs + ir_revs + fused_revs + sparse_revs
|
|
]
|
|
apis = [
|
|
(
|
|
{**api, **{'is_prim': True}}
|
|
if api['name'] in prims
|
|
else {**api, **{'is_prim': False}}
|
|
)
|
|
for api in apis
|
|
]
|
|
|
|
apis = extend_compat_info(apis, compats)
|
|
apis = apis + get_inplace_api(apis)
|
|
process_backward_invoke_info(apis)
|
|
process_optional_inplace_output_info(apis)
|
|
|
|
apis = [
|
|
{**api, **{'class_name': to_pascal_case(api["name"]) + "Op"}}
|
|
for api in apis
|
|
]
|
|
|
|
for item in apis:
|
|
for attr_item in item["attrs"]:
|
|
if attr_item["typename"] not in attr_types_map.keys():
|
|
raise TypeError
|
|
attr_item["mapped_type"] = attr_types_map[attr_item["typename"]]
|
|
for out_item in item["outputs"]:
|
|
if out_item["typename"] not in output_type_map.keys():
|
|
name = out_item["typename"]
|
|
raise TypeError(f"err type {name}")
|
|
if out_item["optional"]:
|
|
out_item["mapped_type"] = (
|
|
"paddle::optional<"
|
|
+ output_type_map[out_item["typename"]]
|
|
+ ">"
|
|
)
|
|
else:
|
|
out_item["mapped_type"] = output_type_map[out_item["typename"]]
|
|
|
|
render(
|
|
templates_dir,
|
|
destination_dir,
|
|
apis=apis,
|
|
backend_black_list=BACKENDS_BLACK_LIST,
|
|
vjp_black_list=VJPS_BLACK_LIST,
|
|
vjp_comp_white_list=VJP_COMPS,
|
|
)
|
|
render_decomp_vjp(
|
|
templates_dir,
|
|
decomp_vjp_destination_dir,
|
|
apis=apis,
|
|
backend_black_list=BACKENDS_BLACK_LIST,
|
|
vjp_black_list=VJPS_BLACK_LIST,
|
|
vjp_comp_white_list=VJP_COMPS,
|
|
decomp_vjp_white_list=decomp_vjp_interface_implementation_gen_op_list,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(
|
|
description='Generate Static Primitive API'
|
|
)
|
|
parser.add_argument(
|
|
'--prim_path',
|
|
type=str,
|
|
help='The primitive API yaml file.',
|
|
)
|
|
parser.add_argument(
|
|
'--fwd_path', type=str, help='The parsed ops yaml file.'
|
|
)
|
|
parser.add_argument(
|
|
'--rev_path', type=str, help='The parsed ops yaml file.'
|
|
)
|
|
parser.add_argument(
|
|
'--compat_path',
|
|
type=str,
|
|
help='The parsed ops compat yaml file.',
|
|
)
|
|
parser.add_argument(
|
|
'--fwd_pd_op_path',
|
|
type=str,
|
|
help='The ir forward ops parsed yaml file.',
|
|
)
|
|
parser.add_argument(
|
|
'--update_fwd_pd_op_path',
|
|
type=str,
|
|
help='The ir update forward ops parsed yaml file.',
|
|
)
|
|
parser.add_argument(
|
|
'--rev_pd_op_path',
|
|
type=str,
|
|
help='The ir backward ops parsed yaml file.',
|
|
)
|
|
parser.add_argument(
|
|
'--fused_op_path',
|
|
type=str,
|
|
help='The parsed fused forward ops yaml file.',
|
|
)
|
|
parser.add_argument(
|
|
'--fused_rev_op_path',
|
|
type=str,
|
|
help='The parsed fused backward ops yaml file.',
|
|
)
|
|
parser.add_argument(
|
|
'--sparse_op_path',
|
|
type=str,
|
|
help='The parsed sparse forward ops yaml file.',
|
|
)
|
|
parser.add_argument(
|
|
'--sparse_rev_op_path',
|
|
type=str,
|
|
help='The parsed sparse backward ops yaml file.',
|
|
)
|
|
parser.add_argument(
|
|
'--templates_dir',
|
|
type=str,
|
|
help='JinJa2 templates base directory.',
|
|
)
|
|
parser.add_argument(
|
|
'--destination_dir',
|
|
type=str,
|
|
help='Destination base directory for generated file.',
|
|
)
|
|
parser.add_argument(
|
|
'--decomp_vjp_destination_dir',
|
|
type=str,
|
|
help='Destination base directory for generated file.',
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
gen(
|
|
pathlib.Path(args.prim_path),
|
|
pathlib.Path(args.fwd_path),
|
|
pathlib.Path(args.rev_path),
|
|
pathlib.Path(args.compat_path),
|
|
pathlib.Path(args.fwd_pd_op_path),
|
|
pathlib.Path(args.update_fwd_pd_op_path),
|
|
pathlib.Path(args.rev_pd_op_path),
|
|
pathlib.Path(args.fused_op_path),
|
|
pathlib.Path(args.fused_rev_op_path),
|
|
pathlib.Path(args.sparse_op_path),
|
|
pathlib.Path(args.sparse_rev_op_path),
|
|
pathlib.Path(args.templates_dir),
|
|
pathlib.Path(args.destination_dir),
|
|
pathlib.Path(args.decomp_vjp_destination_dir),
|
|
)
|