[Mlir-commits] [mlir] [mlir][py][c] Enable setting block arg locations. (PR #169033)
Jacques Pienaar
llvmlistbot at llvm.org
Fri Nov 21 05:22:37 PST 2025
https://github.com/jpienaar updated https://github.com/llvm/llvm-project/pull/169033
>From 5074ad98d1f6260f3bf0f28f9226ab4efd7e6dc1 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jacques+gh at japienaar.info>
Date: Fri, 21 Nov 2025 10:51:02 +0000
Subject: [PATCH 1/2] [mlir][py][c] Enable setting block arg locations.
This enables changing the location of a block argument. Follows the approach for updating type of block arg.
---
mlir/include/mlir-c/IR.h | 4 ++++
mlir/lib/Bindings/Python/IRCore.cpp | 6 +++++
mlir/lib/CAPI/IR/IR.cpp | 5 ++++
mlir/test/python/ir/blocks.py | 37 ++++++++++++++++++++---------
4 files changed, 41 insertions(+), 11 deletions(-)
diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index c464e4da66f17..d2f476286ca69 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -1051,6 +1051,10 @@ MLIR_CAPI_EXPORTED intptr_t mlirBlockArgumentGetArgNumber(MlirValue value);
MLIR_CAPI_EXPORTED void mlirBlockArgumentSetType(MlirValue value,
MlirType type);
+/// Sets the location of the block argument to the given location.
+MLIR_CAPI_EXPORTED void mlirBlockArgumentSetLocation(MlirValue value,
+ MlirLocation loc);
+
/// Returns an operation that produced this value as its result. Asserts if the
/// value is not an op result.
MLIR_CAPI_EXPORTED MlirOperation mlirOpResultGetOwner(MlirValue value);
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 9d5bb9f54e933..03b540de97d4f 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2347,6 +2347,12 @@ class PyBlockArgument : public PyConcreteValue<PyBlockArgument> {
return mlirBlockArgumentSetType(self.get(), type);
},
nb::arg("type"), "Sets the type of this block argument.");
+ c.def(
+ "set_location",
+ [](PyBlockArgument &self, PyLocation loc) {
+ return mlirBlockArgumentSetLocation(self.get(), loc);
+ },
+ nb::arg("loc"), "Sets the location of this block argument.");
}
};
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index 188186598c5c5..ffcbed8b340cd 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -1129,6 +1129,11 @@ void mlirBlockArgumentSetType(MlirValue value, MlirType type) {
blockArg.setType(unwrap(type));
}
+void mlirBlockArgumentSetLocation(MlirValue value, MlirLocation loc) {
+ if (auto blockArg = llvm::dyn_cast<BlockArgument>(unwrap(value)))
+ blockArg.setLoc(unwrap(loc));
+}
+
MlirOperation mlirOpResultGetOwner(MlirValue value) {
return wrap(llvm::dyn_cast<OpResult>(unwrap(value)).getOwner());
}
diff --git a/mlir/test/python/ir/blocks.py b/mlir/test/python/ir/blocks.py
index ced5fce434728..932ac6d793e11 100644
--- a/mlir/test/python/ir/blocks.py
+++ b/mlir/test/python/ir/blocks.py
@@ -180,14 +180,29 @@ def testBlockAddArgs():
# CHECK-LABEL: TEST: testBlockEraseArgs
@run
def testBlockEraseArgs():
- with Context() as ctx, Location.unknown(ctx) as loc:
- ctx.allow_unregistered_dialects = True
- f32 = F32Type.get()
- op = Operation.create("test", regions=1, loc=Location.unknown())
- blocks = op.regions[0].blocks
- blocks.append(f32)
- # CHECK: ^bb0(%{{.+}}: f32 loc(unknown)):
- op.print(enable_debug_info=True)
- blocks[0].erase_argument(0)
- # CHECK: ^bb0:
- op.print(enable_debug_info=True)
+ with Context() as ctx, Location.unknown(ctx) as loc:
+ ctx.allow_unregistered_dialects = True
+ f32 = F32Type.get()
+ op = Operation.create("test", regions=1, loc=Location.unknown())
+ blocks = op.regions[0].blocks
+ blocks.append(f32)
+ # CHECK: ^bb0(%{{.+}}: f32 loc(unknown)):
+ op.print(enable_debug_info=True)
+ blocks[0].erase_argument(0)
+ # CHECK: ^bb0:
+ op.print(enable_debug_info=True)
+
+
+# CHECK-LABEL: TEST: testBlockArgSetLocation
+# CHECK: ^bb0(%{{.+}}: f32 loc("new_loc")):
+ at run
+def testBlockArgSetLocation():
+ with Context() as ctx, Location.unknown(ctx) as loc:
+ ctx.allow_unregistered_dialects = True
+ f32 = F32Type.get()
+ op = Operation.create("test", regions=1, loc=Location.unknown())
+ blocks = op.regions[0].blocks
+ blocks.append(f32)
+ arg = blocks[0].arguments[0]
+ arg.set_location(Location.name("new_loc"))
+ op.print(enable_debug_info=True)
>From 8edc6befc05620d53214844e0c134f7addbee877 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jacques+gh at japienaar.info>
Date: Fri, 21 Nov 2025 13:13:46 +0000
Subject: [PATCH 2/2] Format
---
mlir/test/python/ir/blocks.py | 40 +++++++++++++++++------------------
1 file changed, 20 insertions(+), 20 deletions(-)
diff --git a/mlir/test/python/ir/blocks.py b/mlir/test/python/ir/blocks.py
index 932ac6d793e11..e876c00e0c52d 100644
--- a/mlir/test/python/ir/blocks.py
+++ b/mlir/test/python/ir/blocks.py
@@ -180,29 +180,29 @@ def testBlockAddArgs():
# CHECK-LABEL: TEST: testBlockEraseArgs
@run
def testBlockEraseArgs():
- with Context() as ctx, Location.unknown(ctx) as loc:
- ctx.allow_unregistered_dialects = True
- f32 = F32Type.get()
- op = Operation.create("test", regions=1, loc=Location.unknown())
- blocks = op.regions[0].blocks
- blocks.append(f32)
- # CHECK: ^bb0(%{{.+}}: f32 loc(unknown)):
- op.print(enable_debug_info=True)
- blocks[0].erase_argument(0)
- # CHECK: ^bb0:
- op.print(enable_debug_info=True)
+ with Context() as ctx, Location.unknown(ctx) as loc:
+ ctx.allow_unregistered_dialects = True
+ f32 = F32Type.get()
+ op = Operation.create("test", regions=1, loc=Location.unknown())
+ blocks = op.regions[0].blocks
+ blocks.append(f32)
+ # CHECK: ^bb0(%{{.+}}: f32 loc(unknown)):
+ op.print(enable_debug_info=True)
+ blocks[0].erase_argument(0)
+ # CHECK: ^bb0:
+ op.print(enable_debug_info=True)
# CHECK-LABEL: TEST: testBlockArgSetLocation
# CHECK: ^bb0(%{{.+}}: f32 loc("new_loc")):
@run
def testBlockArgSetLocation():
- with Context() as ctx, Location.unknown(ctx) as loc:
- ctx.allow_unregistered_dialects = True
- f32 = F32Type.get()
- op = Operation.create("test", regions=1, loc=Location.unknown())
- blocks = op.regions[0].blocks
- blocks.append(f32)
- arg = blocks[0].arguments[0]
- arg.set_location(Location.name("new_loc"))
- op.print(enable_debug_info=True)
+ with Context() as ctx, Location.unknown(ctx) as loc:
+ ctx.allow_unregistered_dialects = True
+ f32 = F32Type.get()
+ op = Operation.create("test", regions=1, loc=Location.unknown())
+ blocks = op.regions[0].blocks
+ blocks.append(f32)
+ arg = blocks[0].arguments[0]
+ arg.set_location(Location.name("new_loc"))
+ op.print(enable_debug_info=True)
More information about the Mlir-commits
mailing list