[Mlir-commits] [mlir] [MLIR][DLTI] Make queries visit all ancestors/respect nested scopes (PR #115043)
Rolf Morel
llvmlistbot at llvm.org
Tue Nov 5 10:22:28 PST 2024
https://github.com/rolfmorel created https://github.com/llvm/llvm-project/pull/115043
Change the behaviour of a DLTI `query(op, keys)` from just finding the first `DLTIQueryInterface`-implementing attr of the nearest ancestor of `op` and only looking up `keys` on that attribute to continueing on to further ancestors of `op` in case a lookup doesn't succeed on the first encountered `DLTIQueryInterface`-implementing attributes.
In effect, this gives `query` the expected "scoped" look-up semantics in that DLTI attributes that belong to ops whose regions encompass `op` will now be inspected in case the lookup doesn't succeed in the "closest scope". As usual for scoped lookups we have that names/keys can be shadowed.
Includes a small fix for `dlti.transform.query` documentation.
>From cdd0ef4c8c4d2fb325167f63362d131f9fba95f2 Mon Sep 17 00:00:00 2001
From: Rolf Morel <rolf.morel at intel.com>
Date: Tue, 5 Nov 2024 09:49:23 -0800
Subject: [PATCH] [MLIR][DLTI] Make queries visit all ancestors/respect nested
scopes
Change the behaviour of a DLTI `query(op, keys)` from just finding the
first `DLTIQueryInterface`-implementing attr of the nearest ancestor of
`op` and only looking up `keys` on that attribute to continueing on to
further ancestors of `op` in case a lookup doesn't succeed on the first
encountered `DLTIQueryInterface`-implementing attributes.
In effect, this gives `query` the expected "scoped" look-up semantics
in that DLTI attributes that belong to ops whose regions encompass `op`
will now be inspected in case the lookup doesn't succeed in the "closest
scope". As usual for scoped lookups we have that names/keys can be
shadowed.
Includes a small fix for `dlti.transform.query` documentation.
---
mlir/docs/Dialects/Transform.md | 4 +
mlir/include/mlir/Dialect/DLTI/DLTI.h | 5 +-
mlir/lib/Dialect/DLTI/DLTI.cpp | 120 ++++++++++++------------
mlir/test/Dialect/DLTI/query.mlir | 129 ++++++++++++++++++++++++--
4 files changed, 184 insertions(+), 74 deletions(-)
diff --git a/mlir/docs/Dialects/Transform.md b/mlir/docs/Dialects/Transform.md
index 37fcc0c8880335..5f79116dd00b55 100644
--- a/mlir/docs/Dialects/Transform.md
+++ b/mlir/docs/Dialects/Transform.md
@@ -427,6 +427,10 @@ ops rather than having the methods directly act on the payload IR.
[include "Dialects/DebugExtensionOps.md"]
+## DLTI Transform Operations
+
+[include "Dialects/DLTITransformOps.md"]
+
## IRDL (extension) Transform Operations
[include "Dialects/IRDLExtensionOps.md"]
diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h
index f268fea340a6fb..8ff04927052f21 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTI.h
+++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h
@@ -24,8 +24,9 @@ class DataLayoutEntryAttrStorage;
} // namespace mlir
namespace mlir {
namespace dlti {
-/// Perform a DLTI-query at `op`, recursively querying each key of `keys` on
-/// query interface-implementing attrs, starting from attr obtained from `op`.
+/// Perform a DLTI-query at `op`, by recursively querying each key of `keys` on
+/// `DLTIQueryInterface`-implementing attributes of an op, attempting this query
+/// procedure for all ancestors of `op` in turn, starting with `op` itself.
FailureOr<Attribute> query(Operation *op, ArrayRef<DataLayoutEntryKey> keys,
bool emitError = false);
} // namespace dlti
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 508e50d42e4cf2..ab0ca936099884 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -8,17 +8,13 @@
#include "mlir/Dialect/DLTI/DLTI.h"
#include "mlir/IR/Builders.h"
-#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinDialect.h"
#include "mlir/IR/BuiltinOps.h"
-#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/DialectImplementation.h"
#include "llvm/ADT/TypeSwitch.h"
-#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h"
-#include "llvm/Support/MathExtras.h"
using namespace mlir;
@@ -489,77 +485,77 @@ void TargetSystemSpecAttr::print(AsmPrinter &printer) const {
// DLTIDialect
//===----------------------------------------------------------------------===//
-/// Retrieve the first `DLTIQueryInterface`-implementing attribute that is
-/// attached to `op` or such an attr on as close as possible an ancestor. The
-/// op the attribute is attached to is returned as well.
-static std::pair<DLTIQueryInterface, Operation *>
-getClosestQueryable(Operation *op) {
- DLTIQueryInterface queryable = {};
-
- // Search op and its ancestors for the first attached DLTIQueryInterface attr.
- do {
- for (NamedAttribute attr : op->getAttrs())
- if ((queryable = dyn_cast<DLTIQueryInterface>(attr.getValue())))
- break;
- } while (!queryable && (op = op->getParentOp()));
-
- return std::pair(queryable, op);
-}
-
FailureOr<Attribute>
dlti::query(Operation *op, ArrayRef<DataLayoutEntryKey> keys, bool emitError) {
+ InFlightDiagnostic diag = op->emitError() << "target op of failed DLTI query";
+
if (keys.empty()) {
- if (emitError) {
- auto diag = op->emitError() << "target op of failed DLTI query";
+ if (emitError)
diag.attachNote(op->getLoc()) << "no keys provided to attempt query with";
- }
+ else
+ diag.abandon();
return failure();
}
- auto [queryable, queryOp] = getClosestQueryable(op);
- Operation *reportOp = (queryOp ? queryOp : op);
-
- if (!queryable) {
- if (emitError) {
- auto diag = op->emitError() << "target op of failed DLTI query";
- diag.attachNote(reportOp->getLoc())
- << "no DLTI-queryable attrs on target op or any of its ancestors";
- }
- return failure();
- }
-
- Attribute currentAttr = queryable;
- for (auto &&[idx, key] : llvm::enumerate(keys)) {
- if (auto map = dyn_cast<DLTIQueryInterface>(currentAttr)) {
- auto maybeAttr = map.query(key);
- if (failed(maybeAttr)) {
- if (emitError) {
- auto diag = op->emitError() << "target op of failed DLTI query";
- diag.attachNote(reportOp->getLoc())
- << "key " << keyToStr(key)
- << " has no DLTI-mapping per attr: " << map;
+ auto interleaveComma = [](ArrayRef<DataLayoutEntryKey> keys) {
+ std::string buf;
+ llvm::interleave(
+ keys, [&](auto key) { buf += keyToStr(key); }, [&]() { buf += ","; });
+ return buf;
+ };
+
+ // Recursively replace `currentAttr` by the attribute obtained by querying a
+ // new key on each new `currentAttr` until all `keys` have been exhausted -
+ // `atOp` is only used for error reporting.
+ auto queryKeysOnAttribute = [&](Attribute currentAttr,
+ Operation *atOp) -> FailureOr<Attribute> {
+ for (auto &&[idx, key] : llvm::enumerate(keys)) {
+ if (auto map = dyn_cast<DLTIQueryInterface>(currentAttr)) {
+ auto maybeAttr = map.query(key);
+ if (failed(maybeAttr)) {
+ if (emitError)
+ diag.attachNote(atOp->getLoc())
+ << "key not present - failed at keys: ["
+ << interleaveComma(keys.take_front(idx + 1)) << "]";
+ return failure();
}
+ currentAttr = *maybeAttr;
+ } else {
+ // The previous key, if any, is responsible for the current currentAttr.
+ if (idx > 0 && emitError)
+ diag.attachNote(atOp->getLoc())
+ << "attribute at keys [" << interleaveComma(keys.take_front(idx))
+ << "] is not queryable";
return failure();
}
- currentAttr = *maybeAttr;
- } else {
- if (emitError) {
- std::string commaSeparatedKeys;
- llvm::interleave(
- keys.take_front(idx), // All prior keys.
- [&](auto key) { commaSeparatedKeys += keyToStr(key); },
- [&]() { commaSeparatedKeys += ","; });
-
- auto diag = op->emitError() << "target op of failed DLTI query";
- diag.attachNote(reportOp->getLoc())
- << "got non-DLTI-queryable attribute upon looking up keys ["
- << commaSeparatedKeys << "] at op";
- }
- return failure();
}
+ return currentAttr;
+ };
+
+ // Run over all ancestors of `op`, starting the recursive attribute query for
+ // each ancestor which has an attribute on which we can perform a query.
+ for (Operation *ancestor = op; ancestor; ancestor = ancestor->getParentOp()) {
+ DLTIQueryInterface queryableAttr;
+ // NB: only the op's first DLTI attr will be inspected
+ for (NamedAttribute attr : ancestor->getAttrs())
+ if (auto queryableAttr = dyn_cast<DLTIQueryInterface>(attr.getValue())) {
+ auto maybeAttr = queryKeysOnAttribute(queryableAttr, ancestor);
+ if (succeeded(maybeAttr)) {
+ diag.abandon();
+ return maybeAttr;
+ }
+ }
+ }
+
+ if (emitError) {
+ if (diag.getUnderlyingDiagnostic()->getNotes().empty())
+ diag.attachNote(op->getLoc())
+ << "no DLTI-queryable attrs on target op or any of its ancestors";
+ } else {
+ diag.abandon();
}
- return currentAttr;
+ return failure();
}
constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutAttrName;
diff --git a/mlir/test/Dialect/DLTI/query.mlir b/mlir/test/Dialect/DLTI/query.mlir
index 3825cee6f16164..c5fbf67dc7c2b2 100644
--- a/mlir/test/Dialect/DLTI/query.mlir
+++ b/mlir/test/Dialect/DLTI/query.mlir
@@ -203,16 +203,47 @@ module attributes {transform.with_named_sequence} {
// -----
+// Demonstation of nested lookup by walking ancestors and co-commitant shadowing.
+
+// expected-remark @below {{associated CPU attr at module 42 : i32}}
+// expected-remark @below {{associated GPU attr at module 43 : i32}}
module attributes { test.dlti = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"test.id" = 42 : i32>,
"GPU" = #dlti.target_device_spec<"test.id" = 43 : i32>> } {
- // expected-remark @below {{associated attr 24 : i32}}
+ // expected-remark @below {{associated CPU attr at func 24 : i32}}
+ // expected-remark @below {{associated GPU attr at func 43 : i32}}
func.func private @f() attributes { test.dlti = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"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
+ %module = transform.get_parent_op %func : (!transform.any_op) -> !transform.any_op
+ // First the CPU attributes
+ %cpu_func_param = transform.dlti.query ["CPU","test.id"] at %func : (!transform.any_op) -> !transform.any_param
+ transform.debug.emit_param_as_remark %cpu_func_param, "associated CPU attr at func" at %func : !transform.any_param, !transform.any_op
+ %cpu_module_param = transform.dlti.query ["CPU","test.id"] at %module : (!transform.any_op) -> !transform.any_param
+ transform.debug.emit_param_as_remark %cpu_module_param, "associated CPU attr at module" at %module : !transform.any_param, !transform.any_op
+ // Now the GPU attribute
+ %gpu_func_param = transform.dlti.query ["GPU","test.id"] at %func : (!transform.any_op) -> !transform.any_param
+ transform.debug.emit_param_as_remark %gpu_func_param, "associated GPU attr at func" at %func : !transform.any_param, !transform.any_op
+ %gpu_module_param = transform.dlti.query ["GPU","test.id"] at %func : (!transform.any_op) -> !transform.any_param
+ transform.debug.emit_param_as_remark %gpu_module_param , "associated GPU attr at module" at %module : !transform.any_param, !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+module attributes { test.dlti = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"test.id" = 42 : i32>,
+ "GPU" = #dlti.target_device_spec<"test.id" = 43 : i32>> } {
+ // expected-remark @below {{associated attr 43 : i32}}
+ func.func private @f() attributes { test.dlti = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"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 ["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
}
@@ -220,6 +251,58 @@ module attributes {transform.with_named_sequence} {
// -----
+// Demonstation of nested lookup by walking ancestors and co-commitant shadowing.
+
+// expected-remark @below {{associated CPU attr at module 42 : i32}}
+// expected-remark @below {{associated GPU attr at module 43 : i32}}
+module attributes { test.dlti = #dlti.map<"CPU" = #dlti.map<"test.id" = 42 : i32>,
+ "GPU" = #dlti.map<"test.id" = 43 : i32>> } {
+ // expected-remark @below {{associated CPU attr at func 42 : i32}}
+ // expected-remark @below {{associated GPU attr at func 43 : i32}}
+ func.func @f(%A: tensor<128x128xf32>) {
+ // expected-remark @below {{associated CPU attr at matmul 24 : i32}}
+ // expected-remark @below {{associated GPU attr at matmul 43 : i32}}
+ %0 = linalg.matmul { test.dlti = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"test.id" = 24 : i32>> } ins(%A, %A : tensor<128x128xf32>, tensor<128x128xf32>)
+ outs(%A : tensor<128x128xf32>) -> tensor<128x128xf32>
+ // expected-remark @below {{associated CPU attr at constant 42 : i32}}
+ // expected-remark @below {{associated GPU attr at constant 43 : i32}}
+ arith.constant 0 : i32
+ return
+ }
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg: !transform.any_op) {
+ %constant = transform.structured.match ops{["arith.constant"]} in %arg : (!transform.any_op) -> !transform.any_op
+ %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg : (!transform.any_op) -> !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
+ // First query at the matmul
+ %cpu_matmul_param = transform.dlti.query ["CPU","test.id"] at %matmul : (!transform.any_op) -> !transform.any_param
+ transform.debug.emit_param_as_remark %cpu_matmul_param, "associated CPU attr at matmul" at %matmul : !transform.any_param, !transform.any_op
+ %gpu_matmul_param = transform.dlti.query ["GPU","test.id"] at %matmul : (!transform.any_op) -> !transform.any_param
+ transform.debug.emit_param_as_remark %gpu_matmul_param, "associated GPU attr at matmul" at %matmul : !transform.any_param, !transform.any_op
+ // Now query at the constant
+ %cpu_constant_param = transform.dlti.query ["CPU","test.id"] at %constant : (!transform.any_op) -> !transform.any_param
+ transform.debug.emit_param_as_remark %cpu_constant_param, "associated CPU attr at constant" at %constant : !transform.any_param, !transform.any_op
+ %gpu_constant_param = transform.dlti.query ["GPU","test.id"] at %constant : (!transform.any_op) -> !transform.any_param
+ transform.debug.emit_param_as_remark %gpu_constant_param, "associated GPU attr at constant" at %constant : !transform.any_param, !transform.any_op
+ // Now query at the func
+ %cpu_func_param = transform.dlti.query ["CPU","test.id"] at %func : (!transform.any_op) -> !transform.any_param
+ transform.debug.emit_param_as_remark %cpu_func_param, "associated CPU attr at func" at %func : !transform.any_param, !transform.any_op
+ %gpu_func_param = transform.dlti.query ["GPU","test.id"] at %func : (!transform.any_op) -> !transform.any_param
+ transform.debug.emit_param_as_remark %gpu_func_param, "associated GPU attr at func" at %func : !transform.any_param, !transform.any_op
+ // Now query at the module
+ %cpu_module_param = transform.dlti.query ["CPU","test.id"] at %module : (!transform.any_op) -> !transform.any_param
+ transform.debug.emit_param_as_remark %cpu_module_param, "associated CPU attr at module" at %module : !transform.any_param, !transform.any_op
+ %gpu_module_param = transform.dlti.query ["GPU","test.id"] at %module : (!transform.any_op) -> !transform.any_param
+ transform.debug.emit_param_as_remark %gpu_module_param, "associated GPU attr at module" at %module : !transform.any_param, !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
module attributes { test.dlti = #dlti.target_system_spec<
"CPU" = #dlti.target_device_spec<
"cache::L1::size_in_bytes" = 65536 : i32,
@@ -298,7 +381,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// expected-note @below {{key "NPU" has no DLTI-mapping per attr: #dlti.target_system_spec}}
+// expected-note @below {{key not present - failed at keys: ["NPU"]}}
module attributes { test.dlti = #dlti.target_system_spec<
"CPU" = #dlti.target_device_spec<"test.id" = 42 : i32>,
"GPU" = #dlti.target_device_spec<"test.id" = 43 : i32>> } {
@@ -317,7 +400,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// expected-note @below {{key "unspecified" has no DLTI-mapping per attr: #dlti.target_device_spec}}
+// expected-note @below {{key not present - failed at keys: ["CPU","unspecified"]}}
module attributes { test.dlti = #dlti.target_system_spec<
"CPU" = #dlti.target_device_spec<"test.id" = 42 : i32>,
"GPU" = #dlti.target_device_spec<"test.id" = 43 : i32>> } {
@@ -336,7 +419,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// expected-note @below {{key "test.id" has no DLTI-mapping per attr: #dlti.target_system_spec}}
+// expected-note @below {{key not present - failed at keys: ["test.id"]}}
module attributes { test.dlti = #dlti.target_system_spec<
"CPU" = #dlti.target_device_spec<"test.id" = 42 : i32>,
"GPU" = #dlti.target_device_spec<"test.id" = 43 : i32>> } {
@@ -355,7 +438,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// expected-note @below {{key "CPU" has no DLTI-mapping per attr: #dlti.dl_spec}}
+// expected-note @below {{key not present - failed at keys: ["CPU"]}}
module attributes { test.dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
// expected-error @below {{target op of failed DLTI query}}
func.func private @f()
@@ -372,7 +455,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// expected-note @below {{got non-DLTI-queryable attribute upon looking up keys ["CPU"]}}
+// expected-note @below {{attribute at keys ["CPU"] is not queryable}}
module attributes { test.dlti = #dlti.dl_spec<"CPU" = 42 : i32> } {
// expected-error @below {{target op of failed DLTI query}}
func.func private @f()
@@ -389,7 +472,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// expected-note @below {{got non-DLTI-queryable attribute upon looking up keys [i32]}}
+// expected-note @below {{attribute at keys [i32] is not queryable}}
module attributes { test.dlti = #dlti.dl_spec<i32 = 32 : i32> } {
// expected-error @below {{target op of failed DLTI query}}
func.func private @f()
@@ -423,7 +506,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// expected-note @below {{key i64 has no DLTI-mapping per attr: #dlti.map<i32 = 32 : i64>}}
+// expected-note @below {{key not present - failed at keys: ["width_in_bits",i64]}}
module attributes { test.dlti = #dlti.map<"width_in_bits" = #dlti.map<i32 = 32>>} {
// expected-error @below {{target op of failed DLTI query}}
func.func private @f()
@@ -483,4 +566,30 @@ module attributes {transform.with_named_sequence} {
%param = transform.dlti.query ["test.id"] at %funcs : (!transform.any_op) -> !transform.param<i64>
transform.yield
}
-}
\ No newline at end of file
+}
+
+// -----
+
+// expected-note @below {{attribute at keys ["CPU","test"] is not queryable}}
+module attributes { test.dlti = #dlti.map<"CPU" = #dlti.map<"test" = {"id" = 0}>> } {
+ // expected-note @below {{key not present - failed at keys: ["CPU","test","id"]}}
+ func.func @f(%A: tensor<128x128xf32>) attributes { test.dlti = #dlti.map<"CPU" = #dlti.map<"test" = #dlti.map<"ego" = 0>>> } {
+ scf.execute_region { // NB: No notes/errors on this unannotated ancestor
+ // expected-note @below {{key not present - failed at keys: ["CPU","test"]}}
+ // expected-error @below {{target op of failed DLTI query}}
+ %0 = linalg.matmul { test.dlti = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"test.id" = 24 : i32>> } ins(%A, %A : tensor<128x128xf32>, tensor<128x128xf32>)
+ outs(%A : tensor<128x128xf32>) -> tensor<128x128xf32>
+ scf.yield
+ }
+ return
+ }
+}
+
+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
+ // expected-error @below {{'transform.dlti.query' op failed to apply}}
+ %cpu_matmul_param = transform.dlti.query ["CPU","test","id"] at %matmul : (!transform.any_op) -> !transform.any_param
+ transform.yield
+ }
+}
More information about the Mlir-commits
mailing list