Merge npcomp and mlir python namespaces.

* Now the parts of the MLIR API are directly exported under the npcomp module (i.e. `npcomp.ir`, etc).
* Has required fixes for https://reviews.llvm.org/D108489
* Deletes npcomp.tracing vs fixing it because it was a very early experiment that will not be carried forward.
* This makes the npcomp python distribution completely standalone and separate from an mlir installation.
* Makes most of npcomp itself relocatable for future use as a library.
* Most things are a namespace package now. In the future we can s/torch_mlir/npcomp.frontends.torch/ and have it layer properly.

frontends/pytorch/csrc/CMakeLists.txt

@@ -1,5 +1,9 @@
 include(NpcompPython)
 
+# TODO: Add this to an npcomp header of some kind so it doesn't need to be
+# passed loose.
+add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=npcomp.")
+
 # Sharp edge: Torch extensions need to use the same pybind11 that torch
 # was compiled with, or else there will be issues in cross module exception
 # handling (which will abort instead of raise). We circumvent the possibility

frontends/pytorch/csrc/builder/module_builder.cpp

@@ -21,7 +21,7 @@ namespace py = pybind11;
 using namespace torch_mlir;
 
 static py::object getMlirIrClass(const char *className) {
-  return py::module::import("mlir.ir").attr(className);
+  return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir")).attr(className);
 }
 
 static py::object createPythonContextIfNone(py::object contextObj) {

frontends/pytorch/python/torch_mlir/__init__.py

@@ -7,7 +7,7 @@
 # prior to loading shared objects.
 
 import torch
-import npcomp
+import npcomp._mlir_libs._npcomp
 
 # Our native extension is not self-contained. It references libraries which
 # must come in via the above first.

frontends/pytorch/python/torch_mlir_torchscript_e2e_test_configs/npcomp_backend.py

@@ -10,9 +10,9 @@ import tempfile
 
 import numpy as np
 import torch
-from mlir.passmanager import PassManager
 
 import torch_mlir
+from npcomp.passmanager import PassManager
 from npcomp.compiler.pytorch.backend import refjit
 from npcomp.compiler.pytorch.backend.abc import NpcompBackend
 from torch_mlir_torchscript.e2e_test.framework import TestConfig, Trace, TraceItem

frontends/pytorch/setup.py

@@ -22,6 +22,9 @@ extension_sources = [str(p) for p in this_dir.joinpath("csrc").rglob("*.cpp")]
 include_dirs = npcomp_build.get_include_dirs()
 lib_dirs = npcomp_build.get_lib_dirs()
 npcomp_libs = [npcomp_build.get_capi_link_library_name()]
+# TODO: Export this in some way from an npcomp config file include vs needing
+# it loose here.
+compile_args = ["-DMLIR_PYTHON_PACKAGE_PREFIX=npcomp."]
 
 setup(
   name="npcomp-torch",
@@ -31,7 +34,8 @@ setup(
       sources=extension_sources,
       include_dirs=include_dirs,
       library_dirs=lib_dirs,
-      libraries=npcomp_libs),
+      libraries=npcomp_libs,
+      extra_compile_args=compile_args),
   ],
   cmdclass={
     "build_ext": cpp_extension.BuildExtension

frontends/pytorch/test/extension_coexistence.py

@@ -10,8 +10,7 @@ import sys
 
 print(f"PYTHONPATH={sys.path}")
 
-import mlir
-import npcomp
+import npcomp.ir
 import torch_mlir
 
 print("Extensions all loaded")

python/CMakeLists.txt

@@ -1,6 +1,10 @@
 include(AddMLIRPython)
 include(MLIRDetectPythonEnv)
 
+# Specifies that all MLIR packages are co-located under npcomp.
+# TODO: Add an upstream cmake param for this vs having a global here.
+add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=npcomp.")
+
 ################################################################################
 # Resources that must be packaged into the python tree
 ################################################################################
@@ -28,21 +32,18 @@ declare_mlir_python_sources(NPCOMPPythonSources)
 
 declare_mlir_python_sources(NPCOMPPythonSources.Core
   ADD_TO_PARENT NPCOMPPythonSources
-  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}"
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/npcomp"
   SOURCES
-    npcomp/__init__.py
-    npcomp/build.py
-    npcomp/decorators.py
-    npcomp/exporter.py
-    npcomp/smoketest.py
-    npcomp/types.py
-    npcomp/dialects/_ods_common.py
+    build.py
+    decorators.py
+    exporter.py
+    smoketest.py
+    types.py
   SOURCES_GLOB
-    npcomp/compiler/*.py
-    npcomp/frontends/*.py
-    npcomp/torch/*.py
-    npcomp/tracing/*.py
-    npcomp/utils/*.py
+    compiler/*.py
+    frontends/*.py
+    torch/*.py
+    utils/*.py
 )
 
 declare_mlir_python_sources(NPCOMPPythonSources.Dialects
@@ -83,23 +84,23 @@ declare_mlir_python_extension(NPCOMPPythonExtensions.Core
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT NPCOMPPythonSources.Dialects
-  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}"
-  TD_FILE npcomp/dialects/BasicpyBind.td
-  SOURCES npcomp/dialects/basicpy.py
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/npcomp"
+  TD_FILE dialects/BasicpyBind.td
+  SOURCES dialects/basicpy.py
   DIALECT_NAME basicpy)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT NPCOMPPythonSources.Dialects
-  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}"
-  TD_FILE npcomp/dialects/NumpyBind.td
-  SOURCES npcomp/dialects/numpy.py
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/npcomp"
+  TD_FILE dialects/NumpyBind.td
+  SOURCES dialects/numpy.py
   DIALECT_NAME numpy)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT NPCOMPPythonSources.Dialects
-  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}"
-  TD_FILE npcomp/dialects/TorchBind.td
-  SOURCES npcomp/dialects/torch.py
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/npcomp"
+  TD_FILE dialects/TorchBind.td
+  SOURCES dialects/torch.py
   DIALECT_NAME torch)
 
 ################################################################################
@@ -109,44 +110,30 @@ declare_mlir_dialect_python_bindings(
 # Bundle our own, self-contained CAPI library with all of our deps.
 add_mlir_python_common_capi_library(NPCOMPPythonCAPI
   INSTALL_COMPONENT NPCOMPPythonModules
-  INSTALL_DESTINATION python_packages/npcomp_core/mlir/_mlir_libs
-  # NOTE: When the MLIR API is relocated under npcomp, this would change to
-  # .../npcomp/_mlir_libs
-  OUTPUT_DIRECTORY "${MLIR_NPCOMP_PYTHON_PACKAGES_DIR}/npcomp_core/mlir/_mlir_libs"
+  INSTALL_DESTINATION python_packages/npcomp_core/npcomp/_mlir_libs
+  OUTPUT_DIRECTORY "${MLIR_NPCOMP_PYTHON_PACKAGES_DIR}/npcomp_core/npcomp/_mlir_libs"
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_SOURCES
-    # TODO: This can be chopped down significantly for size.
-    MLIRPythonSources
+    # TODO: Common MLIR deps can be reduced substantially.
+    MLIRPythonSources.Core
+    MLIRPythonSources.Dialects
+    MLIRPythonSources.ExecutionEngine
     MLIRPythonExtension.AllPassesRegistration
     NPCOMPPythonSources
     NPCOMPPythonExtensions
 )
 
-# Bundle the MLIR python sources into our package.
-# The MLIR API is position independent, so we explicitly output it to the mlir/
-# folder as a temporary measure. It will eventually migrate under the npcomp/
-# folder and be accessible under the unified "import npcomp..." namespace.
-add_mlir_python_modules(NPCOMPMLIRPythonModules
-  ROOT_PREFIX "${MLIR_NPCOMP_PYTHON_PACKAGES_DIR}/npcomp_core/mlir"
-  INSTALL_PREFIX "python_packages/npcomp_core/mlir"
+# Bundle MLIR and NPCOMP into the top-level npcomp package.
+add_mlir_python_modules(NPCOMPPythonModules
+  ROOT_PREFIX "${MLIR_NPCOMP_PYTHON_PACKAGES_DIR}/npcomp_core/npcomp"
+  INSTALL_PREFIX "python_packages/npcomp_core/npcomp"
   DECLARED_SOURCES
     MLIRPythonSources
     MLIRPythonExtension.AllPassesRegistration
-    # We need the npcomp extensions co-located with the MLIR extensions. When
-    # the namespace is unified, this moves to the below.
     MLIRPythonCAPIHeaderSources
+    NPCOMPPythonSources
     NPCOMPPythonExtensions
     NPCOMPPythonCAPIHeaderSources
   COMMON_CAPI_LINK_LIBS
     NPCOMPPythonCAPI
 )
-
-# Bundle the NPCOMP python sources into our package.
-add_mlir_python_modules(NPCOMPPythonModules
-  ROOT_PREFIX "${MLIR_NPCOMP_PYTHON_PACKAGES_DIR}/npcomp_core"
-  INSTALL_PREFIX "python_packages/npcomp_core"
-  DECLARED_SOURCES
-    NPCOMPPythonSources
-  COMMON_CAPI_LINK_LIBS
-    NPCOMPPythonCAPI
-)

python/NpcompModule.cpp

@@ -52,10 +52,10 @@ void emitError(MlirLocation loc, std::string message) {
 
 PYBIND11_MODULE(_npcomp, m) {
   m.doc() = "Npcomp native python bindings";
+  ::npcompRegisterAllPasses();
+  ::npcompInitializeLLVMCodegen();
 
   m.def("register_all_dialects", ::npcompRegisterAllDialects);
-  m.def("_register_all_passes", ::npcompRegisterAllPasses);
-  m.def("_initialize_llvm_codegen", ::npcompInitializeLLVMCodegen);
   m.def("shaped_to_ndarray_type", shapedToNdArrayArrayType);
   m.def("ndarray_to_tensor_type", ndarrayToTensorType);
   m.def("slot_object_type", slotObjectType);

python/NpcompPybindUtils.h

@@ -20,153 +20,11 @@
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Pass.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
 #include "llvm/ADT/Optional.h"
 
 namespace py = pybind11;
 
-namespace pybind11 {
-namespace detail {
-
-template <typename T>
-struct type_caster<llvm::Optional<T>> : optional_caster<llvm::Optional<T>> {};
-
-/// Helper to convert a presumed MLIR API object to a capsule, accepting either
-/// an explicit Capsule (which can happen when two C APIs are communicating
-/// directly via Python) or indirectly by querying the MLIR_PYTHON_CAPI_PTR_ATTR
-/// attribute (through which supported MLIR Python API objects export their
-/// contained API pointer as a capsule). This is intended to be used from
-/// type casters, which are invoked with a raw handle (unowned). The returned
-/// object's lifetime may not extend beyond the apiObject handle without
-/// explicitly having its refcount increased (i.e. on return).
-static py::object mlirApiObjectToCapsule(py::handle apiObject) {
-  if (PyCapsule_CheckExact(apiObject.ptr()))
-    return py::reinterpret_borrow<py::object>(apiObject);
-  return apiObject.attr(MLIR_PYTHON_CAPI_PTR_ATTR);
-}
-
-// Note: Currently all of the following support cast from py::object to the
-// Mlir* C-API type, but only a few light-weight, context-bound ones
-// implicitly cast the other way because the use case has not yet emerged and
-// ownership is unclear.
-
-/// Casts object -> MlirAttribute.
-template <> struct type_caster<MlirAttribute> {
-  PYBIND11_TYPE_CASTER(MlirAttribute, _("MlirAttribute"));
-  bool load(handle src, bool) {
-    auto capsule = mlirApiObjectToCapsule(src);
-    value = mlirPythonCapsuleToAttribute(capsule.ptr());
-    if (mlirAttributeIsNull(value)) {
-      return false;
-    }
-    return true;
-  }
-  static handle cast(MlirAttribute v, return_value_policy, handle) {
-    auto capsule =
-        py::reinterpret_steal<py::object>(mlirPythonAttributeToCapsule(v));
-    return py::module::import("mlir.ir")
-        .attr("Attribute")
-        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
-        .release();
-  }
-};
-
-/// Casts object -> MlirContext.
-template <> struct type_caster<MlirContext> {
-  PYBIND11_TYPE_CASTER(MlirContext, _("MlirContext"));
-  bool load(handle src, bool) {
-    auto capsule = mlirApiObjectToCapsule(src);
-    value = mlirPythonCapsuleToContext(capsule.ptr());
-    if (mlirContextIsNull(value)) {
-      return false;
-    }
-    return true;
-  }
-};
-
-/// Casts object -> MlirLocation.
-template <> struct type_caster<MlirLocation> {
-  PYBIND11_TYPE_CASTER(MlirLocation, _("MlirLocation"));
-  bool load(handle src, bool) {
-    auto capsule = mlirApiObjectToCapsule(src);
-    value = mlirPythonCapsuleToLocation(capsule.ptr());
-    if (mlirLocationIsNull(value)) {
-      return false;
-    }
-    return true;
-  }
-  static handle cast(MlirLocation v, return_value_policy, handle) {
-    auto capsule =
-        py::reinterpret_steal<py::object>(mlirPythonLocationToCapsule(v));
-    return py::module::import("mlir.ir")
-        .attr("Location")
-        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
-        .release();
-  }
-};
-
-/// Casts object -> MlirModule.
-template <> struct type_caster<MlirModule> {
-  PYBIND11_TYPE_CASTER(MlirModule, _("MlirModule"));
-  bool load(handle src, bool) {
-    auto capsule = mlirApiObjectToCapsule(src);
-    value = mlirPythonCapsuleToModule(capsule.ptr());
-    if (mlirModuleIsNull(value)) {
-      return false;
-    }
-    return true;
-  }
-};
-
-/// Casts object -> MlirOperation.
-template <> struct type_caster<MlirOperation> {
-  PYBIND11_TYPE_CASTER(MlirOperation, _("MlirOperation"));
-  bool load(handle src, bool) {
-    auto capsule = mlirApiObjectToCapsule(src);
-    value = mlirPythonCapsuleToOperation(capsule.ptr());
-    if (mlirOperationIsNull(value)) {
-      return false;
-    }
-    return true;
-  }
-};
-
-/// Casts object -> MlirPassManager.
-template <> struct type_caster<MlirPassManager> {
-  PYBIND11_TYPE_CASTER(MlirPassManager, _("MlirPassManager"));
-  bool load(handle src, bool) {
-    auto capsule = mlirApiObjectToCapsule(src);
-    value = mlirPythonCapsuleToPassManager(capsule.ptr());
-    if (mlirPassManagerIsNull(value)) {
-      return false;
-    }
-    return true;
-  }
-};
-
-/// Casts object -> MlirType.
-template <> struct type_caster<MlirType> {
-  PYBIND11_TYPE_CASTER(MlirType, _("MlirType"));
-  bool load(handle src, bool) {
-    auto capsule = mlirApiObjectToCapsule(src);
-    value = mlirPythonCapsuleToType(capsule.ptr());
-    if (mlirTypeIsNull(value)) {
-      return false;
-    }
-    return true;
-  }
-  static handle cast(MlirType t, return_value_policy, handle) {
-    auto capsule =
-        py::reinterpret_steal<py::object>(mlirPythonTypeToCapsule(t));
-    return py::module::import("mlir.ir")
-        .attr("Type")
-        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
-        .release();
-  }
-};
-
-} // namespace detail
-} // namespace pybind11
-
 namespace pybind11 {
 
 /// Raises a python exception with the given message.

python/npcomp/__init__.py

@@ -1,13 +0,0 @@
-from mlir import _cext_loader
-_cext_loader._cext.globals.append_dialect_search_prefix("npcomp.dialects")
-
-_cext = _cext_loader._load_extension("_npcomp")
-_cext._register_all_passes()
-_cext._initialize_llvm_codegen()
-
-# Top-level symbols.
-from .exporter import *
-from .types import *
-
-from . import tracing
-from . import utils

python/npcomp/build.py

@@ -7,7 +7,7 @@
 import os
 import platform
 
-from mlir._mlir_libs import get_include_dirs, get_lib_dirs
+from ._mlir_libs import get_include_dirs, get_lib_dirs
 
 __all__ = [
   "get_include_dirs",

python/npcomp/compiler/generic/backend/refjit.py

@@ -18,7 +18,7 @@ def get_refjit():
   global _refjit
   if _refjit is not None:
     return _refjit
-  from .... import _cext
+  from ...._mlir_libs import _npcomp as _cext
   try:
     imported_refjit = _cext.backend.refjit
   except AttributeError:

python/npcomp/compiler/numpy/__init__.py

@@ -1,3 +0,0 @@
-#  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

python/npcomp/compiler/numpy/extensions/numpy/builtin_ops.py

@@ -8,15 +8,15 @@ from typing import Callable, Iterator, Sequence, Tuple
 import functools
 import numpy as np
 
-from mlir import ir as _ir
-
-from npcomp import _cext
-from npcomp.dialects import numpy as numpy_ops
-
-from ....utils import logging
 from ...interfaces import *
 from ...partial_eval_base import *
 
+from ....utils import logging
+
+from ..... import ir as _ir
+from ....._mlir_libs import _npcomp as _cext
+from .....dialects import numpy as numpy_ops
+
 __all__ = [
     "get_ufuncs_from_module",
     "bind_ufuncs",

python/npcomp/compiler/numpy/extensions/numpy/value_coder.py

@@ -6,14 +6,13 @@
 import numpy as np
 from typing import Union
 
-from mlir import ir as _ir
-from mlir.dialects import std as std_ops
-
-from npcomp import _cext
-from npcomp.dialects import numpy as numpy_ops
+from ...interfaces import *
 
 from ....utils import logging
-from ...interfaces import *
+
+from ..... import ir as _ir
+from .....dialects import std as std_ops, numpy as numpy_ops
+from ....._mlir_libs import _npcomp as _cext
 
 __all__ = [
     "CreateNumpyValueCoder",

python/npcomp/compiler/numpy/frontend.py

@@ -10,15 +10,17 @@ import inspect
 import textwrap
 from typing import Optional
 
-from mlir import ir as _ir
 from ..utils import logging
 from .importer import *
 from .interfaces import *
 from .name_resolver_base import *
 from .value_coder_base import *
 from .target import *
+
 from ..utils.mlir_utils import *
 
+from ... import ir as _ir
+
 __all__ = [
     "ImportFrontend",
 ]

python/npcomp/compiler/numpy/importer.py

@@ -8,14 +8,12 @@ import ast
 import sys
 import traceback
 
-from mlir import ir as _ir
-from mlir.dialects import std as std_ops
-
-from npcomp import _cext
-from npcomp.dialects import basicpy as basicpy_ops
-
-from ..utils import logging
 from .interfaces import *
+from ..utils import logging
+
+from ... import ir as _ir
+from ...dialects import std as std_ops, basicpy as basicpy_ops
+from ..._mlir_libs import _npcomp as _cext
 
 __all__ = [
     "FunctionContext",

python/npcomp/compiler/numpy/interfaces.py

@@ -8,10 +8,9 @@ from enum import Enum
 import sys
 from typing import List, Optional, Sequence, Tuple, Union
 
-from mlir import ir as _ir
-
 from .target import *
 from ..utils.mlir_utils import *
+from ... import ir as _ir
 
 __all__ = [
     "Configuration",

python/npcomp/compiler/numpy/name_resolver_base.py

@@ -5,9 +5,10 @@
 
 from typing import Optional
 
-from mlir import ir as _ir
 from .interfaces import *
 
+from ... import ir as _ir
+
 __all__ = [
     "ConstModuleNameResolver",
     "LocalNameResolver",

python/npcomp/compiler/numpy/target.py

@@ -3,10 +3,11 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from typing import *
-from mlir import ir as _ir
 
 from ..utils.mlir_utils import *
 
+from ... import ir as _ir
+
 __all__ = [
     "GenericTarget32",
     "GenericTarget64",

python/npcomp/compiler/numpy/test_config.py

@@ -6,13 +6,13 @@
 import ast
 import functools
 
-from ..utils import logging
 from .frontend import *
 from .interfaces import *
 from .partial_eval_base import *
 from .target import *
 from .value_coder_base import *
 from .extensions import numpy as npc
+from ..utils import logging
 
 
 def create_import_dump_decorator(*,

python/npcomp/compiler/numpy/value_coder_base.py

@@ -5,12 +5,11 @@
 
 from typing import Union
 
-from mlir import ir as _ir
-from mlir.dialects import std as std_ops
-from npcomp.dialects import basicpy as basicpy_ops
-
 from .interfaces import *
 
+from ... import ir as _ir
+from ...dialects import std as std_ops, basicpy as basicpy_ops
+
 __all__ = [
     "BuiltinsValueCoder",
 ]

python/npcomp/compiler/pytorch/backend/abc.py

@@ -7,7 +7,7 @@ from typing import TypeVar
 
 import torch
 
-from mlir.ir import Module
+from npcomp.ir import Module
 
 # A type shared between the result of `NpcompBackend.compile` and the input
 # to `NpcompBackend.load`. Each backend will likely have a different definition
@@ -26,7 +26,7 @@ class NpcompBackend(abc.ABC):
     @abc.abstractmethod
     def compile(self, module: Module) -> CompiledArtifact:
         """Compile the provided MLIR module into a compiled artifact.
-        
+
         The module adheres to the npcomp backend contract
         (see the VerifyBackendContract pass).
 

python/npcomp/compiler/pytorch/backend/refjit.py

@@ -6,8 +6,8 @@ import os
 
 import torch
 
-from mlir.ir import *
-from mlir.passmanager import *
+from npcomp.ir import *
+from npcomp.passmanager import *
 from npcomp.compiler.generic.backend import refjit as refjit_backend
 from npcomp.compiler.utils import logging
 from .abc import NpcompBackend

python/npcomp/compiler/utils/mlir_utils.py

@@ -5,9 +5,9 @@
 
 from typing import Optional, Tuple
 
-from mlir import ir as _ir
-from mlir.dialects import builtin as builtin_ops
-from npcomp import _cext
+from ... import ir as _ir
+from ...dialects import builtin as builtin_ops
+from ..._mlir_libs import _npcomp as _cext
 
 __all__ = [
     "ImportContext",

python/npcomp/dialects/_ods_common.py

@@ -1,7 +0,0 @@
-#  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
-
-# Generated tablegen dialects expect to be able to find some symbols from
-# the mlir.dialects package.
-from mlir.dialects._ods_common import _cext, segmented_accessor, equally_sized_accessor, extend_opview_class, get_default_loc_context

python/npcomp/tracing/__init__.py

@@ -1,3 +0,0 @@
-# Module level symbols.
-from .context import *
-from .mlir_trace import *

python/npcomp/tracing/context.py

@@ -1,189 +0,0 @@
-#  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 typing import Optional
-
-import contextlib
-import os
-import threading
-
-import numpy as np
-
-
-class TracingError(Exception):
-  pass
-
-
-class TraceContext:
-  """Context for intercepting array traces.
-
-  Context manager:
-  ----------------
-  Instances act as context managers, the inner-most of which can be
-  queried with current() or optional_current().
-
-    >>> with TraceContext(desc=1) as tc:
-    ...   print(tc)
-    ...   print(TraceContext.current())
-    <TraceContext 1>
-    <TraceContext 1>
-    >>> print(TraceContext.optional_current())
-    None
-    >>> TraceContext.current()
-    Traceback (most recent call last):
-    ...
-    RuntimeError: No active TraceContext
-
-  Unique ids:
-  -----------
-  Many things in tracing require a context-local id.
-
-    >>> with TraceContext() as tc:
-    ...   print(tc.get_next_id())
-    ...   print(tc.get_next_id())
-    1
-    2
-
-  """
-  _local = threading.local()
-  __slots__ = [
-      "_desc",
-      "_next_id",
-      "active",
-  ]
-
-  def __init__(self, desc=None):
-    _check_numpy_version()
-    self._desc = desc
-    self._next_id = 1
-    self.active = False
-
-  def _handle_array_getitem(self, array, key):
-    """Handles a call to __getitem__ on a traced array."""
-    raise NotImplementedError("Array slicing not implemented")
-
-  def _handle_ufunc(self, ufunc, method, inputs, kwargs):
-    """Handles a ufunc invocation involving at least one TracedArray."""
-    return NotImplemented
-
-  def _handle_array_func(self, func, types, inputs, kwargs):
-    """Handles an __array_func__ hook involving at least on TracedArray."""
-    return NotImplemented
-
-  def get_next_id(self):
-    """Gets the next unique id for the context."""
-    rv = self._next_id
-    self._next_id += 1
-    return rv
-
-  @classmethod
-  def _get_context_stack(cls):
-    try:
-      return cls._local.s
-    except AttributeError:
-      cls._local.s = []
-      return cls._local.s
-
-  @classmethod
-  def optional_current(cls) -> Optional["TraceContext"]:
-    s = cls._get_context_stack()
-    if s:
-      return s[-1]
-    else:
-      return None
-
-  @classmethod
-  def current(cls) -> "TraceContext":
-    c = cls.optional_current()
-    if c is None:
-      raise RuntimeError("No active TraceContext")
-    return c
-
-  def __enter__(self):
-    s = self._get_context_stack()
-    if s:
-      s[-1].active = False
-    s.append(self)
-    self.active = True
-    return self
-
-  def __exit__(self, exc_type, exc_value, traceback):
-    s = self._get_context_stack()
-    s.pop()
-    self.active = False
-    if s:
-      s[-1].active = True
-
-  def __repr__(self):
-    return "<TraceContext %r>" % self._desc
-
-
-def _assert_active(tc: TraceContext):
-  assert tc.active, (
-      "Attempt to trace an action on an inactive trace context: %r" % tc)
-
-
-class TracedArray(np.lib.mixins.NDArrayOperatorsMixin):
-  """An array that traces its operations.
-
-  Unique ids:
-  -----------
-    >>> tc = TraceContext()
-    >>> TracedArray(tc=tc)
-    <TracedArray 1>
-    >>> TracedArray(tc=tc)
-    <TracedArray 2>
-  """
-
-  def __init__(self, tc: Optional[TraceContext] = None):
-    self._tc = tc if tc is not None else TraceContext.current()
-    self._uid = self._tc.get_next_id()
-
-  def __hash__(self):
-    return id(self)
-
-  @property
-  def uid(self):
-    return self._uid
-
-  def __repr__(self):
-    return "<TracedArray %d>" % self._uid
-
-  def __getitem__(self, key):
-    tc = self._tc
-    _assert_active(tc)
-    return tc._handle_array_getitem(self, key)
-
-  def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
-    tc = self._tc
-    _assert_active(tc)
-    return tc._handle_ufunc(ufunc, method, inputs, kwargs)
-
-  def __array_function__(self, func, types, args, kwargs):
-    tc = self._tc
-    _assert_active(tc)
-    return tc._handle_array_func(func, types, args, kwargs)
-
-  @property
-  def T(self):
-    """Shortcut for transpose."""
-    tc = self._tc
-    _assert_active(tc)
-    return tc._handle_array_func(np.transpose, [TracedArray], [self], {})
-
-
-def _check_numpy_version():
-  version = np.lib.NumpyVersion(np.__version__)
-  if version < "1.16.0":
-    raise RuntimeError("Numpy version >= 1.16 is required")
-  if version > "1.17.0":
-    return
-  if os.environ.get("NUMPY_EXPERIMENTAL_ARRAY_FUNCTION") != "1":
-    raise RuntimeError("For numpy 1.16, the environment variable "
-                       "NUMPY_EXPERIMENTAL_ARRAY_FUNCTION must equal 1")
-
-
-if __name__ == "__main__":
-  import doctest
-  doctest.testmod()

python/npcomp/tracing/emitters.py

@@ -1,296 +0,0 @@
-#  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
-
-import numpy as np
-
-from collections import namedtuple
-from enum import Enum
-
-from mlir import ir as _ir
-
-from npcomp.dialects import numpy as numpy_ops
-
-
-class Protocol(Enum):
-  UFUNC = 1
-  ARRAY_FUNC = 2
-
-
-class TraceValueType(Enum):
-  NDARRAY = 1
-
-
-class TraceValue(namedtuple("TraceValue", ["value", "type"])):
-  __slots__ = ()
-  """A Python value and the trace type that it should correspond to."""
-
-
-TraceValue.__new__.__defaults__ = (TraceValueType.NDARRAY,)
-
-
-class TraceInvocation(
-    namedtuple("TraceInvocation", ["inputs", "kwargs", "protocol", "method"])):
-  """An invocation of a single functions.
-
-  This abstracts over both ufuncs and array_funcs, differentiating by the
-  protocol and method.
-  """
-  __slots__ = ()
-
-
-TraceInvocation.__new__.__defaults__ = (Protocol.ARRAY_FUNC, "__call__")
-
-
-class EmissionRequest(
-    namedtuple("EmissionRequest", ["input_ssa_values", "ic", "extra"])):
-  """Represents the result of processing inputs from an invocation.
-
-  The `input_ssa_values` are mlir.ir.Value instances corresponding to
-  input_trace_values in TraceValueMap.
-
-  The `extra` value is only relevant to the producer and can be used as a
-  blackbox mechanism to transfer un-tracked state from an invocation to
-  emission.
-
-  The `dialect_helper` fields correspond to mlir.ir.DialectHelper.
-  """
-  __slots__ = ()
-
-
-EmissionRequest.__new__.__defaults__ = (None,)
-
-
-class TraceValueMap(
-    namedtuple("TraceValueMap",
-               ["input_trace_values", "result_trace_value_types", "extra"])):
-  """The result of mapping an invocation to corresponding op structure.
-
-  This type associates:
-    - Python (object, TraceValueType) representing invocation inputs that
-      correspond to SSA values in the IR.
-    - TraceValueTypes that are the expected logical result types from the
-      invocation.
-    - 'extra' object that is passed to followon Emitter methods.
-  """
-  __slots__ = ()
-
-
-TraceValueMap.__new__.__defaults__ = (None)
-
-
-class FuncEmitter:
-  """An emitter for an op-like function invocation."""
-
-  def map_invocation(self, trace_invocation: TraceInvocation) -> TraceValueMap:
-    """Maps from an invocation to EmissionRequest.
-
-    This hook is also responsible for validating the invocation and should
-    raise appropriate user-visible exceptions (i.e. when invoked with incorrect
-    arguments).
-
-    This hook is used to prepare for emission in a define-by-run scenario.
-    Static emission from an AST needs to be prepared via another mechanism.
-
-    Args:
-      trace_invocation: An Invocation instance to map.
-    Returns:
-      A TraceValueMap describing the structure of the invocation as mapped
-      to/from IR.
-    """
-    raise NotImplementedError()
-
-  def map_results(self, py_results, extra):
-    """Maps a list of python results to actual function return values.
-
-    Args:
-      py_results: List of python results corresponding to the emitted op
-        results.
-      extra: The extra object returned by map_invocation.
-    Returns:
-      Actual function result. Typically this requires special handling to
-      unpack the result of functions that return 1 item.
-    """
-    raise NotImplementedError()
-
-  def emit(self, request: EmissionRequest):
-    """Emits IR using the provided ops and types factories.
-
-    Args:
-      emission_inputs: An EmissionRequest produced by tracing each TraceValue
-        from a previous call to map_invocation and the corresponding extra
-        value.
-    Returns:
-      An iterable of mlir.ir.Value instances representing the outputs of the
-      operation. The `builder` on `ops` must be positioned to consume these
-      values.
-    """
-    raise NotImplementedError()
-
-
-class GenericCallUfuncEmitter(FuncEmitter):
-  """A FuncEmitter for generic ufuncs requiring no special behavior.
-
-  Representation:
-    >>> emitter = GenericCallUfuncEmitter("numpy.add")
-    >>> emitter
-    <ufunc emitter 'numpy.add'>
-    >>> inv = TraceInvocation([1, 2], {}, protocol=Protocol.UFUNC)
-    >>> inputs = emitter.map_invocation(inv)
-    >>> inputs
-    TraceValueMap(input_trace_values=[TraceValue(value=1, type=<TraceValueType.NDARRAY: 1>), TraceValue(value=2, type=<TraceValueType.NDARRAY: 1>)], result_trace_value_types=[<TraceValueType.NDARRAY: 1>], extra=None)
-
-  Error on unsupported kwargs:
-    >>> inv = TraceInvocation([1, 2], {"foobar": 1}, protocol=Protocol.UFUNC)
-    >>> emitter.map_invocation(inv)
-    Traceback (most recent call last):
-    ...
-    ValueError: Unexpected keyword args for ufunc numpy.add: foobar
-
-  """
-  __slots__ = ("_ufunc_name")
-
-  def __init__(self, ufunc_name: str):
-    self._ufunc_name = ufunc_name
-
-  def __repr__(self):
-    return "<ufunc emitter '%s'>" % self._ufunc_name
-
-  def map_invocation(self,
-                     trace_invocation: TraceInvocation) -> EmissionRequest:
-    assert trace_invocation.protocol == Protocol.UFUNC
-    assert trace_invocation.method == "__call__"
-    if trace_invocation.kwargs:
-      raise ValueError(
-          "Unexpected keyword args for ufunc %s: %s" %
-          (self._ufunc_name, ", ".join(trace_invocation.kwargs.keys())))
-    # Without above special cases, any positional args map to emission
-    # inputs.
-    return TraceValueMap([
-        TraceValue(i, TraceValueType.NDARRAY) for i in trace_invocation.inputs
-    ], [TraceValueType.NDARRAY],
-                         extra=None)
-
-  def map_results(self, py_results, extra):
-    # Ufuncs always return one result, so just unpack it.
-    return py_results[0]
-
-  def emit(self, request: EmissionRequest):
-    ic = request.ic
-    name_attr = _ir.StringAttr.get(self._ufunc_name)
-    result_type = ic.unknown_tensor_type
-    call_op = numpy_ops.BuiltinUfuncCallOp(result_type,
-                                           qualified_name=name_attr,
-                                           inputs=request.input_ssa_values,
-                                           loc=ic.loc,
-                                           ip=ic.ip)
-    return call_op.results
-
-
-class GenericArrayFuncEmitter(FuncEmitter):
-  """Emitter for array funcs that don't do anything 'special'."""
-  __slots__ = ("_op_name", "_nresults")
-
-  def __init__(self, op_name: str, nresults: int = 1):
-    self._op_name = op_name
-    self._nresults = nresults
-
-  def __repr__(self):
-    return "<array_func emitter '%s'>" % self._op_name
-
-  def map_invocation(self,
-                     trace_invocation: TraceInvocation) -> EmissionRequest:
-    assert trace_invocation.protocol == Protocol.ARRAY_FUNC
-    if trace_invocation.method != "__call__":
-      raise NotImplementedError("Only __call__ is supported for %s (got '%s')" %
-                                (
-                                    self._op_name,
-                                    trace_invocation.method,
-                                ))
-    if trace_invocation.kwargs:
-      raise ValueError(
-          "Unexpected keyword args for %s: %s" %
-          (self._op_name, ", ".join(trace_invocation.kwargs.keys())))
-    # Without above special cases, any positional args map to emission
-    # inputs.
-    return TraceValueMap([
-        TraceValue(i, TraceValueType.NDARRAY) for i in trace_invocation.inputs
-    ], [TraceValueType.NDARRAY] * self._nresults,
-                         extra=None)
-
-  def map_results(self, py_results, extra):
-    if self._nresults == 1:
-      return py_results[0]
-    else:
-      return tuple(py_results)
-
-  def emit(self, request: EmissionRequest):
-    ic = request.ic
-    op_result_types = [ic.unknown_tensor_type] * self._nresults
-    op = _ir.Operation.create(self._op_name,
-                              results=op_result_types,
-                              operands=request.input_ssa_values,
-                              loc=ic.loc,
-                              ip=ic.ip)
-    return op.results
-
-
-class EmitterRegistry:
-  """Registry of known Emitter instances mapped to source function.
-
-    >>> r = EmitterRegistry.create_default()
-    >>> r.lookup_ufunc(np.add, "__call__")
-    <ufunc emitter 'numpy.add'>
-    >>> r.lookup_array_func(np.dot)
-    <array_func emitter 'numpy.dot'>
-  """
-
-  def __init__(self):
-    self._ufunc_map = {}  # Dictionary of (f, method) -> Emitter
-    self._arrayfunc_map = {}  # Dictionary of f -> Emitter
-
-  @classmethod
-  def create_default(cls):
-    registry = cls()
-    registry.register_defaults()
-    return registry
-
-  def register_ufunc(self, ufunc, method, emitter):
-    # Last registration wins.
-    self._ufunc_map[(ufunc, method)] = emitter
-
-  def register_array_func(self, f, emitter):
-    # Last registration wins.
-    self._arrayfunc_map[f] = emitter
-
-  def lookup_ufunc(self, ufunc, method):
-    return self._ufunc_map.get((ufunc, method))
-
-  def lookup_array_func(self, f):
-    return self._arrayfunc_map.get(f)
-
-  def register_defaults(self):
-    # Find all ufuncs in the numpy module and register by name.
-    for member in sorted(dir(np)):
-      ufunc = getattr(np, member)
-      if isinstance(ufunc, np.ufunc):
-        self.register_ufunc(ufunc, "__call__",
-                            GenericCallUfuncEmitter("numpy." + member))
-    # Register generic 1-result array funcs.
-    GENERIC_FUNCS = (
-        (np.inner, "numpy.inner"),
-        (np.outer, "numpy.outer"),
-        (np.dot, "numpy.dot"),
-        (np.vdot, "numpy.vdot"),
-        (np.linalg.det, "numpy.linalg.det"),
-        # TODO: This needs a custom implementation to differentiate when
-        # axes is specified (this version will fail).
-        (np.transpose, "numpy.transpose"),
-    )
-    for f, op_name in GENERIC_FUNCS:
-      self.register_array_func(f, GenericArrayFuncEmitter(op_name))
-
-
-if __name__ == "__main__":
-  import doctest
-  doctest.testmod()

python/npcomp/tracing/mlir_trace.py

@@ -1,296 +0,0 @@
-#  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
-
-import re
-from typing import Iterable, Optional
-import numpy as np
-
-from mlir import ir as _ir
-from mlir.dialects import std as std_ops
-
-from npcomp import _cext
-from npcomp.dialects import basicpy as basicpy_ops
-from npcomp.dialects import numpy as numpy_ops
-
-from ..exporter import *
-from ..types import *
-from ..compiler.utils.mlir_utils import *
-
-from .context import *
-from .emitters import *
-
-
-class ModuleBuilder:
-  """Builds an MLIR module by tracing functions."""
-
-  __slots__ = [
-      "emitters",
-      "ic",
-  ]
-
-  def __init__(self,
-               mlir_context: Optional[_ir.Context] = None,
-               emitter_registry=None):
-    ic = self.ic = ImportContext(mlir_context)
-    ic.module = _ir.Module.create(loc=ic.loc)
-    self.emitters = (emitter_registry
-                     if emitter_registry else EmitterRegistry.create_default())
-
-  @property
-  def module(self):
-    return self.ic.module
-
-  def trace(self, *export_py_funcs: ExportPyFunction):
-    """Traces exported py functions."""
-    for export_py_func in export_py_funcs:
-      assert isinstance(export_py_func, ExportPyFunction), (
-          "Expected an exported python function (from the Exporter class)")
-      tracer = FunctionTracer(self, export_py_func)
-      with tracer:
-        tracer.trace()
-
-
-class FunctionTracer(TraceContext):
-  """A trace of a single function."""
-  __slots__ = [
-      "module_builder",
-      "epf",
-      "_args_array_params",
-      "_f",
-      "_f_types",
-      "_ic",
-      "_python_args",
-      "_result_array_params",
-      "_traced_arrays",
-      "_external_arrays",
-  ]
-
-  def __init__(self, module_builder: ModuleBuilder, epf: ExportPyFunction):
-    super().__init__(desc="[trace of %s]" % epf.__name__)
-    self.module_builder = module_builder
-    self.epf = epf
-    self._traced_arrays = {}  # Mapping of TracedArray to current consumer value
-    self._external_arrays = {}  # Mapping of id to (ndarray, ir.Value)
-    self._validate()
-
-    # Alias some parent members for convenience.
-    self._ic = module_builder.ic
-    with self._ic.context:
-      # Extract ArrayParams for all args and results.
-      self._args_array_params = [
-          ArrayParams.from_constraints(arg.constraints)
-          for arg in self.epf.sig.args
-      ]
-      self._python_args = [None] * len(self._args_array_params)
-      self._result_array_params = ArrayParams.from_constraints(
-          self.epf.sig.result.constraints)
-
-      # Create the MLIR function.
-      self._f, self._f_types = self._create_mlir_function()
-      self._create_trace_roots()
-
-  @property
-  def entry_block(self) -> _ir.Block:
-    return self._f.regions[0].blocks[0]
-
-  def trace(self):
-    # Invoke the python function with placeholders.
-    # TODO: More sophisticated signature merging
-    # TODO: Multiple results
-    # TODO: Error reporting
-    ic = self._ic
-    ic.insert_end_of_block(self.entry_block)
-    with ic.context:
-      py_results = (self.epf.pyfunc(*self._python_args),)
-      if len(py_results) != len(self._f_types):
-        raise TracingError("Traced function returned != %d results: %r" % (
-            len(self._f_types),
-            py_results,
-        ))
-
-    # Narrow all results to the declared return types.
-    return_operands = []
-    for py_result, mlir_result_type in zip(py_results, self._f_types):
-      mlir_result = self.get_traced_array_value(py_result)
-      # narrow to declared result type.
-      return_operands.extend(
-          numpy_ops.NarrowOp(mlir_result_type,
-                             mlir_result,
-                             loc=ic.loc,
-                             ip=ic.ip).results)
-    std_ops.ReturnOp(return_operands, loc=ic.loc, ip=ic.ip)
-    ic.pop_ip()
-
-  def set_traced_array(self, traced_array, value):
-    """Sets the current SSA value for a traced_array."""
-    assert isinstance(traced_array, TracedArray)
-    self._traced_arrays[traced_array] = value
-
-  def get_traced_array_value(self, traced_array):
-    if not isinstance(traced_array, TracedArray):
-      # Generic import of external value. For now, we just treat these as
-      # local consts.
-      return self._get_external_array_value(traced_array)
-
-    traced_value = self._traced_arrays.get(traced_array)
-    if traced_value is None:
-      raise TracingError("Unregistered traced array: %r", (traced_array,))
-    return traced_value
-
-  def _get_external_array_value(self, external_array):
-    ic = self._ic
-    if not isinstance(external_array, np.ndarray):
-      raise TracingError("Expected ndarray but got: %r" % (external_array,))
-    found_it = self._external_arrays.get(id(external_array))
-    if found_it:
-      return found_it[1]
-    # Import it.
-    dense_attr = _ir.DenseElementsAttr.get(external_array, context=ic.context)
-    const_value = std_ops.ConstantOp(dense_attr.type,
-                                     dense_attr,
-                                     loc=ic.loc,
-                                     ip=ic.ip).result
-    self._external_arrays[id(external_array)] = (external_array, const_value)
-    return const_value
-
-  def _validate(self):
-    if not all(
-        arg.type_class == TypeClass.NdArray for arg in self.epf.sig.args):
-      raise NotImplementedError("Non NdArray args: %r" % (self.epf.sig.args,))
-    if not self.epf.sig.result.type_class == TypeClass.NdArray:
-      raise NotImplementedError("Non NdArray result: %r" %
-                                (self.epf.sig.result,))
-
-  def _create_mlir_function(self):
-    ic = self._ic
-    epf = self.epf
-    f_args = [
-        _ir.Type.parse(ap.mlir_tensor_type_asm)
-        for ap in self._args_array_params
-    ]
-    f_types = [_ir.Type.parse(self._result_array_params.mlir_tensor_type_asm)]
-    ic.insert_end_of_block(ic.module.body)
-    f_type = _ir.FunctionType.get(f_args, f_types)
-    f, _ = ic.FuncOp(epf.__name__, f_type, create_entry_block=True)
-    return f, f_types
-
-  def _create_trace_roots(self):
-    entry_block = self.entry_block
-    for index, ap in enumerate(self._args_array_params):
-      if ap is not None:
-        ta = TracedArray(self)
-        self.set_traced_array(ta, entry_block.arguments[index])
-        self._python_args[index] = ta
-
-  def _resolve_input_ssa_values(self, trace_values: Iterable[TraceValue]):
-    """Resolves input python values to SSA values."""
-    ssa_values = []
-    for tv in trace_values:
-      assert tv.type == TraceValueType.NDARRAY, (
-          "Unsupported TraceValueType: %r" % tv.type)
-      ssa_value = self.get_traced_array_value(tv.value)
-      ssa_values.append(ssa_value)
-    return ssa_values
-
-  def _resolve_result_py_values(self,
-                                trace_value_types: Iterable[TraceValueType],
-                                ssa_values):
-    """Resolves result SSA values to runtime python values."""
-    assert len(trace_value_types) == len(ssa_values), (
-        "Mismatched emitter declared result types and results")
-    py_values = []
-    for trace_value_type, ssa_value in zip(trace_value_types, ssa_values):
-      assert trace_value_type == TraceValueType.NDARRAY, (
-          "Unsupported TraceValueType: %r" % trace_value_type)
-      py_value = TracedArray(self)
-      self.set_traced_array(py_value, ssa_value)
-      py_values.append(py_value)
-    return py_values
-
-  def _emit_invocation(self, emitter: FuncEmitter, invocation: TraceInvocation):
-    tv_map = emitter.map_invocation(invocation)
-    input_ssa_values = self._resolve_input_ssa_values(tv_map.input_trace_values)
-    request = EmissionRequest(input_ssa_values, ic=self._ic, extra=tv_map.extra)
-    result_ssa_values = emitter.emit(request)
-    py_values = self._resolve_result_py_values(tv_map.result_trace_value_types,
-                                               result_ssa_values)
-    return emitter.map_results(py_values, tv_map.extra)
-
-  def _handle_ufunc(self, ufunc, method, inputs, kwargs):
-    emitter = self.module_builder.emitters.lookup_ufunc(ufunc, method)
-    if not emitter:
-      return NotImplemented
-    invocation = TraceInvocation(inputs, kwargs, Protocol.UFUNC, method)
-    return self._emit_invocation(emitter, invocation)
-
-  def _handle_array_func(self, func, types, inputs, kwargs):
-    emitter = self.module_builder.emitters.lookup_array_func(func)
-    if not emitter:
-      return NotImplemented
-    invocation = TraceInvocation(inputs, kwargs, Protocol.ARRAY_FUNC)
-    return self._emit_invocation(emitter, invocation)
-
-  def _emit_slice_value(self, slice_element):
-    ic = self._ic
-    if slice_element == None:
-      return basicpy_ops.SingletonOp(ic.none_type, loc=ic.loc, ip=ic.ip).result
-    elif slice_element == Ellipsis:
-      return basicpy_ops.SingletonOp(ic.ellipsis_type, loc=ic.loc,
-                                     ip=ic.ip).result
-    elif isinstance(slice_element, int):
-      return std_ops.ConstantOp(ic.index_type,
-                                _ir.IntegerAttr.get(ic.index_type,
-                                                    slice_element),
-                                loc=ic.loc,
-                                ip=ic.ip).result
-    elif isinstance(slice_element, slice):
-      return self._emit_slice_object(slice_element)
-    else:
-      # Assume array convertible.
-      raise NotImplementedError(
-          "TODO: Slicing with generic arrays not yet implemented")
-
-  def _emit_slice_object(self, slice_object: slice):
-    ic = self._ic
-
-    def emit_index(index):
-      if index is None:
-        return basicpy_ops.SingletonOp(ic.none_type, loc=ic.loc,
-                                       ip=ic.ip).result
-      else:
-        return std_ops.ConstantOp(ic.index_type,
-                                  _ir.IntegerAttr.get(ic.index_type,
-                                                      int(index)),
-                                  loc=ic.loc,
-                                  ip=ic.ip).result
-
-    start = emit_index(slice_object.start)
-    stop = emit_index(slice_object.stop)
-    step = emit_index(slice_object.step)
-    result_type = _cext.slot_object_type(ic.context, "slice",
-                                         [start.type, stop.type, step.type])
-    return basicpy_ops.SlotObjectMakeOp(result_type, [start, stop, step],
-                                        loc=ic.loc,
-                                        ip=ic.ip).result
-
-  def _handle_array_getitem(self, array, key):
-    ic = self._ic
-    array_value = self.get_traced_array_value(array)
-    # Array slicing is always based on a tuple.
-    slice_tuple = key if isinstance(key, tuple) else (key,)
-    # Resolve and emit each slice element.
-    slice_values = [self._emit_slice_value(elt) for elt in slice_tuple]
-    result_value = numpy_ops.GetSliceOp(ic.unknown_array_type,
-                                        array_value,
-                                        slice_values,
-                                        loc=ic.loc,
-                                        ip=ic.ip).result
-    result_array = TracedArray(self)
-    self.set_traced_array(result_array, result_value)
-    return result_array
-
-
-if __name__ == "__main__":
-  import doctest
-  doctest.testmod()

setup.py

@@ -17,11 +17,21 @@ import shutil
 import subprocess
 import sys
 
-from setuptools import find_packages, setup, Extension
+from distutils.command.build import build as _build
+from setuptools import find_namespace_packages, setup, Extension
 from setuptools.command.build_ext import build_ext
 from setuptools.command.build_py import build_py
 
 
+# Build phase discovery is unreliable. Just tell it what phases to run.
+class CustomBuild(_build):
+
+  def run(self):
+    self.run_command("build_py")
+    self.run_command("build_ext")
+    self.run_command("build_scripts")
+
+
 class CMakeExtension(Extension):
 
   def __init__(self, name, sourcedir=""):
@@ -79,21 +89,25 @@ setup(
     long_description="",
     include_package_data=True,
     ext_modules=[
-        CMakeExtension("mlir._mlir_libs._mlir"),
-        CMakeExtension("mlir._mlir_libs._npcomp"),
+        CMakeExtension("npcomp._mlir_libs._mlir"),
+        CMakeExtension("npcomp._mlir_libs._npcomp"),
         # TODO: We don't really want these but they are along for the ride.
-        CMakeExtension("mlir._mlir_libs._mlirAsyncPasses"),
-        CMakeExtension("mlir._mlir_libs._mlirConversions"),
-        CMakeExtension("mlir._mlir_libs._mlirTransforms"),
-        CMakeExtension("mlir._mlir_libs._mlirSparseTensorPasses"),
-        CMakeExtension("mlir._mlir_libs._mlirAllPassesRegisration"),
-        CMakeExtension("mlir._mlir_libs._mlirLinalgPasses"),
-        CMakeExtension("mlir._mlir_libs._mlirGPUPasses"),
+        CMakeExtension("npcomp._mlir_libs._mlirAsyncPasses"),
+        CMakeExtension("npcomp._mlir_libs._mlirConversions"),
+        CMakeExtension("npcomp._mlir_libs._mlirTransforms"),
+        CMakeExtension("npcomp._mlir_libs._mlirSparseTensorPasses"),
+        CMakeExtension("npcomp._mlir_libs._mlirAllPassesRegisration"),
+        CMakeExtension("npcomp._mlir_libs._mlirLinalgPasses"),
+        CMakeExtension("npcomp._mlir_libs._mlirGPUPasses"),
     ],
     cmdclass={
+        "build": CustomBuild,
         "built_ext": NoopBuildExtension,
         "build_py": CMakeBuild,
     },
     zip_safe=False,
-    packages=find_packages(),
+    packages=find_namespace_packages(include=[
+        "npcomp",
+        "npcomp.*",
+    ]),
 )

test/CMakeLists.txt

@@ -17,7 +17,6 @@ set(NPCOMP_TEST_DEPENDS
         npcomp-opt
         refback-run
         NPCOMPPythonModules
-        NPCOMPMLIRPythonModules
 )
 
 add_lit_testsuite(check-npcomp-lit "Running the npcomp regression tests"

test/Python/Tracing/const.py

@@ -1,34 +0,0 @@
-# RUN: %PYTHON %s | FileCheck %s --dump-input=fail
-
-#  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
-
-import numpy as np
-import npcomp as npc
-from npcomp.types import *
-
-weights = np.random.uniform(size=(16, 4)).astype(np.float32)
-bias = np.random.uniform(size=(4,)).astype(np.float32)
-
-
-def constants(a: np.ndarray) -> np.ndarray:
-  return np.dot(a, weights) + bias
-
-
-# TODO: Implement subclassing and deriving constraints by run
-exp = npc.Exporter()
-exp.constants = constants
-
-mb = npc.tracing.ModuleBuilder()
-mb.trace(exp.constants)
-# CHECK-LABEL:   func @constants(
-# CHECK-SAME:                    %[[VAL_0:.*]]: tensor<*x!numpy.any_dtype>) -> tensor<*x!numpy.any_dtype> {
-# CHECK:           %[[VAL_1:.*]] = constant dense<{{.*}}> : tensor<16x4xf32>
-# CHECK:           %[[VAL_2:.*]] = numpy.dot %[[VAL_0]], %[[VAL_1]] : (tensor<*x!numpy.any_dtype>, tensor<16x4xf32>) -> tensor<*x!basicpy.UnknownType>
-# CHECK:           %[[VAL_3:.*]] = constant dense<{{.*}}> : tensor<4xf32>
-# CHECK:           %[[VAL_4:.*]] = numpy.builtin_ufunc_call<"numpy.add"> (%[[VAL_2]], %[[VAL_3]]) : (tensor<*x!basicpy.UnknownType>, tensor<4xf32>) -> tensor<*x!basicpy.UnknownType>
-# CHECK:           %[[VAL_5:.*]] = numpy.narrow %[[VAL_4]] : (tensor<*x!basicpy.UnknownType>) -> tensor<*x!numpy.any_dtype>
-# CHECK:           return %[[VAL_5]] : tensor<*x!numpy.any_dtype>
-# CHECK:         }
-print(mb.module)

test/Python/Tracing/dot.py

@@ -1,38 +0,0 @@
-# RUN: %PYTHON %s | FileCheck %s --dump-input=fail
-
-#  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
-
-import numpy as np
-import npcomp as npc
-from npcomp.types import *
-
-
-def dot2d(a: np.ndarray, b: np.ndarray) -> np.ndarray:
-  return np.dot(a, b)
-
-
-# TODO: Implement subclassing and deriving constraints by run
-exp = npc.Exporter()
-exp.dot2d = dot2d
-exp.dot2d.sig.args["a"] += Shape(4, 16)
-exp.dot2d.sig.args["a"] += DynamicDim(0)
-exp.dot2d.sig.args["a"] += DType(np.float32)
-exp.dot2d.sig.args["b"] += Shape(16, 32)
-exp.dot2d.sig.args["b"] += DType(np.float32)
-exp.dot2d.sig.result += Shape(4, 32)
-exp.dot2d.sig.result += DynamicDim(0)
-exp.dot2d.sig.result += DType(np.float32)
-
-mb = npc.tracing.ModuleBuilder()
-mb.trace(exp.dot2d)
-
-# CHECK-LABEL:   func @dot2d(
-# CHECK-SAME:                %[[VAL_0:.*]]: tensor<?x16xf32>,
-# CHECK-SAME:                %[[VAL_1:.*]]: tensor<16x32xf32>) -> tensor<?x32xf32> {
-# CHECK:           %[[VAL_2:.*]] = numpy.dot %[[VAL_0]], %[[VAL_1]] : (tensor<?x16xf32>, tensor<16x32xf32>) -> tensor<*x!basicpy.UnknownType>
-# CHECK:           %[[VAL_3:.*]] = numpy.narrow %[[VAL_2]] : (tensor<*x!basicpy.UnknownType>) -> tensor<?x32xf32>
-# CHECK:           return %[[VAL_3]] : tensor<?x32xf32>
-# CHECK:         }
-print(mb.module)

test/Python/Tracing/function_trace.py

@@ -1,41 +0,0 @@
-# RUN: %PYTHON %s | FileCheck %s --dump-input=fail
-
-#  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
-
-import os
-os.environ["NUMPY_EXPERIMENTAL_ARRAY_FUNCTION"] = "1"
-
-from npcomp.types import *
-from npcomp.exporter import *
-from npcomp.tracing.mlir_trace import *
-
-
-def simple_mul(a: np.ndarray, b: np.ndarray) -> np.ndarray:
-  return a * b + a + b
-
-
-# TODO: Implement subclassing and deriving constraints by run
-exp = Exporter()
-exp.simple_mul = simple_mul
-exp.simple_mul.sig.args["a"] += Shape(1, 4)
-exp.simple_mul.sig.args["a"] += DynamicDim(0)
-exp.simple_mul.sig.args["a"] += DType(np.float32)
-exp.simple_mul.sig.args["b"] += Shape(1)
-exp.simple_mul.sig.args["b"] += DType(np.float32)
-exp.simple_mul.sig.result += Shape(1, 4)
-exp.simple_mul.sig.result += DynamicDim(0)
-exp.simple_mul.sig.result += DType(np.float32)
-
-mb = ModuleBuilder()
-mb.trace(exp.simple_mul)
-# This test exercises the full tracing path and incidentally checks the ops.
-# CHECK: func @simple_mul(%arg0: tensor<?x4xf32>, %arg1: tensor<1xf32>) -> tensor<?x4xf32> {
-# CHECK:  %0 = numpy.builtin_ufunc_call<"numpy.multiply"> (%arg0, %arg1) : (tensor<?x4xf32>, tensor<1xf32>) -> tensor<*x!basicpy.UnknownType>
-# CHECK:  %1 = numpy.builtin_ufunc_call<"numpy.add"> (%0, %arg0) : (tensor<*x!basicpy.UnknownType>, tensor<?x4xf32>) -> tensor<*x!basicpy.UnknownType>
-# CHECK:  %2 = numpy.builtin_ufunc_call<"numpy.add"> (%1, %arg1) : (tensor<*x!basicpy.UnknownType>, tensor<1xf32>) -> tensor<*x!basicpy.UnknownType>
-# CHECK:  %3 = numpy.narrow %2 : (tensor<*x!basicpy.UnknownType>) -> tensor<?x4xf32>
-# CHECK:  return %3 : tensor<?x4xf32>
-# CHECK: }
-print(str(mb.module))

test/Python/Tracing/slice.py

@@ -1,48 +0,0 @@
-# RUN: %PYTHON %s | FileCheck %s --dump-input=fail
-
-#  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
-
-import numpy as np
-import npcomp as npc
-from npcomp.types import *
-
-
-def slice_array1(a: np.ndarray) -> np.ndarray:
-  return a[1, 2:10:2, 3:4, ..., :, 0]
-
-
-# TODO: Implement subclassing and deriving constraints by run
-exp = npc.Exporter()
-exp.slice_array1 = slice_array1
-
-mb = npc.tracing.ModuleBuilder()
-mb.trace(exp.slice_array1)
-
-# TODO: The numpy.get_slice op emission should be analyzed: it probably
-# needs to both accept and produce either arrays or tensors and the following
-# narrow should do likewise.
-# CHECK-LABEL:   func @slice_array1(
-# CHECK-SAME:                       %[[VAL_0:.*]]: tensor<*x!numpy.any_dtype>) -> tensor<*x!numpy.any_dtype> {
-# CHECK:           %[[VAL_1:.*]] = constant 1 : index
-# CHECK:           %[[VAL_2:.*]] = constant 2 : index
-# CHECK:           %[[VAL_3:.*]] = constant 10 : index
-# CHECK:           %[[VAL_4:.*]] = constant 2 : index
-# CHECK:           %[[VAL_5:.*]] = basicpy.slot_object_make(%[[VAL_2]], %[[VAL_3]], %[[VAL_4]]) -> !basicpy.SlotObject<slice, index, index, index>
-# CHECK:           %[[VAL_6:.*]] = constant 3 : index
-# CHECK:           %[[VAL_7:.*]] = constant 4 : index
-# CHECK:           %[[VAL_8:.*]] = basicpy.singleton : !basicpy.NoneType
-# CHECK:           %[[VAL_9:.*]] = basicpy.slot_object_make(%[[VAL_6]], %[[VAL_7]], %[[VAL_8]]) -> !basicpy.SlotObject<slice, index, index, !basicpy.NoneType>
-# CHECK:           %[[VAL_10:.*]] = basicpy.singleton : !basicpy.EllipsisType
-# CHECK:           %[[VAL_11:.*]] = basicpy.singleton : !basicpy.NoneType
-# CHECK:           %[[VAL_12:.*]] = basicpy.singleton : !basicpy.NoneType
-# CHECK:           %[[VAL_13:.*]] = basicpy.singleton : !basicpy.NoneType
-# CHECK:           %[[VAL_14:.*]] = basicpy.slot_object_make(%[[VAL_11]], %[[VAL_12]], %[[VAL_13]]) -> !basicpy.SlotObject<slice, !basicpy.NoneType, !basicpy.NoneType, !basicpy.NoneType>
-# CHECK:           %[[VAL_15:.*]] = constant 0 : index
-# CHECK:           %[[VAL_16:.*]] = numpy.get_slice %[[VAL_0]], %[[VAL_1]], %[[VAL_5]], %[[VAL_9]], %[[VAL_10]], %[[VAL_14]], %[[VAL_15]] : (tensor<*x!numpy.any_dtype>, index, !basicpy.SlotObject<slice, index, index, index>, !basicpy.SlotObject<slice, index, index, !basicpy.NoneType>, !basicpy.EllipsisType, !basicpy.SlotObject<slice, !basicpy.NoneType, !basicpy.NoneType, !basicpy.NoneType>, index) -> !numpy.ndarray<*:?>
-# CHECK:           %[[VAL_17:.*]] = numpy.narrow %[[VAL_16]] : (!numpy.ndarray<*:?>) -> tensor<*x!numpy.any_dtype>
-# CHECK:           return %[[VAL_17]] : tensor<*x!numpy.any_dtype>
-# CHECK:         }
-
-print(mb.module)

test/Python/Tracing/transpose.py

@@ -1,42 +0,0 @@
-# RUN: %PYTHON %s | FileCheck %s --dump-input=fail
-
-#  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
-
-import numpy as np
-import npcomp as npc
-from npcomp.types import *
-
-
-def transpose_attribute(a: np.ndarray) -> np.ndarray:
-  return a.T
-
-
-def transpose(a: np.ndarray) -> np.ndarray:
-  return np.transpose(a)
-
-
-# TODO: Implement subclassing and deriving constraints by run
-exp = npc.Exporter()
-exp.transpose_attribute = transpose_attribute
-exp.transpose = transpose
-
-mb = npc.tracing.ModuleBuilder()
-mb.trace(exp.transpose_attribute, exp.transpose)
-
-# TODO: Consolidate any_dtype -> UnknownType.
-# CHECK-LABEL:   func @transpose_attribute(
-# CHECK-SAME:                              %[[VAL_0:.*]]: tensor<*x!numpy.any_dtype>) -> tensor<*x!numpy.any_dtype> {
-# CHECK:           %[[VAL_1:.*]] = numpy.transpose %[[VAL_0]] : (tensor<*x!numpy.any_dtype>) -> tensor<*x!basicpy.UnknownType>
-# CHECK:           %[[VAL_2:.*]] = numpy.narrow %[[VAL_1]] : (tensor<*x!basicpy.UnknownType>) -> tensor<*x!numpy.any_dtype>
-# CHECK:           return %[[VAL_2]] : tensor<*x!numpy.any_dtype>
-# CHECK:         }
-
-# CHECK-LABEL:   func @transpose(
-# CHECK-SAME:                    %[[VAL_0:.*]]: tensor<*x!numpy.any_dtype>) -> tensor<*x!numpy.any_dtype> {
-# CHECK:           %[[VAL_1:.*]] = numpy.transpose %[[VAL_0]] : (tensor<*x!numpy.any_dtype>) -> tensor<*x!basicpy.UnknownType>
-# CHECK:           %[[VAL_2:.*]] = numpy.narrow %[[VAL_1]] : (tensor<*x!basicpy.UnknownType>) -> tensor<*x!numpy.any_dtype>
-# CHECK:           return %[[VAL_2]] : tensor<*x!numpy.any_dtype>
-# CHECK:         }
-print(mb.module)

test/Python/doctests.py

@@ -24,9 +24,6 @@ def run_doctest(mod):
 
 TEST_MODULES = (
     "npcomp.compiler.numpy.py_value_utils",
-    "npcomp.tracing.context",
-    "npcomp.tracing.emitters",
-    "npcomp.tracing.mlir_trace",
     "npcomp.types",
     "npcomp.exporter",
 )

test/Python/extension_coexistence.py

@@ -1,12 +0,0 @@
-# Some checks that we can import the various extensions and libraries and
-# not have symbol collisions or other goings on.
-# RUN: %PYTHON %s
-
-import sys
-
-print(f"PYTHONPATH={sys.path}")
-
-import mlir.ir
-import npcomp
-
-print("Extensions all loaded")