[Mlir-commits] [mlir] 04c66ed - [mlir] Add simple jupyter kernel
Jacques Pienaar
llvmlistbot at llvm.org
Sun Feb 21 18:18:32 PST 2021
Author: Jacques Pienaar
Date: 2021-02-21T18:16:06-08:00
New Revision: 04c66edd3251975fd253e4c4b0e86ea56e8a7f46
URL: https://github.com/llvm/llvm-project/commit/04c66edd3251975fd253e4c4b0e86ea56e8a7f46
DIFF: https://github.com/llvm/llvm-project/commit/04c66edd3251975fd253e4c4b0e86ea56e8a7f46.diff
LOG: [mlir] Add simple jupyter kernel
Simple jupyter kernel using mlir-opt and reproducer to run passes.
Useful for local experimentation & generating examples. The export to
markdown from here is not immediately useful nor did I define a
CodeMirror synax to make the HTML output prettier. It only supports one
level of history (e.g., `_`) as I was mostly using with expanding a
pipeline one pass at a time and so was all I needed.
I placed this in utils directory next to editor & debugger utils.
Differential Revision: https://reviews.llvm.org/D95742
Added:
mlir/utils/jupyter/.gitignore
mlir/utils/jupyter/README.md
mlir/utils/jupyter/mlir_opt_kernel/__init__.py
mlir/utils/jupyter/mlir_opt_kernel/__main__.py
mlir/utils/jupyter/mlir_opt_kernel/assets/kernel.js
mlir/utils/jupyter/mlir_opt_kernel/assets/kernel.json
mlir/utils/jupyter/mlir_opt_kernel/install.py
mlir/utils/jupyter/mlir_opt_kernel/kernel.py
Modified:
Removed:
################################################################################
diff --git a/mlir/utils/jupyter/.gitignore b/mlir/utils/jupyter/.gitignore
new file mode 100644
index 000000000000..3b62e6397412
--- /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
index 000000000000..afed79fb77e6
--- /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
index 000000000000..7ed2bbe0ad49
--- /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
index 000000000000..02582f9b5bfb
--- /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
index 000000000000..55444f9cf568
--- /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
index 000000000000..7ae515b1f7d6
--- /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
index 000000000000..ddb37c87dfea
--- /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
index 000000000000..89947aa80e09
--- /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"<<input>>")
+ errors = errors.replace(fname, b"<<input>>")
+ 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()
More information about the Mlir-commits
mailing list