Revert "[mlir][python] Allow specifying block arg locations"

This reverts commit 4d0d295.

This caused a buildbot failure: https://lab.llvm.org/buildbot/#/builders/61/builds/43479

Author: rkayaith

mlir/lib/Bindings/Python/IRCore.cpp

@@ -193,31 +193,6 @@ 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> argTypes;
-  argTypes.reserve(pyArgTypes.size());
-  for (const auto &pyType : pyArgTypes)
-    argTypes.push_back(pyType.cast<PyType &>());
-
-  SmallVector<MlirLocation> 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); }
@@ -389,10 +364,21 @@ class PyBlockList {
     throw SetPyError(PyExc_IndexError, "attempt to access out of bounds block");
   }
 
-  PyBlock appendBlock(const py::args &pyArgTypes,
-                      const std::optional<py::sequence> &pyArgLocs) {
+  PyBlock appendBlock(const py::args &pyArgTypes) {
     operation->checkValid();
-    MlirBlock block = createBlock(pyArgTypes, pyArgLocs);
+    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());
     mlirRegionAppendOwnedBlock(region, block);
     return PyBlock(operation, block);
   }
@@ -402,8 +388,7 @@ class PyBlockList {
         .def("__getitem__", &PyBlockList::dunderGetItem)
         .def("__iter__", &PyBlockList::dunderIter)
         .def("__len__", &PyBlockList::dunderLen)
-        .def("append", &PyBlockList::appendBlock, kAppendBlockDocstring,
-             py::arg("arg_locs") = std::nullopt);
+        .def("append", &PyBlockList::appendBlock, kAppendBlockDocstring);
   }
 
 private:
@@ -2981,17 +2966,27 @@ void mlir::python::populateIRCore(py::module &m) {
           "Returns a forward-optimized sequence of operations.")
       .def_static(
           "create_at_start",
-          [](PyRegion &parent, const py::list &pyArgTypes,
-             const std::optional<py::sequence> &pyArgLocs) {
+          [](PyRegion &parent, py::list pyArgTypes) {
             parent.checkValid();
-            MlirBlock block = createBlock(pyArgTypes, pyArgLocs);
+            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());
             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 and locations).")
+          "region (with given argument types).")
       .def(
           "append_to",
           [](PyBlock &self, PyRegion &region) {
@@ -3003,30 +2998,50 @@ void mlir::python::populateIRCore(py::module &m) {
           "Append this block to a region, transferring ownership if necessary")
       .def(
           "create_before",
-          [](PyBlock &self, const py::args &pyArgTypes,
-             const std::optional<py::sequence> &pyArgLocs) {
+          [](PyBlock &self, py::args pyArgTypes) {
             self.checkValid();
-            MlirBlock block = createBlock(pyArgTypes, pyArgLocs);
+            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());
             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 and locations).")
+          "(with given argument types).")
       .def(
           "create_after",
-          [](PyBlock &self, const py::args &pyArgTypes,
-             const std::optional<py::sequence> &pyArgLocs) {
+          [](PyBlock &self, py::args pyArgTypes) {
             self.checkValid();
-            MlirBlock block = createBlock(pyArgTypes, pyArgLocs);
+            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());
             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 and locations).")
+          "(with given argument types).")
       .def(
           "__iter__",
           [](PyBlock &self) {

mlir/python/mlir/dialects/_func_ops_ext.py

@@ -90,15 +90,15 @@ def entry_block(self):
       raise IndexError('External function does not have a body')
     return self.regions[0].blocks[0]
 
-  def add_entry_block(self, arg_locs: Optional[Sequence[Location]] = None):
+  def add_entry_block(self):
     """
     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, arg_locs=arg_locs)
+    self.body.blocks.append(*self.type.inputs)
     return self.body.blocks[0]
 
   @property

mlir/test/python/ir/blocks.py

@@ -18,28 +18,28 @@ def run(f):
 
 
 # CHECK-LABEL: TEST: testBlockCreation
-# CHECK: func @test(%[[ARG0:.*]]: i32 loc("arg0"), %[[ARG1:.*]]: i16 loc("arg1"))
+# CHECK: func @test(%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i16)
 # CHECK:   cf.br ^bb1(%[[ARG1]] : i16)
-# CHECK: ^bb1(%[[PHI0:.*]]: i16 loc("middle")):
+# CHECK: ^bb1(%[[PHI0:.*]]: i16):
 # CHECK:   cf.br ^bb2(%[[ARG0]] : i32)
-# CHECK: ^bb2(%[[PHI1:.*]]: i32 loc("successor")):
+# CHECK: ^bb2(%[[PHI1:.*]]: i32):
 # CHECK:   return
 @run
 def testBlockCreation():
   with Context() as ctx, Location.unknown():
-    module = builtin.ModuleOp()
+    module = Module.create()
     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([Location.name("arg0"), Location.name("arg1")])
+      entry_block = f_op.add_entry_block()
       i32_arg, i16_arg = entry_block.arguments
-      successor_block = entry_block.create_after(i32_arg.type, arg_locs=[Location.name("successor")])
+      successor_block = entry_block.create_after(i32_arg.type)
       with InsertionPoint(successor_block) as successor_ip:
         assert successor_ip.block == successor_block
         func.ReturnOp([])
-      middle_block = successor_block.create_before(i16_arg.type, arg_locs=[Location.name("middle")])
+      middle_block = successor_block.create_before(i16_arg.type)
 
       with InsertionPoint(entry_block) as entry_ip:
         assert entry_ip.block == entry_block
@@ -48,57 +48,27 @@ def testBlockCreation():
       with InsertionPoint(middle_block) as middle_ip:
         assert middle_ip.block == middle_block
         cf.BranchOp([i32_arg], dest=successor_block)
-    module.print(enable_debug_info=True)
+    print(module.operation)
     # 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 loc("arg_loc"))
+# CHECK: func @test(%{{.*}}: f32)
 # CHECK:   return
 @run
 def testFirstBlockCreation():
   with Context() as ctx, Location.unknown():
-    module = builtin.ModuleOp()
+    module = Module.create()
     f32 = F32Type.get()
     with InsertionPoint(module.body):
       f = func.FuncOp("test", ([f32], []))
-      entry_block = Block.create_at_start(f.operation.regions[0],
-                                          [f32], [Location.name("arg_loc")])
+      entry_block = Block.create_at_start(f.operation.regions[0], [f32])
       with InsertionPoint(entry_block):
         func.ReturnOp([])
 
-    module.print(enable_debug_info=True)
-    assert module.verify()
+    print(module)
+    assert module.operation.verify()
     assert f.body.blocks[0] == entry_block