[mlir] gen_spirv_dialect.py: Some support for OCL ops generation

[Hardcode84]

Summary

It is not complete and disabled by default, but it can be still useful.

[antiagainst]

Cool, thanks for enabling automatically pulling in OpenCL intrinsics!

mlir/utils/spirv/gen_spirv_dialect.py

33

Can you turn this into a proper command line options? We use argparse for this already. Then use function parameters to pass such options to where you'd need them.

41

Can you make this as a function parameter default?

1042

Can you move this constant URLs to the beginning like the core spec ones?

[Hardcode84]

added proper command line option, moved constants

mlir/utils/spirv/gen_spirv_dialect.py

 SPIRV_HTML_SPEC_URL = 'https://www.khronos.org/registry/spir-v/specs/unified1/SPIRV.html'
 SPIRV_JSON_SPEC_URL = 'https://raw.githubusercontent.com/KhronosGroup/SPIRV-Headers/master/include/spirv/unified1/spirv.core.grammar.json'
 
 AUTOGEN_OP_DEF_SEPARATOR = '\n// -----\n\n'
 AUTOGEN_ENUM_SECTION_MARKER = 'enum section. Generated from SPIR-V spec; DO NOT MODIFY!'
 AUTOGEN_OPCODE_SECTION_MARKER = (
     'opcode section. Generated from SPIR-V spec; DO NOT MODIFY!')
 
+PARSE_OCL_EXT_OPS=False
 
-def get_spirv_doc_from_html_spec():
+def get_spirv_doc_from_html_spec(url):
   """Extracts instruction documentation from SPIR-V HTML spec.
 
   Returns:
     - A dict mapping from instruction opcode to documentation.
   """
-  response = requests.get(SPIRV_HTML_SPEC_URL)
+  if url is None:
+    url = SPIRV_HTML_SPEC_URL
+
+  response = requests.get(url)
   spec = response.content
 
   from bs4 import BeautifulSoup
   spirv = BeautifulSoup(spec, 'html.parser')
 
-  section_anchor = spirv.find('h3', {'id': '_a_id_instructions_a_instructions'})
-
   doc = {}
 
+  if PARSE_OCL_EXT_OPS:
+    section_anchor = spirv.find('h2', {'id': '_a_id_binary_a_binary_form'})
+    for section in section_anchor.parent.find_all('div', {'class': 'sect2'}):
+      for table in section.find_all('table'):
+        inst_html = table.tbody.tr.td
+        opname = inst_html.a['id']
+        # Ignore the first line, which is just the opname.
+        doc[opname] = inst_html.text.split('\n', 1)[1].strip()
+  else:
+    section_anchor = spirv.find('h3', {'id': '_a_id_instructions_a_instructions'})
     for section in section_anchor.parent.find_all('div', {'class': 'sect3'}):
       for table in section.find_all('table'):
         inst_html = table.tbody.tr.td.p
         opname = inst_html.a['id']
         # Ignore the first line, which is just the opname.
         doc[opname] = inst_html.text.split('\n', 1)[1].strip()
 
   return doc
 
 
-def get_spirv_grammar_from_json_spec():
+def get_spirv_grammar_from_json_spec(url):
   """Extracts operand kind and instruction grammar from SPIR-V JSON spec.
 
   Returns:
     - A list containing all operand kinds' grammar
     - A list containing all instructions' grammar
   """
   response = requests.get(SPIRV_JSON_SPEC_URL)
   spec = response.content
 
   import json
   spirv = json.loads(spec)
 
+  if url is None:
     return spirv['operand_kinds'], spirv['instructions']
 
+  response_ext = requests.get(url)
+  spec_ext = response_ext.content
+  spirv_ext = json.loads(spec_ext)
+
+  return spirv['operand_kinds'], spirv_ext['instructions']
+
 
 def split_list_into_sublists(items):
   """Split the list of items into multiple sublists.
 
   This is to make sure the string composed from each sublist won't exceed
   80 characters.
 
   Arguments:
def get_description(text, appendix):
 
   Returns:
     - A string that corresponds to the description of the Tablegen op.
   """
   fmt_str = '{text}\n\n    <!-- End of AutoGen section -->\n{appendix}\n  '
   return fmt_str.format(text=text, appendix=appendix)
 
 
-def get_op_definition(instruction, doc, existing_info, capability_mapping):
+def get_op_definition(instruction, opname, doc, existing_info, capability_mapping):
   """Generates the TableGen op definition for the given SPIR-V instruction.
 
   Arguments:
     - instruction: the instruction's SPIR-V JSON grammar
     - doc: the instruction's SPIR-V HTML doc
     - existing_info: a dict containing potential manually specified sections for
       this instruction
     - capability_mapping: mapping from duplicated capability symbols to the
                    canonicalized symbol chosen for SPIRVBase.td
 
   Returns:
     - A string containing the TableGen op definition
   """
