[mlir][gpu] Add GPU target attribute interface.

[fmorac]

Summary

For an explanation of these patches see D154153.

Commit message:
This patch adds the GPUTargetAttrInterface attribute interface, this interface
is meant to be used as an opaque interface for serializing GPU modules into
binary strings.

[fmorac]

Rebasing.

[krzysz00]

This has memory safety issues.

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

208

StringRef doesn't own its contents, the toolkit path could disappear out from under you at an arbitrary time.

Let's make this either a std::string or a SmallString.

211

Same ownership notes. This should be something like a SmallVector<SmallVector<char>> or SmallVector<std::string> or so on.

[fmorac]

Changed StringRef -> std::string & ArrayRef -> SmallVector.

[mehdi_amini]

Seems fine to me, assuming @krzysz00 comments are addressed

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

35

Seems implied by the fact that there is no default implementation?

[fmorac]

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

35

I'll remove it.

[krzysz00]

Hold removed since my concerns are addressed

[fmorac]

Moved the compilation attributes from GPUDialect.h to CompilationInterfaces.h so they can be used by any dialect without including GPUDialect.h.

[fmorac]

Added a defualt implementation to the interface.

[fmorac]

Changed the error in the default implementation. There's a bug on tablegen, it crashes with:

code ... = [{
#ifdef ...
#endif
}];

because of a lexing error when included from another file.

[fmorac]

Moved the interface method body to DefaultImplementation.

[fmorac]

Added new line.

[fmorac]

Change constraint to isa.

[fmorac]

Sync correct diff.

[fmorac]

Added cmdOptions to TargetOptions for passing raw command line options.

Added a CompilationTarget option for specifying the final represation of the process (LLVM, Assembly, Binary).

[fmorac]

Removed unnecessary tablegen code.

mlir/include/mlir/Dialect/GPU/IR/CMakeLists.txt

 mlir_tablegen(ParallelLoopMapperEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRParallelLoopMapperEnumsGen)
 
 set(LLVM_TARGET_DEFINITIONS GPUOps.td)
 mlir_tablegen(GPUOpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(GPUOpsEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRGPUOpsEnumsGen)
 
+set(LLVM_TARGET_DEFINITIONS CompilationAttrInterfaces.td)
+mlir_tablegen(CompilationAttrInterfaces.h.inc -gen-attr-interface-decls)
+mlir_tablegen(CompilationAttrInterfaces.cpp.inc -gen-attr-interface-defs)
+add_public_tablegen_target(MLIRGPUCompilationAttrInterfacesIncGen)
+
 set(LLVM_TARGET_DEFINITIONS GPUOps.td)
 mlir_tablegen(GPUOpsAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=gpu)
 mlir_tablegen(GPUOpsAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=gpu)
 add_public_tablegen_target(MLIRGPUOpsAttributesIncGen)

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

+//===-- CompilationAttrInterfaces.td - GPU compilation interfaces ---------===//
+//
+// 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 interfaces for GPU target attributes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef GPU_COMPILATIONATTRINTERFACES
+#define GPU_COMPILATIONATTRINTERFACES
+
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// GPU target attribute interface.
+//===----------------------------------------------------------------------===//
+
+def GPUTargetAttrInterface : AttrInterface<"TargetAttrInterface"> {
+  let description = [{
+    Interface for GPU target attributes. Attributes implementing this interface
+    compile GPU modules into binary objects, providing an opaque interface to
+    hide implementation details.
+  }];
+  let cppNamespace = "::mlir::gpu";
+  let methods = [
+    InterfaceMethod<[{
+        Serializes a GPU module to a string containing a representation of the
+        module.
+
+        All attributes implementing this interface must implement this method.
+        If serialization fails then the method should return `std::nullopt`.
+
+        The `options` argument is meant to be used for passing additional
+        options that are not in the attribute.
+      }],
+      "std::optional<SmallVector<char, 0>>", "serializeToObject",
+      (ins "Operation*":$module, "const gpu::TargetOptions&":$options)
+    >
+  ];
+}
+
+def ImplementsTargetAttrInterface : AttrConstraint<
+  CPred<"$_self.hasTrait<::mlir::gpu::TargetAttrInterface::Trait>()">,
+  "Attribute implementing the `TargetAttrInterface` interface."
+>;
+
+def GPUTargetAttr : ConfinedAttr<AnyAttr, [ImplementsTargetAttrInterface]> {
+  let description = [{
+    Generic target attribute implementing the `TargetAttrInterface` interface.
+  }];
+}
+
+def GPUTargetArrayAttr :
+  TypedArrayAttrBase<GPUTargetAttr, "Array of GPU target attributes">;
+
+def GPUNonEmptyTargetArrayAttr :
+  ConfinedAttr<GPUTargetArrayAttr, [ArrayMinCount<1>]>;
+
+#endif // GPU_COMPILATIONATTRINTERFACES

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

public:
   using Base =
       typename Type::TypeBase<SparseHandleType<K>, Type, TypeStorage>::Base;
   using Base::Base;
 };
 
 using SparseDnTensorHandleType = SparseHandleType<SparseHandleKind::DnTensor>;
 using SparseSpMatHandleType = SparseHandleType<SparseHandleKind::SpMat>;
 
