[Mlir-commits] [mlir] [mlir] Add mlirTypeWalk and mlirAttributeWalk to the C API (PR #213182)
Jared Hoberock
llvmlistbot at llvm.org
Thu Jul 30 17:51:16 PDT 2026
https://github.com/jaredhoberock updated https://github.com/llvm/llvm-project/pull/213182
>From 946c3fd90385ebe7ef1333e1120c4e068a5c89e4 Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Thu, 30 Jul 2026 15:58:29 -0500
Subject: [PATCH] [mlir] Add mlirTypeWalk and mlirAttributeWalk to the C API
The C API can walk an operation together with the operations nested in it,
but offers no way to reach the types and attributes nested inside a type or
an attribute. Bindings outside C++ therefore cannot answer questions such as
"does anything nested inside this type come from dialect X", which any
legality check or IR audit written against the C API needs, while the C++
AttrTypeWalker answers them directly.
Add mlirTypeWalk and mlirAttributeWalk, which wrap AttrTypeWalker and report
every type and attribute reachable from the walked element, crossing between
types and attributes in either direction to any depth. They follow
mlirOperationWalk: the caller supplies callbacks, user data, and a traversal
order, and a callback steers the walk with MlirWalkResult. They additionally
return the walk result, so a caller can tell an interrupted walk from one
that ran to completion. MlirWalkResult and MlirWalkOrder now describe walks
in general rather than operation walks alone.
Assisted-by: Claude (Anthropic)
---
mlir/include/mlir-c/IR.h | 44 ++++++++++++-
mlir/lib/CAPI/IR/IR.cpp | 60 ++++++++++++++++++
mlir/test/CAPI/ir.c | 134 +++++++++++++++++++++++++++++++++++++++
3 files changed, 236 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 98ea0dfc00a1b..8bfeed29de6e3 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -856,14 +856,14 @@ MLIR_CAPI_EXPORTED void mlirOperationMoveBefore(MlirOperation op,
/// take O(N) where N is the number of operations within the parent block.
MLIR_CAPI_EXPORTED bool mlirOperationIsBeforeInBlock(MlirOperation op,
MlirOperation other);
-/// Operation walk result.
+/// Walk result.
typedef enum MlirWalkResult {
MlirWalkResultAdvance,
MlirWalkResultInterrupt,
MlirWalkResultSkip
} MlirWalkResult;
-/// Traversal order for operation walk.
+/// Traversal order for a walk.
typedef enum MlirWalkOrder {
MlirWalkPreOrder,
MlirWalkPostOrder
@@ -881,6 +881,46 @@ MLIR_CAPI_EXPORTED
void mlirOperationWalk(MlirOperation op, MlirOperationWalkCallback callback,
void *userData, MlirWalkOrder walkOrder);
+/// Walker type for types. The handler is passed an (opaque) reference to a
+/// type and a pointer to a `userData`.
+typedef MlirWalkResult (*MlirTypeWalkCallback)(MlirType, void *userData);
+
+/// Walker type for attributes. The handler is passed an (opaque) reference to
+/// an attribute and a pointer to a `userData`.
+typedef MlirWalkResult (*MlirAttributeWalkCallback)(MlirAttribute,
+ void *userData);
+
+/// Walks `type` in `walkOrder`, along with every type and attribute nested
+/// inside it, and calls `typeCallback` on each type and `attributeCallback` on
+/// each attribute. The walk crosses between types and attributes in both
+/// directions, so a type parameterized by an attribute reaches the types that
+/// attribute carries, and so on to any depth. Either callback may be null, in
+/// which case elements of that kind are not reported; the walk still descends
+/// through them. `*userData` is passed to the callbacks as well and can be used
+/// to tunnel some context or other data into them.
+///
+/// A type or an attribute that the walked structure mentions more than once is
+/// reported once per walk.
+///
+/// A callback returning `MlirWalkResultInterrupt` ends the walk, which this
+/// function reports by returning `MlirWalkResultInterrupt`. A callback
+/// returning `MlirWalkResultSkip` in `MlirWalkPreOrder` leaves the elements
+/// nested inside the reported element unvisited; in `MlirWalkPostOrder` those
+/// elements have already been visited, so it behaves like
+/// `MlirWalkResultAdvance`. This function returns `MlirWalkResultAdvance`
+/// unless the walk was interrupted.
+MLIR_CAPI_EXPORTED MlirWalkResult
+mlirTypeWalk(MlirType type, MlirTypeWalkCallback typeCallback,
+ MlirAttributeWalkCallback attributeCallback, void *userData,
+ MlirWalkOrder walkOrder);
+
+/// Walks `attribute` in `walkOrder`, along with every attribute and type nested
+/// inside it, with the same contract as `mlirTypeWalk`.
+MLIR_CAPI_EXPORTED MlirWalkResult
+mlirAttributeWalk(MlirAttribute attribute, MlirTypeWalkCallback typeCallback,
+ MlirAttributeWalkCallback attributeCallback, void *userData,
+ MlirWalkOrder walkOrder);
+
/// Replace uses of 'of' value with the 'with' value inside the 'op' operation.
MLIR_CAPI_EXPORTED void
mlirOperationReplaceUsesOfWith(MlirOperation op, MlirValue of, MlirValue with);
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index d5d545a4e1a71..ab9a53f9690e5 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -15,6 +15,7 @@
#include "mlir/CAPI/IRMapping.h"
#include "mlir/CAPI/Support.h"
#include "mlir/CAPI/Utils.h"
+#include "mlir/IR/AttrTypeSubElements.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinOps.h"
@@ -926,6 +927,65 @@ void mlirOperationWalk(MlirOperation op, MlirOperationWalkCallback callback,
}
}
+static MlirWalkResult wrap(mlir::WalkResult result) {
+ if (result.wasInterrupted())
+ return MlirWalkResultInterrupt;
+ if (result.wasSkipped())
+ return MlirWalkResultSkip;
+ return MlirWalkResultAdvance;
+}
+
+/// Registers the caller's callbacks on `walker`. A null callback registers no
+/// walk function for that kind of element, which leaves elements of that kind
+/// unreported while the walk still descends through them.
+static void addSubElementWalks(AttrTypeWalker &walker,
+ MlirTypeWalkCallback typeCallback,
+ MlirAttributeWalkCallback attributeCallback,
+ void *userData) {
+ if (typeCallback) {
+ walker.addWalk([typeCallback, userData](Type type) {
+ return unwrap(typeCallback(wrap(type), userData));
+ });
+ }
+ if (attributeCallback) {
+ walker.addWalk([attributeCallback, userData](Attribute attribute) {
+ return unwrap(attributeCallback(wrap(attribute), userData));
+ });
+ }
+}
+
+/// Runs `walker` over `element`. AttrTypeWalker takes the traversal order as a
+/// template argument, so the run-time order selects the instantiation here.
+template <typename ElementT>
+static MlirWalkResult runSubElementWalk(AttrTypeWalker &walker,
+ ElementT element,
+ MlirWalkOrder walkOrder) {
+ switch (walkOrder) {
+ case MlirWalkPreOrder:
+ return wrap(walker.walk<mlir::WalkOrder::PreOrder>(element));
+ case MlirWalkPostOrder:
+ return wrap(walker.walk<mlir::WalkOrder::PostOrder>(element));
+ }
+ llvm_unreachable("unknown order in runSubElementWalk");
+}
+
+MlirWalkResult mlirTypeWalk(MlirType type, MlirTypeWalkCallback typeCallback,
+ MlirAttributeWalkCallback attributeCallback,
+ void *userData, MlirWalkOrder walkOrder) {
+ AttrTypeWalker walker;
+ addSubElementWalks(walker, typeCallback, attributeCallback, userData);
+ return runSubElementWalk(walker, unwrap(type), walkOrder);
+}
+
+MlirWalkResult mlirAttributeWalk(MlirAttribute attribute,
+ MlirTypeWalkCallback typeCallback,
+ MlirAttributeWalkCallback attributeCallback,
+ void *userData, MlirWalkOrder walkOrder) {
+ AttrTypeWalker walker;
+ addSubElementWalks(walker, typeCallback, attributeCallback, userData);
+ return runSubElementWalk(walker, unwrap(attribute), walkOrder);
+}
+
void mlirOperationReplaceUsesOfWith(MlirOperation op, MlirValue oldValue,
MlirValue newValue) {
unwrap(op)->replaceUsesOfWith(unwrap(oldValue), unwrap(newValue));
diff --git a/mlir/test/CAPI/ir.c b/mlir/test/CAPI/ir.c
index 3c72209ad95c3..923c25d481a82 100644
--- a/mlir/test/CAPI/ir.c
+++ b/mlir/test/CAPI/ir.c
@@ -2943,6 +2943,138 @@ int testDominanceInfo(MlirContext ctx) {
return 0;
}
+typedef struct {
+ const char *label;
+} attrTypeWalkData;
+
+MlirWalkResult attrTypeWalkTypeCallBack(MlirType type, void *dataVoid) {
+ fprintf(stderr, "%s type: ", ((attrTypeWalkData *)(dataVoid))->label);
+ mlirTypePrint(type, printToStderr, NULL);
+ fprintf(stderr, "\n");
+ return MlirWalkResultAdvance;
+}
+
+MlirWalkResult attrTypeWalkAttributeCallBack(MlirAttribute attribute,
+ void *dataVoid) {
+ fprintf(stderr, "%s attr: ", ((attrTypeWalkData *)(dataVoid))->label);
+ mlirAttributePrint(attribute, printToStderr, NULL);
+ fprintf(stderr, "\n");
+ return MlirWalkResultAdvance;
+}
+
+MlirWalkResult attrTypeWalkTypeCallBackTestInterrupt(MlirType type,
+ void *dataVoid) {
+ attrTypeWalkTypeCallBack(type, dataVoid);
+ if (mlirTypeIsAInteger(type))
+ return MlirWalkResultInterrupt;
+ return MlirWalkResultAdvance;
+}
+
+MlirWalkResult attrTypeWalkAttributeCallBackTestSkip(MlirAttribute attribute,
+ void *dataVoid) {
+ attrTypeWalkAttributeCallBack(attribute, dataVoid);
+ if (mlirAttributeIsAArray(attribute))
+ return MlirWalkResultSkip;
+ return MlirWalkResultAdvance;
+}
+
+int testAttrTypeWalk(MlirContext ctx) {
+ // CHECK-LABEL: @testAttrTypeWalk
+ fprintf(stderr, "@testAttrTypeWalk\n");
+
+ // A tensor type encoded by an array attribute that holds a type attribute,
+ // so a type sits two levels down behind an attribute. The string attribute
+ // beside it spells a type name that the walk must not report as a type.
+ MlirAttribute elements[2] = {
+ mlirTypeAttrGet(mlirIntegerTypeGet(ctx, 16)),
+ mlirStringAttrGet(ctx, mlirStringRefCreateFromCString("i64"))};
+ MlirAttribute encoding = mlirArrayAttrGet(ctx, 2, elements);
+ int64_t shape[1] = {4};
+ MlirType tensor =
+ mlirRankedTensorTypeGet(1, shape, mlirF32TypeGet(ctx), encoding);
+
+ attrTypeWalkData data;
+ data.label = "postorder";
+
+ // CHECK-NEXT: postorder type: f32
+ // CHECK-NEXT: postorder type: i16
+ // CHECK-NEXT: postorder attr: i16
+ // CHECK-NEXT: postorder attr: "i64"
+ // CHECK-NEXT: postorder attr: [i16, "i64"]
+ // CHECK-NEXT: postorder type: tensor<4xf32, [i16, "i64"]>
+ mlirTypeWalk(tensor, attrTypeWalkTypeCallBack, attrTypeWalkAttributeCallBack,
+ (void *)(&data), MlirWalkPostOrder);
+
+ data.label = "preorder";
+ // CHECK-NEXT: preorder type: tensor<4xf32, [i16, "i64"]>
+ // CHECK-NEXT: preorder type: f32
+ // CHECK-NEXT: preorder attr: [i16, "i64"]
+ // CHECK-NEXT: preorder attr: i16
+ // CHECK-NEXT: preorder type: i16
+ // CHECK-NEXT: preorder attr: "i64"
+ mlirTypeWalk(tensor, attrTypeWalkTypeCallBack, attrTypeWalkAttributeCallBack,
+ (void *)(&data), MlirWalkPreOrder);
+
+ data.label = "skip preorder";
+ // Skipping the array attribute leaves what it holds unvisited.
+ // CHECK-NEXT: skip preorder type: tensor<4xf32, [i16, "i64"]>
+ // CHECK-NEXT: skip preorder type: f32
+ // CHECK-NEXT: skip preorder attr: [i16, "i64"]
+ MlirWalkResult result = mlirTypeWalk(tensor, attrTypeWalkTypeCallBack,
+ attrTypeWalkAttributeCallBackTestSkip,
+ (void *)(&data), MlirWalkPreOrder);
+ // CHECK-NEXT: result after skip is advance: 1
+ fprintf(stderr, "result after skip is advance: %d\n",
+ result == MlirWalkResultAdvance);
+
+ data.label = "skip postorder";
+ // In postorder the array attribute is reported after what it holds, so
+ // skipping it has nothing left to leave unvisited.
+ // CHECK-NEXT: skip postorder type: f32
+ // CHECK-NEXT: skip postorder type: i16
+ // CHECK-NEXT: skip postorder attr: i16
+ // CHECK-NEXT: skip postorder attr: "i64"
+ // CHECK-NEXT: skip postorder attr: [i16, "i64"]
+ // CHECK-NEXT: skip postorder type: tensor<4xf32, [i16, "i64"]>
+ mlirTypeWalk(tensor, attrTypeWalkTypeCallBack,
+ attrTypeWalkAttributeCallBackTestSkip, (void *)(&data),
+ MlirWalkPostOrder);
+
+ data.label = "interrupt";
+ // Interrupted at `i16`, the type the type attribute holds.
+ // CHECK-NEXT: interrupt type: tensor<4xf32, [i16, "i64"]>
+ // CHECK-NEXT: interrupt type: f32
+ // CHECK-NEXT: interrupt attr: [i16, "i64"]
+ // CHECK-NEXT: interrupt attr: i16
+ // CHECK-NEXT: interrupt type: i16
+ result = mlirTypeWalk(tensor, attrTypeWalkTypeCallBackTestInterrupt,
+ attrTypeWalkAttributeCallBack, (void *)(&data),
+ MlirWalkPreOrder);
+ // CHECK-NEXT: result after interrupt is interrupt: 1
+ fprintf(stderr, "result after interrupt is interrupt: %d\n",
+ result == MlirWalkResultInterrupt);
+
+ data.label = "from attribute";
+ // CHECK-NEXT: from attribute type: i16
+ // CHECK-NEXT: from attribute attr: i16
+ // CHECK-NEXT: from attribute attr: "i64"
+ // CHECK-NEXT: from attribute attr: [i16, "i64"]
+ mlirAttributeWalk(encoding, attrTypeWalkTypeCallBack,
+ attrTypeWalkAttributeCallBack, (void *)(&data),
+ MlirWalkPostOrder);
+
+ data.label = "attributes only";
+ // A null type callback leaves types unreported, and the walk still descends
+ // through them to the attributes they carry.
+ // CHECK-NEXT: attributes only attr: i16
+ // CHECK-NEXT: attributes only attr: "i64"
+ // CHECK-NEXT: attributes only attr: [i16, "i64"]
+ mlirTypeWalk(tensor, NULL, attrTypeWalkAttributeCallBack, (void *)(&data),
+ MlirWalkPostOrder);
+
+ return 0;
+}
+
int main(void) {
MlirContext ctx = mlirContextCreate();
registerAllUpstreamDialects(ctx);
@@ -2998,6 +3130,8 @@ int main(void) {
return 19;
if (testDominanceInfo(ctx))
return 20;
+ if (testAttrTypeWalk(ctx))
+ return 21;
// CHECK: DESTROY MAIN CONTEXT
// CHECK: reportResourceDelete: resource_i64_blob
More information about the Mlir-commits
mailing list