Add SPIR-V Load/Store operations. Currently this only support memory

operands being None, Volatile, Aligned and Nontemporal

PiperOrigin-RevId: 254792353

include/mlir/IR/OpImplementation.h

@@ -244,9 +244,15 @@ class OpAsmParser {
   /// Parse a `[` token.
   virtual ParseResult parseLSquare() = 0;
 
+  /// Parse a `[` token if present.
+  virtual ParseResult parseOptionalLSquare() = 0;
+
   /// Parse a `]` token.
   virtual ParseResult parseRSquare() = 0;
 
+  /// Parse a `]` token if present.
+  virtual ParseResult parseOptionalRSquare() = 0;
+
   //===--------------------------------------------------------------------===//
   // Attribute Parsing
   //===--------------------------------------------------------------------===//

include/mlir/SPIRV/SPIRVBase.td

@@ -190,6 +190,24 @@ def SPV_ImageFormatAttr :
   let underlyingType = "uint32_t";
 }
 
+def SPV_MA_None                    : EnumAttrCase<"None", 0x0000>;
+def SPV_MA_Volatile                : EnumAttrCase<"Volatile", 0x0001>;
+def SPV_MA_Aligned                 : EnumAttrCase<"Aligned", 0x0002>;
+def SPV_MA_Nontemporal             : EnumAttrCase<"Nontemporal", 0x0004>;
+def SPV_MA_MakePointerAvailableKHR : EnumAttrCase<"MakePointerAvailableKHR", 0x0008>;
+def SPV_MA_MakePointerVisibleKHR   : EnumAttrCase<"MakePointerVisibleKHR", 0x0010>;
+def SPV_MA_NonPrivatePointerKHR    : EnumAttrCase<"NonPrivatePointerKHR", 0x0020>;
+
+def SPV_MemoryAccessAttr :
+    EnumAttr<"MemoryAccess", "valid SPIR-V MemoryAccess", [
+      SPV_MA_None, SPV_MA_Volatile, SPV_MA_Aligned, SPV_MA_Nontemporal,
+      SPV_MA_MakePointerAvailableKHR, SPV_MA_MakePointerVisibleKHR,
+      SPV_MA_NonPrivatePointerKHR
+    ]> {
+  let cppNamespace = "::mlir::spirv";
+  let underlyingType = "uint32_t";
+}
+
 def SPV_MM_Simple    : EnumAttrCase<"Simple", 0>;
 def SPV_MM_GLSL450   : EnumAttrCase<"GLSL450", 1>;
 def SPV_MM_OpenCL    : EnumAttrCase<"OpenCL", 2>;

include/mlir/SPIRV/SPIRVOps.td

@@ -65,6 +65,61 @@ def SPV_FMulOp : SPV_Op<"FMul", [NoSideEffect, SameOperandsAndResultType]> {
   let opcode = 133;
 }
 
+def SPV_LoadOp : SPV_Op<"Load"> {
+  let summary = "Load value through a pointer.";
+
+  let description = [{
+    Result Type is the type of the loaded object. It must be a type
+    with fixed size; i.e., it cannot be, nor include, any
+    OpTypeRuntimeArray types.
+
+    Pointer is the pointer to load through. Its type must be an
+    OpTypePointer whose Type operand is the same as Result Type.
+
+    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.
+
+    ### Custom assembly form
+
+    ``` {.ebnf}
+    memory-access ::= `"None"` | `"Volatile"` | `"Aligned"` integer-literal
+                    | `"NonTemporal"`
+
+    load-op ::= ssa-id ` = spv.Load ` storage-class ssa-use
+                (`[` memory-access `]`)? ` : ` spirv-element-type
+    ```
+
+    For example:
+
+    ```
+    %0 = spv.Variable : !spv.ptr<f32, Function>
+    %1 = spv.Load "Function" %0 : f32
+    %2 = spv.Load "Function" %0 ["Volatile"] : f32
+    %3 = spv.Load "Function" %0 ["Aligned", 4] : f32
+    ```
+  }];
+
+  let arguments = (ins
+    SPV_AnyPtr:$ptr,
+    OptionalAttr<SPV_MemoryAccessAttr>:$memory_access,
+    OptionalAttr<I32Attr>:$alignment
+  );
+
+  let results = (outs
+    SPV_Type:$value
+  );
+
+  let extraClassDeclaration = [{
+    static StringRef getMemoryAccessAttrName() {
+      return "memory_access";
+    }
+    static StringRef getAlignmentAttrName() {
+      return "alignment";
+    }
+  }];
+}
+
 def SPV_ReturnOp : SPV_Op<"Return", [Terminator]> {
   let summary = "Return with no value from a function with void return type";
 
@@ -84,6 +139,54 @@ def SPV_ReturnOp : SPV_Op<"Return", [Terminator]> {
   let opcode = 253;
 }
 
+def SPV_StoreOp : SPV_Op<"Store"> {
+  let summary = "Store through a pointer.";
+
+  let description = [{
+    Pointer is the pointer to store through. Its type must be an
+    OpTypePointer whose Type operand is the same as the type of
+    Object.
+
+    Object is the object to store.
+
+    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.
+
+    ### Custom assembly form
+
+    ``` {.ebnf}
+    store-op ::= `spv.Store ` storage-class ssa-use `, ` ssa-use `, `
+                  (memory-access)? : spirv-element-type
+    ```
+
+    For example:
+
+    ```
+    %0 = spv.Variable : !spv.ptr<f32, Function>
+    %1 = spv.FMul ... : f32
+    spv.Store "Function" %0, %1 : f32
+    spv.Store "Function" %0, %1 ["Volatile"] : f32
+    spv.Store "Function" %0, %1 ["Aligned", 4] : f32
+  }];
+
+  let arguments = (ins
+    SPV_AnyPtr:$ptr,
+    SPV_Type:$value,
+    OptionalAttr<SPV_MemoryAccessAttr>:$memory_access,
+    OptionalAttr<APIntAttr>:$alignment
+  );
+
+  let extraClassDeclaration = [{
+    static StringRef getMemoryAccessAttrName() {
+      return "memory_access";
+    }
+    static StringRef getAlignmentAttrName() {
+      return "alignment";
+    }
+  }];
+}
+
 def SPV_VariableOp : SPV_Op<"Variable"> {
   let summary = [{
     Allocate an object in memory, resulting in a pointer to it, which can be

lib/Parser/Parser.cpp

@@ -3104,11 +3104,21 @@ class CustomOpAsmParser : public OpAsmParser {
     return parser.parseToken(Token::l_square, "expected '['");
   }
 
+  /// Parses a '[' if present.
+  ParseResult parseOptionalLSquare() override {
+    return success(parser.consumeIf(Token::l_square));
+  }
+
   /// Parse a `]` token.
   ParseResult parseRSquare() override {
     return parser.parseToken(Token::r_square, "expected ']'");
   }
 
+  /// Parses a ']' if present.
+  ParseResult parseOptionalRSquare() override {
+    return success(parser.consumeIf(Token::r_square));
+  }
+
   //===--------------------------------------------------------------------===//
   // Attribute Parsing
   //===--------------------------------------------------------------------===//