Reimplementing target description concept using DLTI attribute

Author: nhasabni

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

@@ -21,6 +21,8 @@ namespace mlir {
 namespace impl {
 class DataLayoutEntryStorage;
 class DataLayoutSpecStorage;
+class TargetSystemDescSpecAttrStorage;
+class TargetDeviceDescSpecAttrStorage;
 } // namespace impl
 
 //===----------------------------------------------------------------------===//
@@ -124,6 +126,150 @@ class DataLayoutSpecAttr
   static constexpr StringLiteral name = "builtin.data_layout_spec";
 };
 
+//===----------------------------------------------------------------------===//
+// TargetSystemDescSpecAttr
+//===----------------------------------------------------------------------===//
+
+/// A system description attribute is a list of device descriptors, each
+/// having a unique device ID
+class TargetSystemDescSpecAttr
+    : public Attribute::AttrBase<TargetSystemDescSpecAttr, Attribute,
+                                 impl::TargetSystemDescSpecAttrStorage,
+                                 TargetSystemDescSpecInterface::Trait> {
+public:
+  using Base::Base;
+
+  /// The keyword used for this attribute in custom syntax.
+  constexpr const static StringLiteral kAttrKeyword = "tsd_spec";
+
+  /// Returns a system descriptor attribute from the given system descriptor
+  static TargetSystemDescSpecAttr
+  get(MLIRContext *context, ArrayRef<TargetDeviceDescSpecInterface> entries);
+
+  /// Returns the list of entries.
+  TargetDeviceDescSpecListRef getEntries() const;
+
+  /// Return the device descriptor that matches the given device ID
+  TargetDeviceDescSpecInterface getDeviceDescForDeviceID(uint32_t deviceID);
+
+  /// 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 TargetSystemDescSpecAttr
+  getChecked(function_ref<InFlightDiagnostic()> emitError, MLIRContext *context,
+             ArrayRef<TargetDeviceDescSpecInterface> entries);
+
+  /// Checks that the given list of entries does not contain duplicate keys.
+  static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
+                              ArrayRef<TargetDeviceDescSpecInterface> entries);
+
+  /// Parses an instance of this attribute.
+  static TargetSystemDescSpecAttr parse(AsmParser &parser);
+
+  /// Prints this attribute.
+  void print(AsmPrinter &os) const;
+
+  static constexpr StringLiteral name = "builtin.target_system_description";
+};
+
+//===----------------------------------------------------------------------===//
+// TargetDeviceDescSpecAttr
+//===----------------------------------------------------------------------===//
+
+class TargetDeviceDescSpecAttr
+    : public Attribute::AttrBase<TargetDeviceDescSpecAttr, Attribute,
+                                 impl::TargetDeviceDescSpecAttrStorage,
+                                 TargetDeviceDescSpecInterface::Trait> {
+public:
+  using Base::Base;
+
+  /// The keyword used for this attribute in custom syntax.
+  constexpr const static StringLiteral kAttrKeyword = "tdd_spec";
+
+  /// Returns a system descriptor attribute from the given system descriptor
+  static TargetDeviceDescSpecAttr
+  get(MLIRContext *context, 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 TargetDeviceDescSpecAttr
+  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);
+
+  /// Returns the list of entries.
+  DataLayoutEntryListRef getEntries() const;
+
+  /// Parses an instance of this attribute.
+  static TargetDeviceDescSpecAttr parse(AsmParser &parser);
+
+  /// Prints this attribute.
+  void print(AsmPrinter &os) const;
+
+  /// Returns the device ID identifier.
+  StringAttr getDeviceIDIdentifier(MLIRContext *context);
+
+  /// Returns the device type identifier.
+  StringAttr getDeviceTypeIdentifier(MLIRContext *context);
+
+  /// Returns max vector op width identifier.
+  StringAttr getMaxVectorOpWidthIdentifier(MLIRContext *context);
+
+  /// Returns canonicalizer max iterations identifier.
+  StringAttr getCanonicalizerMaxIterationsIdentifier(MLIRContext *context);
+
+  /// Returns canonicalizer max num rewrites identifier.
+  StringAttr getCanonicalizerMaxNumRewritesIdentifier(MLIRContext *context);
+
+  /// Returns L1 cache size identifier
+  StringAttr getL1CacheSizeInBytesIdentifier(MLIRContext *context);
+
+  /// Returns the interface spec for device ID
+  /// Since we verify that the spec contains device ID the function
+  /// will return a valid spec.
+  DataLayoutEntryInterface getSpecForDeviceID(MLIRContext *context);
+
+  /// Returns the interface spec for device type
+  /// Since we verify that the spec contains device type the function
+  /// will return a valid spec.
+  DataLayoutEntryInterface getSpecForDeviceType(MLIRContext *context);
+
+  /// Returns the interface spec for max vector op width
+  /// Since max vector op width is an optional property, this function will
+  /// return a valid spec if the property is defined, otherwise it
+  /// will return an empty spec.
+  DataLayoutEntryInterface getSpecForMaxVectorOpWidth(MLIRContext *context);
+
+  /// Returns the interface spec for L1 cache size
+  /// Since L1 cache size is an optional property, this function will
+  /// return a valid spec if the property is defined, otherwise it
+  /// will return an empty spec.
+  DataLayoutEntryInterface getSpecForL1CacheSizeInBytes(MLIRContext *context);
+
+  /// Returns the interface spec for canonicalizer max iterations.
+  /// Since this is an optional property, this function will
+  /// return a valid spec if the property is defined, otherwise it
+  /// will return an empty spec.
+  DataLayoutEntryInterface
+  getSpecForCanonicalizerMaxIterations(MLIRContext *context);
+
+  /// Returns the interface spec for canonicalizer max num rewrites.
+  /// Since this is an optional property, this function will
+  /// return a valid spec if the property is defined, otherwise it
+  /// will return an empty spec.
+  DataLayoutEntryInterface
+  getSpecForCanonicalizerMaxNumRewrites(MLIRContext *context);
+
+  /// Return the value of device ID
+  uint32_t getDeviceID(MLIRContext *context);
+
+  static constexpr StringLiteral name = "builtin.target_device_description";
+};
+
 } // namespace mlir
 
 #include "mlir/Dialect/DLTI/DLTIDialect.h.inc"

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.tsd_spec";
+
+    constexpr const static ::llvm::StringLiteral
+    kTargetDeviceDescAttrName = "dlti.tdd_spec";
+
     // Constants used in entries.
     constexpr const static ::llvm::StringLiteral
     kDataLayoutEndiannessKey = "dlti.endianness";
