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 32 kB

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

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