Skip to content

[mlir][python] Enable python bindings for Index dialect #85827

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.

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 4 commits into from
Mar 20, 2024

Conversation

stevenvar
Copy link
Contributor

@stevenvar stevenvar commented Mar 19, 2024

This small patch enables python bindings for the index dialect.

Copy link

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be
notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write
permissions for the repository. In which case you can instead tag reviewers by
name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review
by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate
is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

@llvmbot llvmbot added mlir:python MLIR Python bindings mlir labels Mar 19, 2024
@llvmbot
Copy link
Member

llvmbot commented Mar 19, 2024

@llvm/pr-subscribers-mlir

Author: Steven Varoumas (stevenvar)

Changes

This small patch is intended to enable python bindings for the index dialect.


Full diff: https://github.com/llvm/llvm-project/pull/85827.diff

4 Files Affected:

  • (modified) mlir/python/CMakeLists.txt (+9)
  • (added) mlir/python/mlir/dialects/IndexOps.td (+14)
  • (added) mlir/python/mlir/dialects/index.py (+6)
  • (added) mlir/test/python/dialects/index_dialect.py (+335)
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 563d035f155267..c27ee688a04087 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -108,6 +108,15 @@ declare_mlir_dialect_python_bindings(
     dialects/complex.py
   DIALECT_NAME complex)
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/IndexOps.td
+  SOURCES
+    dialects/index.py
+  DIALECT_NAME index
+  GEN_ENUM_BINDINGS)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
diff --git a/mlir/python/mlir/dialects/IndexOps.td b/mlir/python/mlir/dialects/IndexOps.td
new file mode 100644
index 00000000000000..f98e4698f22e12
--- /dev/null
+++ b/mlir/python/mlir/dialects/IndexOps.td
@@ -0,0 +1,14 @@
+//===-- IndexOps.td - Entry point for Index bindings -----*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_INDEX_OPS
+#define PYTHON_BINDINGS_INDEX_OPS
+
+include "mlir/Dialect/Index/IR/IndexOps.td"
+
+#endif
diff --git a/mlir/python/mlir/dialects/index.py b/mlir/python/mlir/dialects/index.py
new file mode 100644
index 00000000000000..73708c7d71a8c8
--- /dev/null
+++ b/mlir/python/mlir/dialects/index.py
@@ -0,0 +1,6 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from ._index_ops_gen import *
+from ._index_enum_gen import *
diff --git a/mlir/test/python/dialects/index_dialect.py b/mlir/test/python/dialects/index_dialect.py
new file mode 100644
index 00000000000000..f2e7071a53f04f
--- /dev/null
+++ b/mlir/test/python/dialects/index_dialect.py
@@ -0,0 +1,335 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import index, arith
+
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+
+
+# CHECK-LABEL: TEST: testConstantOp
+@run
+def testConstantOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+    # CHECK: %[[A:.*]] = index.constant 42
+    print(module)
+
+
+# CHECK-LABEL: TEST: testBoolConstantOp
+@run
+def testBoolConstantOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.BoolConstantOp(value=True)
+    # CHECK: %[[A:.*]] = index.bool.constant true
+    print(module)
+
+
+# CHECK-LABEL: TEST: testAndOp
+@run
+def testAndOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.AndOp(a, a)
+    # CHECK: %[[R:.*]] = index.and {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testOrOp
+@run
+def testOrOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.OrOp(a, a)
+    # CHECK: %[[R:.*]] = index.or {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testXOrOp
+@run
+def testXOrOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.XOrOp(a, a)
+    # CHECK: %[[R:.*]] = index.xor {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testCastSOp
+@run
+def testCastSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
+      c = index.CastSOp(input=a, output=IntegerType.get_signless(32))
+      d = index.CastSOp(input=b, output=IndexType.get())
+    # CHECK: %[[C:.*]] = index.casts {{.*}} : index to i32
+    # CHECK: %[[D:.*]] = index.casts {{.*}} : i64 to index
+    print(module)
+
+
+# CHECK-LABEL: TEST: testCastUOp
+@run
+def testCastUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
+      c = index.CastUOp(input=a, output=IntegerType.get_signless(32))
+      d = index.CastUOp(input=b, output=IndexType.get())
+    # CHECK: %[[C:.*]] = index.castu {{.*}} : index to i32
+    # CHECK: %[[D:.*]] = index.castu {{.*}} : i64 to index
+    print(module)
+
+
+# CHECK-LABEL: TEST: testCeilDivSOp
+@run
+def testCeilDivSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.CeilDivSOp(a, a)
+    # CHECK: %[[R:.*]] = index.ceildivs {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testCeilDivUOp
+@run
+def testCeilDivUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.CeilDivUOp(a, a)
+    # CHECK: %[[R:.*]] = index.ceildivu {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testCmpOp
+@run
+def testCmpOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      pred = AttrBuilder.get('IndexCmpPredicateAttr')("slt", context=ctx)
+      r = index.CmpOp(pred, lhs=a, rhs=b)
+    # CHECK: %[[R:.*]] = index.cmp slt({{.*}}, {{.*}})
+    print(module)
+
+
+# CHECK-LABEL: TEST: testAddOp
+@run
+def testAddOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.AddOp(a, a)
+    # CHECK: %[[R:.*]] = index.add {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testSubOp
+@run
+def testSubOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.SubOp(a, a)
+    # CHECK: %[[R:.*]] = index.sub {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testMulOp
+@run
+def testMulOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.MulOp(a, a)
+    # CHECK: %[[R:.*]] = index.mul {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testDivSOp
+@run
+def testDivSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.DivSOp(a, a)
+    # CHECK: %[[R:.*]] = index.divs {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testDivUOp
+@run
+def testDivUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.DivUOp(a, a)
+    # CHECK: %[[R:.*]] = index.divu {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testFloorDivSOp
+@run
+def testFloorDivSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.FloorDivSOp(a, a)
+    # CHECK: %[[R:.*]] = index.floordivs {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testMaxSOp
+@run
+def testMaxSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      r = index.MaxSOp(a, b)
+    # CHECK: %[[R:.*]] = index.maxs {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testMaxUOp
+@run
+def testMaxUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      r = index.MaxUOp(a, b)
+    # CHECK: %[[R:.*]] = index.maxu {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testMinSOp
+@run
+def testMinSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      r = index.MinSOp(a, b)
+    # CHECK: %[[R:.*]] = index.mins {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testMinUOp
+@run
+def testMinUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      r = index.MinUOp(a, b)
+    # CHECK: %[[R:.*]] = index.minu {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testRemSOp
+@run
+def testRemSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      r = index.RemSOp(a, b)
+    # CHECK: %[[R:.*]] = index.rems {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testRemUOp
+@run
+def testRemUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      r = index.RemUOp(a, b)
+    # CHECK: %[[R:.*]] = index.remu {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testShlOp
+@run
+def testShlOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=3)
+      r = index.ShlOp(a, b)
+    # CHECK: %[[R:.*]] = index.shl {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testShrSOp
+@run
+def testShrSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=3)
+      r = index.ShrSOp(a, b)
+    # CHECK: %[[R:.*]] = index.shrs {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testShrUOp
+@run
+def testShrUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=3)
+      r = index.ShrUOp(a, b)
+    # CHECK: %[[R:.*]] = index.shru {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testSizeOfOp
+@run
+def testSizeOfOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      r = index.SizeOfOp()
+    # CHECK: %[[R:.*]] = index.sizeof
+    print(module)

Copy link

github-actions bot commented Mar 19, 2024

✅ With the latest revision this PR passed the Python code formatter.

@stevenvar stevenvar force-pushed the enable_index_python_bindings branch from db4d850 to 10ce043 Compare March 20, 2024 09:19
@ftynse ftynse merged commit eb861ac into llvm:main Mar 20, 2024
Copy link

@stevenvar Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested
by our build bots. If there is a problem with a build, you may recieve a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as
the builds can include changes from many authors. It is not uncommon for your
change to be included in a build that fails due to someone else's changes, or
infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself.
This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!

chencha3 pushed a commit to chencha3/llvm-project that referenced this pull request Mar 23, 2024
This small patch enables python bindings for the index dialect.

---------

Co-authored-by: Steven Varoumas <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
mlir:python MLIR Python bindings mlir
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants