Update the gen_spirv_dialect.py script to add opcodes from the SPIR-V

JSON spec into the SPIRBase.td file. This is done incrementally to
only import those opcodes that are needed, through use of the script
define_opcode.sh added.

PiperOrigin-RevId: 257517343

2 files changed, 121 insertions, 2 deletions

diff --git a/mlir/utils/spirv/define_opcodes.sh b/mlir/utils/spirv/define_opcodes.sh
new file mode 100755
index 00000000000..e98d55c7128
--- /dev/null
+++ b/mlir/utils/spirv/define_opcodes.sh
@@ -0,0 +1,39 @@
+#!/bin/bash
+
+# Copyright 2019 The MLIR Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Script for defining map for opname to opcode using SPIR-V spec from the
+# Internet
+#
+# Run as:
+# ./define_opcode.sh (<op-name>)*
+#
+# For example:
+# ./define_opcode.sh OpTypeVoid OpTypeFunction
+#
+# If no op-name is specified, the existing opcodes are updated
+#
+# The 'instructions' list of spirv.core.grammar.json contains all instructions
+# in SPIR-V
+
+set -e
+set -x
+
+current_file="$(readlink -f "$0")"
+current_dir="$(dirname "$current_file")"
+
+python3 ${current_dir}/gen_spirv_dialect.py \
+  --base-td-path ${current_dir}/../../include/mlir/SPIRV/SPIRVBase.td \
+  --new-opcode $@
diff --git a/mlir/utils/spirv/gen_spirv_dialect.py b/mlir/utils/spirv/gen_spirv_dialect.py
index 88f640c6ef8..8279f70765b 100755
--- a/mlir/utils/spirv/gen_spirv_dialect.py
+++ b/mlir/utils/spirv/gen_spirv_dialect.py
@@ -32,6 +32,8 @@ SPIRV_HTML_SPEC_URL = 'https://www.khronos.org/registry/spir-v/specs/unified1/SP
 SPIRV_JSON_SPEC_URL = 'https://raw.githubusercontent.com/KhronosGroup/SPIRV-Headers/master/include/spirv/unified1/spirv.core.grammar.json'
 
 AUTOGEN_ENUM_SECTION_MARKER = 'enum section. Generated from SPIR-V spec; DO NOT MODIFY!'
+AUTOGEN_INSTRUCTION_OPCODE_SECTION_MARKER = ('opcode section. Generated from '
+                                             'SPIR-V spec; DO NOT MODIFY!')
 
 
 def get_spirv_grammar_from_json_spec():
@@ -129,6 +131,76 @@ def gen_operand_kind_enum_attr(operand_kind):
   return kind_name, case_defs + '\n\n' + enum_attr
 
 
+def gen_opcode(instructions):
+  """ Generates the TableGen definition to map opname to opcode
+
+   Returns:
+       - A string containing the TableGen SPV_OpCode definition
+   """
+
+  max_len = max([len(inst['opname']) for inst in instructions])
+  def_fmt_str = 'def SPV_OC_{name} {colon:>{offset}} '\
+            'I32EnumAttrCase<"{name}", {value}>;'
+  opcode_defs = [
+      def_fmt_str.format(
+          name=inst['opname'],
+          value=inst['opcode'],
+          colon=':',
+          offset=(max_len + 1 - len(inst['opname']))) for inst in instructions
+  ]
+  opcode_str = '\n'.join(opcode_defs)
+
+  decl_fmt_str = 'SPV_OC_{name}'
+  opcode_list = [
+      decl_fmt_str.format(name=inst['opname']) for inst in instructions
+  ]
+  opcode_list = split_list_into_sublists(opcode_list, 6)
+  opcode_list = [
+      '{:6}'.format('') + ', '.join(sublist) for sublist in opcode_list
+  ]
+  opcode_list = ',\n'.join(opcode_list)
+  enum_attr = 'def SPV_OpcodeAttr :\n'\
+              '    I32EnumAttr<"{name}", "valid SPIR-V instructions", [\n'\
+              '{lst}\n'\
+              '      ]> {{\n'\
+              '    let returnType = "::mlir::spirv::{name}";\n'\
+              '    let convertFromStorage = '\
+              '"static_cast<::mlir::spirv::{name}>($_self.getInt())";\n'\
+              '    let cppNamespace = "::mlir::spirv";\n}}'.format(
+                  name='Opcode', lst=opcode_list)
+  return opcode_str + '\n\n' + enum_attr
+
+
+def update_td_opcodes(path, instructions, filter_list):
+
+  with open(path, 'r') as f:
+    content = f.read()
+
+  content = content.split(AUTOGEN_INSTRUCTION_OPCODE_SECTION_MARKER)
+  assert len(content) == 3
+
+  # Extend opcode list with existing list
+  import re
+  existing_opcodes = [k[11:] for k in re.findall('def SPV_OC_\w+', content[1])]
+  filter_list.extend(existing_opcodes)
+  filter_list = list(set(filter_list))
+
+  # Generate the opcode for all instructions in SPIR-V
+  filter_instrs = list(
+      filter(lambda inst: (inst['opname'] in filter_list), instructions))
+  # Sort instruction based on opcode
+  filter_instrs.sort(key=lambda inst: inst['opcode'])
+  opcode = gen_opcode(filter_instrs)
+
+  # Substitute the opcode
+  content = content[0] + AUTOGEN_INSTRUCTION_OPCODE_SECTION_MARKER + '\n\n' + \
+        opcode + '\n\n// End ' + AUTOGEN_INSTRUCTION_OPCODE_SECTION_MARKER \
+        + content[2]
+
+  with open(path, 'w') as f:
+    f.write(content)
+
+
 def update_td_enum_attrs(path, operand_kinds, filter_list):
   """Updates SPIRBase.td with new generated enum definitions.
 
@@ -136,7 +208,7 @@ def update_td_enum_attrs(path, operand_kinds, filter_list):
         - path: the path to SPIRBase.td
         - operand_kinds: a list containing all operand kinds' grammar
         - filter_list: a list containing new enums to add
-    """
+  """
   with open(path, 'r') as f:
     content = f.read()
 
@@ -175,8 +247,16 @@ if __name__ == '__main__':
                           help='Path to SPIRVBase.td')
   cli_parser.add_argument('--new-enum', dest='new_enum', type=str,
                           help='SPIR-V enum to be added to SPIRVBase.td')
+  cli_parser.add_argument(
+      '--new-opcodes',
+      dest='new_opcodes',
+      type=str,
+      nargs='*',
+      help='update SPIR-V opcodes in SPIRVBase.td')
   args = cli_parser.parse_args()
 
-  operand_kinds, _ = get_spirv_grammar_from_json_spec()
+  operand_kinds, instructions = get_spirv_grammar_from_json_spec()
 
   update_td_enum_attrs(args.base_td_path, operand_kinds, [args.new_enum])
+
+  update_td_opcodes(args.base_td_path, instructions, args.new_opcodes)