[MLIR][SPIRV] Add support for OpCopyMemory.

[ergawy]

Summary

This patch add support for 'spv.CopyMemory'. The following changes are
introduced:

• 'CopyMemory' op is added to SPIRVOps.td.
• Custom parse and print methods are introduced.
• A few Roundtripping tests are added.

[ergawy]

Remove comments added during development.

[ergawy]

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

223

Correct me if I am wrong, but I think we don't need a custom verifier here since in

static ParseResult parseCopyMemoryOp(OpAsmParser &parser,
                                     OperationState &state)

I did the needed verification that target and source types are compatible with each other.

[ergawy]

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

60

Currently trying to come up with more tests. If you have any suggestions, please let me know :).

[ergawy]

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

217–218

According to the spec: https://www.khronos.org/registry/spir-v/specs/unified1/SPIRV.html#OpCopyMemory, the op supports up to 2 memory access operands. However, I still don't know how to model this in ODS.

[rriddle]

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

223

Please do not rely on verification done in the parser unless it is solely related to parsing. The verifier should *always* be complete. The majority of production systems(including MLIR itself via conversions) will never use the parser and instead construct IR directly. If the verification is in the parser and not here, the verification will never happen at all.

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2823

Where are you printing the attribute dictionary?

[ergawy]

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

223

Thanks for your comment @rriddle . I wasn't aware of that point. I added a verifyCopyMemory function and currently trying to find a good test approach. As far as I understand:

• Any test approach that starts with the textual representation won't really test all the possible paths inside the verifier function because we will have to go through the parser first. And, in the parser we have to resolve the SSA values based on the pointer types constructed from the element type attached to the op text for which some verification of type consistency has to be made.
• Therefore, we have to directly construct the IR in order to test the verification logic thoroughly. One way I am considering is to introduce a new unit test under mlir/unittests/Dialect/SPIRV/.

I still need to figure out the pieces to construct the needed IR elements directly. But I just wanted to make sure that such approach would be suitable in that situation.

[rriddle]

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

223

Any test approach that starts with the textual representation won't really test all the possible paths inside the verifier function because we will have to go through the parser first.

This isn't a problem in MLIR, you can always use the generic parser form which won't go through any of the custom parser code paths. With that you could test any part of the verifier that you want.

[ergawy]

• Add verification function for OpCopyMemory.
• Add tests for printing and verifying maa.

[ergawy]

• Rearrange code after rebase.

[ergawy]

• Rename template type arg.

[ergawy]

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

223

Thank you! That's awesome. I added a few tests for verification with the generic syntax.

[ergawy]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2823

I might be wrong, but isn't that done in line 2841?

[ergawy]

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

217–218

Is it ok to add a TODO for that issue and handle it with a later review? I am working on figuring out a way to do that at the moment. But I think it should ok to merge this review as is (after handling any further comments of course).

[ergawy]

• Print attr dictionary.

[ergawy]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2823

Oh, I didn't get what you mean at first. Please, excuse the ignorance of the newbie 😄 .

[antiagainst]

Awesome, this looks very clean. Thanks for the patch @ergawy!!

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2889

We don't need the leading : here. emitOpError will attach the message here with a leading 'xxx' op so we can write the message based on that assumption.

[antiagainst]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2889

I'll take care of this before landing the patch.

[ergawy]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2889

Awesome, thank you!

mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td

 def SPV_OC_OpSpecConstantComposite     : I32EnumAttrCase<"OpSpecConstantComposite", 51>;
 def SPV_OC_OpFunction                  : I32EnumAttrCase<"OpFunction", 54>;
 def SPV_OC_OpFunctionParameter         : I32EnumAttrCase<"OpFunctionParameter", 55>;
 def SPV_OC_OpFunctionEnd               : I32EnumAttrCase<"OpFunctionEnd", 56>;
 def SPV_OC_OpFunctionCall              : I32EnumAttrCase<"OpFunctionCall", 57>;
 def SPV_OC_OpVariable                  : I32EnumAttrCase<"OpVariable", 59>;
 def SPV_OC_OpLoad                      : I32EnumAttrCase<"OpLoad", 61>;
 def SPV_OC_OpStore                     : I32EnumAttrCase<"OpStore", 62>;
