Add pow2 strength reduction (#177)

Replace expensive transcendental pow(x, 2.0) with a simple multiply.
This directly benefits layernorm's variance computation where
`centered_tx .^ 2.0f0` was generating a tile-wide pow instruction.

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: maleadt

README.md

@@ -97,19 +97,23 @@ Benchmarks comparing cuTile.jl against cuTile Python on an RTX 5080 (`tileiras`
 
 | Kernel | Size | Julia | Python | Status |
 |--------|------|-------|--------|--------|
-| Vector Addition | 2^27 f32 | 842 GB/s | 844 GB/s | OK (=) |
-| Matrix Transpose | 8192² f32 | 801 GB/s | 810 GB/s | OK (-1%) |
-| Layer Normalization | 4096² f32 fwd | 687 GB/s | 720 GB/s | -5% |
-| Matrix Multiplication | 4096³ f32 | 47.2 TFLOPS | 43.3 TFLOPS | +9%* |
-| Batch Matrix Multiply | 1024×512×2048 ×8 f32 | 33.5 TFLOPS | 30.8 TFLOPS | +9%* |
-| FFT (3-stage Cooley-Tukey) | 1024-pt ×64 c64 | 3264 μs | 3133 μs | -4% |
-| Mixture of Experts | 256tok 1024h 32e 2048i f16 | 19.5 TFLOPS | 20.1 TFLOPS | -3% |
-| Attention (FMHA) | 8×16×1024² ×64 f16 causal | 87.8 TFLOPS | 61.3 TFLOPS | +43%** |
-
-\* Likely because Julia's `for` loop guards give `tileiras` a guarantee that the
+| Vector Addition | 2^27 f32 | 841 GB/s | 845 GB/s | OK (=) |
+| Matrix Transpose | 8192² f32 | 805 GB/s | 811 GB/s | OK (-1%) |
+| Layer Norm fwd | 4096² f32 | 925 GB/s | 722 GB/s | +28%* |
+| Layer Norm bwd | 4096² f32 | 243 GB/s | 251 GB/s | -3% |
+| Matrix Multiplication | 4096³ f32 | 46.9 TFLOPS | 43.4 TFLOPS | +8%** |
+| Batch Matrix Multiply | 1024×512×2048 ×8 f32 | 33.6 TFLOPS | 30.9 TFLOPS | +9%** |
+| FFT (3-stage Cooley-Tukey) | 1024-pt ×64 c64 | 3263 μs | 3127 μs | -4% |
+| Mixture of Experts | 256tok 1024h 32e 2048i f16 | 19.3 TFLOPS | 20.3 TFLOPS | -5% |
+| Attention (FMHA) | 8×16×1024² ×64 f16 causal | 88.5 TFLOPS | 61.6 TFLOPS | +44%*** |
+
+\* The pow(x, 2) → mulf(x, x) strength reduction eliminates the expensive
+transcendental in the variance computation. Python still emits `pow`.
+
+\*\* Likely because Julia's `for` loop guards give `tileiras` a guarantee that the
 loop body executes at least once, enabling more aggressive warp scheduling.
 
-\*\* Likely due to Python's compiler splitting the causal masking loop into two
+\*\*\* Likely due to Python's compiler splitting the causal masking loop into two
 loops, duplicating the loop body. Julia emits a single loop with a conditional.
 
 

examples/benchmarks.jl

@@ -91,9 +91,10 @@ function run_benchmark(name::String)
     data = @invokelatest mod.prepare(; benchmark=true)
 
     # Get metric info if available
-    metric_total, metric_unit = 0, ""
+    # metric() returns either (total, unit) or Dict("impl" => (total, unit))
+    metric_result = nothing
     if isdefined(mod, :metric)
-        metric_total, metric_unit = @invokelatest mod.metric(data)
+        metric_result = @invokelatest mod.metric(data)
     end
 
     # Run cuTile
@@ -117,7 +118,7 @@ function run_benchmark(name::String)
         merge!(results, others)
     end
 
-    return results, metric_total, metric_unit
+    return results, metric_result
 end
 
 #=============================================================================
@@ -142,14 +143,23 @@ function main()
             continue
         end
 
-        results, metric_total, metric_unit = ret
+        results, metric_result = ret
 
         # Convert to BenchmarkResult for printing
         benchmark_results = BenchmarkResult[]
         for (impl_name, times) in results
             min_t = minimum(times)
             mean_t = sum(times) / length(times)
-            tp = !isempty(metric_unit) ? format_throughput(metric_total, metric_unit, min_t) : ""
+            tp = ""
+            if metric_result isa Dict
+                if haskey(metric_result, impl_name)
+                    mt, mu = metric_result[impl_name]
+                    tp = format_throughput(mt, mu, min_t)
+                end
+            elseif metric_result isa Tuple
+                mt, mu = metric_result
+                tp = format_throughput(mt, mu, min_t)
+            end
             push!(benchmark_results, BenchmarkResult(impl_name, min_t, mean_t, tp))
         end
 

examples/layernorm.jl

@@ -369,8 +369,13 @@ function test_layernorm(M, N, TILE_N; TILE_M::Int=32, eps::Float32=1f-5, name=no
 end
 
 function metric(data)
-    # Forward: 3 reads of X + W + B reads + Y write + Mean/Rstd writes ≈ 4*M*N floats
-    return 4 * data.M * data.N * sizeof(Float32), "GB/s"
+    MN = data.M * data.N * sizeof(Float32)
+    return Dict(
+        # Forward: X read (3 passes: mean, var, normalize) + Y write ≈ 4*M*N floats
+        "cuTile Fwd" => (4 * MN, "GB/s"),
+        # Backward: X read (2 passes) + DY read (2 passes) + DX write ≈ 5*M*N floats
+        "cuTile Bwd" => (5 * MN, "GB/s"),
+    )
 end
 
 # No run_others for layernorm - no simple reference implementation to compare against

src/compiler/passes/pipeline.jl

@@ -189,6 +189,26 @@ const COMPARISON_RULES = RewriteRule[
                               $(ComparisonPredicate.LessThan), $(Signedness.Signed)))
 ]
 
+#=============================================================================
+ Power Strength Reduction (rewrite)
+=============================================================================#
+
+# pow(x, 2) → mulf(x, x): replaces an expensive transcendental with a multiply.
+# The MLIR Tile IR backend has no canonicalization for pow, so this is purely
+# a Julia-level optimization. Applies to the variance computation in layernorm
+# (centered_tx .^ 2.0f0). Uses a guard with == so it matches any float type
+# (Float16, BFloat16, Float32, Float64, TFloat32). Integer-literal exponents
+# (x .^ 2) are already handled by Julia's literal_pow → x*x → mulf(x, x).
+
+function is_pow_two(match, driver)
+    c = const_value(driver.constants, match.bindings[:exp])
+    c !== nothing && c == 2
+end
+
+const POWER_RULES = RewriteRule[
+    @rewrite(Intrinsics.pow(~x, ~exp) => Intrinsics.mulf(~x, ~x), is_pow_two)
+]
+
 #=============================================================================
  Combined Rule Set
 =============================================================================#
@@ -198,6 +218,7 @@ const OPTIMIZATION_RULES = RewriteRule[
     ALGEBRA_RULES...,
     FMA_RULES...,
     COMPARISON_RULES...,
+    POWER_RULES...,
 ]
 
 #=============================================================================

test/codegen/operations.jl

@@ -1390,18 +1390,31 @@ end
             end
         end
 
-        # scalar exponent
+        # scalar exponent (pow(x, 2) is strength-reduced to mulf(x, x))
         @test @filecheck begin
             @check_label "entry"
             code_tiled(Tuple{ct.TileArray{Float32,1,spec1d}}) do a
                 pid = ct.bid(1)
                 tile = ct.load(a, pid, (16,))
-                @check "broadcast"
-                @check "pow"
+                @check "mulf"
+                @check_not "pow"
                 Base.donotdelete(tile .^ 2.0f0)
                 return
             end
         end
+
+        # pow2 strength reduction works for Float64 too
+        @test @filecheck begin
+            @check_label "entry"
+            code_tiled(Tuple{ct.TileArray{Float64,1,spec1d}}) do a
+                pid = ct.bid(1)
+                tile = ct.load(a, pid, (16,))
+                @check "mulf"
+                @check_not "pow"
+                Base.donotdelete(tile .^ 2.0)
+                return
+            end
+        end
     end
 
     @testset "scalar math functions" begin