rebase

Author: makslevental

mlir/python/CMakeLists.txt

@@ -52,7 +52,6 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/AffineOps.td
   SOURCES
     dialects/affine.py
-    dialects/_affine_ops_ext.py
   DIALECT_NAME affine
   GEN_ENUM_BINDINGS)
 

mlir/python/mlir/dialects/_affine_ops_ext.py

@@ -1,5 +1,50 @@
-#  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    
+#  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 *
+from ._affine_ops_gen import _Dialect
+
+try:
+    from ..ir import *
+    from ._ods_common import (
+        get_op_result_or_value as _get_op_result_or_value,
+        get_op_results_or_values as _get_op_results_or_values,
+        _cext as _ods_cext,
+    )
+except ImportError as e:
+    raise RuntimeError("Error loading imports from extension module") from e
+
+from typing import Optional, Sequence, Union
+
+
+@_ods_cext.register_operation(_Dialect, replace=True)
+class AffineStoreOp(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 []
+        indicies = [_get_op_result_or_value(op) for op in map_operands]
+        _ods_successors = None
+        super().__init__(
+            value, memref, indicies, AffineMapAttr.get(map), loc=loc, ip=ip
+        )

mlir/python/mlir/dialects/affine.py

@@ -295,13 +295,17 @@ static bool isODSReserved(StringRef str) {
 /// (does not change the `name` if it already is suitable) and returns the
 /// modified version.
 static std::string sanitizeName(StringRef name) {
-  std::string processed_str = name.str();
+  std::string processedStr = name.str();
+  std::replace_if(
+      processedStr.begin(), processedStr.end(),
+      [](char c) { return !llvm::isAlnum(c); }, '_');
 
-  std::replace(processed_str.begin(), processed_str.end(), '-', '_');
+  if (llvm::isDigit(*processedStr.begin()))
+    return "_" + processedStr;
 
-  if (isPythonReserved(processed_str) || isODSReserved(processed_str))
-    return processed_str + "_";
-  return processed_str;
+  if (isPythonReserved(processedStr) || isODSReserved(processedStr))
+    return processedStr + "_";
+  return processedStr;
 }
 
 static std::string attrSizedTraitForKind(const char *kind) {
@@ -977,7 +981,6 @@ static void emitValueBuilder(const Operator &op,
                              llvm::SmallVector<std::string> functionArgs,
                              raw_ostream &os) {
   auto name = sanitizeName(op.getOperationName());
-  iterator_range<llvm::SplittingIterator> splitName = llvm::split(name, ".");
   // Params with (possibly) default args.
   auto valueBuilderParams =
       llvm::map_range(functionArgs, [](const std::string &argAndMaybeDefault) {
@@ -996,16 +999,16 @@ static void emitValueBuilder(const Operator &op,
         auto lhs = *llvm::split(arg, "=").begin();
         return (lhs + "=" + llvm::convertToSnakeFromCamelCase(lhs)).str();
       });
-  os << llvm::formatv(
-      valueBuilderTemplate,
-      // Drop dialect name and then sanitize again (to catch e.g. func.return).
-      sanitizeName(llvm::join(++splitName.begin(), splitName.end(), "_")),
-      op.getCppClassName(), llvm::join(valueBuilderParams, ", "),
-      llvm::join(opBuilderArgs, ", "),
-      (op.getNumResults() > 1
-           ? "_Sequence[_ods_ir.OpResult]"
-           : (op.getNumResults() > 0 ? "_ods_ir.OpResult"
-                                     : "_ods_ir.Operation")));
+  std::string nameWithoutDialect =
+      op.getOperationName().substr(op.getOperationName().find('.') + 1);
+  os << llvm::formatv(valueBuilderTemplate, sanitizeName(nameWithoutDialect),
+                      op.getCppClassName(),
+                      llvm::join(valueBuilderParams, ", "),
+                      llvm::join(opBuilderArgs, ", "),
+                      (op.getNumResults() > 1
+                           ? "_Sequence[_ods_ir.OpResult]"
+                           : (op.getNumResults() > 0 ? "_ods_ir.OpResult"
+                                                     : "_ods_ir.Operation")));
 }
 
 /// Emits bindings for a specific Op to the given output stream.