Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[MLIR][DLTI][Transform] Introduce transform.dlti.query - 2nd attempt #102652

Merged
merged 1 commit into from
Aug 12, 2024

Conversation

rolfmorel
Copy link
Contributor

This transform op makes it possible to query attributes associated to IR
by means of the DLTI dialect.

The op takes both a key and a target op to perform the query at.
Facility functions automatically find the closest ancestor op which
defines the appropriate DLTI interface or has an attribute implementing
a DLTI interface. By default the lookup uses the data layout interfaces
of DLTI. If the optional device parameter is provided, the lookup
happens with respect to the interfaces for TargetSystemSpec and
TargetDeviceSpec.

This op uses new free-standing functions in the dlti namespace to not
only look up specifications via the DataLayoutSpecOpInterface and on
ModuleOps but also on any ancestor op that has an appropriate DLTI
attribute.

@llvmbot
Copy link
Collaborator

llvmbot commented Aug 9, 2024

@llvm/pr-subscribers-mlir-dlti

@llvm/pr-subscribers-mlir

Author: Rolf Morel (rolfmorel)

Changes

This transform op makes it possible to query attributes associated to IR
by means of the DLTI dialect.

The op takes both a key and a target op to perform the query at.
Facility functions automatically find the closest ancestor op which
defines the appropriate DLTI interface or has an attribute implementing
a DLTI interface. By default the lookup uses the data layout interfaces
of DLTI. If the optional device parameter is provided, the lookup
happens with respect to the interfaces for TargetSystemSpec and
TargetDeviceSpec.

This op uses new free-standing functions in the dlti namespace to not
only look up specifications via the DataLayoutSpecOpInterface and on
ModuleOps but also on any ancestor op that has an appropriate DLTI
attribute.


Patch is 28.11 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/102652.diff

11 Files Affected:

  • (modified) mlir/include/mlir/Dialect/DLTI/CMakeLists.txt (+2)
  • (modified) mlir/include/mlir/Dialect/DLTI/DLTI.h (+13)
  • (added) mlir/include/mlir/Dialect/DLTI/TransformOps/CMakeLists.txt (+6)
  • (added) mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h (+38)
  • (added) mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td (+61)
  • (modified) mlir/include/mlir/InitAllExtensions.h (+2)
  • (modified) mlir/lib/Dialect/DLTI/CMakeLists.txt (+1)
  • (modified) mlir/lib/Dialect/DLTI/DLTI.cpp (+35)
  • (added) mlir/lib/Dialect/DLTI/TransformOps/CMakeLists.txt (+15)
  • (added) mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp (+96)
  • (added) mlir/test/Dialect/DLTI/query.mlir (+296)
diff --git a/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt b/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt
index 5b0cb6c37bfab5..4f8382e8e6e6b3 100644
--- a/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt
@@ -1,3 +1,5 @@
+add_subdirectory(TransformOps)
+
 add_mlir_dialect(DLTI dlti)
 add_mlir_doc(DLTIAttrs DLTIDialect Dialects/ -gen-dialect-doc)
 
diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h
index f50a654f3885d5..6e2623e084be52 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTI.h
+++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h
@@ -22,6 +22,19 @@ namespace detail {
 class DataLayoutEntryAttrStorage;
 } // namespace detail
 } // namespace mlir
