[MLIR][SPIRV] Support two memory access attributes in OpCopyMemory.

[ergawy]

Summary

This commit augments spv.CopyMemory's implementation to support 2 memory
access operands. Hence, more closely following the spec. The following
changes are introduces:

• Customize logic for spv.CopyMemory serialization and deserialization.
• Add 2 additional attributes for source memory access operand.

[ergawy]

• Add TODO to verify the case of 2 two memroy access operands.

[ergawy]

• Better style to test verification failure.

[ergawy]

• Update CopyMemory production rule in docs.

[antiagainst]

Nice! Just some minor issues.

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

194

I think we typically put mandated parameters first?

267

Same here.

mlir/test/Dialect/SPIRV/ops.mlir

1270

We can omit the -> () part actually. (Sorry existing test cases have that due to historical reasons.)

[ergawy]

• Add TODO to verify the case of 2 two memroy access operands.
• Better style to test verification failure.
• Update CopyMemory production rule in docs.
• Handle review comments:

[ergawy]

@antiagainst Thanks for the review. Handled the issues.

[ergawy]

It seems that I still don't have commit access. Just to be sure that I didn't do anything wrong. Here what I tried:
1 - I squashed all the commits in my local dev branch into one. Also in the commit messages cleaned up the unnecessary Arcanist tags except for Differential Revision: ...
2 - Rebased on top of master.
3 - Changed the active branch to master.
4 - Cherry picked the squashed commit from step 1.
5 - Tried to: git push.

What I get is:

remote: Permission to llvm/llvm-project.git denied to KareemErgawy.
fatal: unable to access 'https://github.com/llvm/llvm-project.git/': The requested URL returned error: 403

[bondhugula]

Please update commit title: 'Support 2' -> 'Support two'

[ergawy]

Update commit title.

[antiagainst]

In D82710#2127288, @ergawy wrote:

It seems that I still don't have commit access. Just to be sure that I didn't do anything wrong. Here what I tried:
1 - I squashed all the commits in my local dev branch into one. Also in the commit messages cleaned up the unnecessary Arcanist tags except for Differential Revision: ...
2 - Rebased on top of master.
3 - Changed the active branch to master.
4 - Cherry picked the squashed commit from step 1.
5 - Tried to: git push.

What I get is:

remote: Permission to llvm/llvm-project.git denied to KareemErgawy.
fatal: unable to access 'https://github.com/llvm/llvm-project.git/': The requested URL returned error: 403

Yup. There is a procedure to request direct write access to the mono repo: https://llvm.org/docs/DeveloperPolicy.html#new-contributors. Typically one needs to accumulate a few commits first before asking. I can help landing CLs before that. :)

[antiagainst]

Hey @ergawy, I reverted this in https://reviews.llvm.org/D83075. This causes a VS2017 MSVC internal compiler bug: http://lab.llvm.org:8011/builders/mlir-windows/builds/4147/steps/build-unified-tree/logs/stdio. I don't know the exact reason; maybe due to the templating or decltype? Might want to cherry-pick the original patch and then modify it to avoid using those "advanced" features to give it a try.

I can help to take a look after coming back from the long weekend. Sorry about the back and forth!

[ergawy]

In D82710#2129067, @antiagainst wrote:

Hey @ergawy, I reverted this in https://reviews.llvm.org/D83075. This causes a VS2017 MSVC internal compiler bug: http://lab.llvm.org:8011/builders/mlir-windows/builds/4147/steps/build-unified-tree/logs/stdio. I don't know the exact reason; maybe due to the templating or decltype? Might want to cherry-pick the original patch and then modify it to avoid using those "advanced" features to give it a try.

I can help to take a look after coming back from the long weekend. Sorry about the back and forth!

That's surprising! And here I am thinking I don't know enough "advanced" C++, turns out I know enough to break a compiler :P.

No worries, enjoy your weekend. Will try without the potentially problematic features.

[rriddle]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

205

nit: Drop trivial braces.

2902

Same here and below.

2937

nit: Please do not put user names in TODO comments.

[ergawy]

@rriddle Thanks for your comments. I handled them but in this review: https://reviews.llvm.org/D83241. This one had to be reverted due to some internal error in MSVC.

[rriddle]

In D82710#2135063, @ergawy wrote:

@rriddle Thanks for your comments. I handled them but in this review: https://reviews.llvm.org/D83241. This one had to be reverted due to some internal error in MSVC.

