Propose to improve Puzzles 15, 16, 17

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
 

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]
@@ -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**
 

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
@@ -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
@@ -246,10 +238,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
 
@@ -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={

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"
 

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

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

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