Introduce gpu.func

Introduce a new function-like operation to the GPU dialect to provide a
placeholder for the execution semantic description and to add support for GPU
memory hierarchy.  This aligns with the overall goal of the dialect to expose
the common abstraction layer for GPU devices, in particular by providing an
MLIR unit of semantics (i.e. an operation) for memory modeling.

This proposal has been discussed in the mailing list:
https://groups.google.com/a/tensorflow.org/d/msg/mlir/RfXNP7Hklsc/MBNN7KhjAgAJ
As decided, the "convergence" aspect of the execution model will be factored
out into a new discussion and therefore is not included in this commit. This
commit only introduces the operation but does not hook it up with the remaining
flow. The intention is to develop the new flow while keeping the old flow
operational and do the switch in a simple, separately reversible commit.

PiperOrigin-RevId: 282357599

1 files changed, 61 insertions, 41 deletions

diff --git a/mlir/lib/IR/FunctionSupport.cpp b/mlir/lib/IR/FunctionSupport.cpp
index 1f39575331c..c6f2673ef2a 100644
--- a/mlir/lib/IR/FunctionSupport.cpp
+++ b/mlir/lib/IR/FunctionSupport.cpp
@@ -128,9 +128,11 @@ static ParseResult parseFunctionResultList(
   return parser.parseRParen();
 }
 
