[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

mlir/utils/jupyter/.gitignore

@@ -0,0 +1,5 @@
+__pycache__
+*.pyc
+build/
+dist/
+MANIFEST

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.

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__

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)

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};
+       });

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"
+    }
+}

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()

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()