diff --git a/mlir/utils/jupyter/.gitignore b/mlir/utils/jupyter/.gitignore new file mode 100644 --- /dev/null +++ b/mlir/utils/jupyter/.gitignore @@ -0,0 +1,5 @@ +__pycache__ +*.pyc +build/ +dist/ +MANIFEST diff --git a/mlir/utils/jupyter/README.md b/mlir/utils/jupyter/README.md new file mode 100644 --- /dev/null +++ b/mlir/utils/jupyter/README.md @@ -0,0 +1,19 @@ +A Jupyter kernel for mlir (mlir-opt) + +This is purely for experimentation. This kernel uses the reproducer runner +conventions to run passes. + +To install: + + python3 -m mlir_opt_kernel.install + +To use it, run one of: + +```shell + jupyter notebook + # In the notebook interface, select MlirOpt from the 'New' menu + jupyter console --kernel mlir +``` + +`mlir-opt` is expected to be either in the `PATH` or `MLIR_OPT_EXECUTABLE` is +used to point to the executable directly. diff --git a/mlir/utils/jupyter/mlir_opt_kernel/__init__.py b/mlir/utils/jupyter/mlir_opt_kernel/__init__.py new file mode 100644 --- /dev/null +++ b/mlir/utils/jupyter/mlir_opt_kernel/__init__.py @@ -0,0 +1,6 @@ +# 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 +"""A mlir-opt kernel for Jupyter""" + +from .kernel import __version__ diff --git a/mlir/utils/jupyter/mlir_opt_kernel/__main__.py b/mlir/utils/jupyter/mlir_opt_kernel/__main__.py new file mode 100644 --- /dev/null +++ b/mlir/utils/jupyter/mlir_opt_kernel/__main__.py @@ -0,0 +1,7 @@ +# 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 ipykernel.kernelapp import IPKernelApp +from .kernel import MlirOptKernel +IPKernelApp.launch_instance(kernel_class=MlirOptKernel) diff --git a/mlir/utils/jupyter/mlir_opt_kernel/assets/kernel.js b/mlir/utils/jupyter/mlir_opt_kernel/assets/kernel.js new file mode 100644 --- /dev/null +++ b/mlir/utils/jupyter/mlir_opt_kernel/assets/kernel.js @@ -0,0 +1,9 @@ +define([ "codemirror/lib/codemirror", "base/js/namespace" ], + function(CodeMirror, IPython) { + "use strict"; + var onload = function() { + // TODO: Add syntax highlighting. + console.log("Loading kernel.js from MlirOptKernel"); + }; + return {onload : onload}; + }); diff --git a/mlir/utils/jupyter/mlir_opt_kernel/assets/kernel.json b/mlir/utils/jupyter/mlir_opt_kernel/assets/kernel.json new file mode 100644 --- /dev/null +++ b/mlir/utils/jupyter/mlir_opt_kernel/assets/kernel.json @@ -0,0 +1,15 @@ +{ + "argv": [ + "python3", "-m", "mlir_opt_kernel", "-f", "{connection_file}" + ], + "display_name": "MlirOpt", + "language": "mlir", + "codemirror_mode": "mlir", + "language_info": { + "name": "mlir", + "codemirror_mode": "mlir", + "mimetype": "text/x-mlir", + "file_extension": ".mlir", + "pygments_lexer": "text" + } +} diff --git a/mlir/utils/jupyter/mlir_opt_kernel/install.py b/mlir/utils/jupyter/mlir_opt_kernel/install.py new file mode 100644 --- /dev/null +++ b/mlir/utils/jupyter/mlir_opt_kernel/install.py @@ -0,0 +1,51 @@ +# 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 +import argparse + +from jupyter_client.kernelspec import KernelSpecManager + + +def install_my_kernel_spec(user=True, prefix=None): + """Install the kernel spec for user in given prefix.""" + print('Installing mlir-opt IPython kernel spec') + pkgroot = os.path.dirname(__file__) + KernelSpecManager().install_kernel_spec(os.path.join(pkgroot, 'assets'), + 'mlir', + user=user, + prefix=prefix) + + +def _is_root(): + """Returns whether the current user is root.""" + try: + return os.geteuid() == 0 + except AttributeError: + # Return false wherever unknown. + return False + + +def main(argv=None): + parser = argparse.ArgumentParser( + description='Install KernelSpec for MlirOpt Kernel') + prefix_locations = parser.add_mutually_exclusive_group() + + prefix_locations.add_argument('--user', + help='Install in user home directory', + action='store_true') + prefix_locations.add_argument('--prefix', + help='Install directory prefix', + default=None) + + args = parser.parse_args(argv) + + user = args.user or not _is_root() + prefix = args.prefix + + install_my_kernel_spec(user=user, prefix=prefix) + + +if __name__ == '__main__': + main() diff --git a/mlir/utils/jupyter/mlir_opt_kernel/kernel.py b/mlir/utils/jupyter/mlir_opt_kernel/kernel.py new file mode 100644 --- /dev/null +++ b/mlir/utils/jupyter/mlir_opt_kernel/kernel.py @@ -0,0 +1,197 @@ +# 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 subprocess import Popen +import os +import subprocess +import tempfile +import traceback +from ipykernel.kernelbase import Kernel + +__version__ = '0.0.1' + + +def _get_executable(): + """Find the mlir-opt executable.""" + + def is_exe(fpath): + """Returns whether executable file.""" + return os.path.isfile(fpath) and os.access(fpath, os.X_OK) + + program = os.environ.get('MLIR_OPT_EXECUTABLE', 'mlir-opt') + path, name = os.path.split(program) + # Attempt to get the executable + if path: + if is_exe(program): + return program + else: + for path in os.environ["PATH"].split(os.pathsep): + file = os.path.join(path, name) + if is_exe(file): + return file + raise OSError('mlir-opt not found, please see README') + + +class MlirOptKernel(Kernel): + """Kernel using mlir-opt inside jupyter. + + The reproducer syntax (`// configuration:`) is used to run passes. The + previous result can be referenced to by using `_` (this variable is reset + upon error). E.g., + + ```mlir + // configuration: --pass + func @foo(%tensor: tensor<2x3xf64>) -> tensor<3x2xf64> { ... } + ``` + + ```mlir + // configuration: --next-pass + _ + ``` + """ + + implementation = 'mlir' + implementation_version = __version__ + + language_version = __version__ + language = "mlir" + language_info = { + "name": "mlir", + "codemirror_mode": { + "name": "mlir" + }, + "mimetype": "text/x-mlir", + "file_extension": ".mlir", + "pygments_lexer": "text" + } + + @property + def banner(self): + """Returns kernel banner.""" + # Just a placeholder. + return "mlir-opt kernel %s" % __version__ + + def __init__(self, **kwargs): + Kernel.__init__(self, **kwargs) + self._ = None + self.executable = None + self.silent = False + + def get_executable(self): + """Returns the mlir-opt executable path.""" + if not self.executable: + self.executable = _get_executable() + return self.executable + + def process_output(self, output): + """Reports regular command output.""" + if not self.silent: + # Send standard output + stream_content = {'name': 'stdout', 'text': output} + self.send_response(self.iopub_socket, 'stream', stream_content) + + def process_error(self, output): + """Reports error response.""" + if not self.silent: + # Send standard error + stream_content = {'name': 'stderr', 'text': output} + self.send_response(self.iopub_socket, 'stream', stream_content) + + def do_execute(self, + code, + silent, + store_history=True, + user_expressions=None, + allow_stdin=False): + """Execute user code using mlir-opt binary.""" + + def ok_status(): + """Returns OK status.""" + return { + 'status': 'ok', + 'execution_count': self.execution_count, + 'payload': [], + 'user_expressions': {} + } + + def run(code): + """Run the code by pipeing via filesystem.""" + try: + inputmlir = tempfile.NamedTemporaryFile(delete=False) + command = [ + # Specify input and output file to error out if also + # set as arg. + self.get_executable(), + '--color', + inputmlir.name, + '-o', + '-' + ] + if code.startswith('// configuration:'): + command.append('--run-reproducer') + # Simple handling of repeating last line. + if code.endswith('\n_'): + if not self._: + raise NameError('No previous result set') + code = code[:-1] + self._ + inputmlir.write(code.encode("utf-8")) + inputmlir.close() + pipe = Popen(command, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE) + output, errors = pipe.communicate() + exitcode = pipe.returncode + finally: + os.unlink(inputmlir.name) + +# Replace temporary filename with placeholder. This takes the very +# remote chance where the full input filename (generated above) +# overlaps with something in the dump unrelated to the file. + fname = inputmlir.name.encode("utf-8") + output = output.replace(fname, b"<>") + errors = errors.replace(fname, b"<>") + return output, errors, exitcode + + self.silent = silent + if not code.strip(): + return ok_status() + + try: + output, errors, exitcode = run(code) + + if exitcode: + self._ = None + else: + self._ = output.decode("utf-8") + except KeyboardInterrupt: + return {'status': 'abort', 'execution_count': self.execution_count} + except Exception as error: + # Print traceback for local debugging. + traceback.print_exc() + self._ = None + exitcode = 255 + errors = repr(error).encode("utf-8") + + if exitcode: + content = {'ename': '', 'evalue': str(exitcode), 'traceback': []} + + self.send_response(self.iopub_socket, 'error', content) + self.process_error(errors.decode("utf-8")) + + content['execution_count'] = self.execution_count + content['status'] = 'error' + return content + + if not silent: + data = {} + data['text/x-mlir'] = self._ + content = { + 'execution_count': self.execution_count, + 'data': data, + 'metadata': {} + } + self.send_response(self.iopub_socket, 'execute_result', content) + self.process_output(self._) + self.process_error(errors.decode("utf-8")) + return ok_status()