[Mlir-commits] [mlir] bbb73ad - [mlir][complex] Add Python bindings for complex ops.

[llvmlistbot at llvm.org]

Author: bixia1
Date: 2022-06-16T14:19:11-07:00
New Revision: bbb73ade43a29c90adabca69afb9a4df1a5cfb5e

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

LOG: [mlir][complex] Add Python bindings for complex ops.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D127916

Added: 
    mlir/python/mlir/dialects/ComplexOps.td
    mlir/python/mlir/dialects/complex.py
    mlir/test/python/dialects/complex_dialect.py

Modified: 
    mlir/python/CMakeLists.txt
    utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel

Removed: 
    


################################################################################
diff  --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 13b35f15b8ddc..4581dd6674842 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -81,6 +81,14 @@ declare_mlir_dialect_python_bindings(
     dialects/_builtin_ops_ext.py
   DIALECT_NAME builtin)
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/ComplexOps.td
+  SOURCES
+    dialects/complex.py
+  DIALECT_NAME complex)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"

diff  --git a/mlir/python/mlir/dialects/ComplexOps.td b/mlir/python/mlir/dialects/ComplexOps.td
new file mode 100644
index 0000000000000..6fd846ba6d270
--- /dev/null
+++ b/mlir/python/mlir/dialects/ComplexOps.td
@@ -0,0 +1,15 @@
+//===-- ComplexOps.td - Entry point for ComplexOps bindings ---------------===//
+//
+// 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_COMPLEX_OPS
+#define PYTHON_BINDINGS_COMPLEX_OPS
+
+include "mlir/Bindings/Python/Attributes.td"
+include "mlir/Dialect/Complex/IR/ComplexOps.td"
+
+#endif

diff  --git a/mlir/python/mlir/dialects/complex.py b/mlir/python/mlir/dialects/complex.py
new file mode 100644
index 0000000000000..ca81173cfc970
--- /dev/null
+++ b/mlir/python/mlir/dialects/complex.py
@@ -0,0 +1,5 @@
+#  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 ._complex_ops_gen import *

diff  --git a/mlir/test/python/dialects/complex_dialect.py b/mlir/test/python/dialects/complex_dialect.py
new file mode 100644
index 0000000000000..e724575b5bf5b
--- /dev/null
+++ b/mlir/test/python/dialects/complex_dialect.py
@@ -0,0 +1,32 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+# Naming this file with a `_dialect` suffix to avoid a naming conflict with
+# python package's math module (coming in from random.py).
+
+from mlir.ir import *
+import mlir.dialects.func as func
+import mlir.dialects.complex as mlir_complex
+
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+
+
+# CHECK-LABEL: TEST: testComplexOps
+ at run
+def testComplexOps():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+
+      @func.FuncOp.from_py_func(ComplexType.get(F32Type.get()))
+      def emit_add(arg):
+        return mlir_complex.AddOp(arg, arg)
+
+    # CHECK-LABEL: func @emit_add(
+    # CHECK-SAME:                  %[[ARG:.*]]: complex<f32>) -> complex<f32> {
+    # CHECK:         %[[RES:.*]] = complex.add %[[ARG]], %[[ARG]] : complex<f32>
+    # CHECK:         return %[[RES]] : complex<f32>
+    # CHECK:       }
+    print(module)

diff  --git a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
index c0134920dfe95..ee58abd08f2f2 100644
--- a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
@@ -296,6 +296,48 @@ filegroup(
     ],
 )
 
+##---------------------------------------------------------------------------##
+# Complex dialect.
+##---------------------------------------------------------------------------##
+
+td_library(
+    name = "ComplexOpsPyTdFiles",
+    srcs = [
+        "//mlir:include/mlir/Bindings/Python/Attributes.td",
+    ],
+    includes = ["../include"],
+    deps = [
+        "//mlir:ComplexOpsTdFiles",
+        "//mlir:OpBaseTdFiles",
+    ],
+)
+
+gentbl_filegroup(
+    name = "ComplexOpsPyGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-python-op-bindings",
+                "-bind-dialect=complex",
+            ],
+            "mlir/dialects/_complex_ops_gen.py",
+        ),
+    ],
+    tblgen = "//mlir:mlir-tblgen",
+    td_file = "mlir/dialects/ComplexOps.td",
+    deps = [
+        ":ComplexOpsPyTdFiles",
+    ],
+)
+
+filegroup(
+    name = "ComplexOpsPyFiles",
+    srcs = [
+        "mlir/dialects/complex.py",
+        ":ComplexOpsPyGen",
+    ],
+)
+
 ##---------------------------------------------------------------------------##
 # ControlFlow dialect.
 ##---------------------------------------------------------------------------##