[fir] Update fir.allocmem op

[clementval]

Summary

Updatet the fir.allocmem operation.

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

Co-authored-by: Eric Schweitz <[e-mail]>

[clementval]

Move verifier to the .cpp file

[mehdi_amini]

flang/lib/Optimizer/Dialect/FIROps.cpp

86

Can you document the syntax here?

I'm also curious if this could just use the declarative assembly instead?

[clementval]

Address some review comments

flang/lib/Optimizer/Dialect/FIROps.cpp

86

Let me give it a try with the declarative assembly.

[mehdi_amini]

flang/lib/Optimizer/Dialect/FIROps.cpp

185

I'm not sure about using this in the builder API?
It seems like it'll yield confusing errors at runtime, could this be just an assert here instead?

assert(intype.isa<ReferenceType, HeapType, PointerType, FunctionType>());

((the assert may be already in the HeadType ctor by the way)

[clementval]

flang/lib/Optimizer/Dialect/FIROps.cpp

185

Right. I'll make a separate patch for this.

flang/include/flang/Optimizer/Dialect/FIROps.td

let verifier = allocVerify#[{
     return mlir::success();
   }];
 
   let extraClassDeclaration = extraAllocClassDeclaration#[{
     static mlir::Type wrapResultType(mlir::Type intype);
   }];
 }
 