No worries, thanks for updating the other review.

mlir/include/mlir/Dialect/SPIRV/SPIRVOps.td

let description = [{
     first applies to Target and cannot include MakePointerVisible, and the
     second applies to Source and cannot include MakePointerAvailable.
 
     <!-- End of AutoGen section -->
 
     ```
     copy-memory-op ::= `spv.CopyMemory ` storage-class ssa-use
                        storage-class ssa-use
-                       (`[` memory-access `]`)?
+                       (`[` memory-access `]` (`, [` memory-access `]`)?)?
                        ` : ` spirv-element-type
     ```
 
     #### Example:
 
     ```mlir
     %0 = spv.Variable : !spv.ptr<f32, Function>
     %1 = spv.Variable : !spv.ptr<f32, Function>
     spv.CopyMemory "Function" %0, "Function" %1 : f32
     ```
   }];
 
   let arguments = (ins
     SPV_AnyPtr:$target,
     SPV_AnyPtr:$source,
     OptionalAttr<SPV_MemoryAccessAttr>:$memory_access,
-    OptionalAttr<I32Attr>:$alignment
+    OptionalAttr<I32Attr>:$alignment,
+    OptionalAttr<SPV_MemoryAccessAttr>:$source_memory_access,
+    OptionalAttr<I32Attr>:$source_alignment
   );
 
   let results = (outs);
 
   let verifier = [{ return verifyCopyMemory(*this); }];
