[mlir][gpu] Add `gpu.binary` op and `#gpu.object` attribute.

[fmorac]

Summary

For an explanation of these patches see D154153.

Commit message:
Adds the #gpu.object attribute for holding a binary object and the target
attribute used to create it. Also adds the gpu.binary operation used to
store GPU objects.

Depends on D154108

[fmorac]

Adds D154098 to the commit message as a dependency.

[fmorac]

Rebasing.

[fmorac]

Rebasing.

[mehdi_amini]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1062

What kind of "information of how the object was generated" are stored and how? From the rest of the patches seems like it is "the binary string", period.

1070

Can the target be defined as GPUTargetAttr?

I'm also not sure it is a good idea to use a StringAttr here, it can be large storage that we could store differently than in the MLIRContext. Maybe it's good enough to start with and we can update later (I'm just worried it means we'll never get to it...)

1089

We should document what is the objectManager role here.

1094

the example syntax does not show the objectManager?

1102

Shall we move the attr-dict-with-keyword after the objects?

Also is the keyword needed?

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

1551

Why do you need this custom print/parse here?

Edit: maybe it'll only be needed in a subsequent revision, in which case that's fine...

[fmorac]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1062

It also stores GPUTargetAttr, which has the information on how the object was generated encoded in the method. Also the SelectObjectManager can use the attribute to select the appropriate target.

1070

I think there's an 'egg-chicken' dependency with ObjectManagerAttrInterface if we have it use BinaryOp, but I'll come back with a better answer.

I used StringAttr because the old pipeline uses StringAttr, however we can create a new parameter type that stores strings as smart ptrs?

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

1551

Yes, this is changed in D154147 with the introduction of the SelectObjectManager as the default one.

[fmorac]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1089

I'll add it.

1094

In change that in D154147, I didn't want to introduce a name of something that is not there (SelectObjectManager).

1102

With keywords is not needed.

I put it before $objects, because $objects can be thousands of characters long so if they are before, it's easier to edit them by hand.

[mehdi_amini]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1062

oh I see the "information of how the object was generated" applies to the whole pair and not the binary string. I think I misread this...

1070

We should have a "StringAttrInterface" ;)
(same for all blob-like attributes I guess)

1070

(this is far out of scope to create such interface of course)

1102

makes sense!

[krzysz00]

This all seems reasonable to me

[fmorac]

Switched ObjectManager to BinaryHandler.
Switched the types of the parameters from generic Attributes to the respective interface in gpu.object & gpu.binary.

[mehdi_amini]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1089

It's not documented yet?

[fmorac]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1089

I was thinking on add it on D154147. Because on that patch I expand on the docs because there I have the select_object handler, but I can add it here.

[fmorac]

Improved the docs & moved ObjectAttr to CompilationAttrs.td.

[fmorac]

Update the gpu.binary attributes to use OffloadingTranslationAttrTrait
instead of OffloadingLLVMTranslationAttrInterface.

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

]> {
   let summary = "A pseudo op that marks the end of a gpu.module.";
   let description = [{
     This op terminates the only block inside the only region of a `gpu.module`.
   }];
 
   let assemblyFormat = "attr-dict";
 }
 
+def GPU_ObjectAttr : GPU_Attr<"Object", "object"> {
+  let description = [{
+    A GPU object attribute pairs a GPU target with a binary string,
+    encapsulating the information of how the object was generated with the
+    object itself.
+
+    The target attribute must implement the `TargetAttrInterface` interface.
+
+    ```
+      #gpu.object<#gpu.nvptx, "...">
+    ```
+  }];
+  let parameters = (ins "Attribute":$target, "StringAttr":$object);
+  let assemblyFormat = [{`<` $target `,` $object `>`}];
+  let genVerifyDecl = 1;
+}
+
+def GPUObjectArrayAttr :
+  TypedArrayAttrBase<GPU_ObjectAttr, "An array of GPU object attributes">;
+
+def GPU_BinaryOp : GPU_Op<"binary", [Symbol, GlobalSymbol]>,
+  Arguments<(ins SymbolNameAttr:$sym_name,
+                 GPUObjectManagerAttr:$objectManager,
+                 ConfinedAttr<GPUObjectArrayAttr, [ArrayMinCount<1>]>:$objects)> {
+  let summary = "An op for storing serialized GPU binary objects.";
+  let description = [{
+    GPU binaries provide a semantic mechanism for storing GPU objects,
+    e.g. the result of compiling a GPU module to an object file.
+
+    This operation has 3 arguments:
+     - The name of the binary.
+     - An attribute implementing the `ObjectManagerAttrInterface` interface.
+     - An array of GPU object attributes.
+
+    ```
+      gpu.binary @myobject [#gpu.object<...>, #gpu.object<...>]
+    ```
+  }];
+  let builders = [
+    OpBuilder<(ins "StringRef":$name, "Attribute":$objectManager,
+                   "ArrayAttr":$objects)>
+  ];
+  let skipDefaultBuilders = 1;
+  let assemblyFormat = [{
+    $sym_name custom<ObjectManager>($objectManager) attr-dict-with-keyword $objects
+  }];
+}
+
 def GPU_HostRegisterOp : GPU_Op<"host_register">,
     Arguments<(ins AnyUnrankedMemRef:$value)> {
   let summary = "Registers a memref for access from device.";
   let description = [{
     This op maps the provided host buffer into the device address space.
 
     This operation may not be supported in every environment, there is not yet a
     way to check at runtime whether this feature is supported.

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

p.printOptionalAttrDictWithKeyword(
       (*this)->getAttrs(),
       {mlir::SymbolTable::getSymbolAttrName(), getTargetsAttrName()});
   p << ' ';
   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/false);
 }
 
 //===----------------------------------------------------------------------===//
+// GPUBinaryOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+ObjectAttr::verify(function_ref<mlir::InFlightDiagnostic()> emitError,
+                   Attribute target, StringAttr object) {
+  if (target && target.hasTrait<TargetAttrInterface::Trait>())
+    return success();
+  emitError() << "The target parameter must implement `TargetAttrInterface`.";
+  return failure();
+}
+
+void BinaryOp::build(OpBuilder &builder, OperationState &result, StringRef name,
+                     Attribute manager, ArrayAttr objects) {
+  auto &properties = result.getOrAddProperties<Properties>();
+  result.attributes.push_back(builder.getNamedAttr(
+      SymbolTable::getSymbolAttrName(), builder.getStringAttr(name)));
+  properties.objects = objects;
+  properties.objectManager = manager;
+}
+
+static ParseResult parseObjectManager(OpAsmParser &parser,
+                                      Attribute &objectManager) {
+  if (parser.parseAttribute(objectManager))
+    return failure();
+  return success();
+}
+
+static void printObjectManager(OpAsmPrinter &printer, Operation *op,
+                               Attribute objectManager) {
+  if (objectManager)
+    printer << '<' << objectManager << '>';
+}
+
+//===----------------------------------------------------------------------===//
 // GPUMemcpyOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult MemcpyOp::verify() {
   auto srcType = getSrc().getType();
   auto dstType = getDst().getType();
 
   if (getElementTypeOrSelf(srcType) != getElementTypeOrSelf(dstType))