Add atomic operations to SPIR-V dialect.

Some changes to the dialect generation script to allow specification
of different base class to derive from in ODS.

PiperOrigin-RevId: 285859230

2 files changed, 6 insertions, 5 deletions

diff --git a/mlir/utils/spirv/define_inst.sh b/mlir/utils/spirv/define_inst.sh
index 3508c4f9b4f..f11078a8e76 100755
--- a/mlir/utils/spirv/define_inst.sh
+++ b/mlir/utils/spirv/define_inst.sh
@@ -35,13 +35,13 @@ file_name=$1
 inst_category=$2
 
 case $inst_category in
-  Op | ArithmeticOp | LogicalOp | CastOp | ControlFlowOp | StructureOp)
+  Op | ArithmeticOp | LogicalOp | CastOp | ControlFlowOp | StructureOp | AtomicUpdateOp | AtomicUpdateWithValueOp)
   ;;
   *)
     echo "Usage : " $0 "<filename> <inst_category> (<opname>)*"
     echo "<filename> is the file name of MLIR SPIR-V op definitions spec"
     echo "<inst_category> must be one of " \
-      "(Op|ArithmeticOp|LogicalOp|CastOp|ControlFlowOp|StructureOp)"
+      "(Op|ArithmeticOp|LogicalOp|CastOp|ControlFlowOp|StructureOp|AtomicUpdateOp)"
     exit 1;
   ;;
 esac
diff --git a/mlir/utils/spirv/gen_spirv_dialect.py b/mlir/utils/spirv/gen_spirv_dialect.py
index bf4886dfd51..be7116c211f 100755
--- a/mlir/utils/spirv/gen_spirv_dialect.py
+++ b/mlir/utils/spirv/gen_spirv_dialect.py
@@ -353,7 +353,7 @@ def map_spec_operand_to_ods_argument(operand):
     # and 'IdScope' given that they should be generated from OpConstant.
     assert quantifier == '', ('unexpected to have optional/variadic memory '
                               'semantics or scope <id>')
-    arg_type = 'I32'
+    arg_type = 'SPV_' + kind[2:] + 'Attr'
   elif kind == 'LiteralInteger':
     if quantifier == '':
       arg_type = 'I32Attr'
@@ -651,8 +651,9 @@ def update_td_op_definitions(path, instructions, docs, filter_list,
       instruction = next(
           inst for inst in instructions if inst['opname'] == opname)
       op_defs.append(
-          get_op_definition(instruction, docs[opname],
-                            op_info_dict.get(opname, {})))
+          get_op_definition(
+              instruction, docs[opname],
+              op_info_dict.get(opname, {'inst_category': inst_category})))
     except StopIteration:
       # This is an op added by us; use the existing ODS definition.
       op_defs.append(name_op_map[opname])