[Mlir-commits] [mlir] ad381e3 - [mlir] Provide minimal Python bindings for the math dialect

[Alex Zinenko]

Author: Alex Zinenko
Date: 2021-06-11T13:21:26+02:00
New Revision: ad381e39a52604ba07e1e027e7bdec1c287d9089

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

LOG: [mlir] Provide minimal Python bindings for the math dialect

Reviewed By: ulysseB

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

Added: 
    mlir/python/mlir/dialects/MathOps.td
    mlir/python/mlir/dialects/math.py
    mlir/test/python/dialects/math.py

Modified: 
    mlir/python/mlir/dialects/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/mlir/python/mlir/dialects/CMakeLists.txt b/mlir/python/mlir/dialects/CMakeLists.txt
index 5eeb6d6281208..3c0434475f304 100644
--- a/mlir/python/mlir/dialects/CMakeLists.txt
+++ b/mlir/python/mlir/dialects/CMakeLists.txt
@@ -25,6 +25,11 @@ add_mlir_dialect_python_bindings(MLIRBindingsPythonLinalgOps
   DEPENDS LinalgOdsGen)
 add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonLinalgOps)
 
+add_mlir_dialect_python_bindings(MLIRBindingsPythonMathOps
+  TD_FILE MathOps.td
+  DIALECT_NAME math)
+add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonMathOps)
+
 add_mlir_dialect_python_bindings(MLIRBindingsPythonMemRefOps
   TD_FILE MemRefOps.td
   DIALECT_NAME memref)

diff  --git a/mlir/python/mlir/dialects/MathOps.td b/mlir/python/mlir/dialects/MathOps.td
new file mode 100644
index 0000000000000..03d1fdef0c2e7
--- /dev/null
+++ b/mlir/python/mlir/dialects/MathOps.td
@@ -0,0 +1,15 @@
+//===-- MathOps.td - Entry point for MathOps 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_MATH_OPS
+#define PYTHON_BINDINGS_MATH_OPS
+
+include "mlir/Bindings/Python/Attributes.td"
+include "mlir/Dialect/Math/IR/MathOps.td"
+
+#endif

diff  --git a/mlir/python/mlir/dialects/math.py b/mlir/python/mlir/dialects/math.py
new file mode 100644
index 0000000000000..f082bf4615859
--- /dev/null
+++ b/mlir/python/mlir/dialects/math.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 ._math_ops_gen import *

diff  --git a/mlir/test/python/dialects/math.py b/mlir/test/python/dialects/math.py
new file mode 100644
index 0000000000000..73246e2e60eda
--- /dev/null
+++ b/mlir/test/python/dialects/math.py
@@ -0,0 +1,26 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+import mlir.dialects.builtin as builtin
+import mlir.dialects.math as mlir_math
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+
+# CHECK-LABEL: TEST: testMathOps
+ at run
+def testMathOps():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      @builtin.FuncOp.from_py_func(F32Type.get())
+      def emit_sqrt(arg):
+        return mlir_math.SqrtOp(F32Type.get(), arg)
+
+    # CHECK-LABEL: func @emit_sqrt(
+    # CHECK-SAME:                  %[[ARG:.*]]: f32) {
+    # CHECK:         math.sqrt %[[ARG]] : f32
+    # CHECK:         return
+    # CHECK:       }
+    print(module)