+def SPV_OC_OpCopyMemory                : I32EnumAttrCase<"OpCopyMemory", 63>;
 def SPV_OC_OpAccessChain               : I32EnumAttrCase<"OpAccessChain", 65>;
 def SPV_OC_OpDecorate                  : I32EnumAttrCase<"OpDecorate", 71>;
 def SPV_OC_OpMemberDecorate            : I32EnumAttrCase<"OpMemberDecorate", 72>;
 def SPV_OC_OpCompositeConstruct        : I32EnumAttrCase<"OpCompositeConstruct", 80>;
 def SPV_OC_OpCompositeExtract          : I32EnumAttrCase<"OpCompositeExtract", 81>;
 def SPV_OC_OpCompositeInsert           : I32EnumAttrCase<"OpCompositeInsert", 82>;
 def SPV_OC_OpConvertFToU               : I32EnumAttrCase<"OpConvertFToU", 109>;
 def SPV_OC_OpConvertFToS               : I32EnumAttrCase<"OpConvertFToS", 110>;
SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       SPV_OC_OpCapability, SPV_OC_OpTypeVoid, SPV_OC_OpTypeBool, SPV_OC_OpTypeInt,
       SPV_OC_OpTypeFloat, SPV_OC_OpTypeVector, SPV_OC_OpTypeMatrix,
       SPV_OC_OpTypeArray, SPV_OC_OpTypeRuntimeArray, SPV_OC_OpTypeStruct,
       SPV_OC_OpTypePointer, SPV_OC_OpTypeFunction, SPV_OC_OpConstantTrue,
       SPV_OC_OpConstantFalse, SPV_OC_OpConstant, SPV_OC_OpConstantComposite,
       SPV_OC_OpConstantNull, SPV_OC_OpSpecConstantTrue, SPV_OC_OpSpecConstantFalse,
       SPV_OC_OpSpecConstant, SPV_OC_OpSpecConstantComposite, SPV_OC_OpFunction,
       SPV_OC_OpFunctionParameter, SPV_OC_OpFunctionEnd, SPV_OC_OpFunctionCall,