+/// This class serves as an opaque interface for passing options to the
+/// `TargetAttrInterface` methods. Users of this class must implement the
+/// `classof` method as well as using the macros `MLIR_*_EXPLICIT_TYPE_ID` to
+/// ensure type safeness.
+class TargetOptions {
+public:
+  /// Constructor initializing the toolkit path and the list of bitcode files.
+  TargetOptions(StringRef toolkitPath = {},
+                ArrayRef<std::string> bitcodeFiles = {});
+
+  /// Returns the typeID.
+  TypeID getTypeID() const;
+
+  /// Returns the toolkit path.
+  StringRef getToolkitPath() const;
+
+  /// Returns the bitcode files to link to.
+  ArrayRef<std::string> getBitcodeFiles() const;
+
+protected:
+  /// Derived classes must use this constructor to initialize `typeID` to the
+  /// appropiate value: ie. `TargetOptions(TypeID::get<DerivedClass>())`.
+  TargetOptions(TypeID typeID, StringRef toolkitPath = {},
+                ArrayRef<std::string> bitcodeFiles = {});
+
+  /// Path to the target toolkit.
+  std::string toolkitPath;
+
+  /// List of files to link with the LLVM module.
+  SmallVector<std::string> bitcodeFiles;
+
+private:
+  TypeID typeID;
+};
 } // namespace gpu
 } // namespace mlir
 
+MLIR_DECLARE_EXPLICIT_TYPE_ID(::mlir::gpu::TargetOptions)
+
 #include "mlir/Dialect/GPU/IR/GPUOpsEnums.h.inc"
 
 #include "mlir/Dialect/GPU/IR/GPUOpsDialect.h.inc"
 
 #include "mlir/Dialect/GPU/IR/GPUOpInterfaces.h.inc"
 
 #include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h"
 
+#include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.h.inc"
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/GPU/IR/GPUOpsAttributes.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/IR/GPUOps.h.inc"
 
 #endif // MLIR_DIALECT_GPU_IR_GPUDIALECT_H

mlir/lib/Dialect/GPU/CMakeLists.txt

add_mlir_dialect_library(MLIRGPUDialect
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/GPU
 
   DEPENDS
   MLIRGPUOpsIncGen
   MLIRGPUOpsAttributesIncGen
   MLIRGPUOpsEnumsGen
   MLIRGPUOpInterfacesIncGen
+  MLIRGPUCompilationAttrInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRDLTIDialect
   MLIRInferIntRangeInterface
   MLIRIR
   MLIRMemRefDialect
   MLIRSideEffectInterfaces

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

 
 } // namespace
 
 void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
   results.add<SimplifyDimOfAllocOp>(context);
 }
 
+//===----------------------------------------------------------------------===//
+// GPU Compilation options
+//===----------------------------------------------------------------------===//
+
+TargetOptions::TargetOptions(StringRef toolkitPath,
+                             ArrayRef<std::string> bitcodeFiles)
+    : TargetOptions(TypeID::get<TargetOptions>(), toolkitPath, bitcodeFiles) {}
+
+TargetOptions::TargetOptions(TypeID typeID, StringRef toolkitPath,
+                             ArrayRef<std::string> bitcodeFiles)
+    : toolkitPath(toolkitPath.str()), bitcodeFiles(bitcodeFiles),
+      typeID(typeID) {}
+
+TypeID TargetOptions::getTypeID() const { return typeID; }
+
+StringRef TargetOptions::getToolkitPath() const { return toolkitPath; }
+
+ArrayRef<std::string> TargetOptions::getBitcodeFiles() const {
+  return bitcodeFiles;
+}
+
+MLIR_DEFINE_EXPLICIT_TYPE_ID(::mlir::gpu::TargetOptions)
+
 #include "mlir/Dialect/GPU/IR/GPUOpInterfaces.cpp.inc"
 #include "mlir/Dialect/GPU/IR/GPUOpsEnums.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/GPU/IR/GPUOpsAttributes.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/IR/GPUOps.cpp.inc"
+
+#include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.cpp.inc"