from numba import cuda

from ch08_gpu_offload import GPUBackend
from ch08_gpu_offload import gpu_compiler_config as _ch08_gpu_compiler_config
from demo01_gelu_tanh_approx import *
from utils.report import Report

class GpuUfuncBackend(Backend, GPUBackend):
    # Ufunc + GPU backend
    def __init__(self, compile_only: bool = False):
        GPUBackend.__init__(self, compile_only)

gpu_compiler_config = {
    **_ch08_gpu_compiler_config,
    "converter_class": ExtendEGraphToRVSDG,
    "cost_model": MyCostModel(),
    "backend": GpuUfuncBackend(compile_only=not cuda.is_available()),
}

report = Report("Pipeline execution report", enable_nested_metadata=True)
cuda_vectorized_gelu = ufunc_vectorize(
    input_type=Float32,
    ndim=1,
    compiler_config={**gpu_compiler_config, "pipeline_report": report, "pipeline_debug": True},
    extra_ruleset=additional_rules | optimize_rules,
)(gelu_tanh_forward)
if __name__ == "__main__":
    report.display()

--------------------------------original source---------------------------------
 106|def gelu_tanh_forward(a):
 107|    dt = np.float32
 108|    result = (
 109|        dt(0.5)
 110|        * a
 111|        * (
 112|            dt(1)
 113|            + np.tanh(np.sqrt(dt(2) / dt(np.pi)) * (a + dt(0.044715) * a**3))
 114|        )
 115|    )
 116|    return result
----------------------------------inter source----------------------------------
   1|def transformed_gelu_tanh_forward(a):
   2|    """#file: demo01_gelu_tanh_approx.py"""
   3|    '#loc: 107:8-107:23'
   4|    dt = np.float32
   5|    '#loc: 108:8-115:9'
   6|    result = dt(0.5) * a * (dt(1) + np.tanh(np.sqrt(dt(2) / dt(np.pi)) * (a + dt(0.044715) * a ** 3)))
   7|    '#loc: 116:8-116:21'
   8|    return result

transformed_gelu_tanh_forward = Func (Args (ArgSpec 'a' (PyNone)))
$0 = Region[804] <- !io a
{
  $1 = PyLoadGlobal $0[0] 'np'
  $2 = PyAttr $0[0] $1 'float32'
  $3 = DbgValue 'dt' $2[1]
  $4 = PyFloat 0.5
  $5 = PyCall $3 $2[0] $4
  $6 = PyBinOp * $5[0] $5[1], $0[1]
  $7 = PyInt 1
  $8 = PyCall $3 $6[0] $7
  $9 = PyInt 2
  $10 = PyCall $3 $8[0] $9
  $11 = PyLoadGlobal $10[0] 'np'
  $12 = PyAttr $10[0] $11 'pi'
  $13 = PyCall $3 $12[0] $12[1]
  $14 = PyBinOp / $13[0] $10[1], $13[1]
  $15 = PyLoadGlobal $14[0] 'np'
  $16 = PyAttr $14[0] $15 'sqrt'
  $17 = PyCall $16[1] $16[0] $14[1]
  $18 = PyFloat 0.044715
  $19 = PyCall $3 $17[0] $18
  $20 = PyInt 3
  $21 = PyBinOp ** $19[0] $0[1], $20
  $22 = PyBinOp * $21[0] $19[1], $21[1]
  $23 = PyBinOp + $22[0] $0[1], $22[1]
  $24 = PyBinOp * $23[0] $17[1], $23[1]
  $25 = PyLoadGlobal $24[0] 'np'
  $26 = PyAttr $24[0] $25 'tanh'
  $27 = PyCall $26[1] $26[0] $24[1]
  $28 = PyBinOp + $27[0] $8[1], $27[1]
  $29 = PyBinOp * $28[0] $6[1], $28[1]
  $30 = DbgValue 'result' $29[1]
} [1268] -> !io=$29[0] !ret=$30

time elapsed 16.08ms
timing breakdown:
  11.26ms: Debug Info on RVSDG 
  4.82ms: RVSDG

time elapsed 99.91ms
timing breakdown:
  99.91ms: EGraph

