[mlir][gpu] Add target attribute to GPU modules.

[fmorac]

Summary

For an explanation of these patches see D154153.

Commit message:
Adds support for Target attributes in GPU modules. This change enables attaching
an optional non empty array of GPU target attributes to the module.

Depends on D154104

[fmorac]

Rebasing.

[krzysz00]

Seems reasonable to me, I'd be willing to approve

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

1040

We could also have a version of this that takes ArrayRef<Attribute> and/or ArrayRef<GPUTargetAttr> (not sure of that second name)

[fmorac]

Added a builder taking ArrayRef<Attribute> and a couple of utility methods.

[mehdi_amini]

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

1020

Can you update the description explaining what does it mean to have a target, when is it used in the pipeline, etc.?

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

1517

What about forwarding to the other builder above?

1531

You can just write result.getOrAddProperties<Properties>() I believe

[fmorac]

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

1020

Yes.

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

1517

I'll make the change to fwd it.

[fmorac]

Addressed comments.

[fmorac]

Rebasing.

[fmorac]

Rebasing.

[fmorac]

Rebasing.

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/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 description = [{
 
     Either none or all work items of a workgroup need to execute this op
     in convergence.
   }];
   let assemblyFormat = "attr-dict";
 }
 
 def GPU_GPUModuleOp : GPU_Op<"module", [
     DataLayoutOpInterface, HasDefaultDLTIDataLayout, IsolatedFromAbove,
     SymbolTable, Symbol,
     SingleBlockImplicitTerminator<"ModuleEndOp">
-]> {
+  ]>, Arguments<(ins OptionalAttr<GPUNonEmptyTargetArrayAttr>:$targets)> {
   let summary = "A top level compilation unit containing code to be run on a GPU.";
   let description = [{
     GPU module contains code that is intended to be run on a GPU. A host device
     can launch this code through a gpu.launc_func that creates a fully
     qualified symbol through the gpu.module's symbol and a gpu.func symbol
     contained in the gpu.module.
 
     The module's top-level scope is modeled by a single region with a single
     block. GPU modules are required to have a name that is used for symbol
     resolution by the gpu.launch_func operation.
 
     Using an op with a region to define a GPU module enables "embedding" GPU
     modules with SIMT execution models in other dialects in a clean manner and
     allows filtering of code regions to execute passes on only code intended to
     or not intended to be run on the separate device.
 
+    Modules can contain zero or more target attributes. These attributes encode
+    how to transform modules into binary strings and are used by the
+    `gpu-module-to-binary` pass to transform modules into GPU binaries.
+
     ```
-      gpu.module @symbol_name {
       gpu.func {}
         ...
       gpu.module_end
     }
-
+    gpu.module @symbol_name2 [#nvvm.target, #rocdl.target<chip = "gfx90a">] {
+      gpu.func {}
+        ...
+      gpu.module_end
+    }
     ```
   }];
-  let builders = [OpBuilder<(ins "StringRef":$name)>];
+  let builders = [
+    OpBuilder<(ins "StringRef":$name, CArg<"ArrayAttr", "{}">:$targets)>,
+    OpBuilder<(ins "StringRef":$name, "ArrayRef<Attribute>":$targets)>
+  ];
   let regions = (region SizedRegion<1>:$bodyRegion);
   let hasCustomAssemblyFormat = 1;
 
   // We need to ensure the block inside the region is properly terminated;
   // the auto-generated builders do not guarantee that.
   let skipDefaultBuilders = 1;
+
+  let extraClassDeclaration = [{
+    /// Checks if `target` is in the `targets` list.
+    bool hasTarget(Attribute target);
+
+    /// Sets the targets of the module.
+    void setTargets(ArrayRef<TargetAttrInterface> targets);
+  }];
 }
 
 def GPU_ModuleEndOp : GPU_Op<"module_end", [
   Terminator, HasParent<"GPUModuleOp">
 ]> {
   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`.

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

LogicalResult gpu::ReturnOp::verify() {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // GPUModuleOp
 //===----------------------------------------------------------------------===//
 
 void GPUModuleOp::build(OpBuilder &builder, OperationState &result,
-                        StringRef name) {
+                        StringRef name, ArrayAttr targets) {
   ensureTerminator(*result.addRegion(), builder, result.location);
   result.attributes.push_back(builder.getNamedAttr(
       ::mlir::SymbolTable::getSymbolAttrName(), builder.getStringAttr(name)));
+
+  if (targets)
+    result.getOrAddProperties<Properties>().targets = targets;
+}
+
+void GPUModuleOp::build(OpBuilder &builder, OperationState &result,
+                        StringRef name, ArrayRef<Attribute> targets) {
+  build(builder, result, name,
+        targets.size() > 0 ? builder.getArrayAttr(targets) : ArrayAttr());
 }
 
 ParseResult GPUModuleOp::parse(OpAsmParser &parser, OperationState &result) {
   StringAttr nameAttr;
+  ArrayAttr targetsAttr;
+
   if (parser.parseSymbolName(nameAttr, mlir::SymbolTable::getSymbolAttrName(),
-                             result.attributes) ||
-      // If module attributes are present, parse them.
-      parser.parseOptionalAttrDictWithKeyword(result.attributes))
+                             result.attributes))
+    return failure();
+
+  // Parse the optional array of target attributes.
+  OptionalParseResult targetsAttrResult =
+      parser.parseOptionalAttribute(targetsAttr, Type{});
+  if (targetsAttrResult.has_value()) {
+    if (failed(*targetsAttrResult)) {
+      return failure();
+    }
+    result.getOrAddProperties<Properties>().targets = targetsAttr;
+  }
+
+  if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
     return failure();
 
   // Parse the module body.
   auto *body = result.addRegion();
   if (parser.parseRegion(*body, {}))
     return failure();
 
   // Ensure that this module has a valid terminator.
   GPUModuleOp::ensureTerminator(*body, parser.getBuilder(), result.location);
   return success();
 }
 
 void GPUModuleOp::print(OpAsmPrinter &p) {
   p << ' ';
   p.printSymbolName(getName());
-  p.printOptionalAttrDictWithKeyword((*this)->getAttrs(),
-                                     {mlir::SymbolTable::getSymbolAttrName()});
+
+  if (Attribute attr = getTargetsAttr()) {
+    p << ' ';
+    p.printAttribute(attr);
+    p << ' ';
+  }
+
+  p.printOptionalAttrDictWithKeyword(
+      (*this)->getAttrs(),
+      {mlir::SymbolTable::getSymbolAttrName(), getTargetsAttrName()});
   p << ' ';
   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/false);
 }
 
+bool GPUModuleOp::hasTarget(Attribute target) {
+  if (ArrayAttr targets = getTargetsAttr())
+    return llvm::count(targets.getValue(), target);
+  return false;
+}
+
+void GPUModuleOp::setTargets(ArrayRef<TargetAttrInterface> targets) {
+  ArrayAttr &targetsAttr = getProperties().targets;
+  SmallVector<Attribute> targetsVector(targets);
+  targetsAttr = ArrayAttr::get(getContext(), targetsVector);
+}
+
 //===----------------------------------------------------------------------===//
 // GPUMemcpyOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult MemcpyOp::verify() {
   auto srcType = getSrc().getType();
   auto dstType = getDst().getType();
 

mlir/test/Dialect/GPU/invalid.mlir

 module attributes {gpu.container_module} {
   gpu.module @kernel {
     // expected-error@+1 {{'gpu.func' op gpu.known_block_size must contain exactly 3 elements}}
     gpu.func @kernel() kernel attributes {gpu.known_block_size = array<i32: 2, 1>} {
       gpu.return
     }
   }
 }
+
+// -----
+
+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 [] {
+  }
+}
+
+// -----
+
+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] {
+  }
+}