From 76397782c2452bfe764affb8e56633aa47acbd89 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Wed, 14 Oct 2026 16:14:27 +0300 Subject: [PATCH 1/6] fix: remove the argued race condition --- book/src/puzzle_15/puzzle_15.md | 25 +++---------------------- solutions/p15/p15.mojo | 10 +--------- 2 files changed, 4 insertions(+), 31 deletions(-) diff --git a/book/src/puzzle_15/puzzle_15.md b/book/src/puzzle_15/puzzle_15.md index 5da7eb7c..538958fe 100644 --- a/book/src/puzzle_15/puzzle_15.md +++ b/book/src/puzzle_15/puzzle_15.md @@ -168,24 +168,6 @@ Input Matrix (4×6) with TileTensor: Block Assignment: ```mojo var stride = TPB // 2 - while stride > 0: - var temp_val: Scalar[dtype] = 0 - if local_i < stride: - temp_val = cache[local_i + stride] # Read phase - barrier() - if local_i < stride: - cache[local_i] += temp_val # Write phase - barrier() - stride //= 2 - ``` - - **Note**: The read and write phases are deliberately separated by a - `barrier()`. The simpler form below reads and writes shared memory in the - same iteration, which races: - - ```mojo - # Racy - do not use - var stride = TPB // 2 while stride > 0: if local_i < stride: cache[local_i] += cache[local_i + stride] @@ -213,12 +195,11 @@ Input Matrix (4×6) with TileTensor: Block Assignment: - Efficient parallel reduction pattern 3. **Synchronization**: - - One barrier after the load, then two per reduction step + - One barrier after the load, then one after each reduction step + - The reduction barrier publishes the current stride's writes before the + next stride reads them - Independent processing between rows - No inter-block communication needed - - **Race condition handling**: read-write hazards during the reduction are - avoided by explicit phase separation—each iteration stages its read in - `temp_val`, barriers, then writes ### Complexity analysis diff --git a/solutions/p15/p15.mojo b/solutions/p15/p15.mojo index 443ddff7..e6646d0f 100644 --- a/solutions/p15/p15.mojo +++ b/solutions/p15/p15.mojo @@ -67,16 +67,8 @@ def axis_sum[ # do reduction sum per each block var stride = TPB // 2 while stride > 0: - # Read phase: all threads read the values they need first to avoid race conditions - var temp_val: Scalar[dtype] = 0 if local_i < stride: - temp_val = cache[local_i + stride] - - barrier() - - # Write phase: all threads safely write their computed values - if local_i < stride: - cache[local_i] += temp_val + cache[local_i] += cache[local_i + stride] barrier() stride //= 2 From ddeff79df488c58d79df773623341a2c5f38d9a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Sat, 17 Oct 2026 14:07:51 +0300 Subject: [PATCH 2/6] up style --- book/src/puzzle_16/tiled.md | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/book/src/puzzle_16/tiled.md b/book/src/puzzle_16/tiled.md index e9f9b57f..6a2642c8 100644 --- a/book/src/puzzle_16/tiled.md +++ b/book/src/puzzle_16/tiled.md @@ -119,15 +119,15 @@ Synchronization required: **Example with TPB=3:** ```txt - Block Layout: Global Matrix (9×9): - [B00][B01][B02] [0 1 2 | 3 4 5 | 6 7 8] - [B10][B11][B12] → [9 A B | C D E | F G H] - [B20][B21][B22] [I J K | L M N | O P Q] - —————————————————————— + Block Layout: Global Matrix (9×9): + [B00][B01][B02] [0 1 2 | 3 4 5 | 6 7 8] + [B10][B11][B12] → [9 A B | C D E | F G H] + [B20][B21][B22] [I J K | L M N | O P Q] + ——————————————————————— [R S T | U V W | X Y Z] [a b c | d e f | g h i] [j k l | m n o | p q r] - —————————————————————— + ——————————————————————— [s t u | v w x | y z α] [β γ δ | ε ζ η | θ ι κ] [λ μ ν | ξ ο π | ρ σ τ] @@ -214,12 +214,12 @@ it works: ```txt Input matrices (9×9) - Perfect fit for (3×3) tiling: - A = [0 1 2 3 4 5 6 7 8 ] B = [0 2 4 6 8 10 12 14 16] - [9 10 11 12 13 14 15 16 17] [18 20 22 24 26 28 30 32 34] - [18 19 20 21 22 23 24 25 26] [36 38 40 42 44 46 48 50 52] - [27 28 29 30 31 32 33 34 35] [54 56 58 60 62 64 66 68 70] - [36 37 38 39 40 41 42 43 44] [72 74 76 78 80 82 84 86 88] - [45 46 47 48 49 50 51 52 53] [90 92 94 96 98 100 102 104 106] + A = [0 1 2 3 4 5 6 7 8 ] B = [0 2 4 6 8 10 12 14 16 ] + [9 10 11 12 13 14 15 16 17] [18 20 22 24 26 28 30 32 34 ] + [18 19 20 21 22 23 24 25 26] [36 38 40 42 44 46 48 50 52 ] + [27 28 29 30 31 32 33 34 35] [54 56 58 60 62 64 66 68 70 ] + [36 37 38 39 40 41 42 43 44] [72 74 76 78 80 82 84 86 88 ] + [45 46 47 48 49 50 51 52 53] [90 92 94 96 98 100 102 104 106] [54 55 56 57 58 59 60 61 62] [108 110 112 114 116 118 120 122 124] [63 64 65 66 67 68 69 70 71] [126 128 130 132 134 136 138 140 142] [72 73 74 75 76 77 78 79 80] [144 146 148 150 152 154 156 158 160] From c13606b6fc3c92231eb7038dee5d44b75d3aca7f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Sat, 17 Oct 2026 14:29:47 +0300 Subject: [PATCH 3/6] Improve shared memory bank conflict explanation in puzzle 16 documentation --- book/src/puzzle_16/tiled.md | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/book/src/puzzle_16/tiled.md b/book/src/puzzle_16/tiled.md index 6a2642c8..e5e7f1e4 100644 --- a/book/src/puzzle_16/tiled.md +++ b/book/src/puzzle_16/tiled.md @@ -269,25 +269,28 @@ it works: - Avoids shared memory bank conflicts: ```txt - Bank Conflict Free (Good): Bank Conflicts (Bad): + Bank Conflict Free (Good): Bank Conflicts (Bad): Thread0: a_shared[0,k] b_shared[k,0] Thread0: a_shared[k,0] b_shared[0,k] Thread1: a_shared[0,k] b_shared[k,1] Thread1: a_shared[k,0] b_shared[1,k] Thread2: a_shared[0,k] b_shared[k,2] Thread2: a_shared[k,0] b_shared[2,k] ↓ ↓ Parallel access to different banks Serialized access to same bank of b_shared - (or broadcast for a_shared) if shared memory was column-major + (or broadcast for a_shared) if TPB=32 ``` **Shared memory bank conflicts explained:** - - **Left (Good)**: `b_shared[k,threadIdx.x]` accesses different banks, - `a_shared[0,k]` broadcasts to all threads - - **Right (Bad)**: If b_shared were column-major, threads would access same - bank simultaneously + - **Left (Good)**: `b_shared[k,threadIdx.x]` accesses different banks; + `a_shared[local_row,k]` broadcasts to threads sharing the same row. + - **Right (Bad)**: With row-major storage, varying the first index + gives a stride of `TPB` words. For `TPB=32`, those addresses map + to the same bank; for `TPB=16`, they map to two banks. This is an + illustrative case, not a conflict in this puzzle's `TPB=3` tile. - **Key insight**: This is about shared memory access patterns, not global memory coalescing - **Bank structure**: Shared memory has 32 banks; conflicts occur when - multiple threads access different addresses in the same bank - simultaneously + multiple threads in a warp access different addresses in the same bank + simultaneously. Multiple threads reading the same address are served + by a broadcast, not a bank conflict. 5. **Synchronization points** From fee148988e4c88cb39ff334ef717d6e65ccfcbc3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Sun, 18 Oct 2026 19:08:14 +0300 Subject: [PATCH 4/6] Update links in puzzle 17 documentation for accuracy and clarity, and revise tutorial references for better guidance. --- book/src/puzzle_17/puzzle_17.md | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/book/src/puzzle_17/puzzle_17.md b/book/src/puzzle_17/puzzle_17.md index cda79a18..d50f64fa 100644 --- a/book/src/puzzle_17/puzzle_17.md +++ b/book/src/puzzle_17/puzzle_17.md @@ -133,9 +133,9 @@ The solution is: These two statements do several important things: 1. Calls - [enqueue_function](https://max.modular.com/api/mojo/max/gpu/host/device_context/DeviceContext/#enqueue_function) + [enqueue_function](https://max.modular.com/api/mojo/max/gpu/host/device_context/DeviceContextBase/#enqueue_function) on the GPU context (`gpu_ctx` is of type - [DeviceContext](https://max.modular.com/api/mojo/max/gpu/host/device_context/DeviceContext/)) + [DeviceContext](https://max.modular.com/api/mojo/max/gpu/host/device_context/#devicecontext)) to schedule our kernel execution 2. Binds the layout and size information as **compile-time** parameters through the `comptime kernel = conv1d_kernel[...]` binding @@ -158,12 +158,12 @@ Let's break down how this works in the larger context: 2. **Graph building**: - Defines input and output tensor types with - [TensorType](https://max.modular.com/api/python/graph/type/#max.graph.type.TensorType) + [TensorType](https://max.modular.com/api/python/generated/max.graph.TensorType/#max.graph.TensorType) - Specifies parameters for our operation via `parameters={...}` - Creates a computation graph with - [`Graph("conv_1d_graph", ...)`](https://max.modular.com/api/python/graph/Graph) + [`Graph("conv_1d_graph", ...)`](https://max.modular.com/api/python/generated/max.graph.Graph/#max.graph.Graph) - Calls our operation using - [`ops.custom(name="conv1d", ...)`](https://max.modular.com/api/python/graph/ops#custom) + [`ops.custom(name="conv1d", ...)`](https://max.modular.com/api/python/graph.ops/#max.graph.ops.custom) 3. **Custom op registration**: - The `@extensibility.register("conv1d")` decorator exposes our operation to MAX @@ -246,10 +246,10 @@ functions to create efficient, type-safe, accelerated operations. ## Understanding MAX Graph custom ops -> Check out the follow tutorials for more details: +> Check out these for more details: > -> - [Get started with MAX Graph in Python](https://max.modular.com/tutorials/get-started-with-max-graph-in-python/) -> - [MAX Graph custom op for GPUs](https://max.modular.com/tutorials/build-custom-ops/) +> - [Graph overview](https://max.modular.com/develop/graph/) +> - [Intro to custom ops](https://max.modular.com/develop/custom-ops/) ### Custom op registration From 8fe7a2b3464e45e22f67ccbb1d34b11e5be8b0dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Thu, 22 Oct 2026 15:27:20 +0300 Subject: [PATCH 5/6] Describe puzzle 17's custom op in terms of the current MAX tensor API. Use to_tile_tensor for the kernel arguments, pass device to ops.custom, and show mojo precompile in the packaging example. --- book/src/puzzle_17/puzzle_17.md | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/book/src/puzzle_17/puzzle_17.md b/book/src/puzzle_17/puzzle_17.md index d50f64fa..ee67537e 100644 --- a/book/src/puzzle_17/puzzle_17.md +++ b/book/src/puzzle_17/puzzle_17.md @@ -206,20 +206,12 @@ Let's break down how this works in the larger context: 2. **Tensor Conversion**: ```mojo - comptime out_layout_val = row_major[input_size]() - comptime OutLayout = type_of(out_layout_val) - - var output_tensor = TileTensor[mut=True, dtype, OutLayout, MutAnyOrigin]( - output.unsafe_ptr(), out_layout_val - ) - var input_tensor = TileTensor[mut=True, dtype, OutLayout, MutAnyOrigin]( - input.unsafe_ptr(), out_layout_val - ) + var output_tensor = output.to_tile_tensor().as_unsafe_any_origin() + var input_tensor = input.to_tile_tensor().as_unsafe_any_origin() + var kernel_tensor = kernel.to_tile_tensor().as_unsafe_any_origin() ``` - The layouts are constructed from the op's compile-time parameters and the - tensors are built from raw pointers—they are not extracted from the - `OutputTensor`/`InputTensor` arguments. + `to_tile_tensor()` wraps each argument as a TileTensor over the same memory. - MAX Graph tensors are wrapped as Mojo TileTensors over the same memory - This allows our kernel to work with them directly @@ -283,7 +275,7 @@ Key components of the registration: Before the custom operation can be used from Python, it needs to be packaged: ```bash -mojo package op -o op.mojoc +mojo precompile op -o op.mojoc ``` This command: @@ -315,6 +307,7 @@ with Graph( # Use our custom operation by name output = ops.custom( name="conv1d", # Must match the name in @extensibility.register + device=DeviceRef.from_device(device), values=[input_value, kernel_value], out_types=[...], parameters={ From f43b2f2dc3a2ac0e8ad61ce434e37b1f7c65489f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Thu, 22 Oct 2026 15:27:20 +0300 Subject: [PATCH 6/6] Call mojo precompile instead of the package alias. The puzzle tasks and the skeleton compile script now use the current command name. --- pixi.toml | 6 +++--- pyproject.toml | 6 +++--- scripts/compile_problems.sh | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pixi.toml b/pixi.toml index cb33d183..06fbfeae 100644 --- a/pixi.toml +++ b/pixi.toml @@ -167,13 +167,13 @@ p16 = "bash solutions/gate.sh p16 mojo -I . problems/p16/p16.mojo" viz16 = "cd book/src/puzzle_16 && python puzzle_16_viz.py" rooflineviz = "cd book/src/puzzle_16 && python roofline_viz.py" -p17 = "mojo package problems/p17/op -o problems/p17/op.mojoc && python problems/p17/p17.py" +p17 = "mojo precompile problems/p17/op -o problems/p17/op.mojoc && python problems/p17/p17.py" -p18-package = "mojo package problems/p18/op -o problems/p18/op.mojoc" +p18-package = "mojo precompile problems/p18/op -o problems/p18/op.mojoc" p18-test-kernels = { cmd = "mojo run -I problems/p18 problems/p18/test/test_softmax.mojo", depends-on = ["p18-package"] } p18 = { cmd = "python problems/p18/p18.py", depends-on = ["p18-package"] } -p19 = "bash solutions/gate.sh p19 mojo package problems/p19/op -o problems/p19/op.mojoc && python problems/p19/p19.py" +p19 = "bash solutions/gate.sh p19 mojo precompile problems/p19/op -o problems/p19/op.mojoc && python problems/p19/p19.py" p20 = "python problems/p20/p20.py" diff --git a/pyproject.toml b/pyproject.toml index a3034a6a..366c17af 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,17 +87,17 @@ p15 = "mojo -I . problems/p15/p15.mojo" p16 = "bash solutions/gate.sh p16 mojo -I . problems/p16/p16.mojo" p17 = [ - { cmd = "mojo package problems/p17/op -o problems/p17/op.mojoc" }, + { cmd = "mojo precompile problems/p17/op -o problems/p17/op.mojoc" }, { cmd = "python problems/p17/p17.py" } ] -p18-package = "mojo package problems/p18/op -o problems/p18/op.mojoc" +p18-package = "mojo precompile problems/p18/op -o problems/p18/op.mojoc" p18-test-kernels = { cmd = "mojo run -I problems/p18 problems/p18/test/test_softmax.mojo", deps = ["p18-package"] } p18 = { cmd = "python problems/p18/p18.py", deps = ["p18-package"] } p19 = [ { cmd = "bash solutions/gate.sh p19 true" }, - { cmd = "mojo package problems/p19/op -o problems/p19/op.mojoc" }, + { cmd = "mojo precompile problems/p19/op -o problems/p19/op.mojoc" }, { cmd = "python problems/p19/p19.py" } ] diff --git a/scripts/compile_problems.sh b/scripts/compile_problems.sh index edaa8b57..a739f2c0 100755 --- a/scripts/compile_problems.sh +++ b/scripts/compile_problems.sh @@ -18,7 +18,7 @@ # silently rot: an API migration that updates the solutions but not the # skeletons leaves learners with starter code that no longer compiles (see # issue #254). This script closes that gap by compiling each skeleton: -# - custom-op modules (problems/pNN/op/) -> `mojo package` (no main()) +# - custom-op modules (problems/pNN/op/) -> `mojo precompile` (no main()) # - standalone single-file puzzles -> `mojo build` (have main()) # An unfilled skeleton is expected to compile (it just produces wrong/zero # output until the learner fills the `# FILL ME IN` gaps); it must never fail to @@ -90,7 +90,7 @@ for sol_dir in solutions/p*/; do echo "${puzzle}:" if [ -d "${prob_dir}/op" ]; then # Custom-op module: compile the package (no main()). - compile_target "${prob_dir}/op" package "${prob_dir}/op" -o /tmp/skel.mojoc + compile_target "${prob_dir}/op" precompile "${prob_dir}/op" -o /tmp/skel.mojoc continue fi # Standalone puzzle files: compile each .mojo that defines a main().