+namespace mlir {
+namespace dlti {
+/// Find the first DataLayoutSpec associated to `op`, via either the
+/// DataLayoutOpInterface, a method on ModuleOp, or an attribute implementing
+/// the interface, on `op` and else on `op`'s ancestors in turn.
+DataLayoutSpecInterface getDataLayoutSpec(Operation *op);
+
+/// Find the first TargetSystemSpec associated to `op`, via either the
+/// DataLayoutOpInterface, a method on ModuleOp, or an attribute implementing
+/// the interface, on `op` and else on `op`'s ancestors in turn.
+TargetSystemSpecInterface getTargetSystemSpec(Operation *op);
+} // namespace dlti
+} // namespace mlir
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/DLTI/DLTIAttrs.h.inc"
diff --git a/mlir/include/mlir/Dialect/DLTI/TransformOps/CMakeLists.txt b/mlir/include/mlir/Dialect/DLTI/TransformOps/CMakeLists.txt
new file mode 100644
index 00000000000000..1188d1a9b24d60
--- /dev/null
+++ b/mlir/include/mlir/Dialect/DLTI/TransformOps/CMakeLists.txt
@@ -0,0 +1,6 @@
+set(LLVM_TARGET_DEFINITIONS DLTITransformOps.td)
+mlir_tablegen(DLTITransformOps.h.inc -gen-op-decls)
+mlir_tablegen(DLTITransformOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRDLTITransformOpsIncGen)
+
+add_mlir_doc(DLTITransformOps DLTITransformOps Dialects/ -gen-op-doc)
diff --git a/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h
new file mode 100644
index 00000000000000..970f2597ef7ad3
--- /dev/null
+++ b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h
@@ -0,0 +1,38 @@
+//===- DLTITransformOps.h - DLTI transform ops ------------------*- C++ -*-===//
+//
+// 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_TRANSFORMOPS_DLTITRANSFORMOPS_H
+#define MLIR_DIALECT_DLTI_TRANSFORMOPS_DLTITRANSFORMOPS_H
+
+#include "mlir/Dialect/Transform/IR/TransformAttrs.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
+#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
+
+namespace mlir {
+namespace transform {
+class QueryOp;
+} // namespace transform
+} // namespace mlir
+
+namespace mlir {
+class DialectRegistry;
+
+namespace dlti {
+void registerTransformDialectExtension(DialectRegistry &registry);
+} // namespace dlti
+} // namespace mlir
+
+////===----------------------------------------------------------------------===//
+//// DLTI Transform Operations
+////===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h.inc"
+
+#endif // MLIR_DIALECT_DLTI_TRANSFORMOPS_DLTITRANSFORMOPS_H
diff --git a/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td
new file mode 100644
index 00000000000000..69aacac986ad73
--- /dev/null
+++ b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td
@@ -0,0 +1,61 @@
+//===- DLTITransformOps.td - DLTI transform ops ------------*- 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 DLTI_TRANSFORM_OPS
+#define DLTI_TRANSFORM_OPS
+
+include "mlir/Dialect/Transform/IR/TransformDialect.td"
+include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
+include "mlir/Dialect/Transform/IR/TransformTypes.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/OpBase.td"
+
+def QueryOp : Op<Transform_Dialect, "dlti.query", [
+  TransformOpInterface, TransformEachOpTrait,
+  DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
+]> {
+  let summary = "Return attribute (as param) associated to key via DTLI";
+  let description = [{
+    This op queries data layout and target information associated to payload
+    IR by way of the DLTI dialect. A lookup is performed for the given `key`
+    at the `target` op, with the DLTI dialect determining which interfaces and
+    attributes are consulted - first checking `target` and then its ancestors.
+
+    When only `key` is provided, the lookup occurs with respect to the data
+    layout specification of DLTI. When `device` is provided, the lookup occurs
+    with respect to DLTI's target device specifications associated to a DLTI
+    system device specification.
+
+    #### Return modes
+
+    When succesful, the result, `associated_attr`, associates one attribute as a
+    param for each op in `target`'s payload.
+
+    If the lookup fails - as DLTI specifications or entries with the right
+    names are missing (i.e. the values of `device` and `key`) - a definite
+    failure is returned.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target,
+                       OptionalAttr<StrAttr>:$device,
+                       StrAttr:$key);
+  let results = (outs TransformParamTypeInterface:$associated_attr);
+  let assemblyFormat =
+      "(`:``:` $device^ `:``:`)? $key `at` $target attr-dict `:`"
+      "functional-type(operands, results)";
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
+        ::mlir::Operation *target,
+        ::mlir::transform::ApplyToEachResultList &results,
+        TransformState &state);
+  }];
+}
+
+#endif // DLTI_TRANSFORM_OPS
diff --git a/mlir/include/mlir/InitAllExtensions.h b/mlir/include/mlir/InitAllExtensions.h
index 20a4ab6f18a286..0adc5e52f2a0e5 100644
--- a/mlir/include/mlir/InitAllExtensions.h
+++ b/mlir/include/mlir/InitAllExtensions.h
@@ -25,6 +25,7 @@
 #include "mlir/Conversion/UBToLLVM/UBToLLVM.h"
 #include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
+#include "mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h"
 #include "mlir/Dialect/Func/Extensions/AllExtensions.h"
 #include "mlir/Dialect/Func/TransformOps/FuncTransformOps.h"
 #include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