_Region_1 = Region("804", InPorts(Vec[String]("!io", "a")))
GraphRoot(
    Term.Func(
        "1274",
        "transformed_gelu_tanh_forward",
        Term.RegionEnd(
            _Region_1,
            PortList(
                Vec[Port](
                    Port("!io", _Region_1.get(0)),
                    Port(
                        "!ret",
                        Nb_Mul_Float32(
                            Nb_Mul_Float32(Npy_cast_f64_to_f32(Term.LiteralF64(0.5)), _Region_1.get(1)),
                            Nb_Add_Float32(
                                Npy_cast_i64_to_f32(Term.LiteralI64(1)),
                                Npy_tanh_float32(
                                    Nb_Mul_Float32(
                                        Npy_sqrt_float32(Nb_Div_Float32(Npy_cast_i64_to_f32(Term.LiteralI64(2)), Npy_cast_f64_to_f32(Term.LiteralF64(3.141592653589793)))),
                                        Nb_Add_Float32(
                                            _Region_1.get(1),
                                            Nb_Mul_Float32(Npy_cast_f64_to_f32(Term.LiteralF64(0.044715)), Nb_Pow_Float32_Int64(_Region_1.get(1), Term.LiteralI64(3))),
                                        ),
                                    )
                                ),
                            ),
                        ),
                    ),
                )
            ),
        ),
    )
)

time elapsed 365.70ms
timing breakdown:
  66.69ms: [debug] initial egraph
  258.65ms: [debug] saturated egraph
  40.36ms: [debug] egglog.extract

transformed_gelu_tanh_forward = Func (Args (ArgSpec 'a' (PyNone)))
$0 = Region[1575] <- !io a; #attrs (_, Float32)->(_, Float32)
{
  $1 = PyFloat 0.5
  $2 = NbOp_F64_to_F32 $1
  $3 = NbOp_Mul_Float32 $2 $0[1]
  $4 = PyInt 1
  $5 = NbOp_I64_to_F32 $4
  $6 = PyFloat 10.0
  $7 = NbOp_F64_to_F32 $6
  $8 = PyInt 2
  $9 = NbOp_I64_to_F32 $8
  $10 = PyFloat 3.141592653589793
  $11 = NbOp_F64_to_F32 $10
  $12 = NbOp_Div_Float32 $9 $11
  $13 = NpyOp_Sqrt_Float32 $12
  $14 = PyFloat 0.044715
  $15 = NbOp_F64_to_F32 $14
  $16 = PyFloat 1.0
  $17 = NbOp_F64_to_F32 $16
  $18 = NbOp_Mul_Float32 $0[1] $17
  $19 = NbOp_Mul_Float32 $0[1] $18
  $20 = NbOp_Mul_Float32 $0[1] $19
  $21 = NbOp_Mul_Float32 $15 $20
  $22 = NbOp_Add_Float32 $0[1] $21
  $23 = NbOp_Mul_Float32 $13 $22
  $24 = NbOp_Mul_Float32 $23 $17
  $25 = NbOp_Mul_Float32 $23 $24
  $26 = NbOp_Mul_Float32 $23 $25
  $27 = NbOp_Mul_Float32 $7 $26
  $28 = PyFloat 105.0
  $29 = NbOp_F64_to_F32 $28
  $30 = NbOp_Mul_Float32 $29 $23
  $31 = NbOp_Add_Float32 $27 $30
  $32 = NbOp_Mul_Float32 $23 $26
  $33 = PyFloat 45.0
  $34 = NbOp_F64_to_F32 $33
  $35 = NbOp_Mul_Float32 $34 $25
  $36 = NbOp_Add_Float32 $32 $35
  $37 = NbOp_Add_Float32 $36 $29
  $38 = NbOp_Div_Float32 $31 $37
  $39 = NbOp_Add_Float32 $5 $38
  $40 = NbOp_Mul_Float32 $3 $39
} [1782] -> !io=$0[0] !ret=$40

14747.0

time elapsed 16.36ms
timing breakdown:
  16.35ms: Extracted RVSDG     
  0.01ms: Extracted cost