+  if PARSE_OCL_EXT_OPS:
     fmt_str = ('def SPV_{opname}Op : '
-             'SPV_{inst_category}<"{opname}"{category_args}[{traits}]> '
+               'SPV_{inst_category}<"{opname_src}", {opcode}, <<Insert result type>> > '
                '{{\n  let summary = {summary};\n\n  let description = '
                '[{{\n{description}}}];{availability}\n')
+  else:
+    fmt_str = ('def SPV_{opname_src}Op : '
+               'SPV_{inst_category}<"{opname_src}"{category_args}[{traits}]> '
+               '{{\n  let summary = {summary};\n\n  let description = '
+               '[{{\n{description}}}];{availability}\n')
+
   inst_category = existing_info.get('inst_category', 'Op')
   if inst_category == 'Op':
     fmt_str +='\n  let arguments = (ins{args});\n\n'\
               '  let results = (outs{results});\n'
 
   fmt_str +='{extras}'\
             '}}\n'
 
-  opname = instruction['opname'][2:]
+  opname_src = instruction['opname']
+  if opname.startswith('Op'):
+    opname_src = opname_src[2:]
+
   category_args = existing_info.get('category_args', '')
 
   if '\n' in doc:
     summary, text = doc.split('\n', 1)
   else:
     summary = doc
     text = ''
   wrapper = textwrap.TextWrapper(
assembly = '\n    ```\n'\
                '    #### Example:\n\n'\
                '    ```\n'\
                '    [TODO]\n' \
                '    ```'
     description = get_description(text, assembly)
 
   return fmt_str.format(
       opname=opname,
+      opname_src=opname_src,
+      opcode=instruction['opcode'],
       category_args=category_args,
       inst_category=inst_category,
       traits=existing_info.get('traits', ''),
       summary=summary,
       description=description,
       availability=avail,
       args=arguments,
       results=results,
for op in ops:
     info_dict = extract_td_op_info(op)
     opname = info_dict['opname']
     name_op_map[opname] = op
     op_info_dict[opname] = info_dict
     filter_list.append(opname)
   filter_list = sorted(list(set(filter_list)))
 
   op_defs = []
+
+  if PARSE_OCL_EXT_OPS:
+    fix_opname = lambda src: src.replace('OCL','').lower()
+  else:
+    fix_opname = lambda src: src
+
   for opname in filter_list:
     # Find the grammar spec for this op
     try:
+      fixed_opname = fix_opname(opname)
       instruction = next(
-          inst for inst in instructions if inst['opname'] == opname)
+          inst for inst in instructions if inst['opname'] == fixed_opname)
+
       op_defs.append(
           get_op_definition(
-              instruction, docs[opname],
+              instruction, opname, docs[fixed_opname],
               op_info_dict.get(opname, {'inst_category': inst_category}),
               capability_mapping))
     except StopIteration:
       # This is an op added by us; use the existing ODS definition.
       op_defs.append(name_op_map[opname])
 
   # Substitute the old op definitions
   op_defs = [header] + op_defs + [footer]
cli_parser.add_argument(
       default='Op',
       help='SPIR-V instruction category used for choosing '\
            'the TableGen base class to define this op')
   cli_parser.add_argument('--gen-inst-coverage', dest='gen_inst_coverage', action='store_true')
   cli_parser.set_defaults(gen_inst_coverage=False)
 
   args = cli_parser.parse_args()
 
-  operand_kinds, instructions = get_spirv_grammar_from_json_spec()
+  if PARSE_OCL_EXT_OPS:
+    ext_html_url = 'https://www.khronos.org/registry/SPIR-V/specs/unified1/OpenCL.ExtendedInstructionSet.100.html'
+    ext_json_url = 'https://raw.githubusercontent.com/KhronosGroup/SPIRV-Headers/master/include/spirv/unified1/extinst.opencl.std.100.grammar.json'
+  else:
+    ext_html_url = None
+    ext_json_url = None
+
+  operand_kinds, instructions = get_spirv_grammar_from_json_spec(ext_json_url)
 
   # Define new enum attr
   if args.new_enum is not None:
     assert args.base_td_path is not None
     filter_list = [args.new_enum] if args.new_enum else []
     update_td_enum_attrs(args.base_td_path, operand_kinds, filter_list)
 
   # Define new opcode
   if args.new_opcodes is not None:
     assert args.base_td_path is not None
     update_td_opcodes(args.base_td_path, instructions, args.new_opcodes)
 
   # Define new op
   if args.new_inst is not None:
     assert args.op_td_path is not None
-    docs = get_spirv_doc_from_html_spec()
+    docs = get_spirv_doc_from_html_spec(ext_html_url)
     capability_mapping = get_capability_mapping(operand_kinds)
     update_td_op_definitions(args.op_td_path, instructions, docs, args.new_inst,
                              args.inst_category, capability_mapping)
     print('Done. Note that this script just generates a template; ', end='')
     print('please read the spec and update traits, arguments, and ', end='')
     print('results accordingly.')
 
   if args.gen_inst_coverage:
     gen_instr_coverage_report(args.base_td_path, instructions)