[mlir][python] Allow specifying block arg locations

[rkayaith]

Summary

Currently blocks are always created with UnknownLoc's for their arguments. This
adds an arg_locs argument to all block creation APIs, which takes an optional
sequence of locations to use, one per block argument. If no locations are
supplied, the current Location context is used.

[rkayaith]

update

[ftynse]

mlir/lib/Bindings/Python/IRCore.cpp

196

Nit: let's not specify the number of stack elements unless there is a good reason to.

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

3002–3004

Support for arguments after py::args like this was only added in pybind 2.9.0 (see changelog, https://github.com/pybind/pybind11/pull/3402), but the minimum version we set is 2.8.0. I can probably work around this, but any objections to just raising the min version to 2.9.0? It was released on December 28, 2021 @ftynse @stellaraccident

The mlir-nvidia buildbot is using 2.8, I'm not sure what the process for updating it would be.

[stellaraccident]

mlir/lib/Bindings/Python/IRCore.cpp

3002–3004

I'm fine with upgrading. How to upgrade though has rolled off of my working memory.

[rkayaith]

reopening so I remember to land this again once the bot is updated

mlir/lib/Bindings/Python/IRCore.cpp

createCustomDialectWrapper(const std::string &dialectNamespace,
   // Create the custom implementation.
   return (*dialectClass)(std::move(dialectDescriptor));
 }
 
 static MlirStringRef toMlirStringRef(const std::string &s) {
   return mlirStringRefCreate(s.data(), s.size());
 }
 
+/// Create a block, using the current location context if no locations are
+/// specified.
+static MlirBlock createBlock(const py::sequence &pyArgTypes,
+                             const std::optional<py::sequence> &pyArgLocs) {
+  SmallVector<MlirType, 4> argTypes;
+  argTypes.reserve(pyArgTypes.size());
+  for (const auto &pyType : pyArgTypes)
+    argTypes.push_back(pyType.cast<PyType &>());
+
+  SmallVector<MlirLocation, 4> argLocs;
+  if (pyArgLocs) {
+    argLocs.reserve(pyArgLocs->size());
+    for (const auto &pyLoc : *pyArgLocs)
+      argLocs.push_back(pyLoc.cast<PyLocation &>());
+  } else if (!argTypes.empty()) {
+    argLocs.assign(argTypes.size(), DefaultingPyLocation::resolve());
+  }
+
+  if (argTypes.size() != argLocs.size())
+    throw py::value_error(("Expected " + Twine(argTypes.size()) +
+                           " locations, got: " + Twine(argLocs.size()))
+                              .str());
+  return mlirBlockCreate(argTypes.size(), argTypes.data(), argLocs.data());
+}
+
 /// Wrapper for the global LLVM debugging flag.
 struct PyGlobalDebugFlag {
   static void set(py::object &o, bool enable) { mlirEnableGlobalDebug(enable); }
 
   static bool get(const py::object &) { return mlirIsGlobalDebugEnabled(); }
 
   static void bind(py::module &m) {
     // Debug flags.
while (!mlirBlockIsNull(block)) {
         return PyBlock(operation, block);
       }
       block = mlirBlockGetNextInRegion(block);
       index -= 1;
     }
     throw SetPyError(PyExc_IndexError, "attempt to access out of bounds block");
   }
 
-  PyBlock appendBlock(const py::args &pyArgTypes) {
+  PyBlock appendBlock(const py::args &pyArgTypes,
+                      const std::optional<py::sequence> &pyArgLocs) {
     operation->checkValid();
-    llvm::SmallVector<MlirType, 4> argTypes;
-    llvm::SmallVector<MlirLocation, 4> argLocs;
-    argTypes.reserve(pyArgTypes.size());
-    argLocs.reserve(pyArgTypes.size());
-    for (auto &pyArg : pyArgTypes) {
-      argTypes.push_back(pyArg.cast<PyType &>());
-      // TODO: Pass in a proper location here.
-      argLocs.push_back(
-          mlirLocationUnknownGet(mlirTypeGetContext(argTypes.back())));
-    }
-
-    MlirBlock block =
-        mlirBlockCreate(argTypes.size(), argTypes.data(), argLocs.data());
+    MlirBlock block = createBlock(pyArgTypes, pyArgLocs);
     mlirRegionAppendOwnedBlock(region, block);
     return PyBlock(operation, block);
   }
 
   static void bind(py::module &m) {
     py::class_<PyBlockList>(m, "BlockList", py::module_local())
         .def("__getitem__", &PyBlockList::dunderGetItem)
         .def("__iter__", &PyBlockList::dunderIter)
         .def("__len__", &PyBlockList::dunderLen)
-        .def("append", &PyBlockList::appendBlock, kAppendBlockDocstring);
+        .def("append", &PyBlockList::appendBlock, kAppendBlockDocstring,
+             py::arg("arg_locs") = std::nullopt);
   }
 
 private:
   PyOperationRef operation;
   MlirRegion region;
 };
 
 class PyOperationIterator {
py::class_<PyBlock>(m, "Block", py::module_local())
       .def_property_readonly(
           "operations",
           [](PyBlock &self) {
             return PyOperationList(self.getParentOperation(), self.get());
           },
           "Returns a forward-optimized sequence of operations.")
       .def_static(
           "create_at_start",
-          [](PyRegion &parent, py::list pyArgTypes) {
+          [](PyRegion &parent, const py::list &pyArgTypes,
+             const std::optional<py::sequence> &pyArgLocs) {
             parent.checkValid();
-            llvm::SmallVector<MlirType, 4> argTypes;
-            llvm::SmallVector<MlirLocation, 4> argLocs;
-            argTypes.reserve(pyArgTypes.size());
-            argLocs.reserve(pyArgTypes.size());
-            for (auto &pyArg : pyArgTypes) {
-              argTypes.push_back(pyArg.cast<PyType &>());
-              // TODO: Pass in a proper location here.
-              argLocs.push_back(
-                  mlirLocationUnknownGet(mlirTypeGetContext(argTypes.back())));
-            }
-
-            MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data(),
-                                              argLocs.data());
+            MlirBlock block = createBlock(pyArgTypes, pyArgLocs);
             mlirRegionInsertOwnedBlock(parent, 0, block);
             return PyBlock(parent.getParentOperation(), block);
           },
           py::arg("parent"), py::arg("arg_types") = py::list(),
+          py::arg("arg_locs") = std::nullopt,
           "Creates and returns a new Block at the beginning of the given "
-          "region (with given argument types).")
+          "region (with given argument types and locations).")
       .def(
           "append_to",
           [](PyBlock &self, PyRegion &region) {
             MlirBlock b = self.get();
             if (!mlirRegionIsNull(mlirBlockGetParentRegion(b)))
               mlirBlockDetach(b);
             mlirRegionAppendOwnedBlock(region.get(), b);
           },
           "Append this block to a region, transferring ownership if necessary")
       .def(
           "create_before",
-          [](PyBlock &self, py::args pyArgTypes) {
+          [](PyBlock &self, const py::args &pyArgTypes,
+             const std::optional<py::sequence> &pyArgLocs) {
             self.checkValid();
-            llvm::SmallVector<MlirType, 4> argTypes;
-            llvm::SmallVector<MlirLocation, 4> argLocs;
-            argTypes.reserve(pyArgTypes.size());
-            argLocs.reserve(pyArgTypes.size());
-            for (auto &pyArg : pyArgTypes) {
-              argTypes.push_back(pyArg.cast<PyType &>());
-              // TODO: Pass in a proper location here.
-              argLocs.push_back(
-                  mlirLocationUnknownGet(mlirTypeGetContext(argTypes.back())));
-            }
-
-            MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data(),
-                                              argLocs.data());
+            MlirBlock block = createBlock(pyArgTypes, pyArgLocs);
             MlirRegion region = mlirBlockGetParentRegion(self.get());
             mlirRegionInsertOwnedBlockBefore(region, self.get(), block);
             return PyBlock(self.getParentOperation(), block);
           },
+          py::arg("arg_locs") = std::nullopt,
           "Creates and returns a new Block before this block "
-          "(with given argument types).")
+          "(with given argument types and locations).")
       .def(
           "create_after",
-          [](PyBlock &self, py::args pyArgTypes) {
+          [](PyBlock &self, const py::args &pyArgTypes,
+             const std::optional<py::sequence> &pyArgLocs) {
             self.checkValid();
-            llvm::SmallVector<MlirType, 4> argTypes;
-            llvm::SmallVector<MlirLocation, 4> argLocs;
-            argTypes.reserve(pyArgTypes.size());
-            argLocs.reserve(pyArgTypes.size());
-            for (auto &pyArg : pyArgTypes) {
-              argTypes.push_back(pyArg.cast<PyType &>());
-
-              // TODO: Pass in a proper location here.
-              argLocs.push_back(
-                  mlirLocationUnknownGet(mlirTypeGetContext(argTypes.back())));
-            }
-            MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data(),
-                                              argLocs.data());
+            MlirBlock block = createBlock(pyArgTypes, pyArgLocs);
             MlirRegion region = mlirBlockGetParentRegion(self.get());
             mlirRegionInsertOwnedBlockAfter(region, self.get(), block);
             return PyBlock(self.getParentOperation(), block);
           },
+          py::arg("arg_locs") = std::nullopt,
           "Creates and returns a new Block after this block "
-          "(with given argument types).")
+          "(with given argument types and locations).")
       .def(
           "__iter__",
           [](PyBlock &self) {
             self.checkValid();
             MlirOperation firstOperation =
                 mlirBlockGetFirstOperation(self.get());
             return PyOperationIterator(self.getParentOperation(),
                                        firstOperation);

mlir/python/mlir/dialects/_func_ops_ext.py

def name(self) -> StringAttr:
     return StringAttr(self.attributes["sym_name"])
 
   @property
   def entry_block(self):
     if self.is_external:
       raise IndexError('External function does not have a body')
     return self.regions[0].blocks[0]
 
-  def add_entry_block(self):
+  def add_entry_block(self, arg_locs: Optional[Sequence[Location]] = None):
     """
     Add an entry block to the function body using the function signature to
     infer block arguments.
     Returns the newly created block
     """
     if not self.is_external:
       raise IndexError('The function already has an entry block!')
-    self.body.blocks.append(*self.type.inputs)
+    self.body.blocks.append(*self.type.inputs, arg_locs=arg_locs)
     return self.body.blocks[0]
 
   @property
   def arg_attrs(self):
     return ArrayAttr(self.attributes[ARGUMENT_ATTRIBUTE_NAME])
 
   @arg_attrs.setter
   def arg_attrs(self, attribute: Union[ArrayAttr, list]):

mlir/test/python/ir/blocks.py

def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
   return f
 
 
 # CHECK-LABEL: TEST: testBlockCreation
-# CHECK: func @test(%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i16)
+# CHECK: func @test(%[[ARG0:.*]]: i32 loc("arg0"), %[[ARG1:.*]]: i16 loc("arg1"))
 # CHECK:   cf.br ^bb1(%[[ARG1]] : i16)
-# CHECK: ^bb1(%[[PHI0:.*]]: i16):
+# CHECK: ^bb1(%[[PHI0:.*]]: i16 loc("middle")):
 # CHECK:   cf.br ^bb2(%[[ARG0]] : i32)
-# CHECK: ^bb2(%[[PHI1:.*]]: i32):
+# CHECK: ^bb2(%[[PHI1:.*]]: i32 loc("successor")):
 # CHECK:   return
 @run
 def testBlockCreation():
   with Context() as ctx, Location.unknown():
-    module = Module.create()
+    module = builtin.ModuleOp()
     with InsertionPoint(module.body):
       f_type = FunctionType.get(
           [IntegerType.get_signless(32),
            IntegerType.get_signless(16)], [])
       f_op = func.FuncOp("test", f_type)
-      entry_block = f_op.add_entry_block()
+      entry_block = f_op.add_entry_block([Location.name("arg0"), Location.name("arg1")])
       i32_arg, i16_arg = entry_block.arguments
-      successor_block = entry_block.create_after(i32_arg.type)
+      successor_block = entry_block.create_after(i32_arg.type, arg_locs=[Location.name("successor")])
       with InsertionPoint(successor_block) as successor_ip:
         assert successor_ip.block == successor_block
         func.ReturnOp([])
-      middle_block = successor_block.create_before(i16_arg.type)
+      middle_block = successor_block.create_before(i16_arg.type, arg_locs=[Location.name("middle")])
 
       with InsertionPoint(entry_block) as entry_ip:
         assert entry_ip.block == entry_block
         cf.BranchOp([i16_arg], dest=middle_block)
 
       with InsertionPoint(middle_block) as middle_ip:
         assert middle_ip.block == middle_block
         cf.BranchOp([i32_arg], dest=successor_block)
-    print(module.operation)
+    module.print(enable_debug_info=True)
     # Ensure region back references are coherent.
     assert entry_block.region == middle_block.region == successor_block.region
 
 
+# CHECK-LABEL: TEST: testBlockCreationArgLocs
+@run
+def testBlockCreationArgLocs():
+  with Context() as ctx:
+    ctx.allow_unregistered_dialects = True
+    f32 = F32Type.get()
+    op = Operation.create("test", regions=1, loc=Location.unknown())
+    blocks = op.regions[0].blocks
+
+    with Location.name("default_loc"):
+      blocks.append(f32)
+    blocks.append()
+    # CHECK:      ^bb0(%{{.+}}: f32 loc("default_loc")):
+    # CHECK-NEXT: ^bb1:
+    op.print(enable_debug_info=True)
+
+    try:
+      blocks.append(f32)
+    except RuntimeError as err:
+      # CHECK: Missing loc: An MLIR function requires a Location but none was provided
+      print("Missing loc:", err)
+
+    try:
+      blocks.append(f32, f32, arg_locs=[Location.unknown()])
+    except ValueError as err:
+      # CHECK: Wrong loc count: Expected 2 locations, got: 1
+      print("Wrong loc count:", err)
+
+
 # CHECK-LABEL: TEST: testFirstBlockCreation
-# CHECK: func @test(%{{.*}}: f32)
+# CHECK: func @test(%{{.*}}: f32 loc("arg_loc"))
 # CHECK:   return
 @run
 def testFirstBlockCreation():
   with Context() as ctx, Location.unknown():
-    module = Module.create()
+    module = builtin.ModuleOp()
     f32 = F32Type.get()
     with InsertionPoint(module.body):
       f = func.FuncOp("test", ([f32], []))
-      entry_block = Block.create_at_start(f.operation.regions[0], [f32])
+      entry_block = Block.create_at_start(f.operation.regions[0],
+                                          [f32], [Location.name("arg_loc")])
       with InsertionPoint(entry_block):
         func.ReturnOp([])
 
-    print(module)
-    assert module.operation.verify()
+    module.print(enable_debug_info=True)
+    assert module.verify()
     assert f.body.blocks[0] == entry_block
 
 
 # CHECK-LABEL: TEST: testBlockMove
 # CHECK:  %0 = "realop"() ({
 # CHECK:  ^bb0([[ARG0:%.+]]: f32):
 # CHECK:    "ret"([[ARG0]]) : (f32) -> ()
 # CHECK:  }) : () -> f32