@@ -48,6 +55,25 @@ def DLTI_Dialect : Dialect {
 
     constexpr const static ::llvm::StringLiteral
     kDataLayoutStackAlignmentKey = "dlti.stack_alignment";
+
+    // Constants used in target description part of DLTI
+    constexpr const static ::llvm::StringLiteral
+    kTargetDeviceIDKey = "dlti.device_id";
+
+    constexpr const static ::llvm::StringLiteral
+    kTargetDeviceTypeKey = "dlti.device_type";
+
+    constexpr const static ::llvm::StringLiteral
+    kTargetDeviceMaxVectorOpWidthKey = "dlti.max_vector_op_width";
+
+    constexpr const static ::llvm::StringLiteral
+    kTargetDeviceCanonicalizerMaxIterationsKey = "dlti.canonicalizer_max_iterations";
+
+    constexpr const static ::llvm::StringLiteral
+    kTargetDeviceCanonicalizerMaxNumRewritesKey = "dlti.canonicalizer_max_num_rewrites";
+
+    constexpr const static ::llvm::StringLiteral
+    kTargetDeviceL1CacheSizeInBytesKey = "dlti.L1_cache_size_in_bytes";
   }];
 
   let useDefaultAttributePrinterParser = 1;
@@ -71,6 +97,24 @@ def DLTI_DataLayoutSpecAttr : DialectAttr<
   let convertFromStorage = "$_self";
 }
 
+def DLTI_TargetSystemDescSpecAttr : DialectAttr<
+    DLTI_Dialect,
+    CPred<"::llvm::isa<::mlir::TargetSystemDescSpecAttr>($_self)">,
+    "Target system description part of DLTI"> {
+  let storageType = "::mlir::TargetSystemDescSpecAttr";
+  let returnType = "::mlir::TargetSystemDescSpecAttr";
+  let convertFromStorage = "$_self";
+}
+
+def DLTI_TargetDeviceDescSpecAttr : DialectAttr<
+    DLTI_Dialect,
+    CPred<"::llvm::isa<::mlir::TargetDeviceDescSpecAttr>($_self)">,
+    "Target device description part of DLTI"> {
+  let storageType = "::mlir::TargetDeviceDescSpecAttr";
+  let returnType = "::mlir::TargetDeviceDescSpecAttr";
+  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);
+TargetSystemDescSpecInterface getTargetSystemDescSpec(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
+  TargetSystemDescSpecInterface getTargetSystemDescSpec() {
+    return impl::getTargetSystemDescSpec(this->getOperation());
+  }
 };
 } // namespace mlir
 

mlir/include/mlir/IR/BuiltinOps.td

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

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

