You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

gen_param_defs.py 28 kB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808
  1. #!/usr/bin/env python3
  2. # -*- coding: utf-8 -*-
  3. import argparse
  4. import collections
  5. import textwrap
  6. import os
  7. import hashlib
  8. import struct
  9. class member_defs:
  10. """contain classes to define members of an opr param"""
  11. Dtype = collections.namedtuple('Dtype', ['cname', 'pycvt', 'pyfmt',
  12. 'cppjson', 'cname_attr'])
  13. Dtype.__new__.__defaults__ = ('', )
  14. uint32 = Dtype('uint32_t', 'int', 'I', 'NumberInt')
  15. uint64 = Dtype('uint64_t', 'int', 'Q', 'NumberInt',
  16. 'alignas(sizeof(uint64_t)) ')
  17. int32 = Dtype('int32_t', 'int', 'i', 'NumberInt')
  18. float32 = Dtype('float', 'float', 'f', 'Number')
  19. float64 = Dtype('double', 'float', 'd', 'Number')
  20. dtype = Dtype('DTypeEnum', '_as_dtype_num', 'I', 'Number')
  21. bool = Dtype('bool', 'bool', '?', 'Bool')
  22. class Base:
  23. pass
  24. class Doc:
  25. """wrap an identifier to associate document
  26. note: if the doc starts with a linebreak, it would not be reforamtted.
  27. """
  28. __slots__ = ['id', 'doc']
  29. def __init__(self, id_, doc):
  30. assert isinstance(id_, str) and isinstance(doc, str), (id_, doc)
  31. self.id = id_
  32. self.doc = doc
  33. @property
  34. def no_reformat(self):
  35. """whether reformat is disallowed for this doc string"""
  36. return self.doc.startswith('\n')
  37. @property
  38. def raw_lines(self):
  39. """the doc lines when ``no_format`` is true"""
  40. ret = self.doc.split('\n')
  41. assert not ret[0]
  42. return ret[1:]
  43. @classmethod
  44. def make(cls, v):
  45. """make doc object from str or doc"""
  46. if isinstance(v, cls):
  47. return v
  48. assert isinstance(v, str)
  49. return cls(v, '')
  50. def __str__(self):
  51. return self.id
  52. def __eq__(self, rhs):
  53. if isinstance(rhs, str):
  54. return self.id == rhs
  55. return (isinstance(rhs, Doc) and
  56. (self.id, self.doc) == (rhs.id, rhs.doc))
  57. class Enum(Base):
  58. """define an enum; the result would contain both an enum class def and its
  59. corresponding data field
  60. :param default: index of default member value
  61. :attr name_field: name of the data field of this enum in the param
  62. struct
  63. :attr member_alias: list of (member, alias) pairs
  64. """
  65. __slots__ = ['name', 'name_field', 'members', 'default',
  66. 'member_alias']
  67. all_enums = {}
  68. """(param_name, name) => enum"""
  69. def __init__(self, param_name, name, name_field, members, default,
  70. member_alias):
  71. name = member_defs.Doc.make(name)
  72. assert name.id[0].isupper()
  73. members = tuple(map(member_defs.Doc.make, members))
  74. if isinstance(default, str):
  75. if default not in name_field:
  76. raise ValueError(
  77. "Default value '{}' does not exist.".format(default))
  78. default = name_field.index(default)
  79. assert isinstance(default, int)
  80. self.name = name
  81. self.name_field = self.get_name_field(name.id, name_field)
  82. self.members = members
  83. self.default = default
  84. self.all_enums[(param_name, name.id)] = self
  85. assert isinstance(member_alias, list)
  86. self.member_alias = member_alias
  87. @classmethod
  88. def get_name_field(cls, name, name_field):
  89. if name_field is None:
  90. name_field = name[0].lower() + name[1:]
  91. assert isinstance(name_field, str)
  92. return name_field
  93. class Field(Base):
  94. """define a normal data field"""
  95. __slots__ = ['name', 'dtype', 'default']
  96. def __init__(self, name, dtype, default):
  97. assert isinstance(dtype, member_defs.Dtype)
  98. self.name = member_defs.Doc.make(name)
  99. self.dtype = dtype
  100. self.default = default
  101. class Const(Base):
  102. """define a const data field"""
  103. __slots__ = ['name', 'dtype', 'default']
  104. def __init__(self, name, dtype, default):
  105. assert isinstance(dtype, member_defs.Dtype)
  106. self.name = member_defs.Doc.make(name)
  107. self.dtype = dtype
  108. self.default = default
  109. class EnumAlias(Base):
  110. """alias of enum type from another param"""
  111. __slots__ = ['name', 'name_field', 'src_class', 'src_name', 'default']
  112. def __init__(self, name, name_field, src_class, src_name, default):
  113. self.name = name
  114. self.name_field = member_defs.Enum.get_name_field(name, name_field)
  115. self.src_class = src_class
  116. if src_name is None:
  117. src_name = name
  118. self.src_name = src_name
  119. self.default = default
  120. @property
  121. def src_enum(self):
  122. """source Enum class"""
  123. return member_defs.Enum.all_enums[(self.src_class, self.src_name)]
  124. def get_default(self):
  125. """get default index; fallback to src index if default is not
  126. set"""
  127. if self.default is None:
  128. return self.src_enum.default
  129. return self.default
  130. class ParamDef:
  131. """"""
  132. __all_tags = set()
  133. all_param_defs = []
  134. __slots__ = ['name', 'members', 'tag', 'is_legacy']
  135. def __init__(self, name, doc='', *, version=0, is_legacy=False):
  136. self.members = []
  137. self.all_param_defs.append(self)
  138. h = hashlib.sha256(name.encode('utf-8'))
  139. if version:
  140. h.update(struct.pack('<I', version))
  141. if is_legacy:
  142. name += 'V{}'.format(version)
  143. self.name = member_defs.Doc(name, doc)
  144. self.tag = int(h.hexdigest()[:8], 16)
  145. self.is_legacy = is_legacy
  146. if self.tag < 1024:
  147. self.tag += 1024
  148. assert self.tag not in self.__all_tags, (
  149. 'tag hash confliction: name={} tag={}'.format(name, self.tag))
  150. self.__all_tags.add(self.tag)
  151. def add_fields(self, dtype, *names_defaults):
  152. assert isinstance(dtype, str)
  153. dtype = getattr(member_defs, dtype)
  154. assert len(names_defaults) % 2 == 0
  155. for i, j in zip(names_defaults[::2], names_defaults[1::2]):
  156. self.members.append(member_defs.Field(i, dtype, j))
  157. return self
  158. def add_enum(self, name, *members, default=0, name_field=None,
  159. member_alias=[]):
  160. self.members.append(member_defs.Enum(
  161. self.name.id, name, name_field, members, default, member_alias))
  162. return self
  163. def add_enum_alias(self, name, src_class, src_name=None, name_field=None,
  164. default=None):
  165. self.members.append(member_defs.EnumAlias(
  166. name, name_field, src_class, src_name, default))
  167. return self
  168. def add_const(self, dtype, *names_defaults):
  169. assert isinstance(dtype, str)
  170. dtype = getattr(member_defs, dtype)
  171. assert len(names_defaults) % 2 == 0
  172. for i, j in zip(names_defaults[::2], names_defaults[1::2]):
  173. self.members.append(member_defs.Const(i, dtype, j))
  174. return self
  175. class WriterBase:
  176. """base class for output file writer"""
  177. _fout = None
  178. _input_hash = None
  179. def __call__(self, fout):
  180. self._fout = fout
  181. def set_input_hash(self, h):
  182. self._input_hash = h
  183. return self
  184. def _get_header(self):
  185. return 'generated by {} for {}'.format(
  186. os.path.basename(__file__),
  187. self._input_hash
  188. )
  189. def _process(self, defs):
  190. dispatch = {
  191. member_defs.Enum: self._on_member_enum,
  192. member_defs.EnumAlias: self._on_member_enum_alias,
  193. member_defs.Field: self._on_member_field,
  194. member_defs.Const: self._on_const_field
  195. }
  196. for i in defs:
  197. assert isinstance(i, ParamDef)
  198. self._on_param_begin(i)
  199. for j in i.members:
  200. dispatch[type(j)](j)
  201. self._on_param_end(i)
  202. def _on_param_begin(self, p):
  203. """:type p: :class:`.ParamDef`"""
  204. def _on_param_end(self, p):
  205. """:type p: :class:`.ParamDef`"""
  206. def _on_member_enum(self, e):
  207. """:type p: :class:`.Enum`"""
  208. def _on_member_enum_alias(self, e):
  209. """:type p: :class:`.EnumAlias`"""
  210. def _on_member_field(self, f):
  211. """:type p: :class:`.Field`"""
  212. def _on_const_field(self, f):
  213. """:type p: :class:`.Const`"""
  214. class IndentWriterBase(WriterBase):
  215. _cur_indent = ''
  216. def _indent(self):
  217. self._cur_indent += ' ' * 4
  218. def _unindent(self):
  219. self._cur_indent = self._cur_indent[:-4]
  220. def _write(self, content, *fmt, indent=0):
  221. if indent < 0:
  222. self._unindent()
  223. self._fout.write(self._cur_indent)
  224. if fmt:
  225. content = content % fmt
  226. self._fout.write(content)
  227. self._fout.write('\n')
  228. if indent > 0:
  229. self._indent()
  230. class PyWriter(IndentWriterBase):
  231. FieldDef = collections.namedtuple(
  232. 'FieldDef', ['name', 'cvt', 'fmt', 'default', 'type', 'doc'])
  233. # see _on_param_end() for the use of those fields
  234. _cur_param_name = None
  235. _cur_fields = None
  236. _cur_struct_fmt = None
  237. _enum_member2num = None
  238. def __call__(self, fout, defs):
  239. super().__call__(fout)
  240. self._enum_member2num = []
  241. self._write('# %s', self._get_header())
  242. self._write('import struct')
  243. self._write('from . import enum36 as enum')
  244. self._write(
  245. 'class _ParamDefBase:\n'
  246. ' def serialize(self):\n'
  247. ' tag = struct.pack("I", type(self).TAG)\n'
  248. ' pdata = [getattr(self, i) for i in self.__slots__]\n'
  249. ' for idx, v in enumerate(pdata):\n'
  250. ' if isinstance(v, _EnumBase):\n'
  251. ' pdata[idx] = _enum_member2num[id(v)]\n'
  252. ' return tag + self._packer.pack(*pdata)\n'
  253. '\n'
  254. )
  255. self._write(
  256. 'class _EnumBase(enum.Enum):\n'
  257. ' @classmethod\n'
  258. ' def __normalize(cls, val):\n'
  259. ' if isinstance(val, str):\n'
  260. ' if not hasattr(cls, "__member_upper_dict__"):\n'
  261. ' cls.__member_upper_dict__ = {k.upper(): v\n'
  262. ' for k, v in cls.__members__.items()}\n'
  263. ' val = cls.__member_upper_dict__.get(val.upper(),val)\n'
  264. ' return val\n'
  265. ' @classmethod\n'
  266. ' def convert(cls, val):\n'
  267. ' val = cls.__normalize(val)\n'
  268. ' if isinstance(val, cls):\n'
  269. ' return val\n'
  270. ' return cls(val)\n'
  271. ' @classmethod\n'
  272. ' def _missing_(cls, value):\n'
  273. ' vnorm = cls.__normalize(value)\n'
  274. ' if vnorm is not value:\n'
  275. ' return cls(vnorm)\n'
  276. ' return super()._missing_(value)\n'
  277. '\n'
  278. )
  279. self._write(
  280. 'def _as_dtype_num(dtype):\n'
  281. ' import megengine._internal.mgb as m\n'
  282. ' return m._get_dtype_num(dtype)\n'
  283. '\n'
  284. )
  285. self._write(
  286. '''
  287. def _as_serialized_dtype(dtype):
  288. import megengine._internal.mgb as m
  289. return m._get_serialized_dtype(dtype)
  290. '''
  291. )
  292. self._process(defs)
  293. self._write(
  294. '''
  295. class SerializedDType(_ParamDefBase):
  296. TAG = FakeSerializedDType.TAG
  297. __slots__ = ['dtype']
  298. class IdentityPacker:
  299. def pack(self, *args):
  300. assert all([isinstance(x, bytes) for x in args])
  301. return b''.join(args)
  302. _packer = IdentityPacker()
  303. def __init__(self, dtype):
  304. """
  305. :type dtype: :class:`np.dtype` compatible
  306. """
  307. self.dtype = _as_serialized_dtype(dtype)
  308. '''
  309. )
  310. self._write('_enum_member2num = {\n %s}',
  311. ',\n '.join(self._enum_member2num))
  312. def _write_doc(self, doc):
  313. assert isinstance(doc, member_defs.Doc)
  314. if not doc.doc:
  315. return
  316. if doc.no_reformat:
  317. self._write('"""')
  318. for i in doc.raw_lines:
  319. self._write(i)
  320. self._write('"""')
  321. return
  322. doc = doc.doc.replace('\n', ' ')
  323. textwidth = 80 - len(self._cur_indent)
  324. self._write('"""')
  325. for i in textwrap.wrap(doc, textwidth):
  326. self._write(i)
  327. self._write('"""')
  328. def _on_param_begin(self, p):
  329. self._cur_param_name = str(p.name)
  330. self._cur_fields = []
  331. self._cur_enum_names = []
  332. self._write('class %s(_ParamDefBase):', p.name, indent=1)
  333. self._write_doc(p.name)
  334. self._write('TAG = %d', p.tag)
  335. def _on_param_end(self, p):
  336. # gen slots and packer
  337. self._write('__slots__ = [%s]', ', '.join(
  338. map('"{.name}"'.format, self._cur_fields)))
  339. struct_fmt = ''.join(i.fmt for i in self._cur_fields)
  340. if not struct_fmt:
  341. struct_fmt = 'x'
  342. else:
  343. # add padding at end
  344. max_t = max(struct_fmt, key=struct.calcsize)
  345. struct_fmt += '0{}'.format(max_t)
  346. self._write('_packer = struct.Struct("%s")', struct_fmt)
  347. # gen __init__ signature
  348. self._write('def __init__(%s):',
  349. ', '.join(['self'] +
  350. list('{}={}'.format(i.name, i.default)
  351. for i in self._cur_fields)),
  352. indent=1)
  353. # gen __init__ doc
  354. self._write('"""')
  355. for i in self._cur_fields:
  356. self._write(':type {}: :class:`.{}`'.format(i.name, i.type))
  357. if i.doc:
  358. self._write(':param {}: {}'.format(i.name, i.doc))
  359. self._write('"""')
  360. # gen cvt in __init__
  361. for i in self._cur_fields:
  362. self._write('self.%s = %s', i.name, i.cvt)
  363. self._unindent()
  364. self._unindent()
  365. self._write('')
  366. def _on_member_enum(self, e):
  367. qualname = '{}.{}'.format(self._cur_param_name, e.name)
  368. self._write('class %s(_EnumBase):', e.name, indent=1)
  369. self._write_doc(e.name)
  370. for idx, emem in enumerate(e.members):
  371. self._write('%s = "%s"', emem, emem)
  372. self._write_doc(emem)
  373. self._enum_member2num.append('id({}.{}):{}'.format(
  374. qualname, emem, idx))
  375. for emem, emem_alis in e.member_alias:
  376. self._write('%s = %s', emem_alis, emem)
  377. self._unindent()
  378. self._write('')
  379. self._cur_fields.append(self.FieldDef(
  380. name=e.name_field,
  381. cvt='{}.convert({})'.format(qualname, e.name_field),
  382. fmt='I',
  383. default="'{}'".format(e.members[e.default]),
  384. type=qualname,
  385. doc=None))
  386. def _on_member_enum_alias(self, e):
  387. self._write('%s = %s.%s', e.name, e.src_class, e.src_name)
  388. s = e.src_enum
  389. qualname = '{}.{}'.format(e.src_class, e.src_name)
  390. self._cur_fields.append(self.FieldDef(
  391. name=e.name_field,
  392. cvt='{}.convert({})'.format(qualname, e.name_field),
  393. fmt='I',
  394. default="'{}'".format(s.members[e.get_default()]),
  395. type=qualname,
  396. doc=None))
  397. def _get_py_default(self, cppdefault):
  398. if not isinstance(cppdefault, str):
  399. return cppdefault
  400. d = cppdefault
  401. if d.endswith('f'): # 1.f
  402. return d[:-1]
  403. if d.endswith('ull'):
  404. return d[:-3]
  405. if d == 'false':
  406. return 'False'
  407. if d == 'true':
  408. return 'True'
  409. if d.startswith('DTypeEnum::'):
  410. return '"{}"'.format(d.split(':')[2].lower())
  411. return d
  412. def _on_member_field(self, f):
  413. d = self._get_py_default(f.default)
  414. self._cur_fields.append(self.FieldDef(
  415. name=f.name,
  416. cvt='{}({})'.format(f.dtype.pycvt, f.name),
  417. fmt=f.dtype.pyfmt,
  418. default=d,
  419. type=f.dtype.pycvt,
  420. doc=f.name.doc
  421. ))
  422. def _on_const_field(self, f):
  423. d = self._get_py_default(f.default)
  424. self._write_doc(f.name)
  425. self._write('%s = %s', f.name, d)
  426. class CPPWriter(IndentWriterBase):
  427. _param_namespace = 'param'
  428. _ctor_args = None
  429. """list of (text in func param, var name); func param name must be var name
  430. appended by an underscore"""
  431. _non_static_members = None
  432. def __call__(self, fout, defs):
  433. super().__call__(fout)
  434. self._write('// %s', self._get_header())
  435. self._write('#pragma once')
  436. self._write('#include "megdnn/dtype.h"')
  437. self._write('#include <stdint.h>')
  438. if self._param_namespace == 'param':
  439. self._write('#include <string.h>')
  440. self._write('namespace megdnn {')
  441. self._write('namespace %s {', self._param_namespace)
  442. self._process(defs)
  443. self._write('} // namespace megdnn')
  444. self._write('} // namespace %s', self._param_namespace)
  445. self._write('// vim: syntax=cpp.doxygen')
  446. def _write_doc(self, doc):
  447. assert isinstance(doc, member_defs.Doc)
  448. if not doc.doc:
  449. return
  450. if doc.no_reformat:
  451. self._write('/*')
  452. for i in doc.raw_lines:
  453. self._write('* ' + i)
  454. self._write('*/')
  455. return
  456. doc = doc.doc.replace('\n', ' ')
  457. textwidth = 80 - len(self._cur_indent) - 4
  458. if len(doc) <= textwidth:
  459. self._write('//! ' + doc)
  460. return
  461. self._write('/*!')
  462. for i in textwrap.wrap(doc, textwidth):
  463. self._write(' * ' + i)
  464. self._write(' */')
  465. def _on_param_begin(self, p):
  466. self._write_doc(p.name)
  467. self._write('struct %s {', p.name, indent=1)
  468. self._write('static MEGDNN_CONSTEXPR uint32_t TAG = %du;', p.tag)
  469. self._ctor_args = []
  470. self._non_static_members = []
  471. def _add_ctor_args(self, typename, default, varname):
  472. self._ctor_args.append((
  473. '{} {}_={}'.format(typename, varname, default),
  474. varname))
  475. def _on_param_end(self, p):
  476. '''
  477. MegDNN param structures are not packed and we need to initialize the structure
  478. paddings to zero or it would break MegBrain hash system. We do memset(0) in default
  479. ctor and use a trick, wrapping non-static members in a anonymous union which would
  480. copy the object representation in its default copy/move ctor, for copy/move ctor.
  481. > The implicitly-defined copy/move constructor for a non-union class X performs
  482. > a memberwise copy/move of its bases and members. [class.copy.ctor 14]
  483. > The implicitly-defined copy/move constructor for a union X copies the object
  484. > representation (6.9) of X. [class.copy.ctor 15]
  485. '''
  486. if self._non_static_members:
  487. self._write('union { struct {')
  488. for i in self._non_static_members:
  489. if isinstance(i, member_defs.Field):
  490. self._write_doc(i.name)
  491. self._write('%s%s %s;', i.dtype.cname_attr, i.dtype.cname, i.name)
  492. else:
  493. assert isinstance(i, (member_defs.Enum, member_defs.EnumAlias))
  494. self._write('%s %s;', i.name, i.name_field)
  495. self._write('}; };')
  496. if self._ctor_args:
  497. pdefs, varnames = zip(*self._ctor_args)
  498. self._write('%s(%s) {', p.name, ', '.join(pdefs), indent=1)
  499. self._write('memset(this, 0, sizeof(*this));')
  500. for var in varnames:
  501. self._write('this->%s = %s_;', var, var)
  502. self._write('}', indent=-1)
  503. self._write('};\n', indent=-1)
  504. def _on_member_enum(self, e):
  505. self._write_doc(e.name)
  506. self._write('enum class %s: uint32_t {', e.name, indent=1)
  507. for idx, i in enumerate(e.members):
  508. self._write_doc(i)
  509. v = '{} = {}'.format(i, idx)
  510. if i is not e.members[-1] or e.member_alias:
  511. v += ','
  512. self._write(v)
  513. for mem, alias in e.member_alias:
  514. self._write('%s = %s,', alias, mem)
  515. self._write('};', indent=-1)
  516. self._non_static_members.append(e)
  517. self._write('static MEGDNN_CONSTEXPR uint32_t %s_NR_MEMBER = %d;',
  518. str(e.name).upper(), len(e.members))
  519. self._add_ctor_args(e.name,
  520. '{}::{}'.format(e.name, e.members[e.default]),
  521. e.name_field)
  522. def _on_member_enum_alias(self, e):
  523. s = e.src_enum
  524. self._write('using %s = %s::%s;', e.name, e.src_class, e.src_name)
  525. self._non_static_members.append(e)
  526. self._write('static MEGDNN_CONSTEXPR uint32_t %s_NR_MEMBER = %d;',
  527. str(e.name).upper(), len(s.members))
  528. self._add_ctor_args(e.name,
  529. '{}::{}'.format(e.name,
  530. s.members[e.get_default()]),
  531. e.name_field)
  532. def _on_member_field(self, f):
  533. self._non_static_members.append(f)
  534. self._add_ctor_args(f.dtype.cname, f.default, f.name)
  535. def _on_const_field(self, f):
  536. self._write_doc(f.name)
  537. if 'int' in f.dtype.cname:
  538. self._write('static constexpr %s%s %s = %s;', f.dtype.cname_attr, f.dtype.cname, f.name, f.default)
  539. else:
  540. self._write('static const %s%s %s = %s;', f.dtype.cname_attr, f.dtype.cname, f.name, f.default)
  541. class CPPEnumValueWriter(CPPWriter):
  542. _param_namespace = 'param_enumv'
  543. def _on_member_enum(self, e):
  544. self._write_doc(e.name)
  545. self._write('struct %s {', e.name, indent=1)
  546. for idx, val in enumerate(e.members):
  547. self._write_doc(val)
  548. self._write('static const uint32_t %s = %d;', val, idx)
  549. for mem, alias in e.member_alias:
  550. self._write('static const uint32_t %s = %s;', alias, mem)
  551. self._write('};', indent=-1)
  552. def _on_member_enum_alias(self, e):
  553. s = e.src_enum
  554. self._write('typedef %s::%s %s;', e.src_class, e.src_name, e.name)
  555. def _on_member_field(self, f):
  556. pass
  557. def _on_const_field(self, f):
  558. pass
  559. class CPPEnumItemWriter(WriterBase):
  560. _class_name = None
  561. _enum_name = None
  562. _enable = False
  563. def __init__(self, enum_def):
  564. self._class_name, self._enum_name = enum_def.split(':')
  565. def __call__(self, fout, defs):
  566. super().__call__(fout)
  567. self._process(defs)
  568. def _on_param_begin(self, p):
  569. self._enable = p.name == self._class_name
  570. def _on_member_enum(self, e):
  571. if self._enable and e.name == self._enum_name:
  572. for i in e.members:
  573. self._fout.write('{}\n'.format(i))
  574. class CPPParamJsonFuncWriter(IndentWriterBase):
  575. _param_namespace = 'param'
  576. _param_name = None
  577. _items = None
  578. def _write_json_item(self, json_cls, field):
  579. cls2ctype = {
  580. 'NumberInt': 'int64_t',
  581. 'Number': 'double',
  582. 'Bool': 'bool',
  583. }
  584. self._items.append('{"%s", json::%s::make(static_cast<%s>(p.%s))},' % (
  585. field, json_cls, cls2ctype[json_cls], field))
  586. def __call__(self, fout, defs):
  587. super().__call__(fout)
  588. self._write('// %s', self._get_header())
  589. self._write('// this file can only be included in '
  590. 'megbrain/src/plugin/impl/opr_footprint.cpp\n'
  591. '// please do not include it directly')
  592. self._write('#include "megdnn/opr_param_defs.h"')
  593. self._write('#pragma once')
  594. self._write('using namespace megdnn;')
  595. self._write('namespace mgb {')
  596. self._write('namespace opr {')
  597. self._write('template<class OprParam>')
  598. self._write('std::shared_ptr<mgb::json::Value> opr_param_to_json(const OprParam &param);')
  599. self._process(defs)
  600. self._write('} // namespace opr')
  601. self._write('} // namespace mgb')
  602. self._write('\n// vim: syntax=cpp.doxygen')
  603. def _on_param_begin(self, p):
  604. self._write('template<>', indent=0)
  605. self._write(
  606. 'std::shared_ptr<mgb::json::Value> opr_param_to_json(const param::%s &p) {',
  607. p.name, indent=1)
  608. self._param_name = 'param::{}'.format(p.name)
  609. self._items = []
  610. def _on_param_end(self, p):
  611. self._write('return json::Object::make({', indent=1)
  612. for i in self._items:
  613. self._write(i, indent=0)
  614. self._write('});', indent=-1)
  615. self._write('}', indent=-1)
  616. def _on_member_enum(self, e):
  617. self._write('auto %s2str = [](const %s::%s arg) -> std::string {',
  618. e.name, self._param_name, e.name, indent=1)
  619. self._write('switch (arg) {', indent=1)
  620. enum2str = []
  621. if isinstance(e, member_defs.EnumAlias):
  622. members = e.src_enum.members
  623. else:
  624. members = e.members
  625. for idx, i in enumerate(members):
  626. self._write('case %s::%s::%s: return "%s";',
  627. self._param_name, e.name, i, i, indent=0)
  628. self._write('default: mgb_throw(MegBrainError, "Invalid %s::%s:%%d", static_cast<int>(arg));',
  629. self._param_name, e.name, indent=0)
  630. self._write('}', indent=-1)
  631. self._write('};', indent=-1)
  632. self._items.append('{"%s", json::String::make(%s2str(p.%s))},' % (
  633. e.name_field, e.name, e.name_field))
  634. def _on_member_enum_alias(self, e):
  635. self._on_member_enum(e)
  636. def _on_member_field(self, f):
  637. self._write_json_item(f.dtype.cppjson, f.name)
  638. def _on_const_field(self, f):
  639. pass
  640. def main():
  641. parser = argparse.ArgumentParser(
  642. 'generate opr param defs from description file')
  643. parser.add_argument('--enumv', action='store_true',
  644. help='generate c++03 compatible code which only '
  645. 'contains enum values')
  646. parser.add_argument('-t', '--type', choices=['c++', 'py'], default='c++',
  647. help='output type')
  648. parser.add_argument('--write-enum-items',
  649. help='write enum item names to output file; argument '
  650. 'should be given in the CLASS:ENUM format')
  651. parser.add_argument('--write-cppjson',
  652. help='generate megbrain json serialization implemention'
  653. 'cpp file')
  654. parser.add_argument('input')
  655. parser.add_argument('output')
  656. args = parser.parse_args()
  657. with open(args.input) as fin:
  658. inputs = fin.read()
  659. exec(inputs, {'pdef': ParamDef, 'Doc': member_defs.Doc})
  660. input_hash = hashlib.sha256()
  661. input_hash.update(inputs.encode(encoding='UTF-8'))
  662. input_hash = input_hash.hexdigest()
  663. if args.type == 'py':
  664. writer = PyWriter()
  665. else:
  666. assert args.type == 'c++'
  667. if args.enumv:
  668. writer = CPPEnumValueWriter()
  669. elif args.write_enum_items:
  670. writer = CPPEnumItemWriter(args.write_enum_items)
  671. else:
  672. writer = CPPWriter()
  673. with open(args.output, 'w') as fout:
  674. writer.set_input_hash(input_hash)(fout, ParamDef.all_param_defs)
  675. if args.write_cppjson:
  676. writer = CPPParamJsonFuncWriter()
  677. with open(args.write_cppjson, 'w') as fout:
  678. writer.set_input_hash(input_hash)(fout, ParamDef.all_param_defs)
  679. if __name__ == '__main__':
  680. main()

MegEngine 安装包中集成了使用 GPU 运行代码所需的 CUDA 环境,不用区分 CPU 和 GPU 版。 如果想要运行 GPU 程序,请确保机器本身配有 GPU 硬件设备并安装好驱动。 如果你想体验在云端 GPU 算力平台进行深度学习开发的感觉,欢迎访问 MegStudio 平台

Contributors (1)