Skip to contents

torch.compile: graph-capturing JIT for forward and backward passes

Source: R/geron_w4d_native.R
morie_geron_torch_compile.Rd

Traces the model (morie_geron_trace), fuses consecutive "linear" ops into one matmul (exact, verified by replay), and under mode = "max-autotune" costs the fused product via the matrix-chain DP.

Usage

morie_geron_torch_compile(model, mode = "default", example_inputs)

Arguments

model

List of ops (see morie_geron_trace).

mode

"default", "reduce-overhead", or "max-autotune".

example_inputs

Required input used for graph capture.

Value

list with compiled, graph, output, eager_output, max_diff, n_ops, n_compiled, fused_runs, flops_eager, flops_compiled, speedup, mode, estimate, n, method.

Examples

A2 <- diag(2) * 2
B2 <- diag(2) * 3
C2 <- diag(2) * 5
morie_geron_torch_compile(list(list(kind = "linear", param = A2),
    list(kind = "linear", param = B2), list(kind = "linear",
        param = C2)), example_inputs = matrix(c(1, 1), 1, 2))
#> $compiled
#> $compiled[[1]]
#> $compiled[[1]]$index
#> [1] 0
#> 
#> $compiled[[1]]$kind
#> [1] "linear"
#> 
#> $compiled[[1]]$in_shape
#> [1] 1 2
#> 
#> $compiled[[1]]$out_shape
#> [1] 1 2
#> 
#> $compiled[[1]]$op
#> $compiled[[1]]$op$kind
#> [1] "linear"
#> 
#> $compiled[[1]]$op$param
#>      [,1] [,2]
#> [1,]   30    0
#> [2,]    0   30
#> 
#> 
#> $compiled[[1]]$fused_from
#> [1] 0 1 2
#> 
#> 
#> 
#> $graph
#> $graph[[1]]
#> $graph[[1]]$index
#> [1] 0
#> 
#> $graph[[1]]$kind
#> [1] "linear"
#> 
#> $graph[[1]]$in_shape
#> [1] 1 2
#> 
#> $graph[[1]]$out_shape
#> [1] 1 2
#> 
#> $graph[[1]]$op
#> $graph[[1]]$op$kind
#> [1] "linear"
#> 
#> $graph[[1]]$op$param
#>      [,1] [,2]
#> [1,]    2    0
#> [2,]    0    2
#> 
#> 
#> 
#> $graph[[2]]
#> $graph[[2]]$index
#> [1] 1
#> 
#> $graph[[2]]$kind
#> [1] "linear"
#> 
#> $graph[[2]]$in_shape
#> [1] 1 2
#> 
#> $graph[[2]]$out_shape
#> [1] 1 2
#> 
#> $graph[[2]]$op
#> $graph[[2]]$op$kind
#> [1] "linear"
#> 
#> $graph[[2]]$op$param
#>      [,1] [,2]
#> [1,]    3    0
#> [2,]    0    3
#> 
#> 
#> 
#> $graph[[3]]
#> $graph[[3]]$index
#> [1] 2
#> 
#> $graph[[3]]$kind
#> [1] "linear"
#> 
#> $graph[[3]]$in_shape
#> [1] 1 2
#> 
#> $graph[[3]]$out_shape
#> [1] 1 2
#> 
#> $graph[[3]]$op
#> $graph[[3]]$op$kind
#> [1] "linear"
#> 
#> $graph[[3]]$op$param
#>      [,1] [,2]
#> [1,]    5    0
#> [2,]    0    5
#> 
#> 
#> 
#> 
#> $output
#>      [,1] [,2]
#> [1,]   30   30
#> 
#> $eager_output
#>      [,1] [,2]
#> [1,]   30   30
#> 
#> $max_diff
#> [1] 0
#> 
#> $n_ops
#> [1] 3
#> 
#> $n_compiled
#> [1] 1
#> 
#> $fused_runs
#> [1] 1
#> 
#> $flops_eager
#> [1] 20
#> 
#> $flops_compiled
#> [1] 4
#> 
#> $speedup
#> [1] 5
#> 
#> $mode
#> [1] "default"
#> 
#> $estimate
#> [1] 5
#> 
#> $n
#> [1] 1
#> 
#> $method
#> [1] "Graph capture (hmtsc) + linear-chain fusion, with matrix-chain DP association under max-autotune"
#>