[mlir][spirv] Add memory semantics verify for atomic operations

[xndcn]

Summary

Add memory semantics verify for atomic operations
also fix wrong path of SPIRVBase.td in report_coverage.sh

[antiagainst]

Nice, thanks!

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBarrierOps.td

def SPV_ControlBarrierOp : SPV_Op<"ControlBarrier", []> {
   let arguments = (ins
     SPV_ScopeAttr:$execution_scope,
     SPV_ScopeAttr:$memory_scope,
     SPV_MemorySemanticsAttr:$memory_semantics
   );
 
   let results = (outs);
 
-  let verifier = [{ return verifyMemorySemantics(*this); }];
+  let verifier = [{ return verifyMemorySemantics(getOperation(), memory_semantics()); }];
 
   let autogenSerialization = 0;
 
   let assemblyFormat = [{
     $execution_scope `,` $memory_scope `,` $memory_semantics attr-dict
   }];
 }
 
def SPV_MemoryBarrierOp : SPV_Op<"MemoryBarrier", []> {
 
   let arguments = (ins
     SPV_ScopeAttr:$memory_scope,
     SPV_MemorySemanticsAttr:$memory_semantics
   );
 
   let results = (outs);
 
-  let verifier = [{ return verifyMemorySemantics(*this); }];
+  let verifier = [{ return verifyMemorySemantics(getOperation(), memory_semantics()); }];
 
   let autogenSerialization = 0;
 
   let assemblyFormat = "$memory_scope `,` $memory_semantics attr-dict";
 }
 
 #endif // MLIR_DIALECT_SPIRV_IR_BARRIER_OPS

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

if (op->getAttr(kSourceAlignmentAttrName)) {
       return memoryOp.emitOpError(
           "invalid alignment specification with non-aligned memory access "
           "specification");
     }
   }
   return success();
 }
 
-template <typename BarrierOp>
-static LogicalResult verifyMemorySemantics(BarrierOp op) {
+static LogicalResult
+verifyMemorySemantics(Operation *op, spirv::MemorySemantics memorySemantics) {
   // According to the SPIR-V specification:
   // "Despite being a mask and allowing multiple bits to be combined, it is
   // invalid for more than one of these four bits to be set: Acquire, Release,
   // AcquireRelease, or SequentiallyConsistent. Requesting both Acquire and
   // Release semantics is done by setting the AcquireRelease bit, not by setting
   // two bits."
-  auto memorySemantics = op.memory_semantics();
   auto atMostOneInSet = spirv::MemorySemantics::Acquire |
                         spirv::MemorySemantics::Release |
                         spirv::MemorySemantics::AcquireRelease |
                         spirv::MemorySemantics::SequentiallyConsistent;
 
   auto bitCount = llvm::countPopulation(
       static_cast<uint32_t>(memorySemantics & atMostOneInSet));
   if (bitCount > 1) {
-    return op.emitError("expected at most one of these four memory constraints "
-                        "to be set: `Acquire`, `Release`,"
-                        "`AcquireRelease` or `SequentiallyConsistent`");
+    return op->emitError(
+        "expected at most one of these four memory constraints "
   }
   return success();
 }
 
 template <typename LoadStoreOpTy>
 static LogicalResult verifyLoadStorePtrAndValTypes(LoadStoreOpTy op, Value ptr,
                                                    Value val) {
   // ODS already checks ptr is spirv::PointerType. Just check that the pointee
static LogicalResult verifyAtomicUpdateOp(Operation *op) {
 
   if (op->getNumOperands() > 1) {
     auto valueType = op->getOperand(1).getType();
     if (valueType != elementType)
       return op->emitOpError("expected value to have the same type as the "
                              "pointer operand's pointee type ")
              << elementType << ", but found " << valueType;
   }
+  auto memorySemantics = static_cast<spirv::MemorySemantics>(
+      op->getAttrOfType<IntegerAttr>(kSemanticsAttrName).getInt());
+  if (failed(verifyMemorySemantics(op, memorySemantics))) {
+    return failure();
+  }
   return success();
 }
 
 static ParseResult parseGroupNonUniformArithmeticOp(OpAsmParser &parser,
                                                     OperationState &state) {
   spirv::Scope executionScope;
   spirv::GroupOperation groupOperation;
   OpAsmParser::OperandType valueInfo;

mlir/test/Dialect/SPIRV/IR/atomic-ops.mlir

 func @atomic_and(%ptr : !spv.ptr<i32, StorageBuffer>, %value : i64) -> i64 {
   // expected-error @+1 {{expected value to have the same type as the pointer operand's pointee type 'i32', but found 'i64'}}
   %0 = "spv.AtomicAnd"(%ptr, %value) {memory_scope = 2: i32, semantics = 0x8 : i32} : (!spv.ptr<i32, StorageBuffer>, i64) -> (i64)
   return %0 : i64
 }
 
 // -----
 
+func @atomic_and(%ptr : !spv.ptr<i32, StorageBuffer>, %value : i32) -> i32 {
+  // expected-error @+1 {{expected at most one of these four memory constraints to be set: `Acquire`, `Release`,`AcquireRelease` or `SequentiallyConsistent`}}
+  %0 = spv.AtomicAnd "Device" "Acquire|Release" %ptr, %value : !spv.ptr<i32, StorageBuffer>
+  return %0 : i32
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // spv.AtomicCompareExchangeWeak
 //===----------------------------------------------------------------------===//
 
 func @atomic_compare_exchange_weak(%ptr: !spv.ptr<i32, Workgroup>, %value: i32, %comparator: i32) -> i32 {
   // CHECK: spv.AtomicCompareExchangeWeak "Workgroup" "Release" "Acquire" %{{.*}}, %{{.*}}, %{{.*}} : !spv.ptr<i32, Workgroup>
   %0 = spv.AtomicCompareExchangeWeak "Workgroup" "Release" "Acquire" %ptr, %value, %comparator: !spv.ptr<i32, Workgroup>
   return %0: i32

mlir/utils/spirv/report_coverage.sh

 # ./report_coverage.sh
 
 set -e
 
 current_file="$(readlink -f "$0")"
 current_dir="$(dirname "$current_file")"
 
 python3 ${current_dir}/gen_spirv_dialect.py \
-  --base-td-path ${current_dir}/../../include/mlir/Dialect/SPIRV/SPIRVBase.td \
+  --base-td-path ${current_dir}/../../include/mlir/Dialect/SPIRV/IR/SPIRVBase.td \
   --gen-inst-coverage