+def fir_AllocMemOp : fir_Op<"allocmem",
+    [MemoryEffects<[MemAlloc<DefaultResource>]>, AttrSizedOperandSegments]> {
+  let summary = "allocate storage on the heap for an object of a given type";
+
+  let description = [{
+    Creates a heap memory reference suitable for storing a value of the
+    given type, T.  The heap refernce returned has type `!fir.heap<T>`.
+    The memory object is in an undefined state.  `allocmem` operations must
+    be paired with `freemem` operations to avoid memory leaks.
+
+    ```mlir
+      %0 = fir.allocmem !fir.array<10 x f32>
+      fir.freemem %0 : !fir.heap<!fir.array<10 x f32>>
+    ```
+  }];
+
+  let arguments = (ins
+    TypeAttr:$in_type,
+    OptionalAttr<StrAttr>:$uniq_name,
+    OptionalAttr<StrAttr>:$bindc_name,
+    Variadic<AnyIntegerType>:$typeparams,
+    Variadic<AnyIntegerType>:$shape
+  );
+  let results = (outs fir_HeapType);
+
+  let parser =
+      "return parseAllocatableOp(wrapAllocMemResultType, parser, result);";
+  let printer = "printAllocatableOp(p, (*this));";
+
+  let builders = [
+    OpBuilder<(ins "mlir::Type":$in_type, "llvm::StringRef":$uniq_name,
+      "llvm::StringRef":$bindc_name, CArg<"mlir::ValueRange", "{}">:$typeparams,
+      CArg<"mlir::ValueRange", "{}">:$shape,
+      CArg<"llvm::ArrayRef<mlir::NamedAttribute>", "{}">:$attributes)>,
+    OpBuilder<(ins "mlir::Type":$in_type, "llvm::StringRef":$uniq_name,
+      CArg<"mlir::ValueRange", "{}">:$typeparams,
+      CArg<"mlir::ValueRange", "{}">:$shape,
+      CArg<"llvm::ArrayRef<mlir::NamedAttribute>", "{}">:$attributes)>,
+    OpBuilder<(ins "mlir::Type":$in_type,
+      CArg<"mlir::ValueRange", "{}">:$typeparams,
+      CArg<"mlir::ValueRange", "{}">:$shape,
+      CArg<"llvm::ArrayRef<mlir::NamedAttribute>", "{}">:$attributes)>];
+
+  let verifier = [{ return ::verify(*this); }];
+
+  let extraClassDeclaration = [{
+    mlir::Type getAllocatedType();
+    bool hasLenParams() { return !typeparams().empty(); }
+    bool hasShapeOperands() { return !shape().empty(); }
+    unsigned numLenParams() { return typeparams().size(); }
+    operand_range getLenParams() { return typeparams(); }
+    unsigned numShapeOperands() { return shape().size(); }
+    operand_range getShapeOperands() { return shape(); }
+    static mlir::Type getRefTy(mlir::Type ty);
+    mlir::Type getInType() { return in_type(); }
+  }];
+}
+
 def fir_LoadOp : fir_OneResultOp<"load"> {
   let summary = "load a value from a memory reference";
   let description = [{
     Load a value from a memory reference into an ssa-value (virtual register).
     Produces an immutable ssa-value of the referent type. A memory reference
     has type `!fir.ref<T>`, `!fir.heap<T>`, or `!fir.ptr<T>`.
 
     ```mlir
let description = [{
     The example creates a value of type box where all bits are zero.
   }];
 
   let results = (outs AnyType:$intype);
 
   let assemblyFormat = "type($intype) attr-dict";
 }
 
-def fir_AllocMemOp : fir_AllocatableOp<"allocmem", DefaultResource> {
-  let summary = "allocate storage on the heap for an object of a given type";
-
-  let description = [{
-    Creates a heap memory reference suitable for storing a value of the
-    given type, T.  The heap refernce returned has type `!fir.heap<T>`.
-    The memory object is in an undefined state.  `allocmem` operations must
-    be paired with `freemem` operations to avoid memory leaks.
-
-    ```mlir
-      %0 = fir.allocmem !fir.array<10 x f32>
-      fir.freemem %0 : !fir.heap<!fir.array<10 x f32>>
-    ```
-  }];
-
-  let results = (outs fir_HeapType);
-
-  let verifier = allocVerify#[{
-    mlir::Type outType = getType();
-    if (!outType.dyn_cast<fir::HeapType>())
-      return emitOpError("must be a !fir.heap type");
-    if (fir::isa_unknown_size_box(fir::dyn_cast_ptrEleTy(outType)))
-      return emitOpError("cannot allocate !fir.box of unknown rank or type");
-    return mlir::success();
-  }];
-
-  let extraClassDeclaration = extraAllocClassDeclaration#[{
-    static mlir::Type wrapResultType(mlir::Type intype);
-  }];
-}
-
 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

flang/lib/Optimizer/Dialect/FIROps.cpp

static bool verifyRecordLenParams(mlir::Type inType, unsigned numLenParams) {
   if (numLenParams > 0) {
     if (auto rt = inType.dyn_cast<fir::RecordType>())
       return numLenParams != rt.getNumLenParams();
     return true;
   }
   return false;
 }
 
+static bool verifyTypeParamCount(mlir::Type inType, unsigned numParams) {
+  auto ty = fir::unwrapSequenceType(inType);
+  if (numParams > 0) {
+    if (auto recTy = ty.dyn_cast<fir::RecordType>())
+      return numParams != recTy.getNumLenParams();
+    if (auto chrTy = ty.dyn_cast<fir::CharacterType>())
+      return !(numParams == 1 && chrTy.hasDynamicLen());
+    return true;
+  }
+  if (auto chrTy = ty.dyn_cast<fir::CharacterType>())
+    return !chrTy.hasConstantLen();
+  return false;
+}
+
+/// Parser shared by Alloca and Allocmem
+///
+/// operation ::= %res = (`fir.alloca` | `fir.allocmem`) $in_type
+///                      ( `(` $typeparams `)` )? ( `,` $shape )?
+///                      attr-dict-without-keyword
+template <typename FN>
+static mlir::ParseResult parseAllocatableOp(FN wrapResultType,
+                                            mlir::OpAsmParser &parser,
+                                            mlir::OperationState &result) {
+  mlir::Type intype;
+  if (parser.parseType(intype))
+    return mlir::failure();
+  auto &builder = parser.getBuilder();
+  result.addAttribute("in_type", mlir::TypeAttr::get(intype));
+  llvm::SmallVector<mlir::OpAsmParser::OperandType> operands;
+  llvm::SmallVector<mlir::Type> typeVec;
+  bool hasOperands = false;
+  std::int32_t typeparamsSize = 0;
+  if (!parser.parseOptionalLParen()) {
+    // parse the LEN params of the derived type. (<params> : <types>)
+    if (parser.parseOperandList(operands, mlir::OpAsmParser::Delimiter::None) ||
+        parser.parseColonTypeList(typeVec) || parser.parseRParen())
+      return mlir::failure();
+    typeparamsSize = operands.size();
+    hasOperands = true;
+  }
+  std::int32_t shapeSize = 0;
+  if (!parser.parseOptionalComma()) {
+    // parse size to scale by, vector of n dimensions of type index
+    if (parser.parseOperandList(operands, mlir::OpAsmParser::Delimiter::None))
+      return mlir::failure();
+    shapeSize = operands.size() - typeparamsSize;
+    auto idxTy = builder.getIndexType();
+    for (std::int32_t i = typeparamsSize, end = operands.size(); i != end; ++i)
+      typeVec.push_back(idxTy);
+    hasOperands = true;
+  }
+  if (hasOperands &&
+      parser.resolveOperands(operands, typeVec, parser.getNameLoc(),
+                             result.operands))
+    return mlir::failure();
+  mlir::Type restype = wrapResultType(intype);
+  if (!restype) {
+    parser.emitError(parser.getNameLoc(), "invalid allocate type: ") << intype;
+    return mlir::failure();
+  }
+  result.addAttribute("operand_segment_sizes",
+                      builder.getI32VectorAttr({typeparamsSize, shapeSize}));
+  if (parser.parseOptionalAttrDict(result.attributes) ||
+      parser.addTypeToList(restype, result.types))
+    return mlir::failure();
+  return mlir::success();
+}
+
+template <typename OP>
+static void printAllocatableOp(mlir::OpAsmPrinter &p, OP &op) {
+  p << ' ' << op.in_type();
+  if (!op.typeparams().empty()) {
+    p << '(' << op.typeparams() << " : " << op.typeparams().getTypes() << ')';
+  }
+  // print the shape of the allocation (if any); all must be index type
+  for (auto sh : op.shape()) {
+    p << ", ";
+    p.printOperand(sh);
+  }
+  p.printOptionalAttrDict(op->getAttrs(), {"in_type", "operand_segment_sizes"});
+}
+
 //===----------------------------------------------------------------------===//
 // AllocaOp
 //===----------------------------------------------------------------------===//
 
 mlir::Type fir::AllocaOp::getAllocatedType() {
   return getType().cast<ReferenceType>().getEleTy();
 }
 
 /// Create a legal memory reference as return type
 mlir::Type fir::AllocaOp::wrapResultType(mlir::Type intype) {
   // FIR semantics: memory references to memory references are disallowed
   if (intype.isa<ReferenceType>())
     return {};
   return ReferenceType::get(intype);
 }
 
 mlir::Type fir::AllocaOp::getRefTy(mlir::Type ty) {
   return ReferenceType::get(ty);
 }
 
 //===----------------------------------------------------------------------===//
 // AllocMemOp
 //===----------------------------------------------------------------------===//
 
-mlir::Type fir::AllocMemOp::getAllocatedType() {
-  return getType().cast<HeapType>().getEleTy();
-}
-
-mlir::Type fir::AllocMemOp::getRefTy(mlir::Type ty) {
-  return HeapType::get(ty);
-}
-
 /// Create a legal heap reference as return type
-mlir::Type fir::AllocMemOp::wrapResultType(mlir::Type intype) {
+static mlir::Type wrapAllocMemResultType(mlir::Type intype) {
   // Fortran semantics: C852 an entity cannot be both ALLOCATABLE and POINTER
   // 8.5.3 note 1 prohibits ALLOCATABLE procedures as well
   // FIR semantics: one may not allocate a memory reference value
   if (intype.isa<ReferenceType>() || intype.isa<HeapType>() ||
       intype.isa<PointerType>() || intype.isa<FunctionType>())
     return {};
   return HeapType::get(intype);
 }
 
+mlir::Type fir::AllocMemOp::getAllocatedType() {
+  return getType().cast<HeapType>().getEleTy();
+}
+
+mlir::Type fir::AllocMemOp::getRefTy(mlir::Type ty) {
+  return HeapType::get(ty);
+}
+
+void fir::AllocMemOp::build(mlir::OpBuilder &builder,
+                            mlir::OperationState &result, mlir::Type inType,
+                            llvm::StringRef uniqName,
+                            mlir::ValueRange typeparams, mlir::ValueRange shape,
+                            llvm::ArrayRef<mlir::NamedAttribute> attributes) {
+  auto nameAttr = builder.getStringAttr(uniqName);
+  build(builder, result, wrapAllocMemResultType(inType), inType, nameAttr, {},
+        typeparams, shape);
+  result.addAttributes(attributes);
+}
+
+void fir::AllocMemOp::build(mlir::OpBuilder &builder,
+                            mlir::OperationState &result, mlir::Type inType,
+                            llvm::StringRef uniqName, llvm::StringRef bindcName,
+                            mlir::ValueRange typeparams, mlir::ValueRange shape,
+                            llvm::ArrayRef<mlir::NamedAttribute> attributes) {
+  auto nameAttr = builder.getStringAttr(uniqName);
+  auto bindcAttr = builder.getStringAttr(bindcName);
+  build(builder, result, wrapAllocMemResultType(inType), inType, nameAttr,
+        bindcAttr, typeparams, shape);
+  result.addAttributes(attributes);
+}
+
+void fir::AllocMemOp::build(mlir::OpBuilder &builder,
+                            mlir::OperationState &result, mlir::Type inType,
+                            mlir::ValueRange typeparams, mlir::ValueRange shape,
+                            llvm::ArrayRef<mlir::NamedAttribute> attributes) {
+  build(builder, result, wrapAllocMemResultType(inType), inType, {}, {},
+        typeparams, shape);
+  result.addAttributes(attributes);
+}
+
+static mlir::LogicalResult verify(fir::AllocMemOp op) {
+  llvm::SmallVector<llvm::StringRef> visited;
+  if (verifyInType(op.getInType(), visited, op.numShapeOperands()))
+    return op.emitOpError("invalid type for allocation");
+  if (verifyTypeParamCount(op.getInType(), op.numLenParams()))
+    return op.emitOpError("LEN params do not correspond to type");
+  mlir::Type outType = op.getType();
+  if (!outType.dyn_cast<fir::HeapType>())
+    return op.emitOpError("must be a !fir.heap type");
+  if (fir::isa_unknown_size_box(fir::dyn_cast_ptrEleTy(outType)))
+    return op.emitOpError("cannot allocate !fir.box of unknown rank or type");
+  return mlir::success();
+}
+
 //===----------------------------------------------------------------------===//
 // ArrayCoorOp
 //===----------------------------------------------------------------------===//
 
 static mlir::LogicalResult verify(fir::ArrayCoorOp op) {
   auto eleTy = fir::dyn_cast_ptrOrBoxEleTy(op.memref().getType());
   auto arrTy = eleTy.dyn_cast<fir::SequenceType>();
   if (!arrTy)