[fir][NFC] Update and move fir.freemem and fir.store

Move fir.freemem and fir.store to the Memory SSA operations sections.
Move parser, printer and verifier of fir.store to the .cpp file.

This patch is part of the upstreaming effort from fir-dev branch.

Co-authored-by: Jean Perier <jperier@nvidia.com>
Co-authored-by: Eric Schweitz <eschweitz@nvidia.com>

Reviewed By: jeanPerier

Differential Revision: https://reviews.llvm.org/D110816

GitOrigin-RevId: ccc0f62d1bc96e56d816f9cd7a987666fbd36b5b
diff --git a/include/flang/Optimizer/Dialect/FIROps.td b/include/flang/Optimizer/Dialect/FIROps.td
index 8c392f6..4ae39c4 100644
--- a/include/flang/Optimizer/Dialect/FIROps.td
+++ b/include/flang/Optimizer/Dialect/FIROps.td
@@ -376,6 +376,28 @@
   }];
 }
 
+def fir_FreeMemOp : fir_Op<"freemem", [MemoryEffects<[MemFree]>]> {
+  let summary = "free a heap object";
+
+  let description = [{
+    Deallocates a heap memory reference that was allocated by an `allocmem`.
+    The memory object that is deallocated is placed in an undefined state
+    after `fir.freemem`.  Optimizations may treat the loading of an object
+    in the undefined state as undefined behavior.  This includes aliasing
+    references, such as the result of an `fir.embox`.
+
+    ```mlir
+      %21 = fir.allocmem !fir.type<ZT(p:i32){field:i32}>
+      ...
+      fir.freemem %21 : !fir.heap<!fir.type<ZT>>
+    ```
+  }];
+
+  let arguments = (ins Arg<fir_HeapType, "", [MemFree]>:$heapref);
+
+  let assemblyFormat = "$heapref attr-dict `:` type($heapref)";
+}
+
 def fir_LoadOp : fir_OneResultOp<"load"> {
   let summary = "load a value from a memory reference";
   let description = [{
@@ -405,6 +427,39 @@
   }];
 }
 
