[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

1074

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.

1082

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...)

1101

We should document what is the objectManager role here.

1106

the example syntax does not show the objectManager?

1114

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

Also is the keyword needed?

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

1613

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

1074

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.

1082

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

1613

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

1101

I'll add it.

1106

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

1114

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

1074

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...

1082

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

1082

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

1114

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

1101

It's not documented yet?

[fmorac]

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

1101

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/CompilationAttrs.td

+//===-- CompilationAttrs.td - GPU compilation attributes ---*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines GPU compilation related attributes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef GPU_COMPILATION_ATTRS
+#define GPU_COMPILATION_ATTRS
+
+include "mlir/Dialect/GPU/IR/GPUBase.td"
+include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
+
+//===----------------------------------------------------------------------===//
+// GPU object attribute.
+//===----------------------------------------------------------------------===//
+
+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<#nvvm.target, "...">
+    ```
+  }];
+  let parameters = (ins "TargetAttrInterface":$target, "StringAttr":$object);
+  let assemblyFormat = [{`<` $target `,` $object `>`}];
+}
+
+def GPUObjectArrayAttr :
+  TypedArrayAttrBase<GPU_ObjectAttr, "an array of GPU object attributes">;
+
+#endif // GPU_COMPILATION_ATTRS

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

 //===----------------------------------------------------------------------===//
 
 #ifndef GPU_OPS
 #define GPU_OPS
 
 include "mlir/Dialect/DLTI/DLTIBase.td"
 include "mlir/Dialect/GPU/IR/GPUBase.td"
 include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
+include "mlir/Dialect/GPU/IR/CompilationAttrs.td"
 include "mlir/Dialect/GPU/IR/ParallelLoopMapperAttr.td"
 include "mlir/Dialect/GPU/TransformOps/GPUDeviceMappingAttr.td"
 include "mlir/IR/EnumAttr.td"
 include "mlir/IR/FunctionInterfaces.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
 include "mlir/Interfaces/InferIntRangeInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.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_BinaryOp : GPU_Op<"binary", [Symbol]>, Arguments<(ins
+      SymbolNameAttr:$sym_name,
+      OptionalAttr<GPUOffloadingLLVMTranslationAttrInterface>:$offloadingHandler,
+      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 optional attribute implementing the offloading LLVM translation interface.
+     - An array of GPU object attributes.
+
+    During translation into LLVM, the offloading attribute will be called
+    for translating GPU binary and launch operations into LLVM instructions. If
+    no attribute is provided, the default handler selects the first object from
+    the array and embeds it as a string.
+
+    Examples:
+    ```
+      gpu.binary @myobject [#gpu.object<...>, #gpu.object<...>]
+      gpu.binary @myobject <#foo.my_handler> [#gpu.object<...>, #gpu.object<...>]
+    ```
+  }];
+  let builders = [
+    OpBuilder<(ins "StringRef":$name,
+                   "OffloadingLLVMTranslationAttrInterface":$offloadingHandler,
+                   "ArrayAttr":$objects)>,
+    OpBuilder<(ins "StringRef":$name,
+                   "OffloadingLLVMTranslationAttrInterface":$offloadingHandler,
+                   "ArrayRef<Attribute>":$objects)>
+  ];
+  let skipDefaultBuilders = 1;
+  let assemblyFormat = [{
+    $sym_name (`<` $offloadingHandler ^ `>`)? attr-dict $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.
 
     Writes from the host are guaranteed to be visible to device kernels that are
     launched afterwards. Writes from the device are guaranteed to be visible on
     the host after synchronizing with the device kernel completion.
   }];

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

 
 void GPUModuleOp::setTargets(ArrayRef<TargetAttrInterface> targets) {
   ArrayAttr &targetsAttr = getProperties().targets;
   SmallVector<Attribute> targetsVector(targets);
   targetsAttr = ArrayAttr::get(getContext(), targetsVector);
 }
 
 //===----------------------------------------------------------------------===//
+// GPUBinaryOp
+//===----------------------------------------------------------------------===//
+void BinaryOp::build(OpBuilder &builder, OperationState &result, StringRef name,
+                     OffloadingLLVMTranslationAttrInterface offloadingHandler,
+                     ArrayAttr objects) {
+  auto &properties = result.getOrAddProperties<Properties>();
+  result.attributes.push_back(builder.getNamedAttr(
+      SymbolTable::getSymbolAttrName(), builder.getStringAttr(name)));
+  properties.objects = objects;
+  properties.offloadingHandler = offloadingHandler;
+}
+
+void BinaryOp::build(OpBuilder &builder, OperationState &result, StringRef name,
+                     OffloadingLLVMTranslationAttrInterface offloadingHandler,
+                     ArrayRef<Attribute> objects) {
+  build(builder, result, name, offloadingHandler,
+        objects.size() > 0 ? builder.getArrayAttr(objects) : ArrayAttr());
+}
+
+//===----------------------------------------------------------------------===//
 // GPUMemcpyOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult MemcpyOp::verify() {
   auto srcType = getSrc().getType();
   auto dstType = getDst().getType();
 
   if (getElementTypeOrSelf(srcType) != getElementTypeOrSelf(dstType))
     return emitOpError("arguments have incompatible element type");
 
   if (failed(verifyCompatibleShape(srcType, dstType)))
     return emitOpError("arguments have incompatible shape");
 
   return success();
 }
 
 namespace {
 
 /// Erases a common case of copy ops where a destination value is used only by
 /// the copy op, alloc and dealloc ops.
 struct EraseTrivialCopyOp : public OpRewritePattern<MemcpyOp> {

mlir/test/Dialect/GPU/invalid.mlir

 
 // -----
 
 module {
   // expected-error @+1 {{'gpu.module' op attribute 'targets' failed to satisfy constraint: array of GPU target attributes with at least 1 elements}}
   gpu.module @gpu_funcs [1] {
   }
 }
+
+// -----
+
+module {
+  // expected-error @+1 {{'gpu.binary' op attribute 'objects' failed to satisfy constraint: an array of GPU object attributes with at least 1 elements}}
+  gpu.binary @binary []
+}
+
+// -----
+
+module {
+  // expected-error @+1 {{custom op 'gpu.binary' invalid kind of attribute specified}}
+  gpu.binary @binary <1> [#gpu.object<#nvvm.target, "">]
+}