[mlir][affine] add python bindings for affine ops

[ashay-github]

Summary

Effectively, this patch makes it possible to reference
MLIRPythonSources.Dialects.affine when adding a dependency on the
affine dialect.

[ashay-github]

[nfc] signing commit

[ftynse]

I would prefer us to land op API customizations either in this patch, or very soon after. Generated API for affine ops will be quite raw and low-level, and we will likely grow uses for which the API will be broken when we land the improvements. Look at how SCF loops are supported in https://github.com/llvm/llvm-project/blob/main/mlir/python/mlir/dialects/_scf_ops_ext.py#L13. Also, load/store with non-trivial maps may get nasty and deserve some API work.

mlir/test/python/dialects/affine_dialect.py

17–24

It is preferable to create this structure programmatically than to rely on parsing string IR. API updates are mechanical, embedded string updates -- not so much.

[ftynse]

Note that it is possible to construct operations from any registered dialect using the generic Operation.create which is barely more verbose than the generated APIs in this case.

[ashay-github]

Apologies for the much delayed response. I'm abandoning this revision since I found a workaround, but I will come back to this with a proper fix later.

mlir/python/CMakeLists.txt

 
 ################################################################################
 # Dialect bindings
 ################################################################################
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/AffineOps.td
+  SOURCES
+    dialects/affine.py
+  DIALECT_NAME affine)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/AsyncOps.td
   SOURCES_GLOB dialects/async_dialect/*.py
   DIALECT_NAME async_dialect)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/BufferizationOps.td

mlir/python/mlir/dialects/AffineOps.td

+//===-- AffineOps.td - Entry point for AffineOps 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_AFFINE_OPS
+#define PYTHON_BINDINGS_AFFINE_OPS
+
+include "mlir/Bindings/Python/Attributes.td"
+include "mlir/Dialect/Affine/IR/AffineOps.td"
+
+#endif

mlir/python/mlir/dialects/affine.py

+#  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 ._affine_ops_gen import *

mlir/test/python/dialects/affine_dialect.py

+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import arith
+import mlir.dialects.affine as affine
+
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+
+
+# CHECK-LABEL: TEST: testAffineOp
+@run
+def testAffineOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.parse(r"""
+      func.func @emit_affine(%arg0: memref<f32>) {
+        return
+      }
+    """)
+    f = module.body.operations[0]
+    func_body = f.regions[0].blocks[0]
+    with InsertionPoint.at_block_terminator(func_body):
+      affine.AffineLoadOp(F32Type.get(), f.arguments[0], [])
+
+    # CHECK-LABEL: func @emit_affine(
+    # CHECK-SAME:                  %[[ARG:.*]]: memref<f32>) {
+    # CHECK:         %[[RES:.*]] = affine.load %[[ARG]][] : memref<f32>
+    # CHECK:         return
+    # CHECK:       }
+    print(module)