module {
  func.func @func(%arg0: f32) -> f32 attributes {llvm.emit_c_interface} {
    %cst = arith.constant 5.000000e-01 : f64
    %c1_i64 = arith.constant 1 : i64
    %cst_0 = arith.constant 1.000000e+01 : f64
    %c2_i64 = arith.constant 2 : i64
    %cst_1 = arith.constant 3.1415926535897931 : f64
    %cst_2 = arith.constant 4.471500e-02 : f64
    %cst_3 = arith.constant 1.000000e+00 : f64
    %cst_4 = arith.constant 1.050000e+02 : f64
    %cst_5 = arith.constant 4.500000e+01 : f64
    cf.br ^bb1
  ^bb1:  // pred: ^bb0
    %c0_i32 = arith.constant 0 : i32
    %0 = arith.truncf %cst : f64 to f32
    %1 = arith.mulf %0, %arg0 : f32
    %2 = arith.sitofp %c1_i64 : i64 to f32
    %3 = arith.truncf %cst_0 : f64 to f32
    %4 = arith.sitofp %c2_i64 : i64 to f32
    %5 = arith.truncf %cst_1 : f64 to f32
    %6 = arith.divf %4, %5 : f32
    %7 = math.sqrt %6 : f32
    %8 = arith.truncf %cst_2 : f64 to f32
    %9 = arith.truncf %cst_3 : f64 to f32
    %10 = arith.mulf %arg0, %9 : f32
    %11 = arith.mulf %arg0, %10 : f32
    %12 = arith.mulf %arg0, %11 : f32
    %13 = arith.mulf %8, %12 : f32
    %14 = arith.addf %arg0, %13 : f32
    %15 = arith.mulf %7, %14 : f32
    %16 = arith.mulf %15, %9 : f32
    %17 = arith.mulf %15, %16 : f32
    %18 = arith.mulf %15, %17 : f32
    %19 = arith.mulf %3, %18 : f32
    %20 = arith.truncf %cst_4 : f64 to f32
    %21 = arith.mulf %20, %15 : f32
    %22 = arith.addf %19, %21 : f32
    %23 = arith.mulf %15, %18 : f32
    %24 = arith.truncf %cst_5 : f64 to f32
    %25 = arith.mulf %24, %17 : f32
    %26 = arith.addf %23, %25 : f32
    %27 = arith.addf %26, %20 : f32
    %28 = arith.divf %22, %27 : f32
    %29 = arith.addf %2, %28 : f32
    %30 = arith.mulf %1, %29 : f32
    return %30 : f32
  }
}

if __name__ == "__main__":
    if not cuda.is_available():
        print("SKIPPED. CUDA unavailable")
    else:
        relclose = lambda x, y: np.allclose(x, y, rtol=1e-6)
        input_val = np.random.random(100).astype(np.float32)
        run_test(
            gelu_tanh_forward,
            cuda_vectorized_gelu,
            (input_val,),
            equal=relclose,
            verbose=True,
        )

(array([0.38916567, 0.01010664, 0.20098524, 0.93654305, 0.26475447,
       0.02711666, 0.43591085, 0.7909023 , 0.02640386, 0.04154724,
       0.65611833, 0.47337994, 0.65105635, 0.9289255 , 0.5399822 ,
       0.06580561, 0.64303064, 0.15780859, 0.5683575 , 0.6420226 ,
       0.38611475, 0.31390396, 0.5064093 , 0.15218998, 0.5781999 ,
       0.13255835, 0.45071495, 0.73919684, 0.51800865, 0.3687768 ,
       0.83240014, 0.11023618, 0.14749587, 0.89334303, 0.9199019 ,
       0.8821255 , 0.99873996, 0.2324525 , 0.2587263 , 0.09865194,
       0.91596186, 0.86126864, 0.37732396, 0.6635682 , 0.17678984,
       0.07984433, 0.45669428, 0.60289496, 0.99220383, 0.3832815 ,
       0.08342831, 0.6211122 , 0.91324633, 0.49621433, 0.8436951 ,
       0.00544985, 0.18173647, 0.7941845 , 0.1175717 , 0.32131624,
       0.3600599 , 0.16189952, 0.8056253 , 0.8696681 , 0.8966021 ,
       0.5313608 , 0.980649  , 0.6631159 , 0.5826805 , 0.03971836,
       0.9231509 , 0.7995121 , 0.3719242 , 0.49562123, 0.24867107,
       0.5806849 , 0.39288136, 0.79061097, 0.8867972 , 0.70988667,
       0.17520966, 0.06835352, 0.5016661 , 0.03497642, 0.10607172,
       0.30873862, 0.32629046, 0.03611915, 0.79392153, 0.26159677,
       0.8521278 , 0.30235165, 0.291528  , 0.7643648 , 0.9302362 ,
       0.12348044, 0.64986646, 0.09082226, 0.05749565, 0.6363848 ],
      dtype=float32),)

