(
x,
y,
T: tl.constexpr,
D: tl.constexpr,
BT: tl.constexpr
)
| 63 | ) |
| 64 | @triton.jit |
| 65 | def logsigmoid_fwd_kernel( |
| 66 | x, |
| 67 | y, |
| 68 | T: tl.constexpr, |
| 69 | D: tl.constexpr, |
| 70 | BT: tl.constexpr |
| 71 | ): |
| 72 | i = tl.program_id(0) |
| 73 | o_i = i * BT + tl.arange(0, BT) |
| 74 | |
| 75 | p_x = x + o_i |
| 76 | p_y = y + o_i |
| 77 | mask = o_i < T |
| 78 | |
| 79 | # [D,] |
| 80 | b_x = tl.load(p_x, mask=mask, other=0.).to(tl.float32) |
| 81 | b_m = tl.minimum(0., b_x) |
| 82 | b_z = 1. + tl.exp(-tl.abs(b_x)) |
| 83 | b_y = b_m - tl.log(b_z) |
| 84 | tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=mask) |
| 85 | |
| 86 | |
| 87 | @triton.autotune( |
nothing calls this directly
no outgoing calls
no test coverage detected