From cdd0ef4c8c4d2fb325167f63362d131f9fba95f2 Mon Sep 17 00:00:00 2001 From: Rolf Morel Date: Tue, 5 Nov 2024 09:49:23 -0800 Subject: [PATCH 1/2] [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 37fcc0c888033..5f79116dd00b5 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 f268fea340a6f..8ff04927052f2 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 query(Operation *op, ArrayRef keys, bool emitError = false); } // namespace dlti diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp index 508e50d42e4cf..ab0ca93609988 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 -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(attr.getValue()))) - break; - } while (!queryable && (op = op->getParentOp())); - - return std::pair(queryable, op); -} - FailureOr dlti::query(Operation *op, ArrayRef 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(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 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 { + for (auto &&[idx, key] : llvm::enumerate(keys)) { + if (auto map = dyn_cast(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(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 3825cee6f1616..c5fbf67dc7c2b 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 } { // 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}} +// expected-note @below {{key not present - failed at keys: ["width_in_bits",i64]}} module attributes { test.dlti = #dlti.map<"width_in_bits" = #dlti.map>} { // 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 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 + } +} From 5655f41db9b932b8dfd7a3948a8b671cefcc3aa7 Mon Sep 17 00:00:00 2001 From: Rolf Morel Date: Tue, 5 Nov 2024 11:38:00 -0800 Subject: [PATCH 2/2] Extend docs of transform.dlti.query --- .../DLTI/TransformOps/DLTITransformOps.td | 33 ++++++++++++++----- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td index f25bb383912d4..bce7327812ef0 100644 --- a/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td +++ b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td @@ -24,25 +24,40 @@ def QueryOp : Op>>} { + module attributes {#dlti.map = #dlti.map<"A" = #dlti.map<"B" = 42>>} { func.func private @f() } ``` - and we have that `%func` is a Tranform handle to op `@f`, then + and we have that `%func` is a Tranform handle to func `@f`, then `transform.dlti.query ["A", "B"] at %func` returns 42 as a param and `transform.dlti.query ["A"] at %func` returns the `#dlti.map` attribute containing just the key "B" and its value. Using `["B"]` or `["A","C"]` as `keys` will yield an error. + In the below example we have that querying `["A", "B"]` at `%func` - or any + op that it contains - returns 0, while these same keys yield 42 if the + containing module would be the `target` of the query. ``` + ``` + module attributes {#dlti.map = #dlti.map<"A" = #dlti.map<"B" = 42, + "C" = 1>>} { + func.func private @f() attributes {#dlti.map = #dlti.map<"A" = + #dlti.map<"B" = 0>>} { + ... + } + } + ``` + Querying `["A","C"]` on the module, the func, or any contained op yields 1. + #### Return modes When successful, the result, `associated_attr`, associates one attribute as