gen_opkernel_doc.py

#!/usr/bin/env python
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

importargparse
importio# noqa: F401
importos
importpathlib
fromcollectionsimportdefaultdict

importonnxruntime.capi.onnxruntime_pybind11_stateasrtpy


defformat_version_range(v):
ifv[1] >=2147483647:
returnstr(v[0]) +"+"
else:
ifv[0] ==v[1]:
returnstr(v[0])
else:
return"["+str(v[0]) +", "+str(v[1]) +"]"


defformat_type_constraints(tc):
tcstr=""
firsttcitem=True
fortcitemintc:
iffirsttcitem:
firsttcitem=False
else:
tcstr+=", "
tcstr+=tcitem
returntcstr


defformat_param_strings(params):
firstparam=True
s=""
ifparams:
forparaminsorted(params):
iffirstparam:
firstparam=False
else:
s+="<br><br>or<br><br>"
s+=param
returns


defexpand_providers(provider_filter: [str]):
providers=set()
ifprovider_filter:
forproviderinprovider_filter:
p=provider.lower()
ifnotp.endswith("executionprovider"):
p+="executionprovider"
providers.add(p)

returnproviders


defmain(output_path: pathlib.Path, provider_filter: [str]):
providers=expand_providers(provider_filter)

withopen(output_path, "w", newline="", encoding="utf-8") asfout:
fout.write("## Supported Operators and Data Types\n")
fout.write(
"*This file is automatically generated from the registered kernels by "
"[this script](https://github.com/microsoft/onnxruntime/blob/main/tools/python/gen_opkernel_doc.py).\n"
"Do not modify directly.*\n\n"
 )
opdef=rtpy.get_all_operator_schema()
paramdict= {}
forschemainopdef:
inputs=schema.inputs
domain=schema.domain
ifnotdomain:
domain="ai.onnx"
fullname=domain+"."+schema.name
paramstr=""
firstinput=True
ifinputs:
forinpininputs:
iffirstinput:
firstinput=False
else:
paramstr+="<br> "
paramstr+=f"*in* {inp.name}:**{inp.typeStr}**"

outputs=schema.outputs
ifoutputs:
foroutpinoutputs:
iffirstinput:
firstinput=False
else:
paramstr+="<br> "
paramstr+=f"*out* {outp.name}:**{outp.typeStr}**"

paramstr+=""
paramset=paramdict.get(fullname)
ifparamsetisNone:
paramdict[fullname] =set()

paramdict[fullname].add(paramstr)

index=defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
foropinrtpy.get_all_opkernel_def():
domain=op.domain
ifnotdomain:
domain="ai.onnx"
index[op.provider][domain][op.op_name].append(op)

# TOC
fout.write("## Execution Providers\n\n")
forproviderinsorted(index.keys()):
ifprovidersandprovider.lower() notinproviders:
continue
fout.write(f"- [{provider}](#{provider.lower()})\n")
fout.write("\n---------------")

forprovider, domainmapinsorted(index.items()):
ifprovidersandprovider.lower() notinproviders:
continue

fout.write(f'\n\n<a name="{provider.lower()}"/>\n\n')
fout.write(f"## Operators implemented by {provider}\n\n")
fout.write("| Op Name | Parameters | OpSet Version | Types Supported |\n")
fout.write("|---------|------------|---------------|-----------------|\n")
fordomain, namemapinsorted(domainmap.items()):
fout.write("|**Operator Domain:** *"+domain+"*||||\n")
forname, opsinsorted(namemap.items()):
version_type_index=defaultdict(lambda: defaultdict(set))
foropinops:
fortname, tclistinop.type_constraints.items():
forcintclist:
version_type_index[op.version_range][tname].add(c)

namefirsttime=True
forversion_range, typemapinsorted(version_type_index.items(), key=lambdax: x[0], reverse=True):
ifnamefirsttime:
params=paramdict.get(domain+"."+name, None)
fout.write("|"+name+"|"+format_param_strings(params) +"|")
namefirsttime=False
else:
fout.write("|||")
fout.write(format_version_range(version_range) +"|")
fori, (tname, tcset) inenumerate(sorted(typemap.items())):
tnameindex=i+1
tclist= []
fortcinsorted(tcset):
tclist.append(tc)
fout.write("**"+tname+"** = "+format_type_constraints(tclist))
iftnameindex<len(typemap):
fout.write("<br/> ")
fout.write("|\n")

fout.write("| |\n| |\n")


if__name__=="__main__":
parser=argparse.ArgumentParser(description="ONNX Runtime Operator Kernel Documentation Generator")
parser.add_argument(
"--providers",
nargs="+",
help="Filter to specified execution providers. Case-insensitive. "
"Matches provider names from <ORT>/include/onnxruntime/core/graph/constants.h'. "
"'ExecutionProvider' is automatically appended as needed. "
"e.g. `--providers cpu cuda` will match CPUExecutionProvider and CUDAExecutionProvider.",
 )
parser.add_argument(
"--output_path",
help="output markdown file path",
type=pathlib.Path,
required=True,
default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "OperatorKernels.md"),
 )
args=parser.parse_args()

main(args.output_path, args.providers)