[mlir][py] Enable building ops with raw inputs

[jpienaar]

Summary

For cases where we can automatically construct the Attribute allow for more
user-friendly input. This is consistent with C++ builder generation as well
choice of which single builder to generate here (most
specialized/user-friendly).

Registration of attribute builders from more raw/non-Attribute input is all
Python side.
The downside is that

• same Python code may not work everywhere the same,
• extra checking to see if user provided a custom builder in op builders,
• the ODS attribute name is load bearing,

upside is that

• easily change these/register dialect specific ones in downstream projects,
• we/users can use, for example, numpy where available rather than lowest common denominator,
• adding support/changing to different convenience builders are all along with the rest of the convenience functions in Python (no additional changes to tablegen file or recompilation needed);

Allow for both building with Attributes as well as raw inputs. This change
should therefore be backwards compatible as well as allow for avoiding
recreating Attribute where already available.

[jpienaar]

Have the builders take context additionally

Avoids always using default context.

[jpienaar]

Remove leftover debugging

[jpienaar]

Update test post making context explicit

[jpienaar]

Try moving helpers to common spot

[ftynse]

In general, this feels like a usability improvement. We need user-level documentation on how to use and extend this.

Regarding potential downsides:

• Could we somehow limit the location of AttrBuilder functions so that a downstream integration of bindings can only add functions for attributes from its own dialects (and not core dialects)? I understand this is Python and one can always find a workaround, but making it intentionally hard would send the message. The downstream may want to construct built-in attributes from its custom types, but that's the point where we start getting divergence between downstreams that we'd better avoid if we care for reuse.

• The name of the ODS attribute class being load-bearing can be partially circumvented by having a "mapping" in tablegen, similar to what we do in: https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/Attributes.td. Clients could remap their custom ODS classes to whatever name they want in the builder, and this remapping is more or less visible in the code.

mlir/lib/Bindings/Python/IRCore.cpp

204–205

Could you add a short comment describing what this does? Defining __get__ makes something a property, but this defining __get__ itself as a property, I'm confused.

mlir/python/mlir/ir.py

9–10 ↗

