[Mlir-commits] [mlir] [MLIR][DLTI][Transform] Introduce transform.dlti.query (PR #101561)
Rolf Morel
llvmlistbot at llvm.org
Thu Aug 8 04:09:29 PDT 2024
https://github.com/rolfmorel updated https://github.com/llvm/llvm-project/pull/101561
>From 7448ef1c800c13065bef9ba2d34937c6e581a0d3 Mon Sep 17 00:00:00 2001
From: Rolf Morel <rolf.morel at intel.com>
Date: Thu, 1 Aug 2024 09:16:39 -0700
Subject: [PATCH 1/2] [MLIR][DLTI][Transform] Introduce transform.dlti.query
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.
---
mlir/include/mlir/Dialect/DLTI/CMakeLists.txt | 2 +
mlir/include/mlir/Dialect/DLTI/DLTI.h | 6 +
.../Dialect/DLTI/TransformOps/CMakeLists.txt | 6 +
.../DLTI/TransformOps/DLTITransformOps.h | 38 +++
.../DLTI/TransformOps/DLTITransformOps.td | 61 ++++
mlir/include/mlir/InitAllExtensions.h | 2 +
mlir/lib/Dialect/DLTI/CMakeLists.txt | 1 +
mlir/lib/Dialect/DLTI/DLTI.cpp | 35 +++
.../Dialect/DLTI/TransformOps/CMakeLists.txt | 15 +
.../DLTI/TransformOps/DLTITransformOps.cpp | 94 ++++++
mlir/test/Dialect/DLTI/query.mlir | 296 ++++++++++++++++++
11 files changed, 556 insertions(+)
create mode 100644 mlir/include/mlir/Dialect/DLTI/TransformOps/CMakeLists.txt
create mode 100644 mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.h
create mode 100644 mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td
create mode 100644 mlir/lib/Dialect/DLTI/TransformOps/CMakeLists.txt
create mode 100644 mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp
create mode 100644 mlir/test/Dialect/DLTI/query.mlir
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..4061c98bf9eca4 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTI.h
+++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h
@@ -22,6 +22,12 @@ namespace detail {
class DataLayoutEntryAttrStorage;
} // namespace detail
} // namespace mlir
+namespace mlir {
+namespace dlti {
+DataLayoutSpecInterface getDataLayoutSpec(Operation *op);
+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 ®istry);
+} // 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..438912a6b2623b
--- /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.
+
+ 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 succesfull, 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 ®istry) {
// 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..488c4c97e819d1 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -389,6 +389,41 @@ TargetSystemSpecAttr::verify(function_ref<InFlightDiagnostic()> emitError,
return success();
}
+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;
+}
+
//===----------------------------------------------------------------------===//
// DLTIDialect
//===----------------------------------------------------------------------===//
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..a03e5e885e9dbf
--- /dev/null
+++ b/mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp
@@ -0,0 +1,94 @@
+
+//===- 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"
+
+//===----------------------------------------------------------------------===//
+// FuseOp
+//===----------------------------------------------------------------------===//
+
+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:
+ 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 ®istry) {
+ 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.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 24 : i32}}
+ func.func private @f() attributes { test.dlti = #dlti.target_system_spec<"CPU": #dlti.target_device_spec<#dlti.dl_entry<"test.id", 24 : i32>>> }
+}
+
+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 ::"CPU"::"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.yield
+ }
+}
+
+// -----
+
+module attributes { test.dlti = #dlti.target_system_spec<"CPU": #dlti.target_device_spec<#dlti.dl_entry<"cache::L1::size_in_bytes", 65536 : i32>,
+ #dlti.dl_entry<"cache::L1d::size_in_bytes", 32768 : i32>>> } {
+ // expected-remark @below {{L1::size_in_bytes 65536 : i32}}
+ // expected-remark @below {{L1d::size_in_bytes 32768 : 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
+ %l1_size = transform.dlti.query ::"CPU"::"cache::L1::size_in_bytes" at %func : (!transform.any_op) -> !transform.param<i32>
+ %l1d_size = transform.dlti.query ::"CPU"::"cache::L1d::size_in_bytes" at %func : (!transform.any_op) -> !transform.param<i32>
+ transform.debug.emit_param_as_remark %l1_size, "L1::size_in_bytes" at %func : !transform.param<i32>, !transform.any_op
+ transform.debug.emit_param_as_remark %l1d_size, "L1d::size_in_bytes" at %func : !transform.param<i32>, !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+module attributes { test.dlti = #dlti.target_system_spec<"CPU":
+ #dlti.target_device_spec<#dlti.dl_entry<"inner_most_tile_size", 42 : i32>>>} {
+ // CHECK-LABEL: func @matmul_tensors
+ func.func @matmul_tensors(
+ %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
+ -> tensor<?x?xf32> {
+ // CHECK: scf.for {{.*}} to {{.*}} step {{.*}}42
+ // CHECK: tensor.extract_slice
+ // CHECK: linalg.matmul
+ // CHECK: tensor.insert_slice
+ // CHECK: scf.yield
+ %0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%arg2: tensor<?x?xf32>)
+ -> tensor<?x?xf32>
+ // CHECK: return
+ return %0 : tensor<?x?xf32>
+ }
+}
+
+// Demonstrates obtaining transform op parameters from DLTI attributes and directly putting them to use.
+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
+ %tile_size = transform.dlti.query ::"CPU"::"inner_most_tile_size" at %matmul : (!transform.any_op) -> !transform.param<i32>
+ transform.structured.tile_using_for %matmul tile_sizes [%tile_size] : (!transform.any_op, !transform.param<i32>) -> (!transform.any_op, !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-error @below {{no "NPU" target device spec found}}
+ 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 ::"NPU"::"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.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-error @below {{no DLTI entry for key: "unspecified"}}
+ 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 ::"CPU"::"unspecified" 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.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-error @below {{no data layout spec associated to: }}
+ 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 "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.yield
+ }
+}
+
+// -----
+
+module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} {
+ // expected-error @below {{no target system spec associated to: }}
+ 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 ::"CPU"::"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.yield
+ }
+}
+
+// -----
+
+module {
+ // expected-error @below {{no target system spec associated to: }}
+ 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 ::"CPU"::"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.yield
+ }
+}
+
+// -----
+
+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
+ // expected-error @below {{expected the type of the parameter attribute ('i32') to match the parameter type ('i64')}}
+ %param = transform.dlti.query "test.id" at %funcs : (!transform.any_op) -> !transform.param<i64>
+ transform.debug.emit_param_as_remark %param, "associated attr" at %funcs : !transform.param<i64>, !transform.any_op
+ transform.yield
+ }
+}
>From 7ab0a58d7b1b1460e6f91c5547babfd2de746245 Mon Sep 17 00:00:00 2001
From: Rolf Morel <rolf.morel at intel.com>
Date: Thu, 8 Aug 2024 02:27:52 -0700
Subject: [PATCH 2/2] Update docs based on review
---
mlir/include/mlir/Dialect/DLTI/DLTI.h | 7 +++++++
.../mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td | 6 +++---
mlir/lib/Dialect/DLTI/DLTI.cpp | 8 ++++----
mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp | 2 +-
4 files changed, 15 insertions(+), 8 deletions(-)
diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h
index 4061c98bf9eca4..6e2623e084be52 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTI.h
+++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h
@@ -24,7 +24,14 @@ class DataLayoutEntryAttrStorage;
} // 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
diff --git a/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td
index 438912a6b2623b..69aacac986ad73 100644
--- a/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td
+++ b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td
@@ -24,7 +24,7 @@ def QueryOp : Op<Transform_Dialect, "dlti.query", [
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.
+ 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
@@ -33,8 +33,8 @@ def QueryOp : Op<Transform_Dialect, "dlti.query", [
#### Return modes
- When succesfull, the result, `associated_attr`, associates one attribute
- as a param for each op in `target`'s payload.
+ 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
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 488c4c97e819d1..c995f44c380e57 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -389,6 +389,10 @@ TargetSystemSpecAttr::verify(function_ref<InFlightDiagnostic()> emitError,
return success();
}
+//===----------------------------------------------------------------------===//
+// DLTIDialect
+//===----------------------------------------------------------------------===//
+
DataLayoutSpecInterface dlti::getDataLayoutSpec(Operation *op) {
DataLayoutSpecInterface dlSpec = nullptr;
@@ -424,10 +428,6 @@ TargetSystemSpecInterface dlti::getTargetSystemSpec(Operation *op) {
return sysSpec;
}
-//===----------------------------------------------------------------------===//
-// DLTIDialect
-//===----------------------------------------------------------------------===//
-
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/DLTITransformOps.cpp b/mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp
index a03e5e885e9dbf..259af8857c02f6 100644
--- a/mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp
+++ b/mlir/lib/Dialect/DLTI/TransformOps/DLTITransformOps.cpp
@@ -20,7 +20,7 @@ using namespace mlir::transform;
#define DEBUG_TYPE "dlti-transforms"
//===----------------------------------------------------------------------===//
-// FuseOp
+// QueryOp
//===----------------------------------------------------------------------===//
void transform::QueryOp::getEffects(
More information about the Mlir-commits
mailing list