[mlir][python] Fix generation of Python bindings for async dialect (#75960)

The Python bindings generated for "async" dialect didn't include any of
the "async" dialect ops. This PR fixes issues with generation of Python
bindings for "async" dialect and adds a test case to use them.

Author: adk9

mlir/include/mlir/Dialect/Async/IR/CMakeLists.txt

@@ -1,2 +1,3 @@
+set(LLVM_TARGET_DEFINITIONS AsyncOps.td)
 add_mlir_dialect(AsyncOps async)
 add_mlir_doc(AsyncOps AsyncDialect Dialects/ -gen-dialect-doc)

mlir/python/CMakeLists.txt

@@ -79,7 +79,7 @@ declare_mlir_dialect_python_bindings(
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/AsyncOps.td
   SOURCES_GLOB dialects/async_dialect/*.py
-  DIALECT_NAME async_dialect)
+  DIALECT_NAME async)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -591,7 +591,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
 
 declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
   MODULE_NAME _mlirAsyncPasses
-  ADD_TO_PARENT MLIRPythonSources.Dialects.async_dialect
+  ADD_TO_PARENT MLIRPythonSources.Dialects.async
   ROOT_DIR "${PYTHON_SOURCE_DIR}"
   SOURCES
     AsyncPasses.cpp

mlir/python/mlir/dialects/async_dialect/__init__.py

@@ -2,4 +2,4 @@
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from .._async_dialect_ops_gen import *
+from .._async_ops_gen import *

mlir/test/python/dialects/async_dialect.py

@@ -1,7 +1,8 @@
 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
-import mlir.dialects.async_dialect
+from mlir.dialects import arith
+import mlir.dialects.async_dialect as async_dialect
 import mlir.dialects.async_dialect.passes
 from mlir.passmanager import *
 
@@ -11,6 +12,19 @@ def run(f):
     f()
 
 
+# CHECK-LABEL: TEST: testCreateGroupOp
+@run
+def testCreateGroupOp():
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            i32 = IntegerType.get_signless(32)
+            group_size = arith.ConstantOp(i32, 4)
+            async_dialect.create_group(group_size)
+        # CHECK:         %0 = "arith.constant"() <{value = 4 : i32}> : () -> i32
+        # CHECK:         %1 = "async.create_group"(%0) : (i32) -> !async.group
+        print(module)
+
 def testAsyncPass():
     with Context() as context:
         PassManager.parse("any(async-to-async-runtime)")

utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel

@@ -331,6 +331,53 @@ filegroup(
     ],
 )
 
+##---------------------------------------------------------------------------##
+# Async dialect.
+##---------------------------------------------------------------------------##
+
+gentbl_filegroup(
+    name = "AsyncOpsPyGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-python-enum-bindings",
+                "-bind-dialect=async",
+            ],
+            "mlir/dialects/_async_enum_gen.py",
+        ),
+        (
+            [
+                "-gen-python-op-bindings",
+                "-bind-dialect=async",
+            ],
+            "mlir/dialects/_async_ops_gen.py",
+        ),
+    ],
+    tblgen = "//mlir:mlir-tblgen",
+    td_file = "mlir/dialects/AsyncOps.td",
+    deps = [
+        "//mlir:AsyncOpsTdFiles",
+        "//mlir:OpBaseTdFiles",
+    ],
+)
+
+filegroup(
+    name = "AsyncOpsPyFiles",
+    srcs = [
+        ":AsyncOpsPyGen",
+    ],
+)
+
+filegroup(
+    name = "AsyncOpsPackagePyFiles",
+    srcs = glob(["mlir/dialects/async_dialect/*.py"]),
+)
+
+filegroup(
+    name = "AsyncOpsPackagePassesPyFiles",
+    srcs = glob(["mlir/dialects/async_dialect/passes/*.py"]),
+)
+
 ##---------------------------------------------------------------------------##
 # Arith dialect.
 ##---------------------------------------------------------------------------##