[CIR] Add cir::FuncOp operation

This is necessary to enforce the presence of linkage kind, which
is key to continue implementing methods.

While here add some helpers to apply more reliable parsing for
linkage types.

CodeGen does not use cir::FuncOp just yet, patch a few places that
to keep using mlir::FuncOp for now.

Author: bcardosolopes

mlir/include/mlir/Dialect/CIR/IR/CIROps.td

@@ -17,11 +17,12 @@
 include "mlir/Dialect/CIR/IR/CIRDialect.td"
 include "mlir/Dialect/CIR/IR/CIRTypes.td"
 include "mlir/Dialect/CIR/IR/CIRAttrs.td"
-include "mlir/IR/SymbolInterfaces.td"
+include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/FunctionInterfaces.td"
 include "mlir/IR/SymbolInterfaces.td"
 
 //===----------------------------------------------------------------------===//
@@ -327,7 +328,7 @@ def StoreOp : CIR_Op<"store", [
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
-def ReturnOp : CIR_Op<"return", [HasParent<"FuncOp, ScopeOp, IfOp, SwitchOp, LoopOp">,
+def ReturnOp : CIR_Op<"return", [HasParent<"mlir::FuncOp, ScopeOp, IfOp, SwitchOp, LoopOp">,
                                  Terminator]> {
   let summary = "Return from function";
   let description = [{
@@ -1106,4 +1107,98 @@ def StructElementAddr : CIR_Op<"struct_element_addr"> {
   // FIXME: add verifier.
 }
 
-#endif // MLIR_CIR_DIALECT_CIR_OPS
+//===----------------------------------------------------------------------===//
+// FuncOp
+//===----------------------------------------------------------------------===//
+
+def FuncOp : CIR_Op<"func", [
+  AutomaticAllocationScope, CallableOpInterface, FunctionOpInterface,
+  IsolatedFromAbove, Symbol
+]> {
+  let summary = "Declare or define a function";
+  let description = [{
+
+    Similar to `mlir::FuncOp` built-in:
+    > Operations within the function cannot implicitly capture values defined
+    > outside of the function, i.e. Functions are `IsolatedFromAbove`. All
+    > external references must use function arguments or attributes that establish
+    > a symbolic connection (e.g. symbols referenced by name via a string
+    > attribute like SymbolRefAttr). An external function declaration (used when
+    > referring to a function declared in some other module) has no body. While
+    > the MLIR textual form provides a nice inline syntax for function arguments,
+    > they are internally represented as “block arguments” to the first block in
+    > the region.
+    >
+    > Only dialect attribute names may be specified in the attribute dictionaries
+    > for function arguments, results, or the function itself.
+
+    The function linkage information is specified by `linkage`, as defined by
+    `GlobalLinkageKind` attribute.
+
+    Example:
+
+    ```mlir
+    // External function definitions.
+    func @abort()
+
+    // A function with internal linkage.
+    func internal @count(%x: i64) -> (i64)
+      return %x : i64
+    }
+    ```
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name,
+                       TypeAttr:$type,
+                       DefaultValuedAttr<GlobalLinkageKind,
+                                         "GlobalLinkageKind::ExternalLinkage">:$linkage,
+                       OptionalAttr<StrAttr>:$sym_visibility);
+  let regions = (region AnyRegion:$body);
+  let skipDefaultBuilders = 1;
+
+  let builders = [OpBuilder<(ins
+    "StringRef":$name, "FunctionType":$type,
+    CArg<"GlobalLinkageKind", "GlobalLinkageKind::ExternalLinkage">:$linkage,
+    CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs,
+    CArg<"ArrayRef<DictionaryAttr>", "{}">:$argAttrs)
+  >];
+
+  let extraClassDeclaration = [{
+    // Returns the type of this function.
+    // FIXME: We should derive this via the ODS `type` param.
+    FunctionType getType() {
+      return getTypeAttr().getValue().cast<FunctionType>();
+    }
+
+    /// Returns the region on the current operation that is callable. This may
+    /// return null in the case of an external callable object, e.g. an external
+    /// function.
+    ::mlir::Region *getCallableRegion() {
+      return isExternal() ? nullptr : &getBody();
+    }
+
+    /// Returns the results types that the callable region produces when
+    /// executed.
+    ArrayRef<Type> getCallableResults() { return getType().getResults(); }
+
+    /// This trait needs access to the hooks defined below.
+    friend struct FunctionOpInterfaceTrait<mlir::FuncOp>;
+
+    /// Returns the argument types of this function.
+    ArrayRef<Type> getArgumentTypes() { return getType().getInputs(); }
+
+    /// Returns the result types of this function.
+    ArrayRef<Type> getResultTypes() { return getType().getResults(); }
+
+    /// Hook for OpTrait::FunctionOpInterfaceTrait, called after verifying that
+    /// the 'type' attribute is present and checks if it holds a function type.
+    /// Ensures getType, getNumFuncArguments, and getNumFuncResults can be
+    /// called safely.
+    LogicalResult verifyType();
+  }];
+  let parser = [{ return ::parseCIRFuncOp(parser, result); }];
+  let printer = [{ return ::print(*this, p); }];
+  let verifier = [{ return ::verify(*this); }];
+}
+
+#endif // MLIR_CIR_DIALECT_CIR_OPS

mlir/lib/Dialect/CIR/IR/CIRDialect.cpp

@@ -16,6 +16,7 @@
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TypeUtilities.h"
@@ -58,6 +59,54 @@ void cir::CIRDialect::initialize() {
   addInterfaces<CIROpAsmDialectInterface>();
 }
 
+//===----------------------------------------------------------------------===//
+// Helpers
+//===----------------------------------------------------------------------===//
+
+// Parses one of the keywords provided in the list `keywords` and returns the
+// position of the parsed keyword in the list. If none of the keywords from the
+// list is parsed, returns -1.
+static int parseOptionalKeywordAlternative(OpAsmParser &parser,
+                                           ArrayRef<StringRef> keywords) {
+  for (auto en : llvm::enumerate(keywords)) {
+    if (succeeded(parser.parseOptionalKeyword(en.value())))
+      return en.index();
+  }
+  return -1;
+}
+
+namespace {
+template <typename Ty>
+struct EnumTraits {};
+
+#define REGISTER_ENUM_TYPE(Ty)                                                 \
+  template <>                                                                  \
+  struct EnumTraits<Ty> {                                                      \
+    static StringRef stringify(Ty value) { return stringify##Ty(value); }      \
+    static unsigned getMaxEnumVal() { return getMaxEnumValFor##Ty(); }         \
+  }
+
+REGISTER_ENUM_TYPE(GlobalLinkageKind);
+} // namespace
+
+/// Parse an enum from the keyword, or default to the provided default value.
+/// The return type is the enum type by default, unless overriden with the
+/// second template argument.
+/// TODO: teach other places in this file to use this function.
+template <typename EnumTy, typename RetTy = EnumTy>
+static RetTy parseOptionalCIRKeyword(OpAsmParser &parser,
+                                     OperationState &result,
+                                     EnumTy defaultValue) {
+  SmallVector<StringRef, 10> names;
+  for (unsigned i = 0, e = EnumTraits<EnumTy>::getMaxEnumVal(); i <= e; ++i)
+    names.push_back(EnumTraits<EnumTy>::stringify(static_cast<EnumTy>(i)));
+
+  int index = parseOptionalKeywordAlternative(parser, names);
+  if (index == -1)
+    return static_cast<RetTy>(defaultValue);
+  return static_cast<RetTy>(index);
+}
+
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
@@ -190,7 +239,7 @@ static LogicalResult verify(cir::CastOp castOp) {
 //===----------------------------------------------------------------------===//
 
 static mlir::LogicalResult checkReturnAndFunction(ReturnOp op,
-                                                  FuncOp function) {
+                                                  mlir::FuncOp function) {
   // ReturnOps currently only have a single optional operand.
   if (op.getNumOperands() > 1)
     return op.emitOpError() << "expects at most 1 return operand";
@@ -223,11 +272,11 @@ static mlir::LogicalResult verify(ReturnOp op) {
   // Returns can be present in multiple different scopes, get the
   // wrapping function and start from there.
   auto *fnOp = op->getParentOp();
-  while (!isa<FuncOp>(fnOp))
+  while (!isa<mlir::FuncOp>(fnOp))
     fnOp = fnOp->getParentOp();
 
   // Make sure return types match function return type.
-  if (checkReturnAndFunction(op, cast<FuncOp>(fnOp)).failed())
+  if (checkReturnAndFunction(op, cast<mlir::FuncOp>(fnOp)).failed())
     return failure();
 
   return success();
@@ -495,7 +544,7 @@ static LogicalResult verify(ScopeOp op) {
 
 static mlir::LogicalResult verify(YieldOp op) {
   auto isDominatedByLoopOrSwitch = [](Operation *parentOp) {
-    while (!llvm::isa<FuncOp>(parentOp)) {
+    while (!llvm::isa<mlir::FuncOp>(parentOp)) {
       if (llvm::isa<cir::SwitchOp, cir::LoopOp>(parentOp))
         return true;
       parentOp = parentOp->getParentOp();
@@ -504,7 +553,7 @@ static mlir::LogicalResult verify(YieldOp op) {
   };
 
   auto isDominatedByLoop = [](Operation *parentOp) {
-    while (!llvm::isa<FuncOp>(parentOp)) {
+    while (!llvm::isa<mlir::FuncOp>(parentOp)) {
       if (llvm::isa<cir::LoopOp>(parentOp))
         return true;
       parentOp = parentOp->getParentOp();
@@ -1013,15 +1062,15 @@ static LogicalResult verify(GlobalOp op) {
   }
 
   switch (op.linkage()) {
-  case mlir::cir::GlobalLinkageKind::InternalLinkage:
-  case mlir::cir::GlobalLinkageKind::PrivateLinkage:
+  case GlobalLinkageKind::InternalLinkage:
+  case GlobalLinkageKind::PrivateLinkage:
     if (op.isPublic())
       return op->emitError()
              << "public visibility not allowed with '"
              << stringifyGlobalLinkageKind(op.linkage()) << "' linkage";
     break;
-  case mlir::cir::GlobalLinkageKind::ExternalLinkage:
-  case mlir::cir::GlobalLinkageKind::ExternalWeakLinkage:
+  case GlobalLinkageKind::ExternalLinkage:
+  case GlobalLinkageKind::ExternalWeakLinkage:
     if (op.isPrivate())
       return op->emitError()
              << "private visibility not allowed with '"
@@ -1072,6 +1121,144 @@ GetGlobalOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// FuncOp
+//===----------------------------------------------------------------------===//
+
+/// Returns the name used for the linkage attribute. This *must* correspond to
+/// the name of the attribute in ODS.
+static StringRef getLinkageAttrName() { return "linkage"; }
+
+void cir::FuncOp::build(OpBuilder &builder, OperationState &result,
+                        StringRef name, FunctionType type,
+                        GlobalLinkageKind linkage,
+                        ArrayRef<NamedAttribute> attrs,
+                        ArrayRef<DictionaryAttr> argAttrs) {
+  result.addRegion();
+  result.addAttribute(SymbolTable::getSymbolAttrName(),
+                      builder.getStringAttr(name));
+  result.addAttribute(getTypeAttrName(), TypeAttr::get(type));
+  result.addAttribute(getLinkageAttrName(), GlobalLinkageKindAttr::get(
+                                                builder.getContext(), linkage));
+  result.attributes.append(attrs.begin(), attrs.end());
+  if (argAttrs.empty())
+    return;
+
+  function_interface_impl::addArgAndResultAttrs(builder, result, argAttrs,
+                                                /*resultAttrs=*/llvm::None);
+}
+
+static ParseResult parseCIRFuncOp(OpAsmParser &parser, OperationState &state) {
+  // Default to external linkage if no keyword is provided.
+  state.addAttribute(
+      getLinkageAttrName(),
+      GlobalLinkageKindAttr::get(
+          parser.getContext(),
+          parseOptionalCIRKeyword<GlobalLinkageKind>(
+              parser, state, GlobalLinkageKind::ExternalLinkage)));
+
+  StringAttr nameAttr;
+  SmallVector<OpAsmParser::OperandType, 8> entryArgs;
+  SmallVector<NamedAttrList, 1> argAttrs;
+  SmallVector<Optional<Location>, 1> argLocations;
+  SmallVector<NamedAttrList, 1> resultAttrs;
+  SmallVector<Type, 8> argTypes;
+  SmallVector<Type, 4> resultTypes;
+  auto &builder = parser.getBuilder();
+
+  // Parse the name as a symbol.
+  if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),
+                             state.attributes))
+    return failure();
+
+  // Parse the function signature.
+  bool isVariadic = false;
+  if (function_interface_impl::parseFunctionSignature(
+          parser, /*allowVariadic=*/false, entryArgs, argTypes, argAttrs,
+          argLocations, isVariadic, resultTypes, resultAttrs))
+    return failure();
+
+  auto fnType = builder.getFunctionType(argTypes, resultTypes);
+  state.addAttribute(function_interface_impl::getTypeAttrName(),
+                     TypeAttr::get(fnType));
+
+  // If additional attributes are present, parse them.
+  if (parser.parseOptionalAttrDictWithKeyword(state.attributes))
+    return failure();
+
+  // Add the attributes to the function arguments.
+  assert(argAttrs.size() == argTypes.size());
+  assert(resultAttrs.size() == resultTypes.size());
+  function_interface_impl::addArgAndResultAttrs(builder, state, argAttrs,
+                                                resultAttrs);
+
+  // Parse the optional function body.
+  auto *body = state.addRegion();
+  OptionalParseResult result = parser.parseOptionalRegion(
+      *body, entryArgs, entryArgs.empty() ? ArrayRef<Type>() : argTypes);
+  return failure(result.hasValue() && failed(*result));
+}
+
+static void print(cir::FuncOp op, OpAsmPrinter &p) {
+  p << ' ';
+  if (op.linkage() != GlobalLinkageKind::ExternalLinkage)
+    p << stringifyGlobalLinkageKind(op.linkage()) << ' ';
+
+  // Print function name, signature, and control.
+  p.printSymbolName(op.sym_name());
+  auto fnType = op.getType();
+  function_interface_impl::printFunctionSignature(p, op, fnType.getInputs(),
+                                                  /*isVariadic=*/false,
+                                                  fnType.getResults());
+  function_interface_impl::printFunctionAttributes(p, op, fnType.getNumInputs(),
+                                                   fnType.getNumResults(), {});
+
+  // Print the body if this is not an external function.
+  Region &body = op.body();
+  if (!body.empty())
+    p.printRegion(body, /*printEntryBlockArgs=*/false,
+                  /*printBlockTerminators=*/true);
+}
+
+// Hook for OpTrait::FunctionLike, called after verifying that the 'type'
+// attribute is present.  This can check for preconditions of the
+// getNumArguments hook not failing.
+LogicalResult cir::FuncOp::verifyType() {
+  auto type = getTypeAttr().getValue();
+  if (!type.isa<FunctionType>())
+    return emitOpError("requires '" + getTypeAttrName() +
+                       "' attribute of function type");
+  if (getType().getNumResults() > 1)
+    return emitOpError("cannot have more than one result");
+  return success();
+}
+
+// Verifies linkage types, similar to LLVM:
+// - functions don't have 'common' linkage
+// - external functions have 'external' or 'extern_weak' linkage
+static LogicalResult verify(cir::FuncOp op) {
+  if (op.linkage() == cir::GlobalLinkageKind::CommonLinkage)
+    return op.emitOpError()
+           << "functions cannot have '"
+           << stringifyGlobalLinkageKind(cir::GlobalLinkageKind::CommonLinkage)
+           << "' linkage";
+
+  if (op.isExternal()) {
+    if (op.linkage() != cir::GlobalLinkageKind::ExternalLinkage &&
+        op.linkage() != cir::GlobalLinkageKind::ExternalWeakLinkage)
+      return op.emitOpError()
+             << "external functions must have '"
+             << stringifyGlobalLinkageKind(
+                    cir::GlobalLinkageKind::ExternalLinkage)
+             << "' or '"
+             << stringifyGlobalLinkageKind(
+                    cir::GlobalLinkageKind::ExternalWeakLinkage)
+             << "' linkage";
+    return success();
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // CIR defined traits
 //===----------------------------------------------------------------------===//