[mlir][python] python binding for the affine.store op (#68816)

This PR creates the necessary files to support bindings for operations
in the affine dialect.

This is the first of many PRs which will progressively introduce
affine.load, affine.for, etc operations. I would like to
acknowledge the work by Nelli's author @makslevental :
https://github.com/makslevental/nelli/blob/main/nelli/mlir/affine/affine.py
which jump-starts the work.

mlir/python/CMakeLists.txt

@@ -46,6 +46,16 @@ declare_mlir_python_sources(MLIRPythonCAPI.HeaderSources
 # 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
+    dialects/_affine_ops_ext.py
+  DIALECT_NAME affine
+  GEN_ENUM_BINDINGS)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"

mlir/python/mlir/dialects/AffineOps.td

@@ -0,0 +1,14 @@
+//===-- AffineOps.td - Entry point for Affine 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_AFFINE_OPS
+#define PYTHON_BINDINGS_AFFINE_OPS
+
+include "mlir/Dialect/Affine/IR/AffineOps.td"
+
+#endif

mlir/python/mlir/dialects/_affine_ops_ext.py

@@ -0,0 +1,56 @@
+#  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
+
+try:
+    from ..ir import *
+    from ._ods_common import get_op_result_or_value as _get_op_result_or_value
+    from ._ods_common import get_op_results_or_values as _get_op_results_or_values
+except ImportError as e:
+    raise RuntimeError("Error loading imports from extension module") from e
+
+from typing import Optional, Sequence, Union
+
+
+class AffineStoreOp:
+    """Specialization for the Affine store operation."""
+
+    def __init__(
+        self,
+        value: Union[Operation, OpView, Value],
+        memref: Union[Operation, OpView, Value],
+        map: AffineMap=None,
+        *,
+        map_operands=None,
+        loc=None,
+        ip=None
+    ):
+        """Creates an affine store operation.
+
+        - `value`: the value to store into the memref.
+        - `memref`: the buffer to store into.
+        - `map`: the affine map that maps the map_operands to the index of the 
+          memref.
+        - `map_operands`: the list of arguments to substitute the dimensions, 
+          then symbols in the affine map, in increasing order.
+        """
+        map = map if map is not None else []
+        map_operands = map_operands if map_operands is not None else []
+        operands = [
+            _get_op_result_or_value(value),
+            _get_op_result_or_value(memref),
+            *[_get_op_result_or_value(op) for op in map_operands]
+        ]
+        results = []
+        attributes = {"map": AffineMapAttr.get(map)}
+        regions = None
+        _ods_successors = None
+        super().__init__(self.build_generic(
+            attributes=attributes,
+            results=results,
+            operands=operands,
+            successors=_ods_successors,
+            regions=regions,
+            loc=loc,
+            ip=ip
+        ))

mlir/python/mlir/dialects/affine.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 ._affine_ops_gen import *