-      SPV_OC_OpVariable, SPV_OC_OpLoad, SPV_OC_OpStore, SPV_OC_OpAccessChain,
-      SPV_OC_OpDecorate, SPV_OC_OpMemberDecorate, SPV_OC_OpCompositeConstruct,
-      SPV_OC_OpCompositeExtract, SPV_OC_OpCompositeInsert, SPV_OC_OpConvertFToU,
-      SPV_OC_OpConvertFToS, SPV_OC_OpConvertSToF, SPV_OC_OpConvertUToF,
-      SPV_OC_OpUConvert, SPV_OC_OpSConvert, SPV_OC_OpFConvert, SPV_OC_OpBitcast,
-      SPV_OC_OpFNegate, SPV_OC_OpIAdd, SPV_OC_OpFAdd, SPV_OC_OpISub, SPV_OC_OpFSub,
-      SPV_OC_OpIMul, SPV_OC_OpFMul, SPV_OC_OpUDiv, SPV_OC_OpSDiv, SPV_OC_OpFDiv,
-      SPV_OC_OpUMod, SPV_OC_OpSRem, SPV_OC_OpSMod, SPV_OC_OpFRem, SPV_OC_OpFMod,
+      SPV_OC_OpVariable, SPV_OC_OpLoad, SPV_OC_OpStore, SPV_OC_OpCopyMemory,
+      SPV_OC_OpAccessChain, SPV_OC_OpDecorate, SPV_OC_OpMemberDecorate,
+      SPV_OC_OpCompositeConstruct, SPV_OC_OpCompositeExtract,
+      SPV_OC_OpCompositeInsert, SPV_OC_OpConvertFToU, SPV_OC_OpConvertFToS,
+      SPV_OC_OpConvertSToF, SPV_OC_OpConvertUToF, SPV_OC_OpUConvert,
+      SPV_OC_OpSConvert, SPV_OC_OpFConvert, SPV_OC_OpBitcast, SPV_OC_OpFNegate,
+      SPV_OC_OpIAdd, SPV_OC_OpFAdd, SPV_OC_OpISub, SPV_OC_OpFSub, SPV_OC_OpIMul,
+      SPV_OC_OpFMul, SPV_OC_OpUDiv, SPV_OC_OpSDiv, SPV_OC_OpFDiv, SPV_OC_OpUMod,
+      SPV_OC_OpSRem, SPV_OC_OpSMod, SPV_OC_OpFRem, SPV_OC_OpFMod,
       SPV_OC_OpMatrixTimesScalar, SPV_OC_OpLogicalEqual, SPV_OC_OpLogicalNotEqual,
       SPV_OC_OpLogicalOr, SPV_OC_OpLogicalAnd, SPV_OC_OpLogicalNot, SPV_OC_OpSelect,
       SPV_OC_OpIEqual, SPV_OC_OpINotEqual, SPV_OC_OpUGreaterThan,
       SPV_OC_OpSGreaterThan, SPV_OC_OpUGreaterThanEqual, SPV_OC_OpSGreaterThanEqual,
       SPV_OC_OpULessThan, SPV_OC_OpSLessThan, SPV_OC_OpULessThanEqual,
       SPV_OC_OpSLessThanEqual, SPV_OC_OpFOrdEqual, SPV_OC_OpFUnordEqual,
       SPV_OC_OpFOrdNotEqual, SPV_OC_OpFUnordNotEqual, SPV_OC_OpFOrdLessThan,
       SPV_OC_OpFUnordLessThan, SPV_OC_OpFOrdGreaterThan, SPV_OC_OpFUnordGreaterThan,

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

def SPV_ControlBarrierOp : SPV_Op<"ControlBarrier", []> {
 
   let assemblyFormat = [{
     $execution_scope `,` $memory_scope `,` $memory_semantics attr-dict
   }];
 }
 
 // -----
 
+def SPV_CopyMemoryOp : SPV_Op<"CopyMemory", []> {
+  let summary = [{
+    Copy from the memory pointed to by Source to the memory pointed to by
+    Target. Both operands must be non-void pointers and having the same <id>
+    Type operand in their OpTypePointer type declaration.  Matching Storage
+    Class is not required.  The amount of memory copied is the size of the
+    type pointed to. The copied type must have a fixed size; i.e., it cannot
+    be, nor include, any OpTypeRuntimeArray types.
+  }];
+
+  let description = [{
+    If present, any Memory Operands must begin with a memory operand
+    literal. If not present, it is the same as specifying the memory operand
+    None. Before version 1.4, at most one memory operands mask can be
+    provided. Starting with version 1.4 two masks can be provided, as
+    described in Memory Operands. If no masks or only one mask is present,
+    it applies to both Source and Target. If two masks are present, the
+    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 `]`)?
+                       ` : ` 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
+  );
+
+  let results = (outs);
+
+  let verifier = [{ return success(); }];
+}
+
+// -----
+
 def SPV_ExecutionModeOp : SPV_Op<"ExecutionMode", [InModuleScope]> {
   let summary = "Declare an execution mode for an entry point.";
 
   let description = [{
     Entry Point must be the Entry Point <id> operand of an OpEntryPoint
     instruction.
 
     Mode is the execution mode. See Execution Mode.

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

if (auto resultMatrixColumns =
 
       // Check that the input and result matrices' columns have the same size
       if (inputMatrixColumns.getNumElements() !=
           resultMatrixColumns.getNumElements())
         return op.emitError("input and result matrices' columns must "
                             "have the same size");
     }
   }
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// spv.CopyMemory
+//===----------------------------------------------------------------------===//
+
+static void print(spirv::CopyMemoryOp copyMemory, OpAsmPrinter &printer) {
+  printer << spirv::CopyMemoryOp::getOperationName() << ' ';
+
+  StringRef targetStorageClass =
+      stringifyStorageClass(copyMemory.target()
+                                .getType()
+                                .cast<spirv::PointerType>()
+                                .getStorageClass());
+  printer << " \"" << targetStorageClass << "\" " << copyMemory.target()
+          << ", ";
+
+  StringRef sourceStorageClass =
+      stringifyStorageClass(copyMemory.source()
+                                .getType()
+                                .cast<spirv::PointerType>()
+                                .getStorageClass());
+  printer << " \"" << sourceStorageClass << "\" " << copyMemory.source();
+
+  SmallVector<StringRef, 4> elidedAttrs;
+  printMemoryAccessAttribute(copyMemory, printer, 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)) {
+    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)) {
+    return failure();
+  }
+
   return success();
 }
 
 namespace mlir {
 namespace spirv {
 
 // TableGen'erated operation interfaces for querying versions, extensions, and
 // capabilities.
 #include "mlir/Dialect/SPIRV/SPIRVAvailability.cpp.inc"
 
 // TablenGen'erated operation definitions.
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SPIRV/SPIRVOps.cpp.inc"
 
 // TableGen'erated operation availability interface implementations.
 #include "mlir/Dialect/SPIRV/SPIRVOpAvailabilityImpl.inc"
 
 } // namespace spirv
 } // namespace mlir

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