(On Diff #482693)

We definitely want type annotations on what x can be.
We may also want to raise a TypeError on wrong types, though attribute getters may already be doing so.

11 ↗

(On Diff #482693)

Nit: I'd add vertical whitespace before each new def.

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

540–548

544

677–687

Nit: is it possible to stick to formatv everywhere for uniformity?

[jpienaar]

In D139568#4004298, @ftynse wrote:

In general, this feels like a usability improvement. We need user-level documentation on how to use and extend this.

Regarding potential downsides:

• Could we somehow limit the location of AttrBuilder functions so that a downstream integration of bindings can only add functions for attributes from its own dialects (and not core dialects)? I understand this is Python and one can always find a workaround, but making it intentionally hard would send the message. The downstream may want to construct built-in attributes from its custom types, but that's the point where we start getting divergence between downstreams that we'd better avoid if we care for reuse.

Best I can think of is to make it so that one cannot re-assign and then error if one attempts to ... Mmm, given multiple imports one either would then need to check if function being assigned is the same or add an explicit "register once" method. (I'm not python expert to know of alternatives). I think the way I could do that easiest would be to have an explicit set call which then takes a string as input and verified C++ side.

AttrBuilder.SymbolNameAttr = _symbolNameAttr

becomes

AttrBuilder.set_builder("SymbolNameAt", _symbolNameAttr)

(at least I couldn't figure out via pybind11 how to enable AttrBuilder['SymbolNameAttr"] = _symbolNameAttr but I can try a bit more)

• The name of the ODS attribute class being load-bearing can be partially circumvented by having a "mapping" in tablegen, similar to what we do in: https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/Attributes.td. Clients could remap their custom ODS classes to whatever name they want in the builder, and this remapping is more or less visible in the code.

One could, but I like this without having any side data structure. I see this as an enhancement we could add if it is cumbersome (e.g.,

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

677–687

You want the entire thing as one string? It means having two string constants and switching between them with the two string constants sharing values that need to be kept consistent between them.

[ftynse]

In D139568#4005166, @jpienaar wrote:

In D139568#4004298, @ftynse wrote:

In general, this feels like a usability improvement. We need user-level documentation on how to use and extend this.

Regarding potential downsides:

• Could we somehow limit the location of AttrBuilder functions so that a downstream integration of bindings can only add functions for attributes from its own dialects (and not core dialects)? I understand this is Python and one can always find a workaround, but making it intentionally hard would send the message. The downstream may want to construct built-in attributes from its custom types, but that's the point where we start getting divergence between downstreams that we'd better avoid if we care for reuse.

Best I can think of is to make it so that one cannot re-assign and then error if one attempts to ... Mmm, given multiple imports one either would then need to check if function being assigned is the same or add an explicit "register once" method. (I'm not python expert to know of alternatives). I think the way I could do that easiest would be to have an explicit set call which then takes a string as input and verified C++ side.

Works for me.

AttrBuilder.SymbolNameAttr = _symbolNameAttr

becomes

AttrBuilder.set_builder("SymbolNameAt", _symbolNameAttr)

(at least I couldn't figure out via pybind11 how to enable AttrBuilder['SymbolNameAttr"] = _symbolNameAttr but I can try a bit more)

Usually __setitem__ works, but you may need an instance of the builder class, not the class itself.

• The name of the ODS attribute class being load-bearing can be partially circumvented by having a "mapping" in tablegen, similar to what we do in: https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Bindings/Python/Attributes.td. Clients could remap their custom ODS classes to whatever name they want in the builder, and this remapping is more or less visible in the code.

One could, but I like this without having any side data structure. I see this as an enhancement we could add if it is cumbersome (e.g.,

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

677–687

We already have similar cases. I find it easier to guess what the generated code should look like if there is less string manipulation.

[jpienaar]

Switch to explicit methods and avoid redefining.

[jpienaar]

Make 2 self-contained strings

[jpienaar]

Update doc

[ftynse]

Nice!

mlir/lib/Bindings/Python/Globals.h

61–63

Something went wrong with the merge/rebase and now the documentation is duplicated,, and wrong for registerAttributeBuilder

mlir/lib/Bindings/Python/IRModule.cpp

70–72

Nit: throw std::runtime_error(message) should do what you expect without calling into low-level API.
Also nit: I'm always suspicious about usage of llvm::Twine in places that expect something else, please double-check that it doesn't get destroyed too early.

97

Nit: is *not* defined? Also, this is a py::function.

[jpienaar]

mlir/lib/Bindings/Python/IRModule.cpp

70–72

Yeah passing Twine unless as && often an issue. Here it converts to std::string before call.

97

The style followed here (taken from below in file) is to say what is being asserted (that the function is defined). Kept it as is to be consistent, I can follow up to flip all.

mlir/lib/Bindings/Python/Globals.h

public:
   /// contains calls to decorators and helpers that register the salient
   /// entities.
   void loadDialectModule(llvm::StringRef dialectNamespace);
 
   /// Decorator for registering a custom Dialect class. The class object must
   /// have a DIALECT_NAMESPACE attribute.
   pybind11::object registerDialectDecorator(pybind11::object pyClass);
 
   /// Adds a concrete implementation dialect class.
   /// Raises an exception if the mapping already exists.
   /// This is intended to be called by implementation code.
   void registerDialectImpl(const std::string &dialectNamespace,
                            pybind11::object pyClass);
 
   /// Adds a concrete implementation operation class.
   /// Raises an exception if the mapping already exists.
   /// This is intended to be called by implementation code.
   void registerOperationImpl(const std::string &operationName,
                              pybind11::object pyClass,
                              pybind11::object rawOpViewClass);
 
+  /// Returns the custom Attribute builder dictionary.
+  pybind11::object attributeBuilderMap();
+
   /// Looks up a registered dialect class by namespace. Note that this may
   /// trigger loading of the defining module and can arbitrarily re-enter.
   llvm::Optional<pybind11::object>
   lookupDialectClass(const std::string &dialectNamespace);
 
   /// Looks up a registered raw OpView class by operation name. Note that this
   /// may trigger a load of the dialect, which can arbitrarily re-enter.
   llvm::Optional<pybind11::object>
   lookupRawOpViewClass(llvm::StringRef operationName);
 
 private:
   static PyGlobals *instance;
   /// Module name prefixes to search under for dialect implementation modules.
   std::vector<std::string> dialectSearchPrefixes;
   /// Map of dialect namespace to external dialect class object.
   llvm::StringMap<pybind11::object> dialectClassMap;
   /// Map of full operation name to external operation class object.
   llvm::StringMap<pybind11::object> operationClassMap;
   /// Map of operation name to custom subclass that directly initializes
   /// the OpView base class (bypassing the user class constructor).
   llvm::StringMap<pybind11::object> rawOpViewClassMap;
+  /// Map of attribute ODS name to custom builder.
+  pybind11::object attributeBuilder;
 
   /// Set of dialect namespaces that we have attempted to import implementation
   /// modules for.
   llvm::StringSet<> loadedDialectModulesCache;
   /// Cache of operation name to custom OpView subclass that directly
   /// initializes the OpView base class (or an undefined object for negative
   /// lookup). This is maintained on loopup as a shadow of rawOpViewClassMap
   /// in order for repeat lookups of the OpView classes to only incur the cost
   /// of one hashtable lookup.
   llvm::StringMap<pybind11::object> rawOpViewClassMapCache;
 };
 
 } // namespace python
 } // namespace mlir
 
 #endif // MLIR_BINDINGS_PYTHON_GLOBALS_H

mlir/lib/Bindings/Python/IRCore.cpp

struct PyGlobalDebugFlag {
   static void bind(py::module &m) {
     // Debug flags.
     py::class_<PyGlobalDebugFlag>(m, "_GlobalDebug", py::module_local())
         .def_property_static("flag", &PyGlobalDebugFlag::get,
                              &PyGlobalDebugFlag::set, "LLVM-wide debug flag");
   }
 };
 
+struct PyAttrBuilder {
+  static py::function get(const py::object &) {
+    return PyGlobals::get().attributeBuilderMap();
+  }
+
+  static void bind(py::module &m) {
+    py::class_<PyAttrBuilder>(m, "AttrBuilder", py::module_local())
+        .def_property_readonly_static("__get__", &PyAttrBuilder::get,
+                                      "User-friendly attribute builders");
+  }
+};
+
 //------------------------------------------------------------------------------
 // Collections.
 //------------------------------------------------------------------------------
 
 namespace {
 
 class PyRegionIterator {
 public:

   PyOpAttributeMap::bind(m);
   PyOpOperandList::bind(m);
   PyOpResultList::bind(m);
   PyRegionIterator::bind(m);
   PyRegionList::bind(m);
 
   // Debug bindings.
   PyGlobalDebugFlag::bind(m);
+
+  // Attribute builder getter.
+  PyAttrBuilder::bind(m);
 }

mlir/lib/Bindings/Python/IRModule.cpp

 PyGlobals *PyGlobals::instance = nullptr;
 
 PyGlobals::PyGlobals() {
   assert(!instance && "PyGlobals already constructed");
   instance = this;
   // The default search path include {mlir.}dialects, where {mlir.} is the
   // package prefix configured at compile time.
   dialectSearchPrefixes.emplace_back(MAKE_MLIR_PYTHON_QUALNAME("dialects"));
+  py::object SimpleNamespace =
+      py::module_::import("types").attr("SimpleNamespace");
+  attributeBuilder = SimpleNamespace();
 }
 
 PyGlobals::~PyGlobals() { instance = nullptr; }
 
 void PyGlobals::loadDialectModule(llvm::StringRef dialectNamespace) {
   if (loadedDialectModulesCache.contains(dialectNamespace))
     return;
   // Since re-entrancy is possible, make a copy of the search prefixes.
void PyGlobals::loadDialectModule(llvm::StringRef dialectNamespace) {
   // may have occurred, which may do anything.
   loadedDialectModulesCache.insert(dialectNamespace);
 }
 
 void PyGlobals::registerDialectImpl(const std::string &dialectNamespace,
                                     py::object pyClass) {
   py::object &found = dialectClassMap[dialectNamespace];
   if (found) {
     throw SetPyError(PyExc_RuntimeError, llvm::Twine("Dialect namespace '") +
                                              dialectNamespace +
                                              "' is already registered.");
   }
   found = std::move(pyClass);
 }
 
 void PyGlobals::registerOperationImpl(const std::string &operationName,
                                       py::object pyClass,
                                       py::object rawOpViewClass) {
   py::object &found = operationClassMap[operationName];
   if (found) {
     throw SetPyError(PyExc_RuntimeError, llvm::Twine("Operation '") +
                                              operationName +
                                              "' is already registered.");
   }
   found = std::move(pyClass);
   rawOpViewClassMap[operationName] = std::move(rawOpViewClass);
 }
 
+py::object PyGlobals::attributeBuilderMap() { return attributeBuilder; }
+
 llvm::Optional<py::object>
 PyGlobals::lookupDialectClass(const std::string &dialectNamespace) {
   loadDialectModule(dialectNamespace);
   // Fast match against the class map first (common case).
   const auto foundIt = dialectClassMap.find(dialectNamespace);
   if (foundIt != dialectClassMap.end()) {
     if (foundIt->second.is_none())
       return std::nullopt;
     assert(foundIt->second && "py::object is defined");
     return foundIt->second;
   }
 
   // Not found and loading did not yield a registration. Negative cache.
   dialectClassMap[dialectNamespace] = py::none();

mlir/test/mlir-tblgen/op-python-bindings.td

def DefaultValuedAttrsOp : TestOp<"default_valued_attrs"> {
 
   let arguments = (ins DefaultValuedAttr<I64ArrayAttr, "{}">:$arr,
                        DefaultValuedAttr<I64ArrayAttr, "dummy_func()">:$unsupported);
   let results = (outs);
 }
 
 // CHECK-LABEL: OPERATION_NAME = "test.derive_result_types_op"
 def DeriveResultTypesOp : TestOp<"derive_result_types_op", [FirstAttrDerivedResultType]> {
-  // CHECK: def __init__(self, type, *, loc=None, ip=None):
+  // CHECK: def __init__(self, type_, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   _ods_result_type_source_attr = attributes["type"]
   // CHECK:   _ods_derived_result_type = (
   // CHECK:       _ods_ir.TypeAttr(_ods_result_type_source_attr).value
   // CHECK:       if _ods_ir.TypeAttr.isinstance(_ods_result_type_source_attr) else
   // CHECK:       _ods_result_type_source_attr.type)
   // CHECK:   results.extend([_ods_derived_result_type] * 2)
   let arguments = (ins TypeAttr:$type);
   let results = (outs AnyType:$res, AnyType);
 }
 
 // CHECK-LABEL: OPERATION_NAME = "test.derive_result_types_variadic_op"
 def DeriveResultTypesVariadicOp : TestOp<"derive_result_types_variadic_op", [FirstAttrDerivedResultType]> {
-  // CHECK: def __init__(self, res, _gen_res_1, type, *, loc=None, ip=None):
+  // CHECK: def __init__(self, res, _gen_res_1, type_, *, loc=None, ip=None):
   let arguments = (ins TypeAttr:$type);
   let results = (outs AnyType:$res, Variadic<AnyType>);
 }
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class EmptyOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.empty"
 def EmptyOp : TestOp<"empty">;

mlir/test/python/dialects/shape.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 import numpy as np
 import mlir.dialects.func as func
 import mlir.dialects.shape as shape
 
 
+# FIXME: These should be in common spot, here for convenience.
+AttrBuilder.BoolAttr = lambda x: BoolAttr.get(x)
+AttrBuilder.IndexAttr = lambda x: IntegerAttr.get(IndexType.get(), x)
+AttrBuilder.I32Attr = lambda x: IntegerAttr.get(IntegerType.get_signless(32), x)
+AttrBuilder.I64Attr = lambda x: IntegerAttr.get(IntegerType.get_signless(64), x)
+AttrBuilder.SymbolNameAttr = lambda x: StringAttr.get(x)
+
+AttrBuilder.IndexElementsAttr = lambda x: DenseElementsAttr.get(
+    np.array(x, dtype=np.int64), type=IndexType.get())
+
+
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   return f
 
 
 # CHECK-LABEL: TEST: testConstShape
 @run
 def testConstShape():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     with InsertionPoint(module.body):
       @func.FuncOp.from_py_func(
           RankedTensorType.get((12, ShapedType.get_dynamic_size()), f32))
       def const_shape_tensor(arg):
+        shape.ConstWitnessOp(False)
+        shape.ConstSizeOp(30)
+        shape.ConstSizeOp(IntegerAttr.get(IndexType.get(), 40))
+        shape.ConstShapeOp([1, 2])
         return shape.ConstShapeOp(
-          DenseElementsAttr.get(np.array([10, 20], dtype=np.int64), type=IndexType.get()))
+            DenseElementsAttr.get(
+                np.array([3, 4], dtype=np.int64), type=IndexType.get()))
 
     # CHECK-LABEL: func @const_shape_tensor(%arg0: tensor<12x?xf32>)
-    # CHECK: shape.const_shape [10, 20] : tensor<2xindex>
+    # CHECK-DAG: shape.const_witness false
+    # CHECK-DAG: shape.const_size 30
+    # CHECK-DAG: shape.const_size 40
+    # CHECK-DAG: shape.const_shape [1, 2] : tensor<2xindex>
+    # CHECK-DAG: shape.const_shape [3, 4] : tensor<2xindex>
     print(module)
+    print("bye")

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

clDialectName("bind-dialect",
                   llvm::cl::init(""), llvm::cl::cat(clOpPythonBindingCat));
 
 static llvm::cl::opt<std::string> clDialectExtensionName(
     "dialect-extension", llvm::cl::desc("The prefix of the dialect extension"),
     llvm::cl::init(""), llvm::cl::cat(clOpPythonBindingCat));
 
 using AttributeClasses = DenseMap<StringRef, StringRef>;
 
-/// Checks whether `str` is a Python keyword.
-static bool isPythonKeyword(StringRef str) {
-  static llvm::StringSet<> keywords(
-      {"and",   "as",     "assert",   "break", "class",  "continue",
-       "def",   "del",    "elif",     "else",  "except", "finally",
-       "for",   "from",   "global",   "if",    "import", "in",
-       "is",    "lambda", "nonlocal", "not",   "or",     "pass",
-       "raise", "return", "try",      "while", "with",   "yield"});
-  return keywords.contains(str);
+/// Checks whether `str` is a Python keyword or would shadow builtin function.
+static bool isPythonReserved(StringRef str) {
+  static llvm::StringSet<> reserved(
+      {"and",      "as",    "assert", "break",      "callable", "class",
+       "continue", "def",   "del",    "elif",       "else",     "except",
+       "finally",  "for",   "from",   "global",     "if",       "import",
+       "in",       "is",    "lambda", "nonlocal",   "not",      "or",
+       "pass",     "raise", "return", "issubclass", "try",      "type",
+       "while",    "with",  "yield"});
+  return reserved.contains(str);
 }
 
 /// Checks whether `str` would shadow a generated variable or attribute
 /// part of the OpView API.
 static bool isODSReserved(StringRef str) {
   static llvm::StringSet<> reserved(
       {"attributes", "create", "context", "ip", "operands", "print", "get_asm",
        "loc", "verify", "regions", "results", "self", "operation",
        "DIALECT_NAMESPACE", "OPERATION_NAME"});
   return str.startswith("_ods_") || str.endswith("_ods") ||
          reserved.contains(str);
 }
 
 /// Modifies the `name` in a way that it becomes suitable for Python bindings
 /// (does not change the `name` if it already is suitable) and returns the
 /// modified version.
 static std::string sanitizeName(StringRef name) {
-  if (isPythonKeyword(name) || isODSReserved(name))
+  if (isPythonReserved(name) || isODSReserved(name))
     return (name + "_").str();
   return name.str();
 }
 
 static std::string attrSizedTraitForKind(const char *kind) {
   return llvm::formatv("::mlir::OpTrait::AttrSized{0}{1}Segments",
                        llvm::StringRef(kind).take_front().upper(),
                        llvm::StringRef(kind).drop_front());
constexpr const char *multiOperandAppendPackTemplate =
     "operands.append(_get_op_results_or_values({0}))";
 constexpr const char *multiResultAppendTemplate = "results.extend({0})";
 
 /// Template for setting an attribute in the operation builder.
 ///   {0} is the attribute name;
 ///   {1} is the builder argument name.
 constexpr const char *initAttributeTemplate = R"Py(attributes["{0}"] = {1})Py";
 
-/// Template for setting an optional attribute in the operation builder.
-///   {0} is the attribute name;
-///   {1} is the builder argument name.
-constexpr const char *initOptionalAttributeTemplate =
-    R"Py(if {1} is not None: attributes["{0}"] = {1})Py";
+/// Template for builder attribute from raw input in the operation builder.
+///   {0} is the builder argument name;
+///   {1} is the attribute builder from raw.
+/// Use the value the user passed in if either it is already an Attribute or
+/// there is not a method registered to make it an Attribute.
+// TODO: Should we just blindly try AttrBuilder post checking if Attribute?
+// Different error message if neither condition holds?
+constexpr const char *initAttributeWithBuilderTemplate =
+    R"Py({0} if (issubclass(type({0}), _ods_ir.Attribute) or not callable(getattr(_ods_ir.AttrBuilder, '{1}', None))) else _ods_ir.AttrBuilder.{1}({0}))Py";
 
 constexpr const char *initUnitAttributeTemplate =
     R"Py(if bool({1}): attributes["{0}"] = _ods_ir.UnitAttr.get(
       _ods_get_default_loc_context(loc)))Py";
 
 /// Template to initialize the successors list in the builder if there are any
 /// successors.
 ///   {0} is the value to initialize the successors list to.
for (int i = 0, e = op.getNumArgs(); i < e; ++i) {
 
     // Unit attributes are handled specially.
     if (attribute->attr.getStorageType().trim().equals("::mlir::UnitAttr")) {
       builderLines.push_back(llvm::formatv(initUnitAttributeTemplate,
                                            attribute->name, argNames[i]));
       continue;
     }
 
-    builderLines.push_back(llvm::formatv(
-        (attribute->attr.isOptional() || attribute->attr.hasDefaultValue())
-            ? initOptionalAttributeTemplate
-            : initAttributeTemplate,
-        attribute->name, argNames[i]));
+    std::string builderLine;
+    if (attribute->attr.isOptional() || attribute->attr.hasDefaultValue()) {
+      builderLine = llvm::formatv(R"Py(if {0} is not None: )Py", argNames[i]);
+    }
+    std::string val =
+        llvm::formatv(initAttributeWithBuilderTemplate, argNames[i],
+                      attribute->attr.getAttrDefName());
+    builderLine.append(
+        llvm::formatv(initAttributeTemplate, attribute->name, val));
+    builderLines.push_back(builderLine);
   }
 }
 
 /// Populates `builderLines` with additional lines that are required in the
 /// builder to set up successors. successorArgNames is expected to correspond
 /// to the Python argument name for each successor on the op.
 static void populateBuilderLinesSuccessors(
     const Operator &op, llvm::ArrayRef<std::string> successorArgNames,
     llvm::SmallVectorImpl<std::string> &builderLines) {