+def fir_StoreOp : fir_Op<"store", []> {
+  let summary = "store an SSA-value to a memory location";
+
+  let description = [{
+    Store an ssa-value (virtual register) to a memory reference.  The stored
+    value must be of the same type as the referent type of the memory
+    reference.
+
+    ```mlir
+      %v = ... : f64
+      %p = ... : !fir.ptr<f64>
+      fir.store %v to %p : !fir.ptr<f64>
+    ```
+
+    The above store changes the value to which the pointer is pointing and not
+    the pointer itself. The operation is undefined if the memory reference,
+    `%p`, is undefined or null.
+  }];
+
+  let arguments = (ins AnyType:$value,
+                   Arg<AnyReferenceLike, "", [MemWrite]>:$memref);
+
+  let parser = "return parseStoreOp(parser, result);";
+
+  let printer = "::print(p, *this);";
+
+  let verifier = "return ::verify(*this);";
+
+  let extraClassDeclaration = [{
+    static mlir::Type elementType(mlir::Type refType);
+  }];
+}
+
 def fir_SaveResultOp : fir_Op<"save_result", [AttrSizedOperandSegments]> {
   let summary = [{
     save an array, box, or record function result SSA-value to a memory location
@@ -454,66 +509,6 @@
   let verifier = [{ return ::verify(*this); }];
 }
 
-def fir_StoreOp : fir_Op<"store", []> {
-  let summary = "store an SSA-value to a memory location";
-
-  let description = [{
-    Store an ssa-value (virtual register) to a memory reference.  The stored
-    value must be of the same type as the referent type of the memory
-    reference.
-
-    ```mlir
-      %v = ... : f64
-      %p = ... : !fir.ptr<f64>
-      fir.store %v to %p : !fir.ptr<f64>
-    ```
-
-    The above store changes the value to which the pointer is pointing and not
-    the pointer itself. The operation is undefined if the memory reference,
-    `%p`, is undefined or null.
-  }];
-
-  let arguments = (ins AnyType:$value,
-                   Arg<AnyReferenceLike, "", [MemWrite]>:$memref);
-
-  let parser = [{
-    mlir::Type type;
-    mlir::OpAsmParser::OperandType oper;
-    mlir::OpAsmParser::OperandType store;
-    if (parser.parseOperand(oper) ||
-        parser.parseKeyword("to") ||
-        parser.parseOperand(store) ||
-        parser.parseOptionalAttrDict(result.attributes) ||
-        parser.parseColonType(type) ||
-        parser.resolveOperand(oper, elementType(type),
-          result.operands) ||
-        parser.resolveOperand(store, type, result.operands))
-       return mlir::failure();
-    return mlir::success();
-  }];
-
-  let printer = [{
-    p << ' ';
-    p.printOperand(value());
-    p << " to ";
-    p.printOperand(memref());
-    p.printOptionalAttrDict((*this)->getAttrs(), {});
-    p << " : " << memref().getType();
-  }];
-
-  let verifier = [{
-    if (value().getType() != fir::dyn_cast_ptrEleTy(memref().getType()))
-      return emitOpError("store value type must match memory reference type");
-    if (fir::isa_unknown_size_box(value().getType()))
-      return emitOpError("cannot store !fir.box of unknown rank or type");
-    return mlir::success();
-  }];
-
-  let extraClassDeclaration = [{
-    static mlir::Type elementType(mlir::Type refType);
-  }];
-}
-
 def fir_UndefOp : fir_OneResultOp<"undefined", [NoSideEffect]> {
   let summary = "explicit undefined value of some type";
   let description = [{
@@ -557,28 +552,6 @@
   let assemblyFormat = "type($intype) attr-dict";
 }
 
-def fir_FreeMemOp : fir_Op<"freemem", [MemoryEffects<[MemFree]>]> {
-  let summary = "free a heap object";
-
-  let description = [{
-    Deallocates a heap memory reference that was allocated by an `allocmem`.
-    The memory object that is deallocated is placed in an undefined state
-    after `fir.freemem`.  Optimizations may treat the loading of an object
-    in the undefined state as undefined behavior.  This includes aliasing
-    references, such as the result of an `fir.embox`.
-
-    ```mlir
-      %21 = fir.allocmem !fir.type<ZT(p:i32){field:i32}>
-      ...
-      fir.freemem %21 : !fir.heap<!fir.type<ZT>>
-    ```
-  }];
-
-  let arguments = (ins Arg<fir_HeapType, "", [MemFree]>:$heapref);
-
-  let assemblyFormat = "$heapref attr-dict `:` type($heapref)";
-}
-
 //===----------------------------------------------------------------------===//
 // Terminator operations
 //===----------------------------------------------------------------------===//
diff --git a/lib/Optimizer/Dialect/FIROps.cpp b/lib/Optimizer/Dialect/FIROps.cpp
index 7a73e42..e3757b6 100644
--- a/lib/Optimizer/Dialect/FIROps.cpp
+++ b/lib/Optimizer/Dialect/FIROps.cpp
@@ -2673,13 +2673,40 @@
 //===----------------------------------------------------------------------===//
 
 mlir::Type fir::StoreOp::elementType(mlir::Type refType) {
-  if (auto ref = refType.dyn_cast<ReferenceType>())
-    return ref.getEleTy();
-  if (auto ref = refType.dyn_cast<PointerType>())
-    return ref.getEleTy();
-  if (auto ref = refType.dyn_cast<HeapType>())
-    return ref.getEleTy();
-  return {};
+  return fir::dyn_cast_ptrEleTy(refType);
+}
+
+static mlir::ParseResult parseStoreOp(mlir::OpAsmParser &parser,
+                                      mlir::OperationState &result) {
+  mlir::Type type;
+  mlir::OpAsmParser::OperandType oper;
+  mlir::OpAsmParser::OperandType store;
+  if (parser.parseOperand(oper) || parser.parseKeyword("to") ||
+      parser.parseOperand(store) ||
+      parser.parseOptionalAttrDict(result.attributes) ||
+      parser.parseColonType(type) ||
+      parser.resolveOperand(oper, fir::StoreOp::elementType(type),
+                            result.operands) ||
+      parser.resolveOperand(store, type, result.operands))
+    return mlir::failure();
+  return mlir::success();
+}
+
+static void print(mlir::OpAsmPrinter &p, fir::StoreOp &op) {
+  p << ' ';
+  p.printOperand(op.value());
+  p << " to ";
+  p.printOperand(op.memref());
+  p.printOptionalAttrDict(op.getOperation()->getAttrs(), {});
+  p << " : " << op.memref().getType();
+}
+
+static mlir::LogicalResult verify(fir::StoreOp &op) {
+  if (op.value().getType() != fir::dyn_cast_ptrEleTy(op.memref().getType()))
+    return op.emitOpError("store value type must match memory reference type");
+  if (fir::isa_unknown_size_box(op.value().getType()))
+    return op.emitOpError("cannot store !fir.box of unknown rank or type");
+  return mlir::success();
 }
 
 //===----------------------------------------------------------------------===//