@@ -69,6 +70,7 @@ inline void registerAllExtensions(DialectRegistry &registry) {
   // Register all transform dialect extensions.
   affine::registerTransformDialectExtension(registry);
   bufferization::registerTransformDialectExtension(registry);
+  dlti::registerTransformDialectExtension(registry);
   func::registerTransformDialectExtension(registry);
   gpu::registerTransformDialectExtension(registry);
   linalg::registerTransformDialectExtension(registry);
diff --git a/mlir/lib/Dialect/DLTI/CMakeLists.txt b/mlir/lib/Dialect/DLTI/CMakeLists.txt
index f74f5f80130908..7691a8a10a37ad 100644
--- a/mlir/lib/Dialect/DLTI/CMakeLists.txt
+++ b/mlir/lib/Dialect/DLTI/CMakeLists.txt
@@ -1,3 +1,4 @@
+add_subdirectory(TransformOps)
 add_mlir_dialect_library(MLIRDLTIDialect
   DLTI.cpp
   Traits.cpp
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 420c605d1a19b2..c995f44c380e57 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -393,6 +393,41 @@ TargetSystemSpecAttr::verify(function_ref<InFlightDiagnostic()> emitError,
 // DLTIDialect
 //===----------------------------------------------------------------------===//
 
+DataLayoutSpecInterface dlti::getDataLayoutSpec(Operation *op) {
+  DataLayoutSpecInterface dlSpec = nullptr;
+
+  for (Operation *cur = op; cur && !dlSpec; cur = cur->getParentOp()) {
+    if (auto dataLayoutOp = dyn_cast<DataLayoutOpInterface>(cur))
+      dlSpec = dataLayoutOp.getDataLayoutSpec();
+    else if (auto moduleOp = dyn_cast<ModuleOp>(cur))
+      dlSpec = moduleOp.getDataLayoutSpec();
+    else
+      for (NamedAttribute attr : cur->getAttrs())
+        if ((dlSpec = llvm::dyn_cast<DataLayoutSpecInterface>(attr.getValue())))
+          break;
+  }
+
+  return dlSpec;
+}
+
+TargetSystemSpecInterface dlti::getTargetSystemSpec(Operation *op) {
+  TargetSystemSpecInterface sysSpec = nullptr;
+
+  for (Operation *cur = op; cur && !sysSpec; cur = cur->getParentOp()) {
+    if (auto dataLayoutOp = dyn_cast<DataLayoutOpInterface>(cur))
+      sysSpec = dataLayoutOp.getTargetSystemSpec();
+    else if (auto moduleOp = dyn_cast<ModuleOp>(cur))
+      sysSpec = moduleOp.getTargetSystemSpec();
+    else
+      for (NamedAttribute attr : cur->getAttrs())
+        if ((sysSpec =
+                 llvm::dyn_cast<TargetSystemSpecInterface>(attr.getValue())))
+          break;
+  }
+
+  return sysSpec;
+}
+
 constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutAttrName;
 constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessKey;
 constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessBig;
diff --git a/mlir/lib/Dialect/DLTI/TransformOps/CMakeLists.txt b/mlir/lib/Dialect/DLTI/TransformOps/CMakeLists.txt
new file mode 100644
index 00000000000000..05734353f9d11d
--- /dev/null
+++ b/mlir/lib/Dialect/DLTI/TransformOps/CMakeLists.txt
@@ -0,0 +1,15 @@
+add_mlir_dialect_library(MLIRDLTITransformOps
+  DLTITransformOps.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/DLTI/TransformOps
+
+  DEPENDS
+  MLIRDLTITransformOpsIncGen
+  MLIRDLTIDialect
+
+  LINK_LIBS PUBLIC
+  MLIRDLTIDialect
+  MLIRSideEffectInterfaces
+  MLIRTransformDialect
+  )
diff --git a/mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp b/mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp
new file mode 100644
index 00000000000000..032228b2f05ca2
--- /dev/null
+++ b/mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp
@@ -0,0 +1,96 @@
+
+//===- DLTITransformOps.cpp - Implementation of DLTI transform ops --------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h"
+
+#include "mlir/Dialect/DLTI/DLTI.h"
+#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
+#include "mlir/Dialect/Transform/Utils/Utils.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
+
+using namespace mlir;
+using namespace mlir::transform;
+
+#define DEBUG_TYPE "dlti-transforms"
+
+//===----------------------------------------------------------------------===//
+// QueryOp
+//===----------------------------------------------------------------------===//
+
+void transform::QueryOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  onlyReadsHandle(getTargetMutable(), effects);
+  producesHandle(getOperation()->getOpResults(), effects);
+  onlyReadsPayload(effects);
+}
+
+DiagnosedSilenceableFailure transform::QueryOp::applyToOne(
+    transform::TransformRewriter &rewriter, Operation *target,
+    transform::ApplyToEachResultList &results, TransformState &state) {
+  StringAttr deviceId = getDeviceAttr();
+  StringAttr key = getKeyAttr();
+
+  DataLayoutEntryInterface entry;
+  if (deviceId) {
+    TargetSystemSpecInterface sysSpec = dlti::getTargetSystemSpec(target);
+    if (!sysSpec)
+      return mlir::emitDefiniteFailure(target->getLoc())
+             << "no target system spec associated to: " << target;
+
+    if (auto targetSpec = sysSpec.getDeviceSpecForDeviceID(deviceId))
+      entry = targetSpec->getSpecForIdentifier(key);
+    else
+      return mlir::emitDefiniteFailure(target->getLoc())
+             << "no " << deviceId << " target device spec found";
+  } else {
+    DataLayoutSpecInterface dlSpec = dlti::getDataLayoutSpec(target);
+    if (!dlSpec)
+      return mlir::emitDefiniteFailure(target->getLoc())
+             << "no data layout spec associated to: " << target;
+
+    entry = dlSpec.getSpecForIdentifier(key);
+  }
+
+  if (!entry)
+    return mlir::emitDefiniteFailure(target->getLoc())
+           << "no DLTI entry for key: " << key;
+
+  results.push_back(entry.getValue());
+
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
+// Transform op registration
+//===----------------------------------------------------------------------===//
+
+namespace {
+class DLTITransformDialectExtension
+    : public transform::TransformDialectExtension<
+          DLTITransformDialectExtension> {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DLTITransformDialectExtension)
+
+  using Base::Base;
+
+  void init() {
+    registerTransformOps<
+#define GET_OP_LIST
+#include "mlir/Dialect/DLTI/TransformOps/DLTITransformOps.cpp.inc"
+        >();
+  }
+};
+} // namespace
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/DLTI/TransformOps/DLTITransformOps.cpp.inc"
+
+void mlir::dlti::registerTransformDialectExtension(DialectRegistry &registry) {
+  registry.addExtensions<DLTITransformDialectExtension>();
+}
diff --git a/mlir/test/Dialect/DLTI/query.mlir b/mlir/test/Dialect/DLTI/query.mlir
new file mode 100644
index 00000000000000..2504958de09289
--- /dev/null
+++ b/mlir/test/Dialect/DLTI/query.mlir
@@ -0,0 +1,296 @@
+// RUN: mlir-opt -transform-interpreter -canonicalize -split-input-file -verify-diagnostics %s | FileCheck %s
+
+// expected-remark @below {{associated attr 42 : i32}}
+module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} {
+  func.func private @f() 
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg: !transform.any_op) {
+    %funcs = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op
+    %module = transform.get_parent_op %funcs : (!transform.any_op) -> !transform.any_op
+    %param = transform.dlti.query "test.id" at %module : (!transform.any_op) -> !transform.any_param
+    transform.debug.emit_param_as_remark %param, "associated attr" at %module : !transform.any_param, !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} {
+  // expected-remark @below {{associated attr 24 : i32}}
+  func.func private @f() attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 24 : i32>>}
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg: !transform.any_op) {
+    %funcs = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op
+    %param = transform.dlti.query "test.id" at %funcs : (!transform.any_op) -> !transform.any_param
+    transform.debug.emit_param_as_remark %param, "associated attr" at %funcs : !transform.any_param, !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+// expected-remark @below {{associated attr 42 : i32}}
+module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} {
+  func.func private @f() attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 24 : i32>>}
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg: !transform.any_op) {
+    %funcs = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op
+    %module = transform.get_parent_op %funcs : (!transform.any_op) -> !transform.any_op
+    %param = transform.dlti.query "test.id" at %module : (!transform.any_op) -> !transform.any_param
+    transform.debug.emit_param_as_remark %param, "associated attr" at %module : !transform.any_param, !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} {
+  func.func @matmul_tensors(
+    %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
+      -> tensor<?x?xf32> {
+    // expected-remark @below {{associated attr 42 : i32}}
+    %0 = linalg.matmul  ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
+                       outs(%arg2: tensor<?x?xf32>)
+      -> tensor<?x?xf32>
+    return %0 : tensor<?x?xf32>
+  }
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg: !transform.any_op) {
+    %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg : (!transform.any_op) -> !transform.any_op
+    %param = transform.dlti.query "test.id" at %matmul : (!transform.any_op) -> !transform.any_param
+    transform.debug.emit_param_as_remark %param, "associated attr" at %matmul : !transform.any_param, !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} {
+  func.func @matmul_tensors(
+    %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
+      -> tensor<?x?xf32> attributes {test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 24 : i32>>} {
+    // expected-remark @below {{associated attr 24 : i32}}
+    %0 = linalg.matmul  ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
+                       outs(%arg2: tensor<?x?xf32>)
+      -> tensor<?x?xf32>
+    return %0 : tensor<?x?xf32>
+  }
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg: !transform.any_op) {
+    %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg : (!transform.any_op) -> !transform.any_op
+    %param = transform.dlti.query "test.id" at %matmul : (!transform.any_op) -> !transform.any_param
+    transform.debug.emit_param_as_remark %param, "associated attr" at %matmul : !transform.any_param, !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+// expected-remark @below {{associated attr 42 : i32}}
+module attributes { test.dlti = #dlti.target_system_spec<"CPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 42 : i32>>>} {
+  func.func private @f()
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg: !transform.any_op) {
+    %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op
+    %module = transform.get_parent_op %func : (!transform.any_op) -> !transform.any_op
+    %param = transform.dlti.query ::"CPU"::"test.id" at %module : (!transform.any_op) -> !transform.any_param
+    transform.debug.emit_param_as_remark %param, "associated attr" at %module : !transform.any_param, !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+module attributes { test.dlti = #dlti.target_system_spec<"CPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 42 : i32>>,
+                                                         "GPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 43 : i32>>>} {
+  // expected-remark @below {{associated attr 43 : i32}}
+  func.func private @f() 
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg: !transform.any_op) {
+    %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op
+    %param = transform.dlti.query ::"GPU"::"test.id" at %func : (!transform.any_op) -> !transform.any_param
+    transform.debug.emit_param_as_remark %param, "associated attr" at %func : !transform.any_param, !transform.any_op
+    transform.yi...
[truncated]

@rolfmorel
Copy link
Contributor Author

rolfmorel commented Aug 9, 2024

This is the 2nd attempt at #101561. The first attempt had to be reverted - thanks @rengolin - as it did not use the new mechanism for registering dialect extensions (see #101119). Mea culpa.

@banach-space and @rengolin, I will rebase again on Monday and ask for your help with merging, if you don't mind. Thank you!

@rolfmorel
Copy link
Contributor Author

rolfmorel commented Aug 12, 2024

Good morning @rengolin and @banach-space,

I just now rebased on latest main and squashed the fix-up commit. mlir-check succeeds for me locally.

If you could help with trying to merge again, that would be much appreciated!

The fix-up commit amounted to the following one line change:

--- a/mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp
+++ b/mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp
@@ -75,6 +75,9 @@ class DLTITransformDialectExtension
     : public transform::TransformDialectExtension<
           DLTITransformDialectExtension> {
 public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DLTITransformDialectExtension)
+
   using Base::Base;

   void init()

This transform op makes it possible to query attributes associated to IR
by means of the DLTI dialect.

The op takes both a `key` and a target `op` to perform the query at.
Facility functions automatically find the closest ancestor op which
defines the appropriate DLTI interface or has an attribute implementing
a DLTI interface. By default the lookup uses the data layout interfaces
of DLTI. If the optional `device` parameter is provided, the lookup
happens with respect to the interfaces for TargetSystemSpec and
TargetDeviceSpec.

This op uses new free-standing functions in the `dlti` namespace to not
only look up specifications via the `DataLayoutSpecOpInterface` and on
`ModuleOp`s but also on any ancestor op that has an appropriate DLTI
attribute.

(The first attempt that got merged was out-of-date w.r.t. a new
 mechanism for registering dialect extensions.)
@rengolin rengolin merged commit 2ad3bcd into llvm:main Aug 12, 2024
8 checks passed
bwendling pushed a commit to bwendling/llvm-project that referenced this pull request Aug 15, 2024
…lvm#102652)

This transform op makes it possible to query attributes associated to IR
by means of the DLTI dialect.

The op takes both a `key` and a target `op` to perform the query at.
Facility functions automatically find the closest ancestor op which
defines the appropriate DLTI interface or has an attribute implementing
a DLTI interface. By default the lookup uses the data layout interfaces
of DLTI. If the optional `device` parameter is provided, the lookup
happens with respect to the interfaces for TargetSystemSpec and
TargetDeviceSpec.

This op uses new free-standing functions in the `dlti` namespace to not
only look up specifications via the `DataLayoutSpecOpInterface` and on
`ModuleOp`s but also on any ancestor op that has an appropriate DLTI
attribute.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants