Reimplementing target description concept using DLTI attribute

mlir/include/mlir/Dialect/DLTI/CMakeLists.txt

@@ -1,2 +1,8 @@
 add_mlir_dialect(DLTI dlti)
 add_mlir_doc(DLTI DLTIDialect Dialects/ -gen-dialect-doc)
+
+set(LLVM_TARGET_DEFINITIONS DLTIAttrs.td)
+mlir_tablegen(DLTIAttrs.h.inc -gen-attrdef-decls -attrdefs-dialect=dlti)
+mlir_tablegen(DLTIAttrs.cpp.inc -gen-attrdef-defs -attrdefs-dialect=dlti)
+add_public_tablegen_target(MLIRDLTIAttrsIncGen)
+add_dependencies(mlir-headers MLIRDLTIAttrsIncGen)

mlir/include/mlir/Dialect/DLTI/DLTI.h

@@ -18,114 +18,13 @@
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 
 namespace mlir {
-namespace impl {
-class DataLayoutEntryStorage;
-class DataLayoutSpecStorage;
-} // namespace impl
-
-//===----------------------------------------------------------------------===//
-// DataLayoutEntryAttr
-//===----------------------------------------------------------------------===//
-
-/// A data layout entry attribute is a key-value pair where the key is a type or
-/// an identifier and the value is another attribute. These entries form a data
-/// layout specification.
-class DataLayoutEntryAttr
-    : public Attribute::AttrBase<DataLayoutEntryAttr, Attribute,
-                                 impl::DataLayoutEntryStorage,
-                                 DataLayoutEntryInterface::Trait> {
-public:
-  using Base::Base;
-
-  /// The keyword used for this attribute in custom syntax.
-  constexpr const static llvm::StringLiteral kAttrKeyword = "dl_entry";
-
-  /// Returns the entry with the given key and value.
-  static DataLayoutEntryAttr get(StringAttr key, Attribute value);
-  static DataLayoutEntryAttr get(Type key, Attribute value);
-
-  /// Returns the key of this entry.
-  DataLayoutEntryKey getKey() const;
-
-  /// Returns the value of this entry.
-  Attribute getValue() const;
-
-  /// Parses an instance of this attribute.
-  static DataLayoutEntryAttr parse(AsmParser &parser);
-
-  /// Prints this attribute.
-  void print(AsmPrinter &os) const;
-
-  static constexpr StringLiteral name = "builtin.data_layout_entry";
-};
-
-//===----------------------------------------------------------------------===//
-// DataLayoutSpecAttr
-//===----------------------------------------------------------------------===//
-
-/// A data layout specification is a list of entries that specify (partial) data
-/// layout information. It is expected to be attached to operations that serve
-/// as scopes for data layout requests.
-class DataLayoutSpecAttr
-    : public Attribute::AttrBase<DataLayoutSpecAttr, Attribute,
-                                 impl::DataLayoutSpecStorage,
-                                 DataLayoutSpecInterface::Trait> {
-public:
-  using Base::Base;
-
-  /// The keyword used for this attribute in custom syntax.
-  constexpr const static StringLiteral kAttrKeyword = "dl_spec";
-
-  /// Returns the specification containing the given list of keys.
-  static DataLayoutSpecAttr get(MLIRContext *ctx,
-                                ArrayRef<DataLayoutEntryInterface> entries);
-
-  /// Returns the specification containing the given list of keys. If the list
-  /// contains duplicate keys or is otherwise invalid, reports errors using the
-  /// given callback and returns null.
-  static DataLayoutSpecAttr
-  getChecked(function_ref<InFlightDiagnostic()> emitError, MLIRContext *context,
-             ArrayRef<DataLayoutEntryInterface> entries);
-
-  /// Checks that the given list of entries does not contain duplicate keys.
-  static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
-                              ArrayRef<DataLayoutEntryInterface> entries);
-
-  /// Combines this specification with `specs`, enclosing specifications listed
-  /// from outermost to innermost. This overwrites the older entries with the
-  /// same key as the newer entries if the entries are compatible. Returns null
-  /// if the specifications are not compatible.
-  DataLayoutSpecAttr combineWith(ArrayRef<DataLayoutSpecInterface> specs) const;
-
-  /// Returns the list of entries.
-  DataLayoutEntryListRef getEntries() const;
-
-  /// Returns the endiannes identifier.
-  StringAttr getEndiannessIdentifier(MLIRContext *context) const;
-
-  /// Returns the alloca memory space identifier.
-  StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const;
-
-  /// Returns the program memory space identifier.
-  StringAttr getProgramMemorySpaceIdentifier(MLIRContext *context) const;
-
-  /// Returns the global memory space identifier.
-  StringAttr getGlobalMemorySpaceIdentifier(MLIRContext *context) const;
-
-  /// Returns the stack alignment identifier.
-  StringAttr getStackAlignmentIdentifier(MLIRContext *context) const;
-
-  /// Parses an instance of this attribute.
-  static DataLayoutSpecAttr parse(AsmParser &parser);
-
-  /// Prints this attribute.
-  void print(AsmPrinter &os) const;
-
-  static constexpr StringLiteral name = "builtin.data_layout_spec";
-};
-
+namespace detail {
+class DataLayoutEntryAttrStorage;
+} // namespace detail
 } // namespace mlir
 
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/DLTI/DLTIAttrs.h.inc"
 #include "mlir/Dialect/DLTI/DLTIDialect.h.inc"
 
 #endif // MLIR_DIALECT_DLTI_DLTI_H

mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td

@@ -0,0 +1,188 @@
+//===- DLTIAttrs.td - DLTI dialect attributes definition --*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_DLTI_DLTIATTRS_TD
+#define MLIR_DIALECT_DLTI_DLTIATTRS_TD
+
+include "mlir/Dialect/DLTI/DLTI.td"
+include "mlir/IR/AttrTypeBase.td"
+
+class DLTIAttr<string name, list<Trait> traits = [],
+                string baseCppClass = "::mlir::Attribute">
+    : AttrDef<DLTI_Dialect, name, traits, baseCppClass> { }
+
+//===----------------------------------------------------------------------===//
+// DataLayoutEntryAttr
+//===----------------------------------------------------------------------===//
+
+def DataLayoutEntryTrait
+    : NativeAttrTrait<"DataLayoutEntryInterface::Trait"> {
+  let cppNamespace = "::mlir";
+}
+
+def DLTI_DataLayoutEntryAttr :
+    DLTIAttr<"DataLayoutEntry", [DataLayoutEntryTrait]> {
+  let summary = [{
+    An attribute to represent an entry of a data layout specification.
+  }];
+  let description = [{
+    A data layout entry attribute is a key-value pair where the key is a type or
+    an identifier and the value is another attribute. These entries form a data
+    layout specification.
+  }];
+  let parameters = (ins
+    "DataLayoutEntryKey":$key, "Attribute":$value
+  );
+  // TODO: We do not generate storage class because llvm::PointerUnion
+  // does not work with hash_key method.
+  let genStorageClass = 0;
+  let mnemonic = "dl_entry";
+  let genVerifyDecl = 0;
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDeclaration = [{
+    /// Returns the entry with the given key and value.
+    static DataLayoutEntryAttr get(StringAttr key, Attribute value);
+    static DataLayoutEntryAttr get(MLIRContext *context, Type key, Attribute value);
+    static DataLayoutEntryAttr get(Type key, Attribute value);
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// DataLayoutSpecAttr
+//===----------------------------------------------------------------------===//
+def DataLayoutSpecTrait
+    : NativeAttrTrait<"DataLayoutSpecInterface::Trait"> {
+  let cppNamespace = "::mlir";
+}
+
+def DLTI_DataLayoutSpecAttr :
+    DLTIAttr<"DataLayoutSpec", [DataLayoutSpecTrait]> {
+  let summary = [{
+    An attribute to represent a data layout specification.
+  }];
+  let description = [{
+    A data layout specification is a list of entries that specify (partial) data
+    layout information. It is expected to be attached to operations that serve
+    as scopes for data layout requests.
+  }];
+  let parameters = (ins
+    ArrayRefParameter<"DataLayoutEntryInterface", "">:$entries
+  );
+  let mnemonic = "dl_spec";
+  let genVerifyDecl = 1;
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDeclaration = [{
+    /// Combines this specification with `specs`, enclosing specifications listed
+    /// from outermost to innermost. This overwrites the older entries with the
+    /// same key as the newer entries if the entries are compatible. Returns null
+    /// if the specifications are not compatible.
+    DataLayoutSpecAttr combineWith(ArrayRef<DataLayoutSpecInterface> specs) const;
+
+    /// Returns the endiannes identifier.
+    StringAttr getEndiannessIdentifier(MLIRContext *context) const;
+
+    /// Returns the alloca memory space identifier.
+    StringAttr getAllocaMemorySpaceIdentifier(MLIRContext *context) const;
+
+    /// Returns the program memory space identifier.
+    StringAttr getProgramMemorySpaceIdentifier(MLIRContext *context) const;
+
+    /// Returns the global memory space identifier.
+    StringAttr getGlobalMemorySpaceIdentifier(MLIRContext *context) const;
+
+    /// Returns the stack alignment identifier.
+    StringAttr getStackAlignmentIdentifier(MLIRContext *context) const;
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// TargetSystemSpecAttr
+//===----------------------------------------------------------------------===//
+
+def TargetSystemSpecTrait
+    : NativeAttrTrait<"TargetSystemSpecInterface::Trait"> {
+  let cppNamespace = "::mlir";
+}
+
+def DLTI_TargetSystemSpecAttr :
+    DLTIAttr<"TargetSystemSpec", [TargetSystemSpecTrait]> {
+  let summary = [{
+    An attribute to represent target system specification.
+  }];
+  let description = [{
+    A system specification describes the overall system containing
+    multiple devices, with each device having a unique ID (string)
+    and its corresponding TargetDeviceSpec object.
+
+    Example:
+          dlti.target_system_spec =
+           #dlti.target_system_spec<
+            "CPU": #dlti.target_device_spec<
+                    #dlti.dl_entry<"dlti.L1_cache_size_in_bytes", 4096: ui32>>,
+            "GPU": #dlti.target_device_spec<
+                    #dlti.dl_entry<"dlti.max_vector_op_width", 64 : ui32>>,
+            "XPU": #dlti.target_device_spec<
+                    #dlti.dl_entry<"dlti.max_vector_op_width", 4096 : ui32>>>
+  }];
+  let parameters = (ins
+    ArrayRefParameter<"DeviceIDTargetDeviceSpecPair", "">:$entries
+  );
+  let mnemonic = "target_system_spec";
+  let genVerifyDecl = 1;
+  let assemblyFormat = "`<` $entries `>`";
+  let extraClassDeclaration = [{
+    /// Return the device specification that matches the given device ID
+    std::optional<TargetDeviceSpecInterface>
+    getDeviceSpecForDeviceID(
+      TargetSystemSpecInterface::DeviceID deviceID);
+  }];
+  let extraClassDefinition = [{
+    std::optional<TargetDeviceSpecInterface>
+    $cppClass::getDeviceSpecForDeviceID(
+        TargetSystemSpecInterface::DeviceID deviceID) {
+      for (const auto& entry : getEntries()) {
+        if (entry.first == deviceID)
+          return entry.second;
+      }
+      return std::nullopt;
+    }
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// TargetDeviceSpecAttr
+//===----------------------------------------------------------------------===//
+
+def TargetDeviceSpecTrait
+    : NativeAttrTrait<"TargetDeviceSpecInterface::Trait"> {
+  let cppNamespace = "::mlir";
+}
+
+def DLTI_TargetDeviceSpecAttr :
+    DLTIAttr<"TargetDeviceSpec", [TargetDeviceSpecTrait]> {
+  let summary = [{
+    An attribute to represent target device specification.
+  }];
+  let description = [{
+    Each device specification describes a single device and its
+    hardware properties. Each device specification can contain any number
+    of optional hardware properties (e.g., max_vector_op_width below).
+
+    Example:
+      #dlti.target_device_spec<
+        #dlti.dl_entry<"dlti.max_vector_op_width", 64 : ui32>>
+  }];
+  let parameters = (ins
+    ArrayRefParameter<"DataLayoutEntryInterface", "">:$entries
+  );
+  let mnemonic = "target_device_spec";
+  let genVerifyDecl = 1;
+  let assemblyFormat = "`<` $entries `>`";
+}
+
+#endif  // MLIR_DIALECT_DLTI_DLTIATTRS_TD

mlir/include/mlir/Dialect/DLTI/DLTIBase.td

@@ -27,6 +27,13 @@ def DLTI_Dialect : Dialect {
     constexpr const static ::llvm::StringLiteral
     kDataLayoutAttrName = "dlti.dl_spec";
 
+    // Top level attribute name for target system description
+    constexpr const static ::llvm::StringLiteral
+    kTargetSystemDescAttrName = "dlti.target_system_spec";
+
+    constexpr const static ::llvm::StringLiteral
+    kTargetDeviceDescAttrName = "dlti.target_device_spec";
+
     // Constants used in entries.
     constexpr const static ::llvm::StringLiteral
     kDataLayoutEndiannessKey = "dlti.endianness";
@@ -53,24 +60,6 @@ def DLTI_Dialect : Dialect {
   let useDefaultAttributePrinterParser = 1;
 }
 
-def DLTI_DataLayoutEntryAttr : DialectAttr<
-    DLTI_Dialect,
-    CPred<"::llvm::isa<::mlir::DataLayoutEntryAttr>($_self)">,
-    "Target data layout entry"> {
-  let storageType = "::mlir::DataLayoutEntryAttr";
-  let returnType = "::mlir::DataLayoutEntryAttr";
-  let convertFromStorage = "$_self";
-}
-
-def DLTI_DataLayoutSpecAttr : DialectAttr<
-    DLTI_Dialect,
-    CPred<"::llvm::isa<::mlir::DataLayoutSpecAttr>($_self)">,
-    "Target data layout specification"> {
-  let storageType = "::mlir::DataLayoutSpecAttr";
-  let returnType = "::mlir::DataLayoutSpecAttr";
-  let convertFromStorage = "$_self";
-}
-
 def HasDefaultDLTIDataLayout : NativeOpTrait<"HasDefaultDLTIDataLayout"> {
   let cppNamespace = "::mlir";
 }

mlir/include/mlir/Dialect/DLTI/Traits.h

@@ -18,6 +18,7 @@ class DataLayoutSpecAttr;
 namespace impl {
 LogicalResult verifyHasDefaultDLTIDataLayoutTrait(Operation *op);
 DataLayoutSpecInterface getDataLayoutSpec(Operation *op);
+TargetSystemSpecInterface getTargetSystemSpec(Operation *op);
 } // namespace impl
 
 /// Trait to be used by operations willing to use the implementation of the
@@ -37,6 +38,12 @@ class HasDefaultDLTIDataLayout
   DataLayoutSpecInterface getDataLayoutSpec() {
     return impl::getDataLayoutSpec(this->getOperation());
   }
+
+  /// Returns the target system description specification as provided by DLTI
+  /// dialect
+  TargetSystemSpecInterface getTargetSystemSpec() {
+    return impl::getTargetSystemSpec(this->getOperation());
+  }
 };
 } // namespace mlir
 

mlir/include/mlir/IR/BuiltinOps.td

@@ -78,6 +78,7 @@ def ModuleOp : Builtin_Op<"module", [
     //===------------------------------------------------------------------===//
 
     DataLayoutSpecInterface getDataLayoutSpec();
+    TargetSystemSpecInterface getTargetSystemSpec();
 
     //===------------------------------------------------------------------===//
     // OpAsmOpInterface Methods