Pybind bazel experiments

mlir/include/mlir/IR/DialectRegistry.h

@@ -18,6 +18,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 
+#include <iostream>
 #include <map>
 #include <tuple>
 
@@ -242,6 +243,9 @@ class DialectRegistry {
       ~Extension() override = default;
 
       void apply(MLIRContext *context, DialectsT *...dialects) const final {
+        for (auto *dialect : {dialects...})
+          std::cerr << "Applying extension to dialect: "
+                    << dialect->getNamespace().str() << "\n";
         extensionFn(context, dialects...);
       }
       ExtensionFnT extensionFn;

mlir/include/mlir/IR/OpDefinition.h

@@ -24,6 +24,7 @@
 #include "mlir/IR/Operation.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
+#include <iostream>
 #include <optional>
 #include <type_traits>
 
@@ -1735,6 +1736,7 @@ class Op : public OpState, public Traits<ConcreteType>... {
       llvm::report_fatal_error(
           "Attempting to attach an interface to an unregistered operation " +
           ConcreteType::getOperationName() + ".");
+    std::cerr << "attachInterface: " << info->getStringRef().str() << "\n";
     (checkInterfaceTarget<Models>(), ...);
     info->attachInterface<Models...>();
   }

mlir/include/mlir/IR/StorageUniquerSupport.h

@@ -13,6 +13,8 @@
 #ifndef MLIR_IR_STORAGEUNIQUERSUPPORT_H
 #define MLIR_IR_STORAGEUNIQUERSUPPORT_H
 
+#include <iostream>
+
 #include "mlir/IR/AttrTypeSubElements.h"
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/Support/InterfaceSupport.h"
@@ -154,12 +156,14 @@ class StorageUserBase : public BaseT, public Traits<ConcreteT>... {
   /// call will abort otherwise.
   template <typename... IfaceModels>
   static void attachInterface(MLIRContext &context) {
+    std::cerr << "attachInterface\n";
     typename ConcreteT::AbstractTy *abstract =
         ConcreteT::AbstractTy::lookupMutable(TypeID::get<ConcreteT>(),
                                              &context);
-    if (!abstract)
-      llvm::report_fatal_error("Registering an interface for an attribute/type "
+    if (!abstract) {
+      llvm::report_fatal_error("HERE! Registering an interface for an attribute/type "
                                "that is not itself registered.");
+    }
 
     // Handle the case where the models resolve a promised interface.
     (dialect_extension_detail::handleAdditionOfUndefinedPromisedInterface(

mlir/include/mlir/InitAllDialects.h

@@ -14,6 +14,8 @@
 #ifndef MLIR_INITALLDIALECTS_H_
 #define MLIR_INITALLDIALECTS_H_
 
+#include <iostream>
+
 #include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
 #include "mlir/Dialect/AMX/AMXDialect.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
@@ -98,6 +100,7 @@ namespace mlir {
 
 /// Add all the MLIR dialects to the provided registry.
 inline void registerAllDialects(DialectRegistry &registry) {
+  std::cerr << "Starting to register dialects!\n";
   // clang-format off
   registry.insert<acc::OpenACCDialect,
                   affine::AffineDialect,
@@ -179,6 +182,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
   NVVM::registerNVVMTargetInterfaceExternalModels(registry);
   ROCDL::registerROCDLTargetInterfaceExternalModels(registry);
   spirv::registerSPIRVTargetInterfaceExternalModels(registry);
+  std::cerr << "All dialects registered!\n";
 }
 
 /// Append all the MLIR dialects to the registry contained in the given context.

mlir/lib/IR/MLIRContext.cpp

@@ -6,7 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/IR/MLIRContext.h"
+#include <iostream>
+
 #include "AffineExprDetail.h"
 #include "AffineMapDetail.h"
 #include "AttributeDetail.h"
@@ -23,6 +24,7 @@
 #include "mlir/IR/ExtensibleDialect.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/Types.h"
@@ -749,6 +751,11 @@ const AbstractAttribute &AbstractAttribute::lookup(TypeID typeID,
 AbstractAttribute *AbstractAttribute::lookupMutable(TypeID typeID,
                                                     MLIRContext *context) {
   auto &impl = context->getImpl();
+  llvm::errs() << "REGISTERED Attributes:\n";
+  for (auto &attr : impl.registeredAttributes) {
+    llvm::errs() << attr.second->getDialect().getNamespace() << "."
+                 << attr.second->getName() << "\n";
+  }
   return impl.registeredAttributes.lookup(typeID);
 }
 
@@ -973,6 +980,12 @@ const AbstractType &AbstractType::lookup(TypeID typeID, MLIRContext *context) {
 
 AbstractType *AbstractType::lookupMutable(TypeID typeID, MLIRContext *context) {
   auto &impl = context->getImpl();
+  // Print out all register dialect names
+  llvm::errs() << "REGISTERED Types:\n";
+  for (auto &ty : impl.registeredTypes) {
+    llvm::errs() << ty.second->getDialect().getNamespace() << "."
+                 << ty.second->getName() << "\n";
+  }
   return impl.registeredTypes.lookup(typeID);
 }
 

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

@@ -870,19 +870,6 @@ exports_files(
     glob(["lib/Bindings/Python/**/*.cpp"]),
 )
 
-# In the targets related to Python bindings, the projects @pybind11 and
-# @local_config_python are defined by @pybind11_bazel. The latter contains
-# python headers, and can be configured in an out-of-tree bazel project via
-#
-#   load("@pybind11_bazel//:python_configure.bzl", "python_configure")
-#   python_configure(name = "local_config_python")
-#
-# For more up-to-date instructions, see
-# https://github.com/pybind/pybind11_bazel
-#
-# Some out-of-tree projects alias @python_runtime//:headers to
-# @local_config_python//:python_headers.
-
 MLIR_BINDINGS_PYTHON_HEADERS = [
     "lib/Bindings/Python/*.h",
     "include/mlir-c/Bindings/Python/*.h",
@@ -1004,121 +991,6 @@ cc_library(
     ],
 )
 
-# Dynamic library with the MLIR Python extension.
-cc_binary(
-    name = "_mlir.so",
-    srcs = ["lib/Bindings/Python/MainModule.cpp"],
-    # These flags are needed for pybind11 to work.
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
-    linkshared = 1,
-    linkstatic = 0,
-    tags = [
-        "manual",  # External dependency
-        "nobuildkite",  # TODO(gcmn): Add support for this target
-    ],
-    deps = [
-        ":MLIRBindingsPythonCore",
-        ":MLIRBindingsPythonHeadersAndDeps",
-    ],
-)
-
-cc_binary(
-    name = "_mlirDialectsLinalg.so",
-    srcs = ["lib/Bindings/Python/DialectLinalg.cpp"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
-    linkshared = 1,
-    linkstatic = 0,
-    tags = [
-        "manual",  # External dependency
-        "nobuildkite",  # TODO(gcmn): Add support for this target
-    ],
-    deps = [
-        ":CAPIIR",
-        ":CAPILinalg",
-        ":MLIRBindingsPythonHeadersAndDeps",
-    ],
-)
-
-cc_binary(
-    name = "_mlirDialectsQuant.so",
-    srcs = ["lib/Bindings/Python/DialectQuant.cpp"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
-    linkshared = 1,
-    linkstatic = 0,
-    tags = [
-        "manual",  # External dependency
-        "nobuildkite",  # TODO(gcmn): Add support for this target
-    ],
-    deps = [
-        ":CAPIIR",
-        ":CAPIQuant",
-        ":MLIRBindingsPythonHeadersAndDeps",
-        "@pybind11",
-    ],
-)
-
-cc_binary(
-    name = "_mlirDialectsSparseTensor.so",
-    srcs = ["lib/Bindings/Python/DialectSparseTensor.cpp"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
-    linkshared = 1,
-    linkstatic = 0,
-    tags = [
-        "manual",  # External dependency
-        "nobuildkite",  # TODO(gcmn): Add support for this target
-    ],
-    deps = [
-        ":CAPIIR",
-        ":CAPISparseTensor",
-        ":MLIRBindingsPythonHeadersAndDeps",
-        "@pybind11",
-    ],
-)
-
-# Dynamic library with the MLIR Conversions Python extension.
-cc_binary(
-    name = "_mlirExecutionEngine.so",
-    srcs = ["lib/Bindings/Python/ExecutionEngineModule.cpp"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
-    linkshared = 1,
-    linkstatic = 0,
-    tags = [
-        "manual",  # External dependency
-        "nobuildkite",  # TODO(gcmn): Add support for this target
-    ],
-    deps = [
-        ":CAPIExecutionEngine",
-        ":MLIRBindingsPythonHeadersAndDeps",
-        "@local_config_python//:python_headers",
-        "@pybind11",
-    ],
-)
-
-# Dynamic library with the MLIR Linalg dialect+passes Python extension.
-cc_binary(
-    name = "_mlirLinalgPasses.so",
-    srcs = ["lib/Bindings/Python/LinalgPasses.cpp"],
-    copts = PYBIND11_COPTS,
-    features = PYBIND11_FEATURES,
-    linkshared = 1,
-    linkstatic = 0,
-    tags = [
-        "manual",  # External dependency
-        "nobuildkite",  # TODO(gcmn): Add support for this target
-    ],
-    deps = [
-        ":CAPILinalg",
-        ":MLIRBindingsPythonHeadersAndDeps",
-        "@local_config_python//:python_headers",
-        "@pybind11",
-    ],
-)
-
 ##---------------------------------------------------------------------------##
 
 td_library(