From d5eb51f6f351f079b757c90411a86e0777b88f9a Mon Sep 17 00:00:00 2001 From: Sasha Lopoukhine Date: Wed, 26 Jun 2024 23:46:13 +0100 Subject: [PATCH 01/11] docs: allow marimo notebooks as documentation --- .pre-commit-config.yaml | 2 +- Makefile | 9 ++++++- docs/marimo/IR_Chapter_0_Attributes.py | 33 ++++++++++++++++++++++++++ pyproject.toml | 3 +++ 4 files changed, 45 insertions(+), 2 deletions(-) create mode 100644 docs/marimo/IR_Chapter_0_Attributes.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8f593cbf8c..72ba463eef 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -19,4 +19,4 @@ repos: rev: 24.2.0 hooks: - id: black - exclude: "/(__pycache__|.asv|venv|build|tests/filecheck)/|versioneer.py$|xdsl/_version.py$" + exclude: "/(__pycache__|.asv|venv|build|tests/filecheck)/|versioneer.py$|xdsl/_version.py$|docs/marimo" diff --git a/Makefile b/Makefile index 4bc2b26b01..bb2f759ce9 100644 --- a/Makefile +++ b/Makefile @@ -51,8 +51,15 @@ pytest-toy: tests-toy: filecheck-toy pytest-toy +tests-marimo: + @for file in docs/marimo/*.py; do \ + echo "Running $$file"; \ + python3 "$$file" || exit 1; \ + done + @echo "All marimo tests passed successfully." + # run all tests -tests: pytest tests-toy filecheck pytest-nb pyright +tests: pytest tests-toy filecheck pytest-nb tests-marimo pyright @echo All tests done. # re-generate the output from all jupyter notebooks in the docs directory diff --git a/docs/marimo/IR_Chapter_0_Attributes.py b/docs/marimo/IR_Chapter_0_Attributes.py new file mode 100644 index 0000000000..8d7cba3ed7 --- /dev/null +++ b/docs/marimo/IR_Chapter_0_Attributes.py @@ -0,0 +1,33 @@ +import marimo + +__generated_with = "0.6.10" +app = marimo.App() + + +@app.cell +def __(): + import marimo as mo + return mo, + + +@app.cell +def __(mo): + from xdsl.dialects.builtin import IntAttr + + mo.md(f""" + # IR Chapter 1: Attributes + + In xDSL, Attributes carry information known at compile time. + For example, `{IntAttr.__name__}` is an attribute representing an integer. + """) + return IntAttr, + + +@app.cell +def __(): + raise Exception("a") + return + + +if __name__ == "__main__": + app.run() diff --git a/pyproject.toml b/pyproject.toml index 1fd8f26318..ca6b57c4b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ dev = [ "nbval<0.12", "filecheck<0.0.25", "lit<19.0.0", + "marimo==0.6.10", "pre-commit==3.7.1", "ruff==0.4.10", "asv<0.7", @@ -81,6 +82,7 @@ typeCheckingMode = "strict" "xdsl/interpreters/onnx.py", ] "ignore" = [ + "docs/marimo", "tests/filecheck/frontend/dialects/builtin.py", "tests/filecheck/frontend/dialects/invalid.py", "tests/filecheck/frontend/dialects/arith.py", @@ -136,6 +138,7 @@ exclude = """ /(__pycache__|.asv|venv|build|tests/filecheck)/ |versioneer.py$ |xdsl/_version.py$ +|docs/marimo """ [tool.pytest.ini_options] From 70a87a873e4e62461c2a19eddd221f6d30d4be94 Mon Sep 17 00:00:00 2001 From: Sasha Lopoukhine Date: Wed, 26 Jun 2024 23:53:58 +0100 Subject: [PATCH 02/11] add marimo to CI --- .github/workflows/ci-core.yml | 5 +++++ .github/workflows/ci-mlir.yml | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/.github/workflows/ci-core.yml b/.github/workflows/ci-core.yml index 8bf9da0348..eae3999e2f 100644 --- a/.github/workflows/ci-core.yml +++ b/.github/workflows/ci-core.yml @@ -49,3 +49,8 @@ jobs: export PYTHONPATH=$(pwd) lit -v tests/filecheck/ lit -v docs/Toy/examples/ + + - name: Test marimo notebooks + run: | + cd xdsl + make tests-marimo diff --git a/.github/workflows/ci-mlir.yml b/.github/workflows/ci-mlir.yml index 2d93edcf8c..b068d4bb72 100644 --- a/.github/workflows/ci-mlir.yml +++ b/.github/workflows/ci-mlir.yml @@ -112,6 +112,11 @@ jobs: export PATH=$PATH:${GITHUB_WORKSPACE}/llvm-project/build/bin/ pytest --nbval docs/mlir_interoperation.ipynb --maxfail 1 -vv + - name: Test marimo notebooks + run: | + cd xdsl + make tests-marimo + - name: Combine coverage data run: | cd xdsl From 54b971084b75f0aacf1cf2b9bac763965085886d Mon Sep 17 00:00:00 2001 From: Sasha Lopoukhine Date: Wed, 26 Jun 2024 23:54:27 +0100 Subject: [PATCH 03/11] Revert "add marimo to CI" This reverts commit 70a87a873e4e62461c2a19eddd221f6d30d4be94. --- .github/workflows/ci-core.yml | 5 ----- .github/workflows/ci-mlir.yml | 5 ----- 2 files changed, 10 deletions(-) diff --git a/.github/workflows/ci-core.yml b/.github/workflows/ci-core.yml index eae3999e2f..8bf9da0348 100644 --- a/.github/workflows/ci-core.yml +++ b/.github/workflows/ci-core.yml @@ -49,8 +49,3 @@ jobs: export PYTHONPATH=$(pwd) lit -v tests/filecheck/ lit -v docs/Toy/examples/ - - - name: Test marimo notebooks - run: | - cd xdsl - make tests-marimo diff --git a/.github/workflows/ci-mlir.yml b/.github/workflows/ci-mlir.yml index b068d4bb72..2d93edcf8c 100644 --- a/.github/workflows/ci-mlir.yml +++ b/.github/workflows/ci-mlir.yml @@ -112,11 +112,6 @@ jobs: export PATH=$PATH:${GITHUB_WORKSPACE}/llvm-project/build/bin/ pytest --nbval docs/mlir_interoperation.ipynb --maxfail 1 -vv - - name: Test marimo notebooks - run: | - cd xdsl - make tests-marimo - - name: Combine coverage data run: | cd xdsl From c1ed347885c7df060a4bd7cddf13af1c8746a2b6 Mon Sep 17 00:00:00 2001 From: Sasha Lopoukhine Date: Wed, 26 Jun 2024 23:54:33 +0100 Subject: [PATCH 04/11] use correct CI --- .github/workflows/ci-notebooks.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/ci-notebooks.yml b/.github/workflows/ci-notebooks.yml index 879e9e0ad5..d8abff1146 100644 --- a/.github/workflows/ci-notebooks.yml +++ b/.github/workflows/ci-notebooks.yml @@ -34,3 +34,7 @@ jobs: run: | # mlir_interoperation.ipynb is dependent on MLIR, and is tested in the MLIR-enabled workflow. pytest -W error --nbval -vv docs --ignore=docs/mlir_interoperation.ipynb + - name: Test marimo notebooks + run: | + cd xdsl + make tests-marimo From 4798ce0c9ff761f168aeb8ec8601b58f0f3454b8 Mon Sep 17 00:00:00 2001 From: Sasha Lopoukhine Date: Wed, 26 Jun 2024 23:56:42 +0100 Subject: [PATCH 05/11] try again --- .github/workflows/ci-notebooks.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/ci-notebooks.yml b/.github/workflows/ci-notebooks.yml index d8abff1146..1386ac9da5 100644 --- a/.github/workflows/ci-notebooks.yml +++ b/.github/workflows/ci-notebooks.yml @@ -36,5 +36,4 @@ jobs: pytest -W error --nbval -vv docs --ignore=docs/mlir_interoperation.ipynb - name: Test marimo notebooks run: | - cd xdsl make tests-marimo From d5b5864d7d2b062aead144ab3bbc8721c011a67a Mon Sep 17 00:00:00 2001 From: Sasha Lopoukhine Date: Wed, 26 Jun 2024 23:58:55 +0100 Subject: [PATCH 06/11] remove exception --- docs/marimo/IR_Chapter_0_Attributes.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/docs/marimo/IR_Chapter_0_Attributes.py b/docs/marimo/IR_Chapter_0_Attributes.py index 8d7cba3ed7..38b88036be 100644 --- a/docs/marimo/IR_Chapter_0_Attributes.py +++ b/docs/marimo/IR_Chapter_0_Attributes.py @@ -23,11 +23,5 @@ def __(mo): return IntAttr, -@app.cell -def __(): - raise Exception("a") - return - - if __name__ == "__main__": app.run() From b6d60888b95a1f93ec749869ded048f66d534dd9 Mon Sep 17 00:00:00 2001 From: Sasha Lopoukhine Date: Thu, 27 Jun 2024 00:01:41 +0100 Subject: [PATCH 07/11] rename file --- .../{IR_Chapter_0_Attributes.py => IR_Chapter_1_Attributes.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename docs/marimo/{IR_Chapter_0_Attributes.py => IR_Chapter_1_Attributes.py} (100%) diff --git a/docs/marimo/IR_Chapter_0_Attributes.py b/docs/marimo/IR_Chapter_1_Attributes.py similarity index 100% rename from docs/marimo/IR_Chapter_0_Attributes.py rename to docs/marimo/IR_Chapter_1_Attributes.py From b3c576869230479b635de2a9c10ed35c639f4afe Mon Sep 17 00:00:00 2001 From: Sasha Lopoukhine Date: Thu, 27 Jun 2024 09:03:05 +0100 Subject: [PATCH 08/11] swap notebook for ONNX --- docs/marimo/IR_Chapter_1_Attributes.py | 27 -- docs/marimo/README.md | 7 + docs/marimo/onnx_demo.py | 421 +++++++++++++++++++++++++ pyproject.toml | 1 + 4 files changed, 429 insertions(+), 27 deletions(-) delete mode 100644 docs/marimo/IR_Chapter_1_Attributes.py create mode 100644 docs/marimo/README.md create mode 100644 docs/marimo/onnx_demo.py diff --git a/docs/marimo/IR_Chapter_1_Attributes.py b/docs/marimo/IR_Chapter_1_Attributes.py deleted file mode 100644 index 38b88036be..0000000000 --- a/docs/marimo/IR_Chapter_1_Attributes.py +++ /dev/null @@ -1,27 +0,0 @@ -import marimo - -__generated_with = "0.6.10" -app = marimo.App() - - -@app.cell -def __(): - import marimo as mo - return mo, - - -@app.cell -def __(mo): - from xdsl.dialects.builtin import IntAttr - - mo.md(f""" - # IR Chapter 1: Attributes - - In xDSL, Attributes carry information known at compile time. - For example, `{IntAttr.__name__}` is an attribute representing an integer. - """) - return IntAttr, - - -if __name__ == "__main__": - app.run() diff --git a/docs/marimo/README.md b/docs/marimo/README.md new file mode 100644 index 0000000000..9c9aec6f4f --- /dev/null +++ b/docs/marimo/README.md @@ -0,0 +1,7 @@ +# Marimo Notebooks + +This folder is for Python notebooks created with the [Marimo](https://marimo.io/) tool. + +These can be run as apps with `marimo run path/to/notebook.py`, as interactive development environments with `marimo edit path/to/notebook.py`, or as headless scripts with `python path/to/notebook.py`. + +With the Marimo VSCode extension, they can also be launched within the IDE in either "run" or "edit" modes. diff --git a/docs/marimo/onnx_demo.py b/docs/marimo/onnx_demo.py new file mode 100644 index 0000000000..d2253a39e2 --- /dev/null +++ b/docs/marimo/onnx_demo.py @@ -0,0 +1,421 @@ +import marimo + +__generated_with = "0.6.10" +app = marimo.App() + + +@app.cell +def __(): + import marimo as mo + + mo.md( + """ + # ONNX to Snitch + + This notebook uses Marimo, a Jupyter-like notebook with interactive UI elements and reactive state. + """ + ) + return mo, + + +@app.cell +def __(mo): + rank = mo.ui.slider(1, 4, value=2, label="Rank") + + mo.md( + f""" + For example, here is a slider, which can take on values from 1 to 4. + + {rank} + """ + ) + return rank, + + +@app.cell +def __(mo, rank): + shape = list(range(2, 2 + rank.value)) + + mo.md( + f""" + We use the slider to determine the shape of our inputs and outputs: + + {shape} + """ + ) + return shape, + + +@app.cell +def __(mo, shape): + import onnx + from onnx import AttributeProto, GraphProto, TensorProto, ValueInfoProto, helper + + # Create one input (ValueInfoProto) + X1 = helper.make_tensor_value_info("X1", TensorProto.DOUBLE, shape) + X2 = helper.make_tensor_value_info("X2", TensorProto.DOUBLE, shape) + + # Create one output (ValueInfoProto) + Y = helper.make_tensor_value_info("Y", TensorProto.DOUBLE, shape) + + # Create a node (NodeProto) - This is based on Pad-11 + node_def = helper.make_node( + "Sub", # node name + ["X1", "X2"], # inputs + ["Y"], # outputs + ) + + # Create the graph (GraphProto) + graph_def = helper.make_graph( + [node_def], + "main_graph", + [X1, X2], + [Y], + ) + + # Set opset version to 18 + opset_import = [helper.make_operatorsetid("", 18)] + + # Create the model (ModelProto) without using helper.make_model + model_def = helper.make_model( + graph_def, producer_name="onnx-example", opset_imports=opset_import + ) + + print(f"The model is:\n{model_def}") + onnx.checker.check_model(model_def) + # onnx.save(model_def, "add.onnx") + print("The model is checked!") + + mo.md( + f""" + ### The ONNX model + + We use the ONNX API to build a simple function, one that returns the elementwise sum of two arrays of shape {shape} + """ + ) + return ( + AttributeProto, + GraphProto, + TensorProto, + ValueInfoProto, + X1, + X2, + Y, + graph_def, + helper, + model_def, + node_def, + onnx, + opset_import, + ) + + +@app.cell +def __(mo): + from xdsl.ir import Attribute, SSAValue + + mo.md( + """ + We then convert the ONNX Graph to the xDSL representation, in the onnx dialect. + """ + ) + return Attribute, SSAValue + + +@app.cell +def __(mo, model_def): + from xdsl.frontend.onnx.ir_builder import build_module + + init_module = build_module(model_def.graph).clone() + + print(init_module) + + mo.md( + """ + Here is the same function, it takes two `tensor` values of our chosen shape, passes them as operands to the `onnx.Add` operation, and returns it. + """ + ) + return build_module, init_module + + +@app.cell +def __(init_module, mo): + from xdsl.context import MLContext + from xdsl.tools.command_line_tool import get_all_dialects + from xdsl.transforms.convert_onnx_to_linalg import ConvertOnnxToLinalgPass + + ctx = MLContext() + + for dialect_name, dialect_factory in get_all_dialects().items(): + ctx.register_dialect(dialect_name, dialect_factory) + + linalg_module = init_module.clone() + + ConvertOnnxToLinalgPass().apply(ctx, linalg_module) + + print(linalg_module) + + mo.md( + """ + We can use a pass implemented in xDSL to convert the ONNX operations to builtin operations, here we can use the `tensor.empty` op to create our output buffer, and `linalg.add` to represent the addition in destination-passing style. + """ + ) + return ( + ConvertOnnxToLinalgPass, + MLContext, + ctx, + dialect_factory, + dialect_name, + get_all_dialects, + linalg_module, + ) + + +@app.cell +def __(ctx, linalg_module, mo): + from xdsl.transforms.mlir_opt import MLIROptPass + + generalized_module = linalg_module.clone() + + MLIROptPass(generic=False, arguments=["--linalg-generalize-named-ops"]).apply( + ctx, generalized_module + ) + + print(generalized_module) + + mo.md( + """ + We can also call into MLIR, here to convert `linalg.add` to `linalg.generic`, a representation of Einstein summation. + """ + ) + return MLIROptPass, generalized_module + + +@app.cell +def __(MLIROptPass, ctx, generalized_module, mo): + bufferized_module = generalized_module.clone() + + MLIROptPass( + arguments=[ + "--empty-tensor-to-alloc-tensor", + "--one-shot-bufferize=bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map", + ] + ).apply(ctx, bufferized_module) + + print(bufferized_module) + + mo.md( + """ + We then use MLIR to bufferize our function. + """ + ) + return bufferized_module, + + +@app.cell +def __(MLIROptPass, bufferized_module, ctx): + scf_module = bufferized_module.clone() + + MLIROptPass( + arguments=["--convert-linalg-to-loops", "--lower-affine", "--canonicalize"] + ).apply(ctx, scf_module) + + print(scf_module) + return scf_module, + + +@app.cell +def __(ctx, scf_module): + from xdsl.backend.riscv.lowering import ( + convert_arith_to_riscv, + convert_func_to_riscv_func, + convert_memref_to_riscv, + convert_scf_to_riscv_scf, + ) + from xdsl.passes import PipelinePass + from xdsl.transforms import reconcile_unrealized_casts + + riscv_module = scf_module.clone() + + lower_to_riscv = PipelinePass( + [ + convert_func_to_riscv_func.ConvertFuncToRiscvFuncPass(), + convert_memref_to_riscv.ConvertMemrefToRiscvPass(), + convert_arith_to_riscv.ConvertArithToRiscvPass(), + convert_scf_to_riscv_scf.ConvertScfToRiscvPass(), + reconcile_unrealized_casts.ReconcileUnrealizedCastsPass(), + ] + ).apply(ctx, riscv_module) + + print(riscv_module) + return ( + PipelinePass, + convert_arith_to_riscv, + convert_func_to_riscv_func, + convert_memref_to_riscv, + convert_scf_to_riscv_scf, + lower_to_riscv, + reconcile_unrealized_casts, + riscv_module, + ) + + +@app.cell +def __(PipelinePass, ctx, riscv_module): + from xdsl.backend.riscv.lowering.convert_snitch_stream_to_snitch import ( + ConvertSnitchStreamToSnitch, + ) + from xdsl.transforms.canonicalize import CanonicalizePass + from xdsl.transforms.lower_snitch import LowerSnitchPass + from xdsl.transforms.riscv_register_allocation import RISCVRegisterAllocation + from xdsl.transforms.riscv_scf_loop_range_folding import ( + RiscvScfLoopRangeFoldingPass, + ) + from xdsl.transforms.snitch_register_allocation import SnitchRegisterAllocation + + regalloc_module = riscv_module.clone() + + PipelinePass( + [ + RISCVRegisterAllocation(), + CanonicalizePass(), + ] + ).apply(ctx, regalloc_module) + + print(regalloc_module) + return ( + CanonicalizePass, + ConvertSnitchStreamToSnitch, + LowerSnitchPass, + RISCVRegisterAllocation, + RiscvScfLoopRangeFoldingPass, + SnitchRegisterAllocation, + regalloc_module, + ) + + +@app.cell +def __(CanonicalizePass, ctx, regalloc_module): + from xdsl.backend.riscv.lowering.convert_riscv_scf_to_riscv_cf import ( + ConvertRiscvScfToRiscvCfPass, + ) + from xdsl.dialects.riscv import riscv_code + + assembly_module = regalloc_module.clone() + + ConvertRiscvScfToRiscvCfPass().apply(ctx, assembly_module) + CanonicalizePass().apply(ctx, assembly_module) + + print(assembly_module) + return ConvertRiscvScfToRiscvCfPass, assembly_module, riscv_code + + +@app.cell +def __(assembly_module, mo, riscv_code): + assembly = riscv_code(assembly_module) + + print(assembly) + + mo.md( + """ + This representation of the program in xDSL corresponds ~1:1 to RISC-V assembly, and we can use a helper function to print that out. + """ + ) + return assembly, + + +@app.cell +def __( + CanonicalizePass, + PipelinePass, + bufferized_module, + convert_arith_to_riscv, + convert_func_to_riscv_func, + convert_memref_to_riscv, + convert_scf_to_riscv_scf, + ctx, + mo, + reconcile_unrealized_casts, +): + from xdsl.transforms import ( + arith_add_fastmath, + convert_linalg_to_memref_stream, + convert_memref_stream_to_loops, + convert_memref_stream_to_snitch_stream, + convert_riscv_scf_for_to_frep, + dead_code_elimination, + loop_hoist_memref, + lower_affine, + memref_streamify, + ) + + snitch_stream_module = bufferized_module.clone() + + pass_pipeline = PipelinePass( + [ + convert_linalg_to_memref_stream.ConvertLinalgToMemrefStreamPass(), + memref_streamify.MemrefStreamifyPass(), + convert_memref_stream_to_loops.ConvertMemrefStreamToLoopsPass(), + convert_memref_stream_to_snitch_stream.ConvertMemrefStreamToSnitch(), + arith_add_fastmath.AddArithFastMathFlagsPass(), + loop_hoist_memref.LoopHoistMemrefPass(), + lower_affine.LowerAffinePass(), + convert_func_to_riscv_func.ConvertFuncToRiscvFuncPass(), + convert_memref_to_riscv.ConvertMemrefToRiscvPass(), + convert_arith_to_riscv.ConvertArithToRiscvPass(), + CanonicalizePass(), + convert_scf_to_riscv_scf.ConvertScfToRiscvPass(), + dead_code_elimination.DeadCodeElimination(), + reconcile_unrealized_casts.ReconcileUnrealizedCastsPass(), + convert_riscv_scf_for_to_frep.ConvertRiscvScfForToFrepPass(), + ] + ) + + pass_pipeline.apply(ctx, snitch_stream_module) + + print(snitch_stream_module) + + mo.md( + """ + ### Compiling to Snitch + + xDSL is also capable of targeting Snitch, and making use of its streaming registers and fixed-repetition loop. We use a different lowering flow from the linalg.generic representation to represent a high-level, structured, but Snitch-specific representation of the code: + """ + ) + return ( + arith_add_fastmath, + convert_linalg_to_memref_stream, + convert_memref_stream_to_loops, + convert_memref_stream_to_snitch_stream, + convert_riscv_scf_for_to_frep, + dead_code_elimination, + loop_hoist_memref, + lower_affine, + memref_streamify, + pass_pipeline, + snitch_stream_module, + ) + + +@app.cell +def __(ctx, mo, riscv_code, snitch_stream_module): + from xdsl.transforms import test_lower_snitch_stream_to_asm + + snitch_asm_module = snitch_stream_module.clone() + + test_lower_snitch_stream_to_asm.TestLowerSnitchStreamToAsm().apply( + ctx, snitch_asm_module + ) + + print(riscv_code(snitch_asm_module)) + + mo.md( + """ + We can then lower this to assembly that includes assembly instructions from the Snitch-extended ISA: + """ + ) + return snitch_asm_module, test_lower_snitch_stream_to_asm + + +if __name__ == "__main__": + app.run() diff --git a/pyproject.toml b/pyproject.toml index ca6b57c4b1..29e2d623dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -128,6 +128,7 @@ target-version = "py310" "tests/test_declarative_assembly_format.py" = ["F811"] "versioneer.py" = ["ALL"] "_version.py" = ["ALL"] +"**/{docs/marimo}/*" = ["E501"] [tool.ruff.mccabe] max-complexity = 10 From 466709a3c8105e92fb8eb594f72d4312774595bd Mon Sep 17 00:00:00 2001 From: Sasha Lopoukhine Date: Thu, 27 Jun 2024 09:13:04 +0100 Subject: [PATCH 09/11] move marimo test to MLIR one --- .github/workflows/ci-mlir.yml | 4 ++++ .github/workflows/ci-notebooks.yml | 3 --- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci-mlir.yml b/.github/workflows/ci-mlir.yml index 2d93edcf8c..4c3115e9fa 100644 --- a/.github/workflows/ci-mlir.yml +++ b/.github/workflows/ci-mlir.yml @@ -112,6 +112,10 @@ jobs: export PATH=$PATH:${GITHUB_WORKSPACE}/llvm-project/build/bin/ pytest --nbval docs/mlir_interoperation.ipynb --maxfail 1 -vv + - name: Test marimo notebooks + run: | + make tests-marimo + - name: Combine coverage data run: | cd xdsl diff --git a/.github/workflows/ci-notebooks.yml b/.github/workflows/ci-notebooks.yml index 1386ac9da5..879e9e0ad5 100644 --- a/.github/workflows/ci-notebooks.yml +++ b/.github/workflows/ci-notebooks.yml @@ -34,6 +34,3 @@ jobs: run: | # mlir_interoperation.ipynb is dependent on MLIR, and is tested in the MLIR-enabled workflow. pytest -W error --nbval -vv docs --ignore=docs/mlir_interoperation.ipynb - - name: Test marimo notebooks - run: | - make tests-marimo From 09b488fd444aad878e927998cedd895d78e6f78b Mon Sep 17 00:00:00 2001 From: Sasha Lopoukhine Date: Thu, 27 Jun 2024 09:24:15 +0100 Subject: [PATCH 10/11] Revert "swap notebook for ONNX" This reverts commit b3c576869230479b635de2a9c10ed35c639f4afe. --- .github/workflows/ci-mlir.yml | 4 - .github/workflows/ci-notebooks.yml | 3 + docs/marimo/IR_Chapter_1_Attributes.py | 27 ++ docs/marimo/README.md | 7 - docs/marimo/onnx_demo.py | 421 ------------------------- pyproject.toml | 1 - 6 files changed, 30 insertions(+), 433 deletions(-) create mode 100644 docs/marimo/IR_Chapter_1_Attributes.py delete mode 100644 docs/marimo/README.md delete mode 100644 docs/marimo/onnx_demo.py diff --git a/.github/workflows/ci-mlir.yml b/.github/workflows/ci-mlir.yml index 4c3115e9fa..2d93edcf8c 100644 --- a/.github/workflows/ci-mlir.yml +++ b/.github/workflows/ci-mlir.yml @@ -112,10 +112,6 @@ jobs: export PATH=$PATH:${GITHUB_WORKSPACE}/llvm-project/build/bin/ pytest --nbval docs/mlir_interoperation.ipynb --maxfail 1 -vv - - name: Test marimo notebooks - run: | - make tests-marimo - - name: Combine coverage data run: | cd xdsl diff --git a/.github/workflows/ci-notebooks.yml b/.github/workflows/ci-notebooks.yml index 879e9e0ad5..1386ac9da5 100644 --- a/.github/workflows/ci-notebooks.yml +++ b/.github/workflows/ci-notebooks.yml @@ -34,3 +34,6 @@ jobs: run: | # mlir_interoperation.ipynb is dependent on MLIR, and is tested in the MLIR-enabled workflow. pytest -W error --nbval -vv docs --ignore=docs/mlir_interoperation.ipynb + - name: Test marimo notebooks + run: | + make tests-marimo diff --git a/docs/marimo/IR_Chapter_1_Attributes.py b/docs/marimo/IR_Chapter_1_Attributes.py new file mode 100644 index 0000000000..38b88036be --- /dev/null +++ b/docs/marimo/IR_Chapter_1_Attributes.py @@ -0,0 +1,27 @@ +import marimo + +__generated_with = "0.6.10" +app = marimo.App() + + +@app.cell +def __(): + import marimo as mo + return mo, + + +@app.cell +def __(mo): + from xdsl.dialects.builtin import IntAttr + + mo.md(f""" + # IR Chapter 1: Attributes + + In xDSL, Attributes carry information known at compile time. + For example, `{IntAttr.__name__}` is an attribute representing an integer. + """) + return IntAttr, + + +if __name__ == "__main__": + app.run() diff --git a/docs/marimo/README.md b/docs/marimo/README.md deleted file mode 100644 index 9c9aec6f4f..0000000000 --- a/docs/marimo/README.md +++ /dev/null @@ -1,7 +0,0 @@ -# Marimo Notebooks - -This folder is for Python notebooks created with the [Marimo](https://marimo.io/) tool. - -These can be run as apps with `marimo run path/to/notebook.py`, as interactive development environments with `marimo edit path/to/notebook.py`, or as headless scripts with `python path/to/notebook.py`. - -With the Marimo VSCode extension, they can also be launched within the IDE in either "run" or "edit" modes. diff --git a/docs/marimo/onnx_demo.py b/docs/marimo/onnx_demo.py deleted file mode 100644 index d2253a39e2..0000000000 --- a/docs/marimo/onnx_demo.py +++ /dev/null @@ -1,421 +0,0 @@ -import marimo - -__generated_with = "0.6.10" -app = marimo.App() - - -@app.cell -def __(): - import marimo as mo - - mo.md( - """ - # ONNX to Snitch - - This notebook uses Marimo, a Jupyter-like notebook with interactive UI elements and reactive state. - """ - ) - return mo, - - -@app.cell -def __(mo): - rank = mo.ui.slider(1, 4, value=2, label="Rank") - - mo.md( - f""" - For example, here is a slider, which can take on values from 1 to 4. - - {rank} - """ - ) - return rank, - - -@app.cell -def __(mo, rank): - shape = list(range(2, 2 + rank.value)) - - mo.md( - f""" - We use the slider to determine the shape of our inputs and outputs: - - {shape} - """ - ) - return shape, - - -@app.cell -def __(mo, shape): - import onnx - from onnx import AttributeProto, GraphProto, TensorProto, ValueInfoProto, helper - - # Create one input (ValueInfoProto) - X1 = helper.make_tensor_value_info("X1", TensorProto.DOUBLE, shape) - X2 = helper.make_tensor_value_info("X2", TensorProto.DOUBLE, shape) - - # Create one output (ValueInfoProto) - Y = helper.make_tensor_value_info("Y", TensorProto.DOUBLE, shape) - - # Create a node (NodeProto) - This is based on Pad-11 - node_def = helper.make_node( - "Sub", # node name - ["X1", "X2"], # inputs - ["Y"], # outputs - ) - - # Create the graph (GraphProto) - graph_def = helper.make_graph( - [node_def], - "main_graph", - [X1, X2], - [Y], - ) - - # Set opset version to 18 - opset_import = [helper.make_operatorsetid("", 18)] - - # Create the model (ModelProto) without using helper.make_model - model_def = helper.make_model( - graph_def, producer_name="onnx-example", opset_imports=opset_import - ) - - print(f"The model is:\n{model_def}") - onnx.checker.check_model(model_def) - # onnx.save(model_def, "add.onnx") - print("The model is checked!") - - mo.md( - f""" - ### The ONNX model - - We use the ONNX API to build a simple function, one that returns the elementwise sum of two arrays of shape {shape} - """ - ) - return ( - AttributeProto, - GraphProto, - TensorProto, - ValueInfoProto, - X1, - X2, - Y, - graph_def, - helper, - model_def, - node_def, - onnx, - opset_import, - ) - - -@app.cell -def __(mo): - from xdsl.ir import Attribute, SSAValue - - mo.md( - """ - We then convert the ONNX Graph to the xDSL representation, in the onnx dialect. - """ - ) - return Attribute, SSAValue - - -@app.cell -def __(mo, model_def): - from xdsl.frontend.onnx.ir_builder import build_module - - init_module = build_module(model_def.graph).clone() - - print(init_module) - - mo.md( - """ - Here is the same function, it takes two `tensor` values of our chosen shape, passes them as operands to the `onnx.Add` operation, and returns it. - """ - ) - return build_module, init_module - - -@app.cell -def __(init_module, mo): - from xdsl.context import MLContext - from xdsl.tools.command_line_tool import get_all_dialects - from xdsl.transforms.convert_onnx_to_linalg import ConvertOnnxToLinalgPass - - ctx = MLContext() - - for dialect_name, dialect_factory in get_all_dialects().items(): - ctx.register_dialect(dialect_name, dialect_factory) - - linalg_module = init_module.clone() - - ConvertOnnxToLinalgPass().apply(ctx, linalg_module) - - print(linalg_module) - - mo.md( - """ - We can use a pass implemented in xDSL to convert the ONNX operations to builtin operations, here we can use the `tensor.empty` op to create our output buffer, and `linalg.add` to represent the addition in destination-passing style. - """ - ) - return ( - ConvertOnnxToLinalgPass, - MLContext, - ctx, - dialect_factory, - dialect_name, - get_all_dialects, - linalg_module, - ) - - -@app.cell -def __(ctx, linalg_module, mo): - from xdsl.transforms.mlir_opt import MLIROptPass - - generalized_module = linalg_module.clone() - - MLIROptPass(generic=False, arguments=["--linalg-generalize-named-ops"]).apply( - ctx, generalized_module - ) - - print(generalized_module) - - mo.md( - """ - We can also call into MLIR, here to convert `linalg.add` to `linalg.generic`, a representation of Einstein summation. - """ - ) - return MLIROptPass, generalized_module - - -@app.cell -def __(MLIROptPass, ctx, generalized_module, mo): - bufferized_module = generalized_module.clone() - - MLIROptPass( - arguments=[ - "--empty-tensor-to-alloc-tensor", - "--one-shot-bufferize=bufferize-function-boundaries function-boundary-type-conversion=identity-layout-map", - ] - ).apply(ctx, bufferized_module) - - print(bufferized_module) - - mo.md( - """ - We then use MLIR to bufferize our function. - """ - ) - return bufferized_module, - - -@app.cell -def __(MLIROptPass, bufferized_module, ctx): - scf_module = bufferized_module.clone() - - MLIROptPass( - arguments=["--convert-linalg-to-loops", "--lower-affine", "--canonicalize"] - ).apply(ctx, scf_module) - - print(scf_module) - return scf_module, - - -@app.cell -def __(ctx, scf_module): - from xdsl.backend.riscv.lowering import ( - convert_arith_to_riscv, - convert_func_to_riscv_func, - convert_memref_to_riscv, - convert_scf_to_riscv_scf, - ) - from xdsl.passes import PipelinePass - from xdsl.transforms import reconcile_unrealized_casts - - riscv_module = scf_module.clone() - - lower_to_riscv = PipelinePass( - [ - convert_func_to_riscv_func.ConvertFuncToRiscvFuncPass(), - convert_memref_to_riscv.ConvertMemrefToRiscvPass(), - convert_arith_to_riscv.ConvertArithToRiscvPass(), - convert_scf_to_riscv_scf.ConvertScfToRiscvPass(), - reconcile_unrealized_casts.ReconcileUnrealizedCastsPass(), - ] - ).apply(ctx, riscv_module) - - print(riscv_module) - return ( - PipelinePass, - convert_arith_to_riscv, - convert_func_to_riscv_func, - convert_memref_to_riscv, - convert_scf_to_riscv_scf, - lower_to_riscv, - reconcile_unrealized_casts, - riscv_module, - ) - - -@app.cell -def __(PipelinePass, ctx, riscv_module): - from xdsl.backend.riscv.lowering.convert_snitch_stream_to_snitch import ( - ConvertSnitchStreamToSnitch, - ) - from xdsl.transforms.canonicalize import CanonicalizePass - from xdsl.transforms.lower_snitch import LowerSnitchPass - from xdsl.transforms.riscv_register_allocation import RISCVRegisterAllocation - from xdsl.transforms.riscv_scf_loop_range_folding import ( - RiscvScfLoopRangeFoldingPass, - ) - from xdsl.transforms.snitch_register_allocation import SnitchRegisterAllocation - - regalloc_module = riscv_module.clone() - - PipelinePass( - [ - RISCVRegisterAllocation(), - CanonicalizePass(), - ] - ).apply(ctx, regalloc_module) - - print(regalloc_module) - return ( - CanonicalizePass, - ConvertSnitchStreamToSnitch, - LowerSnitchPass, - RISCVRegisterAllocation, - RiscvScfLoopRangeFoldingPass, - SnitchRegisterAllocation, - regalloc_module, - ) - - -@app.cell -def __(CanonicalizePass, ctx, regalloc_module): - from xdsl.backend.riscv.lowering.convert_riscv_scf_to_riscv_cf import ( - ConvertRiscvScfToRiscvCfPass, - ) - from xdsl.dialects.riscv import riscv_code - - assembly_module = regalloc_module.clone() - - ConvertRiscvScfToRiscvCfPass().apply(ctx, assembly_module) - CanonicalizePass().apply(ctx, assembly_module) - - print(assembly_module) - return ConvertRiscvScfToRiscvCfPass, assembly_module, riscv_code - - -@app.cell -def __(assembly_module, mo, riscv_code): - assembly = riscv_code(assembly_module) - - print(assembly) - - mo.md( - """ - This representation of the program in xDSL corresponds ~1:1 to RISC-V assembly, and we can use a helper function to print that out. - """ - ) - return assembly, - - -@app.cell -def __( - CanonicalizePass, - PipelinePass, - bufferized_module, - convert_arith_to_riscv, - convert_func_to_riscv_func, - convert_memref_to_riscv, - convert_scf_to_riscv_scf, - ctx, - mo, - reconcile_unrealized_casts, -): - from xdsl.transforms import ( - arith_add_fastmath, - convert_linalg_to_memref_stream, - convert_memref_stream_to_loops, - convert_memref_stream_to_snitch_stream, - convert_riscv_scf_for_to_frep, - dead_code_elimination, - loop_hoist_memref, - lower_affine, - memref_streamify, - ) - - snitch_stream_module = bufferized_module.clone() - - pass_pipeline = PipelinePass( - [ - convert_linalg_to_memref_stream.ConvertLinalgToMemrefStreamPass(), - memref_streamify.MemrefStreamifyPass(), - convert_memref_stream_to_loops.ConvertMemrefStreamToLoopsPass(), - convert_memref_stream_to_snitch_stream.ConvertMemrefStreamToSnitch(), - arith_add_fastmath.AddArithFastMathFlagsPass(), - loop_hoist_memref.LoopHoistMemrefPass(), - lower_affine.LowerAffinePass(), - convert_func_to_riscv_func.ConvertFuncToRiscvFuncPass(), - convert_memref_to_riscv.ConvertMemrefToRiscvPass(), - convert_arith_to_riscv.ConvertArithToRiscvPass(), - CanonicalizePass(), - convert_scf_to_riscv_scf.ConvertScfToRiscvPass(), - dead_code_elimination.DeadCodeElimination(), - reconcile_unrealized_casts.ReconcileUnrealizedCastsPass(), - convert_riscv_scf_for_to_frep.ConvertRiscvScfForToFrepPass(), - ] - ) - - pass_pipeline.apply(ctx, snitch_stream_module) - - print(snitch_stream_module) - - mo.md( - """ - ### Compiling to Snitch - - xDSL is also capable of targeting Snitch, and making use of its streaming registers and fixed-repetition loop. We use a different lowering flow from the linalg.generic representation to represent a high-level, structured, but Snitch-specific representation of the code: - """ - ) - return ( - arith_add_fastmath, - convert_linalg_to_memref_stream, - convert_memref_stream_to_loops, - convert_memref_stream_to_snitch_stream, - convert_riscv_scf_for_to_frep, - dead_code_elimination, - loop_hoist_memref, - lower_affine, - memref_streamify, - pass_pipeline, - snitch_stream_module, - ) - - -@app.cell -def __(ctx, mo, riscv_code, snitch_stream_module): - from xdsl.transforms import test_lower_snitch_stream_to_asm - - snitch_asm_module = snitch_stream_module.clone() - - test_lower_snitch_stream_to_asm.TestLowerSnitchStreamToAsm().apply( - ctx, snitch_asm_module - ) - - print(riscv_code(snitch_asm_module)) - - mo.md( - """ - We can then lower this to assembly that includes assembly instructions from the Snitch-extended ISA: - """ - ) - return snitch_asm_module, test_lower_snitch_stream_to_asm - - -if __name__ == "__main__": - app.run() diff --git a/pyproject.toml b/pyproject.toml index 29e2d623dc..ca6b57c4b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -128,7 +128,6 @@ target-version = "py310" "tests/test_declarative_assembly_format.py" = ["F811"] "versioneer.py" = ["ALL"] "_version.py" = ["ALL"] -"**/{docs/marimo}/*" = ["E501"] [tool.ruff.mccabe] max-complexity = 10 From 9148f4f6961738930477c3e8418a4e1d05dc7b36 Mon Sep 17 00:00:00 2001 From: Sasha Lopoukhine Date: Thu, 27 Jun 2024 10:15:11 +0100 Subject: [PATCH 11/11] tests-marimo PHONY --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index bb2f759ce9..d368a94ea8 100644 --- a/Makefile +++ b/Makefile @@ -14,7 +14,7 @@ TESTS_COVERAGE_FILE = ${COVERAGE_FILE}.tests .ONESHELL: # these targets don't produce files: -.PHONY: ${VENV_DIR}/ venv clean filecheck pytest pytest-nb tests-toy tests rerun-notebooks precommit-install precommit black pyright +.PHONY: ${VENV_DIR}/ venv clean filecheck pytest pytest-nb tests-toy tests rerun-notebooks precommit-install precommit black pyright tests-marimo .PHONY: coverage coverage-tests coverage-filecheck-tests coverage-report-html coverage-report-md # set up the venv with all dependencies for development