Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

documentation: Fix nits on the first marimo notebook #3933

Merged
merged 4 commits into from
Feb 18, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 71 additions & 58 deletions docs/marimo/mlir_ir.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import marimo

__generated_with = "0.11.0"
__generated_with = "0.11.5"
app = marimo.App(width="medium")


Expand Down Expand Up @@ -30,7 +30,8 @@ def _(mo):

@app.cell(hide_code=True)
def _(mo, triangle_text):
mo.md(fr"""
mo.md(
rf"""
MLIR and xDSL use a textual encoding of the IR for debugging, testing, and storing intermediate representations of programs.
It can be very useful to take a program at some stage of compilation, and inspect it.
The textual format makes this easy to do.
Expand All @@ -44,7 +45,7 @@ def _(mo, triangle_text):

@app.cell(hide_code=True)
def _(mo):
mo.md(r"""This notebook explains all the components of the above snippet of code. The sections below are structured by _dialect_, a namespace for related constructs.""")
mo.md(r"""This notebook explains all the components of the above snippet of code. The sections below are structured by _dialect_, which represents a namespace for related abstractions and constructs.""")
return


Expand All @@ -66,7 +67,7 @@ def _():
@app.cell(hide_code=True)
def _(mo, swap_text, xmo):
mo.md(
fr"""
rf"""
The [func dialect](https://mlir.llvm.org/docs/Dialects/Func/) contains building blocks to model function definitions and calls.

{xmo.module_html(swap_text)}
Expand Down Expand Up @@ -103,7 +104,12 @@ def _(first_text, mo):

@app.cell(hide_code=True)
def _(exercise_text, first_text_area, mo):
_first_info_text = exercise_text(first_text_area.value, "first", ((1, 2), (3, 4)), ("first(1, 2) = ", "first(3, 4) = "))
_first_info_text = exercise_text(
first_text_area.value,
"first",
((1, 2), (3, 4)),
("first(1, 2) = ", "first(3, 4) = "),
)
Comment on lines +107 to +112
Copy link
Collaborator Author

@compor compor Feb 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Our ruff formatting excludes notebooks, so it's my local settings I assume.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are you editing the raw python locally or using marimo edit?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Raw python locally, do you want me to run with the editor? I'll see if it can reformat these.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah it would be good to run it in the editor, before merging because they have their own formatter that will likely conflict with your changes.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, there is no "format all cells", I have to go over each cell and do a ctrl-b

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"Run all cells" seems to do it

mo.vstack((first_text_area, mo.md(_first_info_text)))
return

Expand All @@ -116,15 +122,13 @@ def _(mo, xmo):
}\
"""

mo.accordion({
"Solution": xmo.module_html(second_text)
})
mo.accordion({"Solution": xmo.module_html(second_text)})
return (second_text,)


@app.cell(hide_code=True)
def _(mo):
mo.md(r"""## The Arith Dialect""")
mo.md(r"""## The `arith` Dialect""")
return


Expand All @@ -142,7 +146,7 @@ def _():
@app.cell(hide_code=True)
def _(add_one_text, mo, xmo):
mo.md(
fr"""
rf"""
The [arith dialect](https://mlir.llvm.org/docs/Dialects/ArithOps/) contains arithmetic operations on integers, floating-point values, and other numeric constructs. To start with, here is a function that adds one to its only argument:

{xmo.module_html(add_one_text)}
Expand Down Expand Up @@ -170,8 +174,8 @@ def _():
@app.cell(hide_code=True)
def _(add_one_text, mo, xmo):
mo.md(
fr"""
The `arith` dialect also contains operations for comparisons. The function below returns the value True if a is less than b when the 32-bit values passed in are interpreted as signed integers. Note that the signedness is communicated by the operation itself, not the types of the operands:
rf"""
The `arith` dialect also contains operations for comparisons. The function below returns the value `true` if a is less than b when the 32-bit values passed in are interpreted as signed integers. Note that the signedness is communicated by the operation itself, not the types of the operands:

{xmo.module_html(add_one_text)}
"""
Expand All @@ -181,7 +185,7 @@ def _(add_one_text, mo, xmo):

@app.cell(hide_code=True)
def _(mo):
mo.md(r"""MLIR IR is in [SSA form](https://en.wikipedia.org/wiki/Static_single-assignment_form), meaning that each value can only be assigned to once. This property helps reason about the possible runtime data these values can hold, such as whether they are constant, as `%one` in the snippet above.""")
mo.md(r"""MLIR IR is in [SSA form](https://en.wikipedia.org/wiki/Static_single-assignment_form), meaning that each value can only be assigned to _once_. This property helps reason about the possible runtime data these values can hold, such as whether they are constant, like the SSA value `%one` in the snippet above.""")
return


Expand Down Expand Up @@ -209,7 +213,12 @@ def _(fma_text, mo):

@app.cell(hide_code=True)
def _(exercise_text, fma_text_area, mo):
_fma_info_text = exercise_text(fma_text_area.value, "multiply_and_add", ((1, 2, 3), (4, 5, 6)), ("first(1, 2, 3) = ", "first(4, 5, 6) = "))
_fma_info_text = exercise_text(
fma_text_area.value,
"multiply_and_add",
((1, 2, 3), (4, 5, 6)),
("first(1, 2, 3) = ", "first(4, 5, 6) = "),
)
mo.vstack((fma_text_area, mo.md(_fma_info_text)))
return

Expand All @@ -223,9 +232,7 @@ def _(mo, xmo):
func.return %res : i32
}"""

mo.accordion({
"Solution": xmo.module_html(fma_impl)
})
mo.accordion({"Solution": xmo.module_html(fma_impl)})
return (fma_impl,)


Expand All @@ -235,6 +242,12 @@ def _(mo):
return


@app.cell(hide_code=True)
def _(mo):
mo.md(r"""The [`scf` dialect](https://mlir.llvm.org/docs/Dialects/Scf/) contains operations for structured control flow.""")
return


@app.cell(hide_code=True)
def _(mo):
mo.md(r"""### `scf.if`""")
Expand All @@ -258,14 +271,13 @@ def _():
@app.cell(hide_code=True)
def _(mo, select_text, xmo):
mo.md(
fr"""
The [`scf` dialect](https://mlir.llvm.org/docs/Dialects/Scf/) contains operations for control flow.
rf"""
Here is a function that returns the second argument if the first argument is `true`, and the third argument otherwise:

{xmo.module_html(select_text)}

Note that we did not put early returns in the branches of the if operation.
This is because MLIR's SSA blocks have a contract, which is that all the operations are executed from top to bottom, and operations are guaranteed to yield to the outer block.
Note that we did not put early returns in the branches of the `scf.if` operation.
This is due to MLIR's SSA blocks adhering to a specific contract: operations within a block are executed sequentially from top to bottom, and each operation is guaranteed to complete and yield control back to the outer block.
"""
)
return
Expand All @@ -280,9 +292,7 @@ def _(mo):
@app.cell(hide_code=True)
def _(abs_function_text, mo):
abs_input_text = mo.ui.text("-5")
abs_text_area = mo.ui.code_editor(
abs_function_text, language="javascript"
)
abs_text_area = mo.ui.code_editor(abs_function_text, language="javascript")
return abs_input_text, abs_text_area


Expand Down Expand Up @@ -347,7 +357,7 @@ def _(abs_info_text, abs_input_text, abs_text_area, mo):
(
mo.md(f"""Input: {abs_input_text} {mo.ui.button(label="run")}"""),
abs_text_area,
mo.md(abs_info_text)
mo.md(abs_info_text),
)
)
return
Expand All @@ -368,10 +378,7 @@ def _(mo, xmo):
func.return %res : i32
}"""


mo.accordion({
"Solution": xmo.module_html(abs_impl)
})
mo.accordion({"Solution": xmo.module_html(abs_impl)})
return (abs_impl,)


Expand All @@ -384,7 +391,7 @@ def _(mo):
@app.cell(hide_code=True)
def _(mo, triangle_text, xmo):
mo.md(
fr"""
rf"""
The `scf` dialect also contains abstractions to represent for loops, allowing us to implement the triangle function 1 + 2 + 3 + ... + n.

{xmo.module_html(triangle_text)}
Expand Down Expand Up @@ -435,7 +442,7 @@ def _(Parser, ctx, run_func, second_input_text, second_text_area):
"""
else:
second_info_text = f"""\
Change the definition of `second` to compute the factorial of the input.
Change the definition of `factorial` to compute the factorial of the input, instead of the triangle.
Assume that the input is non-negative.

```
Expand All @@ -459,7 +466,7 @@ def _(mo, second_info_text, second_input_text, second_text_area):
(
mo.md(f"""Input: {second_input_text} {mo.ui.button(label="run")}"""),
second_text_area,
mo.md(second_info_text)
mo.md(second_info_text),
)
)
return
Expand All @@ -479,26 +486,25 @@ def _(mo, xmo):
func.return %res : i32
}"""

mo.accordion({
"Solution": xmo.module_html(fact_impl)
})
mo.accordion({"Solution": xmo.module_html(fact_impl)})
return (fact_impl,)


@app.cell(hide_code=True)
def _(mo):
mo.md(r"""## `triangle` revisited""")
mo.md(r"""## The `triangle` revisited""")
return


@app.cell(hide_code=True)
def _(mo, triangle_text):
mo.md(fr"""
mo.md(
rf"""
This notebook contains a very light overview of the most commonly used dialects and operations in MLIR and xDSL, as well as the key concepts of SSA and structured control flow.

{mo.ui.code_editor(triangle_text, language="javascript", disabled=True)}

The sections below are a deeper dive into some of the structures that were implicit in the IR snippets we looked at.
The sections below are a deeper dive into some of the structures that were implicit in the IR snippets we looked at so far, reusing the `triangle` function from earlier.
"""
)
return
Expand All @@ -514,19 +520,17 @@ def _(mo):
def _(mo):
mo.md(
r"""
The [`builtin` dialect](https://mlir.llvm.org/docs/Dialects/Builtin/) contains the most commonly-used operation in MLIR/xDSL: `builtin.module`

A module is a unit of code in xDSL and MLIR.
It holds a single region.
The [`builtin` dialect](https://mlir.llvm.org/docs/Dialects/Builtin/) contains the most commonly-used operation in MLIR and xDSL: the `builtin.module` operation.

A module is a unit of code which holds a single region.
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Regions are not explained, so my main issue here is adding a definition earlier or omitting all together.
I'm inclined to do that later.

The smallest possible piece of code in MLIR IR is an empty module:

```
builtin.module {
}
```

When the first operation in a file is not a `builtin.module`, it is assumed, as is the case for all the snippets above.
When the first operation in a file is not a `builtin.module`, it is implicitly assumed and can be omitted, as is the case for all the snippets above.
"""
)
return
Expand All @@ -540,16 +544,18 @@ def _(mo):

@app.cell(hide_code=True)
def _(builtin, mo):
mo.md(fr"""
mo.md(
rf"""
Attributes hold compile-time data, such as constants, types, and other information.
The IR above contains four attributes: `@triangle`, `0`, `1` and `i32`.
`i32` is the type of integer values that fit in a register on the target.
As the IR here is independent of the machine on which it will run, we don't yet specify the bitwidth of the integer.
In MLIR, the common way to represent integers of 16, 32, 64, or other bitwidths is `i16`, `i32`, `i64`, etc.
`@triangle` is a symbol name, denoting the name of the function.

There are many other attributes, and many of them are in the `builtin` dialect.
We will look into defining those in a later notebook, and will only be using attributes from the builtin dialect in this one.
The IR for the `triangle` snippet contains four attributes: `@triangle`, `0`, `1` and `i32`.
The `i32` attribute is the type of integer values that fit in a register on the target platform.
As the IR presented here is independent of the machine on which it will be executed on, we do not yet specify the bitwidth of the integer.
In MLIR, the common way to represent integers of 16, 32, 64, or other bitwidths is, respectively, with the types of `i16`, `i32`, `i64`, etc.
The `@triangle` attribute is a symbol name, denoting the name of the function.

There are many other attributes, and a lot of them are part of the `builtin` dialect.
We will look into defining those in a later notebook, and will only be using attributes from the `builtin` dialect here.

One important attribute is the dictionary attribute, which looks like this:

Expand All @@ -562,7 +568,7 @@ def _(builtin, mo):
})}
```

The last entry denotes a key-value pair where the value is the unit attribute, which is omitted.
The last entry denotes a key-value pair (i.e., `a_unit_attr: unit`) where the omitted value is the `unit` attribute.
"""
)
return
Expand All @@ -578,8 +584,8 @@ def _(mo):
def _(mo):
mo.md(
r"""
The IRs above are in what's called the _custom format_, a format that allows functions to specify a pretty and concise representation.
The _generic format_ is a more uniform and verbose representation that unambiguously shows the structure of an operation.
All the snippets presented so far are in IRs that follow the _custom format_, a format that allows operations to specify a pretty and concise representation.
On the other hand, the _generic format_ is a more uniform and verbose representation that unambiguously shows the structure of an operation.
Here is the above `triangle` function in generic format:
"""
)
Expand All @@ -591,17 +597,19 @@ def _(Parser, Printer, StringIO, ctx, mo, triangle_text):
_triangle_module = Parser(ctx, triangle_text).parse_module()
_file = StringIO()
Printer(print_generic_format=True, stream=_file).print(_triangle_module)
mo.md(f"""
mo.md(
f"""
```
{_file.getvalue()}
```
""")
"""
)
return


@app.cell(hide_code=True)
def _(mo):
mo.md(r"""In the next notebooks, we'll take a deeper dive into the APIs used to process and construct MLIR IR.""")
mo.md(r"""In the next notebooks, we will take a deeper dive into the APIs used to process and construct MLIR IR.""")
return


Expand Down Expand Up @@ -676,7 +684,12 @@ def run_func(module: ModuleOp, name: str, args: tuple[Any, ...]):

@app.cell(hide_code=True)
def _(Any, Parser, ctx, run_func):
def exercise_text(module_text: str, function_name: str, inputs: tuple[Any, ...], formats: tuple[str, ...]) -> str:
def exercise_text(
module_text: str,
function_name: str,
inputs: tuple[Any, ...],
formats: tuple[str, ...],
) -> str:
error_text = ""
results_text = ""
try:
Expand Down
Loading