Skip to content

[mlir][python]Python Bindings for select edit operations on Block arguments #94305

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jun 5, 2024
Merged

[mlir][python]Python Bindings for select edit operations on Block arguments #94305

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
31cbfc9
Python Bindings for select edit operations on block arguments
sdasgup3 Jun 4, 2024
2257918
Address feedback: remove erase_arguments as erase_argument should be …
sdasgup3 Jun 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions mlir/include/mlir-c/IR.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -858,6 +858,9 @@ MLIR_CAPI_EXPORTED MlirValue mlirBlockAddArgument(MlirBlock block,
MlirType type,
MlirLocation loc);

/// Erase the argument at 'index' and remove it from the argument list.
MLIR_CAPI_EXPORTED void mlirBlockEraseArgument(MlirBlock block, unsigned index);

/// Inserts an argument of the specified type at a specified index to the block.
/// Returns the newly added argument.
MLIR_CAPI_EXPORTED MlirValue mlirBlockInsertArgument(MlirBlock block,
Expand Down
13 changes: 13 additions & 0 deletions mlir/lib/Bindings/Python/IRCore.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3238,6 +3238,19 @@ void mlir::python::populateIRCore(py::module &m) {
return PyBlockArgumentList(self.getParentOperation(), self.get());
},
"Returns a list of block arguments.")
.def(
"add_argument",
[](PyBlock &self, const PyType &type, const PyLocation &loc) {
return mlirBlockAddArgument(self.get(), type, loc);
},
"Append an argument of the specified type to the block and returns "
"the newly added argument.")
.def(
"erase_argument",
[](PyBlock &self, unsigned index) {
return mlirBlockEraseArgument(self.get(), index);
},
"Erase the argument at 'index' and remove it from the argument list.")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: This sounds like it actually erases the value but it doesn't right? It just removes from the argument list? If so, can you actually change the comment to say "but does not erase the value" or something like that.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can see you point and agree. I actually I borrowed the comment from the corresponding cpp APIs (

llvm-project/mlir/include/mlir/IR/Block.h

Line 116 in 539b72f

void eraseArgument(unsigned index);
). I believe the word erase is used there to mean arguments.erase(index), which also looks good to me.

Not sure if we update both and leave them as is.
wdyt?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nah in that case just leave both

.def_property_readonly(
"operations",
[](PyBlock &self) {
Expand Down
4 changes: 4 additions & 0 deletions mlir/lib/CAPI/IR/IR.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -906,6 +906,10 @@ MlirValue mlirBlockAddArgument(MlirBlock block, MlirType type,
return wrap(unwrap(block)->addArgument(unwrap(type), unwrap(loc)));
}

void mlirBlockEraseArgument(MlirBlock block, unsigned index) {
return unwrap(block)->eraseArgument(index);
}

MlirValue mlirBlockInsertArgument(MlirBlock block, intptr_t pos, MlirType type,
MlirLocation loc) {
return wrap(unwrap(block)->insertArgument(pos, unwrap(type), unwrap(loc)));
Expand Down
32 changes: 32 additions & 0 deletions mlir/test/python/ir/blocks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,3 +145,35 @@ def testBlockHash():
block1 = Block.create_at_start(dummy.operation.regions[0], [f32])
block2 = Block.create_at_start(dummy.operation.regions[0], [f32])
assert hash(block1) != hash(block2)


# CHECK-LABEL: TEST: testBlockAddArgs
@run
def testBlockAddArgs():
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()
# CHECK: ^bb0:
op.print(enable_debug_info=True)
blocks[0].add_argument(f32, loc)
# CHECK: ^bb0(%{{.+}}: f32 loc(unknown)):
op.print(enable_debug_info=True)


# 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)
Loading