+
+  let autogenSerialization = 0;
 }
 
 // -----
 
 def SPV_ExecutionModeOp : SPV_Op<"ExecutionMode", [InModuleScope]> {
   let summary = "Declare an execution mode for an entry point.";
 
   let description = [{

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/bit.h"
 
 using namespace mlir;
 
 // TODO(antiagainst): generate these strings using ODS.
+static constexpr const char kMemoryAccessAttrName[] = "memory_access";
+static constexpr const char kSourceMemoryAccessAttrName[] =
+    "source_memory_access";
 static constexpr const char kAlignmentAttrName[] = "alignment";
+static constexpr const char kSourceAlignmentAttrName[] = "source_alignment";
 static constexpr const char kBranchWeightAttrName[] = "branch_weights";
 static constexpr const char kCallee[] = "callee";
 static constexpr const char kClusterSize[] = "cluster_size";
 static constexpr const char kDefaultValueAttrName[] = "default_value";
 static constexpr const char kExecutionScopeAttrName[] = "execution_scope";
 static constexpr const char kEqualSemanticsAttrName[] = "equal_semantics";
 static constexpr const char kFnNameAttrName[] = "fn";
 static constexpr const char kGroupOperationAttrName[] = "group_operation";
parseEnumKeywordAttr(EnumClass &value, OpAsmParser &parser,
   if (parseEnumKeywordAttr(value, parser)) {
     return failure();
   }
   state.addAttribute(attrName, parser.getBuilder().getI32IntegerAttr(
                                    llvm::bit_cast<int32_t>(value)));
   return success();
 }
 
+template <const char memoryAccessAttrName[] = kMemoryAccessAttrName,
+          const char alignmentAttrName[] = kAlignmentAttrName>
 static ParseResult parseMemoryAccessAttributes(OpAsmParser &parser,
                                                OperationState &state) {
   // Parse an optional list of attributes staring with '['
   if (parser.parseOptionalLSquare()) {
     // Nothing to do
     return success();
   }
 
   spirv::MemoryAccess memoryAccessAttr;
-  if (parseEnumStrAttr(memoryAccessAttr, parser, state)) {
+  if (parseEnumStrAttr(memoryAccessAttr, parser, state, memoryAccessAttrName)) {
     return failure();
   }
 
   if (spirv::bitEnumContains(memoryAccessAttr, spirv::MemoryAccess::Aligned)) {
     // Parse integer attribute for alignment.
     Attribute alignmentAttr;
     Type i32Type = parser.getBuilder().getIntegerType(32);
     if (parser.parseComma() ||
-        parser.parseAttribute(alignmentAttr, i32Type, kAlignmentAttrName,
+        parser.parseAttribute(alignmentAttr, i32Type, alignmentAttrName,
                               state.attributes)) {
       return failure();
     }
   }
   return parser.parseRSquare();
 }
 
-template <typename MemoryOpTy>
-static void
-printMemoryAccessAttribute(MemoryOpTy memoryOp, OpAsmPrinter &printer,
-                           SmallVectorImpl<StringRef> &elidedAttrs) {
+template <const char memoryAccessAttrName[] = kMemoryAccessAttrName,
+          const char alignmentAttrName[] = kAlignmentAttrName,
+          bool first = true, typename MemoryOpTy>
+static void printMemoryAccessAttribute(
+    MemoryOpTy memoryOp, OpAsmPrinter &printer,
+    SmallVectorImpl<StringRef> &elidedAttrs,
+    Optional<spirv::MemoryAccess> memoryAccessAtrrValue = None,
+    Optional<llvm::APInt> alignmentAttrValue = None) {
   // Print optional memory access attribute.
-  if (auto memAccess = memoryOp.memory_access()) {
-    elidedAttrs.push_back(spirv::attributeName<spirv::MemoryAccess>());
+  if (auto memAccess = (memoryAccessAtrrValue ? memoryAccessAtrrValue
+                                              : memoryOp.memory_access())) {
+    elidedAttrs.push_back(memoryAccessAttrName);
+
+    if (!first) {
+      printer << ", ";
+    }
+
     printer << " [\"" << stringifyMemoryAccess(*memAccess) << "\"";
 
+    if (spirv::bitEnumContains(*memAccess, spirv::MemoryAccess::Aligned)) {
       // Print integer alignment attribute.
-    if (auto alignment = memoryOp.alignment()) {
-      elidedAttrs.push_back(kAlignmentAttrName);
+      if (auto alignment = (alignmentAttrValue ? alignmentAttrValue
+                                               : memoryOp.alignment())) {
+        elidedAttrs.push_back(alignmentAttrName);
         printer << ", " << alignment;
       }
+    }
     printer << "]";
   }
   elidedAttrs.push_back(spirv::attributeName<spirv::StorageClass>());
 }
 
 static LogicalResult verifyCastOp(Operation *op,
                                   bool requireSameBitWidth = true) {
   Type operandType = op->getOperand(0).getType();
if (isSameBitWidth) {
     return op->emitOpError(
                "expected the different bit widths for operand type and result "
                "type, but provided ")
            << operandType << " and " << resultType;
   }
   return success();
 }
 
-template <typename MemoryOpTy>
+template <const char memoryAccessAttrName[] = kMemoryAccessAttrName,
+          const char alignmentAttrName[] = kAlignmentAttrName,
+          typename MemoryOpTy>
 static LogicalResult verifyMemoryAccessAttribute(MemoryOpTy memoryOp) {
   // ODS checks for attributes values. Just need to verify that if the
   // memory-access attribute is Aligned, then the alignment attribute must be
   // present.
   auto *op = memoryOp.getOperation();
-  auto memAccessAttr = op->getAttr(spirv::attributeName<spirv::MemoryAccess>());
+  auto memAccessAttr = op->getAttr(memoryAccessAttrName);
   if (!memAccessAttr) {
     // Alignment attribute shouldn't be present if memory access attribute is
     // not present.
-    if (op->getAttr(kAlignmentAttrName)) {
+    if (op->getAttr(alignmentAttrName)) {
       return memoryOp.emitOpError(
           "invalid alignment specification without aligned memory access "
           "specification");
     }
     return success();
   }
 
   auto memAccessVal = memAccessAttr.template cast<IntegerAttr>();
   auto memAccess = spirv::symbolizeMemoryAccess(memAccessVal.getInt());
 
   if (!memAccess) {
     return memoryOp.emitOpError("invalid memory access specifier: ")
            << memAccessVal;
   }
 
   if (spirv::bitEnumContains(*memAccess, spirv::MemoryAccess::Aligned)) {
-    if (!op->getAttr(kAlignmentAttrName)) {
+    if (!op->getAttr(alignmentAttrName)) {
       return memoryOp.emitOpError("missing alignment value");
     }
   } else {
-    if (op->getAttr(kAlignmentAttrName)) {
+    if (op->getAttr(alignmentAttrName)) {
       return memoryOp.emitOpError(
           "invalid alignment specification with non-aligned memory access "
           "specification");
     }
   }
   return success();
 }
 
StringRef sourceStorageClass =
       stringifyStorageClass(copyMemory.source()
                                 .getType()
                                 .cast<spirv::PointerType>()
                                 .getStorageClass());
   printer << " \"" << sourceStorageClass << "\" " << copyMemory.source();
 
   SmallVector<StringRef, 4> elidedAttrs;
   printMemoryAccessAttribute(copyMemory, printer, elidedAttrs);
+  printMemoryAccessAttribute<kSourceMemoryAccessAttrName,
+                             kSourceAlignmentAttrName, false>(
+      copyMemory, printer, elidedAttrs, copyMemory.source_memory_access(),
+      copyMemory.source_alignment());
 
   printer.printOptionalAttrDict(op->getAttrs(), elidedAttrs);
 
   Type pointeeType =
       copyMemory.target().getType().cast<spirv::PointerType>().getPointeeType();
   printer << " : " << pointeeType;
 }
 
 static ParseResult parseCopyMemoryOp(OpAsmParser &parser,
                                      OperationState &state) {
   spirv::StorageClass targetStorageClass;
   OpAsmParser::OperandType targetPtrInfo;
 
   spirv::StorageClass sourceStorageClass;
   OpAsmParser::OperandType sourcePtrInfo;
 
   Type elementType;
 
   if (parseEnumStrAttr(targetStorageClass, parser) ||
       parser.parseOperand(targetPtrInfo) || parser.parseComma() ||
       parseEnumStrAttr(sourceStorageClass, parser) ||
       parser.parseOperand(sourcePtrInfo) ||
-      parseMemoryAccessAttributes(parser, state) ||
-      parser.parseOptionalAttrDict(state.attributes) || parser.parseColon() ||
-      parser.parseType(elementType)) {
+      parseMemoryAccessAttributes(parser, state)) {
+    return failure();
+  }
+
+  if (!parser.parseOptionalComma()) {
+    // Parse 2nd memory access attributes.
+    if (parseMemoryAccessAttributes<kSourceMemoryAccessAttrName,
+                                    kSourceAlignmentAttrName>(parser, state)) {
+      return failure();
+    }
+  }
+
+  if (parser.parseColon() || parser.parseType(elementType)) {
+    return failure();
+  }
+
+  if (parser.parseOptionalAttrDict(state.attributes)) {
     return failure();
   }
 
   auto targetPtrType = spirv::PointerType::get(elementType, targetStorageClass);
   auto sourcePtrType = spirv::PointerType::get(elementType, sourceStorageClass);
 
   if (parser.resolveOperand(targetPtrInfo, targetPtrType, state.operands) ||
       parser.resolveOperand(sourcePtrInfo, sourcePtrType, state.operands)) {
static LogicalResult verifyCopyMemory(spirv::CopyMemoryOp copyMemory) {
   Type sourceType =
       copyMemory.source().getType().cast<spirv::PointerType>().getPointeeType();
 
   if (targetType != sourceType) {
     return copyMemory.emitOpError(
         "both operands must be pointers to the same type");
   }
 
-  return verifyMemoryAccessAttribute(copyMemory);
+  if (failed(verifyMemoryAccessAttribute(copyMemory))) {
+    return failure();
+  }
+
+  // TODO (ergawy): According to the spec:
+  //
+  // If two masks are present, the first applies to Target and cannot include
+  // MakePointerVisible, and the second applies to Source and cannot include
+  // MakePointerAvailable.
+  //
+  // Add such verification here.
+
+  return verifyMemoryAccessAttribute<kSourceMemoryAccessAttrName,
+                                     kSourceAlignmentAttrName>(copyMemory);
 }
 
 //===----------------------------------------------------------------------===//
 // spv.Transpose
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verifyTranspose(spirv::TransposeOp op) {
   auto inputMatrix = op.matrix().getType().cast<spirv::MatrixType>();

mlir/lib/Dialect/SPIRV/Serialization/Deserializer.cpp

private:
   LogicalResult
   dispatchToExtensionSetAutogenDeserialization(StringRef extensionSetName,
                                                uint32_t instructionID,
                                                ArrayRef<uint32_t> words);
 
   /// Method to deserialize an operation in the SPIR-V dialect that is a mirror
   /// of an instruction in the SPIR-V spec. This is auto generated if hasOpcode
   /// == 1 and autogenSerialization == 1 in ODS.
-  template <typename OpTy> LogicalResult processOp(ArrayRef<uint32_t> words) {
+  template <typename OpTy>
+  LogicalResult processOp(ArrayRef<uint32_t> words) {
     return emitError(unknownLoc, "unsupported deserialization for ")
            << OpTy::getOperationName() << " op";
   }
 
 private:
   /// The SPIR-V binary module.
   ArrayRef<uint32_t> binary;
 
LogicalResult Deserializer::processConstantNull(ArrayRef<uint32_t> operands) {
   if (resultType.isIntOrFloat() || resultType.isa<VectorType>()) {
     auto attr = opBuilder.getZeroAttr(resultType);
     // For normal constants, we just record the attribute (and its type) for
     // later materialization at use sites.
     constantMap.try_emplace(resultID, attr, resultType);
     return success();
   }
 
-    return emitError(unknownLoc, "unsupported OpConstantNull type: ")
-           << resultType;
 }
 
 //===----------------------------------------------------------------------===//
 // Control flow
 //===----------------------------------------------------------------------===//
 
 Block *Deserializer::getOrCreateBlock(uint32_t id) {
   if (auto *block = getBlock(id)) {
for (auto operand : operands) {
     argAttrs.push_back(argAttr);
   }
 
   opBuilder.create<spirv::MemoryBarrierOp>(unknownLoc, argAttrs[0],
                                            argAttrs[1]);
   return success();
 }
 
+template <>
+LogicalResult
+Deserializer::processOp<spirv::CopyMemoryOp>(ArrayRef<uint32_t> words) {
+  SmallVector<Type, 1> resultTypes;
+  size_t wordIndex = 0;
+  SmallVector<Value, 4> operands;
+  SmallVector<NamedAttribute, 4> attributes;
+
+  if (wordIndex < words.size()) {
+    auto arg = getValue(words[wordIndex]);
+
+    if (!arg) {
+      return emitError(unknownLoc, "unknown result <id> : ")
+             << words[wordIndex];
+    }
+
+    operands.push_back(arg);
+    wordIndex++;
+  }
+
+  if (wordIndex < words.size()) {
+    auto arg = getValue(words[wordIndex]);
+
+    if (!arg) {
+      return emitError(unknownLoc, "unknown result <id> : ")
+             << words[wordIndex];
+    }
+
+    operands.push_back(arg);
+    wordIndex++;
+  }
+
+  bool isAlignedAttr = false;
+
+  if (wordIndex < words.size()) {
+    auto attrValue = words[wordIndex++];
+    attributes.push_back(opBuilder.getNamedAttr(
+        "memory_access", opBuilder.getI32IntegerAttr(attrValue)));
+    isAlignedAttr = (attrValue == 2);
+  }
+
+  if (isAlignedAttr && wordIndex < words.size()) {
+    attributes.push_back(opBuilder.getNamedAttr(
+        "alignment", opBuilder.getI32IntegerAttr(words[wordIndex++])));
+  }
+
+  if (wordIndex < words.size()) {
+    attributes.push_back(opBuilder.getNamedAttr(
+        "source_memory_access",
+        opBuilder.getI32IntegerAttr(words[wordIndex++])));
+  }
+
+  if (wordIndex < words.size()) {
+    attributes.push_back(opBuilder.getNamedAttr(
+        "source_alignment", opBuilder.getI32IntegerAttr(words[wordIndex++])));
+  }
+
+  if (wordIndex != words.size()) {
+    return emitError(unknownLoc,
+                     "found more operands than expected when deserializing "
+                     "spirv::CopyMemoryOp, only ")
+           << wordIndex << " of " << words.size() << " processed";
+  }
+
+  Location loc = createFileLineColLoc(opBuilder);
+  opBuilder.create<spirv::CopyMemoryOp>(loc, resultTypes, operands, attributes);
+
+  return success();
+}
+
 // Pull in auto-generated Deserializer::dispatchToAutogenDeserialization() and
 // various Deserializer::processOp<...>() specializations.
 #define GET_DESERIALIZATION_FNS
 #include "mlir/Dialect/SPIRV/SPIRVSerialization.inc"
 } // namespace
 
 Optional<spirv::ModuleOp> spirv::deserialize(ArrayRef<uint32_t> binary,
                                              MLIRContext *context) {
   Deserializer deserializer(binary, context);
 
   if (failed(deserializer.deserialize()))
     return llvm::None;
 
   return deserializer.collect();
 }

mlir/lib/Dialect/SPIRV/Serialization/Serializer.cpp

private:
   /// SPIR-V dialect that is a mirror of an instruction in the SPIR-V spec.
   /// This is auto-generated from ODS. Dispatch is handled for all operations
   /// in SPIR-V dialect that have hasOpcode == 1.
   LogicalResult dispatchToAutogenSerialization(Operation *op);
 
   /// Method to serialize an operation in the SPIR-V dialect that is a mirror of
   /// an instruction in the SPIR-V spec. This is auto generated if hasOpcode ==
   /// 1 and autogenSerialization == 1 in ODS.
-  template <typename OpTy> LogicalResult processOp(OpTy op) {
+  template <typename OpTy>
+  LogicalResult processOp(OpTy op) {
     return op.emitError("unsupported op serialization");
   }
 
   //===--------------------------------------------------------------------===//
   // Utilities
   //===--------------------------------------------------------------------===//
 
   /// Emits an OpDecorate instruction to decorate the given `target` with the
Serializer::processOp<spirv::FunctionCallOp>(spirv::FunctionCallOp op) {
 
   if (!resultTy.isa<NoneType>())
     valueIDMap[op.getResult(0)] = funcCallID;
 
   return encodeInstructionInto(functionBody, spirv::Opcode::OpFunctionCall,
                                operands);
 }
 
+template <>
+LogicalResult
+Serializer::processOp<spirv::CopyMemoryOp>(spirv::CopyMemoryOp op) {
+  SmallVector<uint32_t, 4> operands;
+  SmallVector<StringRef, 2> elidedAttrs;
+
+  for (Value operand : op.getOperation()->getOperands()) {
+    auto id = getValueID(operand);
+    assert(id && "use before def!");
+    operands.push_back(id);
+  }
+
+  if (auto attr = op.getAttr("memory_access")) {
+    operands.push_back(static_cast<uint32_t>(
+        attr.cast<IntegerAttr>().getValue().getZExtValue()));
+  }
+
+  elidedAttrs.push_back("memory_access");
+
+  if (auto attr = op.getAttr("alignment")) {
+    operands.push_back(static_cast<uint32_t>(
+        attr.cast<IntegerAttr>().getValue().getZExtValue()));
+  }
+
+  elidedAttrs.push_back("alignment");
+
+  if (auto attr = op.getAttr("source_memory_access")) {
+    operands.push_back(static_cast<uint32_t>(
+        attr.cast<IntegerAttr>().getValue().getZExtValue()));
+  }
+
+  elidedAttrs.push_back("source_memory_access");
+
+  if (auto attr = op.getAttr("source_alignment")) {
+    operands.push_back(static_cast<uint32_t>(
+        attr.cast<IntegerAttr>().getValue().getZExtValue()));
+  }
+
+  elidedAttrs.push_back("source_alignment");
+  emitDebugLine(functionBody, op.getLoc());
+  encodeInstructionInto(functionBody, spirv::Opcode::OpCopyMemory, operands);
+
+  return success();
+}
+
 // Pull in auto-generated Serializer::dispatchToAutogenSerialization() and
 // various Serializer::processOp<...>() specializations.
 #define GET_SERIALIZATION_FNS
 #include "mlir/Dialect/SPIRV/SPIRVSerialization.inc"
 } // namespace
 
 LogicalResult Serializer::emitDecoration(uint32_t target,
                                          spirv::Decoration decoration,

mlir/test/Dialect/SPIRV/Serialization/memory-ops.mlir

spv.func @copy_memory_with_access_operands() "None" {
     %0 = spv.Variable : !spv.ptr<f32, Function>
     %1 = spv.Variable : !spv.ptr<f32, Function>
     // CHECK: spv.CopyMemory "Function" %{{.*}}, "Function" %{{.*}} ["Aligned", 4] : f32
     spv.CopyMemory "Function" %0, "Function" %1 ["Aligned", 4] : f32
 
     // CHECK: spv.CopyMemory "Function" %{{.*}}, "Function" %{{.*}} ["Volatile"] : f32
     spv.CopyMemory "Function" %0, "Function" %1 ["Volatile"] : f32
 
+    // CHECK: spv.CopyMemory "Function" %{{.*}}, "Function" %{{.*}} ["Volatile"], ["Volatile"] : f32
+    spv.CopyMemory "Function" %0, "Function" %1 ["Volatile"], ["Volatile"] : f32
+
+    // CHECK: spv.CopyMemory "Function" %{{.*}}, "Function" %{{.*}} ["Aligned", 4], ["Volatile"] : f32
+    spv.CopyMemory "Function" %0, "Function" %1 ["Aligned", 4], ["Volatile"] : f32
+
+    // CHECK: spv.CopyMemory "Function" %{{.*}}, "Function" %{{.*}} ["Volatile"], ["Aligned", 4] : f32
+    spv.CopyMemory "Function" %0, "Function" %1 ["Volatile"], ["Aligned", 4] : f32
+
+    // CHECK: spv.CopyMemory "Function" %{{.*}}, "Function" %{{.*}} ["Aligned", 8], ["Aligned", 4] : f32
+    spv.CopyMemory "Function" %0, "Function" %1 ["Aligned", 8], ["Aligned", 4] : f32
+
     spv.Return
   }
 }
 

mlir/test/Dialect/SPIRV/ops.mlir

func @copy_memory_invalid_maa() -> () {
   %1 = spv.Variable : !spv.ptr<f32, Function>
   // expected-error @+1 {{missing alignment value}}
   "spv.CopyMemory"(%0, %1) {memory_access=0x0002 : i32} : (!spv.ptr<f32, Function>, !spv.ptr<f32, Function>) -> ()
   spv.Return
 }
 
 // -----
 
+func @copy_memory_invalid_source_maa() -> () {
+  %0 = spv.Variable : !spv.ptr<f32, Function>
+  %1 = spv.Variable : !spv.ptr<f32, Function>
+  // expected-error @+1 {{invalid alignment specification with non-aligned memory access specification}}
+  "spv.CopyMemory"(%0, %1) {source_memory_access=0x0001 : i32, memory_access=0x0002 : i32, source_alignment=8 : i32, alignment=4 : i32} : (!spv.ptr<f32, Function>, !spv.ptr<f32, Function>) -> ()
+  spv.Return
+}
+
+// -----
+
+func @copy_memory_invalid_source_maa2() -> () {
+  %0 = spv.Variable : !spv.ptr<f32, Function>
+  %1 = spv.Variable : !spv.ptr<f32, Function>
+  // expected-error @+1 {{missing alignment value}}
+  "spv.CopyMemory"(%0, %1) {source_memory_access=0x0002 : i32, memory_access=0x0002 : i32, alignment=4 : i32} : (!spv.ptr<f32, Function>, !spv.ptr<f32, Function>) -> ()
+  spv.Return
+}
+
+// -----
+
 func @copy_memory_print_maa() -> () {
   %0 = spv.Variable : !spv.ptr<f32, Function>
   %1 = spv.Variable : !spv.ptr<f32, Function>
 
   // CHECK: spv.CopyMemory "Function" %{{.*}}, "Function" %{{.*}} ["Volatile"] : f32
   "spv.CopyMemory"(%0, %1) {memory_access=0x0001 : i32} : (!spv.ptr<f32, Function>, !spv.ptr<f32, Function>) -> ()
 
   // CHECK: spv.CopyMemory "Function" %{{.*}}, "Function" %{{.*}} ["Aligned", 4] : f32
   "spv.CopyMemory"(%0, %1) {memory_access=0x0002 : i32, alignment=4 : i32} : (!spv.ptr<f32, Function>, !spv.ptr<f32, Function>) -> ()
 
+  // CHECK: spv.CopyMemory "Function" %{{.*}}, "Function" %{{.*}} ["Aligned", 4], ["Volatile"] : f32
+  "spv.CopyMemory"(%0, %1) {source_memory_access=0x0001 : i32, memory_access=0x0002 : i32, alignment=4 : i32} : (!spv.ptr<f32, Function>, !spv.ptr<f32, Function>) -> ()
+
+  // CHECK: spv.CopyMemory "Function" %{{.*}}, "Function" %{{.*}} ["Aligned", 4], ["Aligned", 8] : f32
+  "spv.CopyMemory"(%0, %1) {source_memory_access=0x0002 : i32, memory_access=0x0002 : i32, source_alignment=8 : i32, alignment=4 : i32} : (!spv.ptr<f32, Function>, !spv.ptr<f32, Function>) -> ()
+
   spv.Return
 }