ptx.py 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120
  1. """
  2. pygments.lexers.ptx
  3. ~~~~~~~~~~~~~~~~~~~
  4. Lexer for other PTX language.
  5. :copyright: Copyright 2006-2023 by the Pygments team, see AUTHORS.
  6. :license: BSD, see LICENSE for details.
  7. """
  8. from pygments.lexer import RegexLexer, include, words
  9. from pygments.token import Comment, Keyword, Name, String, Number, \
  10. Punctuation, Whitespace, Operator
  11. __all__ = ["PtxLexer"]
  12. class PtxLexer(RegexLexer):
  13. """
  14. For NVIDIA `PTX <https://docs.nvidia.com/cuda/parallel-thread-execution/>`_
  15. source.
  16. .. versionadded:: 2.16
  17. """
  18. name = 'PTX'
  19. url = "https://docs.nvidia.com/cuda/parallel-thread-execution/"
  20. filenames = ['*.ptx']
  21. aliases = ['ptx']
  22. mimetypes = ['text/x-ptx']
  23. #: optional Comment or Whitespace
  24. string = r'"[^"]*?"'
  25. followsym = r'[a-zA-Z0-9_$]'
  26. identifier = r'([-a-zA-Z$._][\w\-$.]*|' + string + ')'
  27. block_label = r'(' + identifier + r'|(\d+))'
  28. tokens = {
  29. 'root': [
  30. include('whitespace'),
  31. (block_label + r'\s*:', Name.Label),
  32. include('keyword'),
  33. (r'%' + identifier, Name.Variable),
  34. (r'%\d+', Name.Variable.Anonymous),
  35. (r'c?' + string, String),
  36. (identifier, Name.Variable),
  37. (r';', Punctuation),
  38. (r'[*+-/]', Operator),
  39. (r'0[xX][a-fA-F0-9]+', Number),
  40. (r'-?\d+(?:[.]\d+)?(?:[eE][-+]?\d+(?:[.]\d+)?)?', Number),
  41. (r'[=<>{}\[\]()*.,!]|x\b', Punctuation)
  42. ],
  43. 'whitespace': [
  44. (r'(\n|\s+)+', Whitespace),
  45. (r'//.*?\n', Comment)
  46. ],
  47. 'keyword': [
  48. # Instruction keywords
  49. (words((
  50. 'abs', 'discard', 'min', 'shf', 'vadd',
  51. 'activemask', 'div', 'mma', 'shfl', 'vadd2',
  52. 'add', 'dp2a', 'mov', 'shl', 'vadd4',
  53. 'addc', 'dp4a', 'movmatrix', 'shr', 'vavrg2',
  54. 'alloca', 'elect', 'mul', 'sin', 'vavrg4',
  55. 'and', 'ex2', 'mul24', 'slct', 'vmad',
  56. 'applypriority', 'exit', 'multimem', 'sqrt', 'vmax',
  57. 'atom', 'fence', 'nanosleep', 'st', 'vmax2',
  58. 'bar', 'fma', 'neg', 'stackrestore', 'vmax4',
  59. 'barrier', 'fns', 'not', 'stacksave', 'vmin',
  60. 'bfe', 'getctarank', 'or', 'stmatrix', 'vmin2',
  61. 'bfi', 'griddepcontrol', 'pmevent', 'sub', 'vmin4',
  62. 'bfind', 'isspacep', 'popc', 'subc', 'vote',
  63. 'bmsk', 'istypep', 'prefetch', 'suld', 'vset',
  64. 'bra', 'ld', 'prefetchu', 'suq', 'vset2',
  65. 'brev', 'ldmatrix', 'prmt', 'sured', 'vset4',
  66. 'brkpt', 'ldu', 'rcp', 'sust', 'vshl',
  67. 'brx', 'lg2', 'red', 'szext', 'vshr',
  68. 'call', 'lop3', 'redux', 'tanh', 'vsub',
  69. 'clz', 'mad', 'rem', 'testp', 'vsub2',
  70. 'cnot', 'mad24', 'ret', 'tex', 'vsub4',
  71. 'copysign', 'madc', 'rsqrt', 'tld4', 'wgmma',
  72. 'cos', 'mapa', 'sad', 'trap', 'wmma',
  73. 'cp', 'match', 'selp', 'txq', 'xor',
  74. 'createpolicy', 'max', 'set', 'vabsdiff', 'cvt',
  75. 'mbarrier', 'setmaxnreg', 'vabsdiff2', 'cvta',
  76. 'membar', 'setp', 'vabsdiff4')), Keyword),
  77. # State Spaces and Suffixes
  78. (words((
  79. 'reg', '.sreg', '.const', '.global',
  80. '.local', '.param', '.shared', '.tex',
  81. '.wide', '.loc'
  82. )), Keyword.Pseudo),
  83. # PTX Directives
  84. (words((
  85. '.address_size', '.explicitcluster', '.maxnreg', '.section',
  86. '.alias', '.extern', '.maxntid', '.shared',
  87. '.align', '.file', '.minnctapersm', '.sreg',
  88. '.branchtargets', '.func', '.noreturn', '.target',
  89. '.callprototype', '.global', '.param', '.tex',
  90. '.calltargets', '.loc', '.pragma', '.version',
  91. '.common', '.local', '.reg', '.visible',
  92. '.const', '.maxclusterrank', '.reqnctapercluster', '.weak',
  93. '.entry', '.maxnctapersm', '.reqntid')), Keyword.Reserved),
  94. # Fundamental Types
  95. (words((
  96. '.s8', '.s16', '.s32', '.s64',
  97. '.u8', '.u16', '.u32', '.u64',
  98. '.f16', '.f16x2', '.f32', '.f64',
  99. '.b8', '.b16', '.b32', '.b64',
  100. '.pred'
  101. )), Keyword.Type)
  102. ],
  103. }