@@ -23,11 +23,17 @@
 namespace mlir {
 class DataLayout;
 class DataLayoutEntryInterface;
+class TargetDeviceDescSpecInterface;
+class TargetSystemDescSpecInterface;
 using DataLayoutEntryKey = llvm::PointerUnion<Type, StringAttr>;
 // Using explicit SmallVector size because we cannot infer the size from the
 // forward declaration, and we need the typedef in the actual declaration.
 using DataLayoutEntryList = llvm::SmallVector<DataLayoutEntryInterface, 4>;
 using DataLayoutEntryListRef = llvm::ArrayRef<DataLayoutEntryInterface>;
+// using TargetDeviceDescSpecList =
+// llvm::SmallVector<TargetDeviceDescSpecInterface, 4>;
+using TargetDeviceDescSpecListRef =
+    llvm::ArrayRef<TargetDeviceDescSpecInterface>;
 class DataLayoutOpInterface;
 class DataLayoutSpecInterface;
 class ModuleOp;
@@ -84,6 +90,24 @@ Attribute getDefaultGlobalMemorySpace(DataLayoutEntryInterface entry);
 /// DataLayoutInterface if specified, otherwise returns the default.
 uint64_t getDefaultStackAlignment(DataLayoutEntryInterface entry);
 
+/// return max vector op width from the specified DataLayoutEntry. If the
+/// property is missing from the entry, then return std::nullopt.
+std::optional<uint32_t> getMaxVectorOpWidth(DataLayoutEntryInterface entry);
+
+/// return L1 cache size in bytes from the specified DataLayoutEntry. If the
+/// property is missing from the entry, then return std::nullopt.
+std::optional<uint32_t> getL1CacheSizeInBytes(DataLayoutEntryInterface entry);
+
+/// return canonicalizer max iterations from the specified DataLayoutEntry.
+/// If the property is missing from the entry, then return std::nullopt.
+std::optional<int64_t>
+getCanonicalizerMaxIterations(DataLayoutEntryInterface entry);
+
+/// returncanonicalizer max num rewrites from the specified DataLayoutEntry.
+/// If the property is missing from the entry, then return std::nullopt.
+std::optional<int64_t>
+getCanonicalizerMaxNumRewrites(DataLayoutEntryInterface entry);
+
 /// Given a list of data layout entries, returns a new list containing the
 /// entries with keys having the given type ID, i.e. belonging to the same type
 /// class.
@@ -95,6 +119,11 @@ DataLayoutEntryList filterEntriesForType(DataLayoutEntryListRef entries,
 DataLayoutEntryInterface
 filterEntryForIdentifier(DataLayoutEntryListRef entries, StringAttr id);
 
+/// Given a list of target device entries, returns the entry that has the given
+/// identifier as key, if such an entry exists in the list.
+TargetDeviceDescSpecInterface
+filterEntryForIdentifier(TargetDeviceDescSpecListRef entries, StringAttr id);
+
 /// Verifies that the operation implementing the data layout interface, or a
 /// module operation, is valid. This calls the verifier of the spec attribute
 /// and checks if the layout is compatible with specs attached to the enclosing
@@ -106,6 +135,12 @@ LogicalResult verifyDataLayoutOp(Operation *op);
 /// and dialect interfaces for type and identifier keys respectively.
 LogicalResult verifyDataLayoutSpec(DataLayoutSpecInterface spec, Location loc);
 
+/// Verifies that a target system desc spec is valid. This dispatches to
+/// individual entry verifiers, and then to the verifiers implemented by the
+/// relevant dialect interfaces for identifier keys.
+LogicalResult verifyTargetSystemDescSpec(TargetSystemDescSpecInterface spec,
+                                         Location loc);
+
 /// Divides the known min value of the numerator by the denominator and rounds
 /// the result up to the next integer. Preserves the scalable flag.
 llvm::TypeSize divideCeil(llvm::TypeSize numerator, uint64_t denominator);
@@ -137,6 +172,13 @@ class DataLayoutDialectInterface
     return success();
   }
 
+  /// Checks whether the given data layout entry is valid and reports any errors
+  /// at the provided location. Derived classes should override this.
+  virtual LogicalResult verifyEntry(TargetDeviceDescSpecInterface entry,
+                                    Location loc) const {
+    return success();
+  }
+
   /// Default implementation of entry combination that combines identical
   /// entries and returns null otherwise.
   static DataLayoutEntryInterface
@@ -214,10 +256,33 @@ class DataLayout {
   /// unspecified.
   uint64_t getStackAlignment() const;
 
+  /// Returns for max vector op width if the property is defined for the given
+  /// device ID, otherwise return std::nullopt.
+  std::optional<uint32_t>
+      getMaxVectorOpWidth(TargetDeviceDescSpecInterface::DeviceID) const;
+
+  /// Returns for L1 cache size if the property is defined for the given
+  /// device ID, otherwise return std::nullopt.
+  std::optional<uint32_t>
+      getL1CacheSizeInBytes(TargetDeviceDescSpecInterface::DeviceID) const;
+
+  /// Returns for canonicalizer max iterations if the property is defined for
+  /// the given device ID, otherwise return std::nullopt.
+  std::optional<int64_t> getCanonicalizerMaxIterations(
+      TargetDeviceDescSpecInterface::DeviceID) const;
+
+  /// Returns for canonicalizer max rewrites if the property is defined for
+  /// the given device ID, otherwise return std::nullopt.
+  std::optional<int64_t> getCanonicalizerMaxNumRewrites(
+      TargetDeviceDescSpecInterface::DeviceID) const;
+
 private:
   /// Combined layout spec at the given scope.
   const DataLayoutSpecInterface originalLayout;
 
+  /// Combined target system desc spec at the given scope.
+  const TargetSystemDescSpecInterface originalTargetSystemDesc;
+
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
   /// List of enclosing layout specs.
   SmallVector<DataLayoutSpecInterface, 2> layoutStack;