[0.25350475 0.00509407 0.11649956 0.7729887  0.16001624 0.01385164
 0.2914175  0.62117213 0.01348003 0.02146206 0.48819003 0.3228447
 0.4833625  0.7648759  0.3808782  0.03462912 0.47573528 0.08879809
 0.40640742 0.47477958 0.2510816  0.1956229  0.35128585 0.08529948
 0.41537035 0.0732687  0.30372617 0.56919426 0.36143333 0.2374325
 0.6636684  0.05995619 0.08239541 0.7272118  0.7552877  0.7154208
 0.83982706 0.13758926 0.15577169 0.05320223 0.7511089  0.6936089
 0.24413522 0.49531832 0.10079875 0.04246275 0.3087379  0.43809676
 0.83275586 0.24883701 0.04448767 0.45507485 0.7482315  0.34243485
 0.6753472  0.00273677 0.10397204 0.6245089  0.06428774 0.2011454
 0.23064888 0.0913609  0.6361733  0.7023753  0.7306451  0.37321466
 0.82028127 0.49488473 0.41946867 0.02048836 0.7587372  0.6299342
 0.23989482 0.3419219  0.14875194 0.4176419  0.25646454 0.6208762
 0.7203263  0.5402459  0.09978905 0.03603924 0.34716    0.01797615
 0.057516   0.19179793 0.20487356 0.01857992 0.6242414  0.15778947
 0.6840965  0.18709503 0.17919308 0.5943542  0.7662706  0.06780756
 0.48222965 0.04869734 0.0300659  0.4694444 ]

[0.25350475 0.00509407 0.11649956 0.77298903 0.16001624 0.01385164
 0.2914175  0.6211722  0.01348003 0.02146206 0.48819003 0.3228447
 0.48336256 0.7648762  0.3808782  0.03462912 0.47573528 0.08879809
 0.40640742 0.4747796  0.25108156 0.1956229  0.35128585 0.08529948
 0.41537035 0.0732687  0.30372617 0.56919426 0.36143333 0.2374325
 0.66366845 0.05995619 0.08239541 0.727212   0.755288   0.71542096
 0.83982766 0.13758926 0.15577169 0.05320224 0.7511091  0.693609
 0.24413522 0.49531832 0.10079875 0.04246275 0.3087379  0.43809676
 0.83275646 0.24883698 0.04448767 0.4550749  0.7482317  0.34243485
 0.6753474  0.00273677 0.10397203 0.624509   0.06428774 0.2011454
 0.2306489  0.0913609  0.6361733  0.7023754  0.7306453  0.37321466
 0.82028174 0.4948848  0.41946867 0.02048836 0.7587375  0.62993425
 0.23989482 0.3419219  0.14875193 0.4176419  0.2564645  0.62087625
 0.7203265  0.54024595 0.09978905 0.03603924 0.34716004 0.01797615
 0.057516   0.19179791 0.20487356 0.01857992 0.6242415  0.15778947
 0.6840967  0.18709503 0.17919308 0.5943543  0.7662709  0.06780756
 0.48222965 0.04869734 0.0300659  0.46944442]

if __name__ == "__main__":
    input_val = np.random.random(50000000).astype(np.float32)
    out = np.zeros_like(input_val)
    
    print("original")
    t_original = %timeit -o gelu_tanh_forward(input_val)
    print("superoptimized")
    t_superopt = %timeit -o cuda_vectorized_gelu(input_val, out=out)
    print("t_original / t_superopt", t_original.best / t_superopt.best)

original
532 ms ± 1.01 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
superoptimized
198 ms ± 1.89 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
t_original / t_superopt 2.7060199654441752

Demo 2: CUDA Backend for Tanh Approximation in GELU Activation Layer¶

Setup GPU Backend¶

Configure the CUDA Ufunc Pipeline¶

Test GELU Ufunc on CUDA¶

Benchmark¶