spv.func @load_store_zero_rank_int(%arg0: !spv.ptr<!spv.struct<!spv.array<1 x i32, stride=4> [0]>, StorageBuffer>, %arg1: !spv.ptr<!spv.struct<!spv.array<1 x i32, stride=4> [0]>, StorageBuffer>) "None" {
     // CHECK: [[STORE_PTR:%.*]] = spv.AccessChain {{%.*}}[{{%.*}}, {{%.*}}] : !spv.ptr<!spv.struct<!spv.array<1 x i32, stride=4> [0]>
     // CHECK-NEXT: spv.Store "StorageBuffer" [[STORE_PTR]], [[VAL]] : i32
     %3 = spv.constant 0 : i32
     %4 = spv.AccessChain %arg1[%3, %3] : !spv.ptr<!spv.struct<!spv.array<1 x i32, stride=4> [0]>, StorageBuffer>, i32, i32
     spv.Store "StorageBuffer" %4, %2 : i32
     spv.Return
   }
 }
+
+// -----
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+  spv.func @copy_memory_simple() "None" {
+    %0 = spv.Variable : !spv.ptr<f32, Function>
+    %1 = spv.Variable : !spv.ptr<f32, Function>
+    // CHECK: spv.CopyMemory "Function" %{{.*}}, "Function" %{{.*}} : f32
+    spv.CopyMemory "Function" %0, "Function" %1 : f32
+    spv.Return
+  }
+}
+
+// -----
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+  spv.func @copy_memory_different_storage_classes(%in : !spv.ptr<!spv.array<4xf32>, Input>, %out : !spv.ptr<!spv.array<4xf32>, Output>) "None" {
+    // CHECK: spv.CopyMemory "Output" %{{.*}}, "Input" %{{.*}} : !spv.array<4 x f32>
+    spv.CopyMemory "Output" %out, "Input" %in : !spv.array<4xf32>
+    spv.Return
+  }
+}
+
+
+// -----
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+  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
+
+    %2 = spv.Variable : !spv.ptr<f32, Function>
+    %3 = spv.Variable : !spv.ptr<f32, Function>
+    // CHECK: spv.CopyMemory "Function" %{{.*}}, "Function" %{{.*}} ["Volatile"] : f32
+    spv.CopyMemory "Function" %2, "Function" %3 ["Volatile"] : f32
+    spv.Return
+  }
+}
+