[Mlir-commits] [mlir] 5ab49ed - [mlir][py][c] Enable setting block arg locations. (#169033)

[llvmlistbot at llvm.org]

Author: Jacques Pienaar
Date: 2025-11-21T13:31:46Z
New Revision: 5ab49edde282814f41b90431194afaff694deba7

URL: https://github.com/llvm/llvm-project/commit/5ab49edde282814f41b90431194afaff694deba7
DIFF: https://github.com/llvm/llvm-project/commit/5ab49edde282814f41b90431194afaff694deba7.diff

LOG: [mlir][py][c] Enable setting block arg locations. (#169033)

This enables changing the location of a block argument. Follows the
approach for updating type of block arg.

Added: 
    

Modified: 
    mlir/include/mlir-c/IR.h
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/lib/CAPI/IR/IR.cpp
    mlir/test/python/ir/blocks.py

Removed: 
    


################################################################################
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..e876c00e0c52d 100644
--- a/mlir/test/python/ir/blocks.py
+++ b/mlir/test/python/ir/blocks.py
@@ -191,3 +191,18 @@ def testBlockEraseArgs():
         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)