Skip to content

Commit 0b9b014

Browse files
authored
[mlir][dlti] Query by strings (#126716)
Adds DLTI utility to query using strings directly as keys.
1 parent 46f1bab commit 0b9b014

File tree

2 files changed

+22
-0
lines changed

2 files changed

+22
-0
lines changed

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,12 @@ namespace dlti {
2828
/// query interface-implementing attrs, starting from attr obtained from `op`.
2929
FailureOr<Attribute> query(Operation *op, ArrayRef<DataLayoutEntryKey> keys,
3030
bool emitError = false);
31+
32+
/// Perform a DLTI-query at `op` using each string in `keys` as a separate DLTI
33+
/// entry key, recursively querying on query interface-implementing attrs,
34+
/// starting from attr obtained from `op`.
35+
FailureOr<Attribute> query(Operation *op, ArrayRef<StringRef> keys,
36+
bool emitError = false);
3137
} // namespace dlti
3238
} // namespace mlir
3339

mlir/lib/Dialect/DLTI/DLTI.cpp

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -508,6 +508,9 @@ getClosestQueryable(Operation *op) {
508508

509509
FailureOr<Attribute>
510510
dlti::query(Operation *op, ArrayRef<DataLayoutEntryKey> keys, bool emitError) {
511+
if (!op)
512+
return failure();
513+
511514
if (keys.empty()) {
512515
if (emitError) {
513516
auto diag = op->emitError() << "target op of failed DLTI query";
@@ -562,6 +565,19 @@ dlti::query(Operation *op, ArrayRef<DataLayoutEntryKey> keys, bool emitError) {
562565
return currentAttr;
563566
}
564567

568+
FailureOr<Attribute> dlti::query(Operation *op, ArrayRef<StringRef> keys,
569+
bool emitError) {
570+
if (!op)
571+
return failure();
572+
573+
MLIRContext *ctx = op->getContext();
574+
SmallVector<DataLayoutEntryKey> entryKeys(keys.size());
575+
for (StringRef key : keys)
576+
entryKeys.push_back(StringAttr::get(ctx, key));
577+
578+
return dlti::query(op, entryKeys, emitError);
579+
}
580+
565581
constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutAttrName;
566582
constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessKey;
567583
constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessBig;

0 commit comments

Comments
 (0)