-/// Parse a function signature, starting with a name and including the
-/// parameter list.
-static ParseResult parseFunctionSignature(
+/// Parses a function signature using `parser`. The `allowVariadic` argument
+/// indicates whether functions with variadic arguments are supported. The
+/// trailing arguments are populated by this function with names, types and
+/// attributes of the arguments and those of the results.
+ParseResult mlir::impl::parseFunctionSignature(
     OpAsmParser &parser, bool allowVariadic,
     SmallVectorImpl<OpAsmParser::OperandType> &argNames,
     SmallVectorImpl<Type> &argTypes,
@@ -145,6 +147,24 @@ static ParseResult parseFunctionSignature(
   return success();
 }
 
+void mlir::impl::addArgAndResultAttrs(
+    Builder &builder, OperationState &result,
+    ArrayRef<SmallVector<NamedAttribute, 2>> argAttrs,
+    ArrayRef<SmallVector<NamedAttribute, 2>> resultAttrs) {
+  // Add the attributes to the function arguments.
+  SmallString<8> attrNameBuf;
+  for (unsigned i = 0, e = argAttrs.size(); i != e; ++i)
+    if (!argAttrs[i].empty())
+      result.addAttribute(getArgAttrName(i, attrNameBuf),
+                          builder.getDictionaryAttr(argAttrs[i]));
+
+  // Add the attributes to the function results.
+  for (unsigned i = 0, e = resultAttrs.size(); i != e; ++i)
+    if (!resultAttrs[i].empty())
+      result.addAttribute(getResultAttrName(i, attrNameBuf),
+                          builder.getDictionaryAttr(resultAttrs[i]));
+}
+
 /// Parser implementation for function-like operations.  Uses `funcTypeBuilder`
 /// to construct the custom function type given lists of input and output types.
 ParseResult
@@ -158,7 +178,7 @@ mlir::impl::parseFunctionLikeOp(OpAsmParser &parser, OperationState &result,
   SmallVector<Type, 4> resultTypes;
   auto &builder = parser.getBuilder();
 
-  // Parse the name as a symbol reference attribute.
+  // Parse the name as a symbol.
   StringAttr nameAttr;
   if (parser.parseSymbolName(nameAttr, ::mlir::SymbolTable::getSymbolAttrName(),
                              result.attributes))
@@ -185,26 +205,14 @@ mlir::impl::parseFunctionLikeOp(OpAsmParser &parser, OperationState &result,
     return failure();
 
   // Add the attributes to the function arguments.
-  SmallString<8> attrNameBuf;
-  for (unsigned i = 0, e = argTypes.size(); i != e; ++i)
-    if (!argAttrs[i].empty())
-      result.addAttribute(getArgAttrName(i, attrNameBuf),
-                          builder.getDictionaryAttr(argAttrs[i]));
-
-  // Add the attributes to the function results.
-  for (unsigned i = 0, e = resultTypes.size(); i != e; ++i)
-    if (!resultAttrs[i].empty())
-      result.addAttribute(getResultAttrName(i, attrNameBuf),
-                          builder.getDictionaryAttr(resultAttrs[i]));
+  assert(argAttrs.size() == argTypes.size());
+  assert(resultAttrs.size() == resultTypes.size());
+  addArgAndResultAttrs(builder, result, argAttrs, resultAttrs);
 
   // Parse the optional function body.
   auto *body = result.addRegion();
-  if (parser.parseOptionalRegion(*body, entryArgs,
-                                 entryArgs.empty() ? llvm::ArrayRef<Type>()
-                                                   : argTypes))
-    return failure();
-
-  return success();
+  return parser.parseOptionalRegion(
+      *body, entryArgs, entryArgs.empty() ? llvm::ArrayRef<Type>() : argTypes);
 }
 
 // Print a function result list.
@@ -227,9 +235,10 @@ static void printFunctionResultList(OpAsmPrinter &p, ArrayRef<Type> types,
 
 /// Print the signature of the function-like operation `op`.  Assumes `op` has
 /// the FunctionLike trait and passed the verification.
-static void printSignature(OpAsmPrinter &p, Operation *op,
-                           ArrayRef<Type> argTypes, bool isVariadic,
-                           ArrayRef<Type> resultTypes) {
+void mlir::impl::printFunctionSignature(OpAsmPrinter &p, Operation *op,
+                                        ArrayRef<Type> argTypes,
+                                        bool isVariadic,
+                                        ArrayRef<Type> resultTypes) {
   Region &body = op->getRegion(0);
   bool isExternal = body.empty();
 
@@ -264,42 +273,53 @@ static void printSignature(OpAsmPrinter &p, Operation *op,
   }
 }
 
-/// Printer implementation for function-like operations.  Accepts lists of
-/// argument and result types to use while printing.
-void mlir::impl::printFunctionLikeOp(OpAsmPrinter &p, Operation *op,
-                                     ArrayRef<Type> argTypes, bool isVariadic,
-                                     ArrayRef<Type> resultTypes) {
-  // Print the operation and the function name.
-  auto funcName =
-      op->getAttrOfType<StringAttr>(::mlir::SymbolTable::getSymbolAttrName())
-          .getValue();
-  p << op->getName() << ' ';
-  p.printSymbolName(funcName);
-
-  // Print the signature.
-  printSignature(p, op, argTypes, isVariadic, resultTypes);
-
+/// Prints the list of function prefixed with the "attributes" keyword. The
+/// attributes with names listed in "elided" as well as those used by the
+/// function-like operation internally are not printed. Nothing is printed
+/// if all attributes are elided. Assumes `op` has the `FunctionLike` trait and
+/// passed the verification.
+void mlir::impl::printFunctionAttributes(OpAsmPrinter &p, Operation *op,
+                                         unsigned numInputs,
+                                         unsigned numResults,
+                                         ArrayRef<StringRef> elided) {
   // Print out function attributes, if present.
   SmallVector<StringRef, 2> ignoredAttrs = {
       ::mlir::SymbolTable::getSymbolAttrName(), getTypeAttrName()};
+  ignoredAttrs.append(elided.begin(), elided.end());
 
   SmallString<8> attrNameBuf;
 
   // Ignore any argument attributes.
   std::vector<SmallString<8>> argAttrStorage;
-  for (unsigned i = 0, e = argTypes.size(); i != e; ++i)
+  for (unsigned i = 0; i != numInputs; ++i)
     if (op->getAttr(getArgAttrName(i, attrNameBuf)))
       argAttrStorage.emplace_back(attrNameBuf);
   ignoredAttrs.append(argAttrStorage.begin(), argAttrStorage.end());
 
   // Ignore any result attributes.
   std::vector<SmallString<8>> resultAttrStorage;
-  for (unsigned i = 0, e = resultTypes.size(); i != e; ++i)
+  for (unsigned i = 0; i != numResults; ++i)
     if (op->getAttr(getResultAttrName(i, attrNameBuf)))
       resultAttrStorage.emplace_back(attrNameBuf);
   ignoredAttrs.append(resultAttrStorage.begin(), resultAttrStorage.end());
 
   p.printOptionalAttrDictWithKeyword(op->getAttrs(), ignoredAttrs);
+}
+
+/// Printer implementation for function-like operations.  Accepts lists of
+/// argument and result types to use while printing.
+void mlir::impl::printFunctionLikeOp(OpAsmPrinter &p, Operation *op,
+                                     ArrayRef<Type> argTypes, bool isVariadic,
+                                     ArrayRef<Type> resultTypes) {
+  // Print the operation and the function name.
+  auto funcName =
+      op->getAttrOfType<StringAttr>(::mlir::SymbolTable::getSymbolAttrName())
+          .getValue();
+  p << op->getName() << ' ';
+  p.printSymbolName(funcName);
+
+  printFunctionSignature(p, op, argTypes, isVariadic, resultTypes);
+  printFunctionAttributes(p, op, argTypes.size(), resultTypes.size());
 
   // Print the body if this is not an external function.
   Region &body = op->getRegion(0);