stdplus-ext 0.0.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- stdplus_ext/__init__.py +4 -0
- stdplus_ext/log_stdplus_triton.py +173 -0
- stdplus_ext/stdplus.py +246 -0
- stdplus_ext/stdplus_triton.py +205 -0
- stdplus_ext-0.0.1.dist-info/METADATA +42 -0
- stdplus_ext-0.0.1.dist-info/RECORD +9 -0
- stdplus_ext-0.0.1.dist-info/WHEEL +5 -0
- stdplus_ext-0.0.1.dist-info/licenses/LICENSE +21 -0
- stdplus_ext-0.0.1.dist-info/top_level.txt +1 -0
stdplus_ext/__init__.py
ADDED
|
@@ -0,0 +1,173 @@
|
|
|
1
|
+
import math
|
|
2
|
+
|
|
3
|
+
import torch
|
|
4
|
+
from torch.utils._triton import has_triton
|
|
5
|
+
|
|
6
|
+
import triton
|
|
7
|
+
from triton import language as tl
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
@triton.autotune(
|
|
11
|
+
configs=[
|
|
12
|
+
triton.Config(kwargs={"BLOCK_SIZE": 64}, num_warps=8, num_stages=1),
|
|
13
|
+
triton.Config(kwargs={"BLOCK_SIZE": 64}, num_warps=16, num_stages=1),
|
|
14
|
+
triton.Config(kwargs={"BLOCK_SIZE": 64}, num_warps=32, num_stages=1),
|
|
15
|
+
triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=8, num_stages=1),
|
|
16
|
+
triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=16, num_stages=1),
|
|
17
|
+
triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=32, num_stages=1),
|
|
18
|
+
triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=8, num_stages=1),
|
|
19
|
+
triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=16, num_stages=1),
|
|
20
|
+
triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=32, num_stages=1),
|
|
21
|
+
],
|
|
22
|
+
key=[],
|
|
23
|
+
)
|
|
24
|
+
@triton.jit
|
|
25
|
+
def log_stdplus_kernel(
|
|
26
|
+
in_ptr0,
|
|
27
|
+
out_ptr,
|
|
28
|
+
out_d_ptr,
|
|
29
|
+
n_elements,
|
|
30
|
+
n_iterations,
|
|
31
|
+
BLOCK_SIZE: "tl.constexpr",
|
|
32
|
+
):
|
|
33
|
+
pid = tl.program_id(axis=0)
|
|
34
|
+
block_start = pid * BLOCK_SIZE
|
|
35
|
+
offsets = block_start + tl.arange(0, BLOCK_SIZE)
|
|
36
|
+
|
|
37
|
+
mask = offsets < n_elements
|
|
38
|
+
|
|
39
|
+
x = tl.load(in_ptr0 + offsets, mask=mask)
|
|
40
|
+
|
|
41
|
+
small_mask = (x < 0.04) & (x > -0.04)
|
|
42
|
+
|
|
43
|
+
b1 = 0.7071067811865476 # M_SQRT1_2
|
|
44
|
+
b2 = 0.25 # 1.0/4
|
|
45
|
+
b3 = 0.007856742013183984 # M_SQRT1_2/90.0
|
|
46
|
+
a1 = 0.5892556509887896 # 5.0*M_SQRT1_2/6.0
|
|
47
|
+
a2 = 0.09444444444444444 # 17.0/180
|
|
48
|
+
|
|
49
|
+
x_small = tl.where(small_mask, x, 0.0)
|
|
50
|
+
fx = x_small * (b1 + x_small * (b2 + x_small * b3))
|
|
51
|
+
d_fx = b1 + x_small * (2.0 * b2 + 3.0 * x_small * b3)
|
|
52
|
+
gx = 1.0 + x_small * (a1 + x_small * (a2))
|
|
53
|
+
d_gx = a1 + 2.0 * x_small * (a2)
|
|
54
|
+
|
|
55
|
+
result = fx / gx
|
|
56
|
+
|
|
57
|
+
result_d = (d_fx * gx - fx * d_gx) / (gx * gx)
|
|
58
|
+
|
|
59
|
+
# branch2
|
|
60
|
+
regular_mask = ~small_mask
|
|
61
|
+
x = tl.where(regular_mask, x, 1.0)
|
|
62
|
+
|
|
63
|
+
xSqr_p_1 = x * x + 1
|
|
64
|
+
logy = tl.math.log((x + tl.math.sqrt(x * x + 4.0)) / 2.0) * 2.0
|
|
65
|
+
|
|
66
|
+
is_positive = x >= 0
|
|
67
|
+
|
|
68
|
+
for i in range(n_iterations):
|
|
69
|
+
exp_nabslogy = tl.exp(-tl.abs(logy))
|
|
70
|
+
|
|
71
|
+
exp_r_a = tl.where(is_positive, 1.0, exp_nabslogy)
|
|
72
|
+
exp_a = tl.where(is_positive, exp_nabslogy, 1.0)
|
|
73
|
+
|
|
74
|
+
logy = ((logy - 1) * exp_r_a + xSqr_p_1 * exp_a) / (exp_r_a - exp_a)
|
|
75
|
+
|
|
76
|
+
d_logy = x / (tl.math.exp(logy) - 1)
|
|
77
|
+
logy = logy / 2
|
|
78
|
+
|
|
79
|
+
result = tl.where(regular_mask, logy, result)
|
|
80
|
+
result_d = tl.where(regular_mask, d_logy, result_d)
|
|
81
|
+
|
|
82
|
+
tl.store(out_ptr + offsets, result, mask=mask)
|
|
83
|
+
tl.store(out_d_ptr + offsets, result_d, mask=mask)
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def _log_stdplus_triton(x, eps: float):
|
|
87
|
+
|
|
88
|
+
n_iterations = int(
|
|
89
|
+
math.ceil(math.log(-math.log(eps) / math.log(2)) / math.log(2) - 1)
|
|
90
|
+
)
|
|
91
|
+
|
|
92
|
+
output = torch.empty_like(x)
|
|
93
|
+
output_d = torch.empty_like(x)
|
|
94
|
+
|
|
95
|
+
n_elements = output.numel()
|
|
96
|
+
|
|
97
|
+
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
|
|
98
|
+
|
|
99
|
+
log_stdplus_kernel[grid](x, output, output_d, n_elements, n_iterations)
|
|
100
|
+
|
|
101
|
+
return output, output_d
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
class _LOG_STDPLUS_TRITON(torch.autograd.Function):
|
|
105
|
+
|
|
106
|
+
@staticmethod
|
|
107
|
+
def forward(ctx, x):
|
|
108
|
+
eps = torch.finfo(x.dtype).eps
|
|
109
|
+
|
|
110
|
+
logy, d_logy = _log_stdplus_triton(x, eps)
|
|
111
|
+
|
|
112
|
+
ctx.save_for_backward(d_logy)
|
|
113
|
+
|
|
114
|
+
return logy
|
|
115
|
+
|
|
116
|
+
@staticmethod
|
|
117
|
+
def backward(ctx, grad_output):
|
|
118
|
+
(derivative,) = ctx.saved_tensors
|
|
119
|
+
|
|
120
|
+
return grad_output * derivative
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
log_stdplus_triton = _LOG_STDPLUS_TRITON.apply
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
if __name__ == "__main__":
|
|
127
|
+
|
|
128
|
+
import matplotlib.pyplot as plt
|
|
129
|
+
import time
|
|
130
|
+
import torch.nn.functional as F
|
|
131
|
+
import pytorch_stdplus
|
|
132
|
+
|
|
133
|
+
# x = torch.linspace(-0.08,0.08,65535, device="cuda")
|
|
134
|
+
# x = torch.linspace(-1.00,1.00,65535, device="cuda")
|
|
135
|
+
x = torch.linspace(-5.00, 5.00, 65535, device="cuda")
|
|
136
|
+
|
|
137
|
+
x.requires_grad_()
|
|
138
|
+
|
|
139
|
+
out = log_stdplus_triton(x).exp()
|
|
140
|
+
|
|
141
|
+
t1 = time.time()
|
|
142
|
+
for i in range(1000):
|
|
143
|
+
out = F.softplus(x)
|
|
144
|
+
# out = torch.log(torch.exp(x)+1)
|
|
145
|
+
t2 = time.time()
|
|
146
|
+
print(t2 - t1)
|
|
147
|
+
|
|
148
|
+
t1 = time.time()
|
|
149
|
+
for i in range(1000):
|
|
150
|
+
yGT = log_stdplus_triton(x) # .exp()
|
|
151
|
+
t2 = time.time()
|
|
152
|
+
print(t2 - t1)
|
|
153
|
+
|
|
154
|
+
t1 = time.time()
|
|
155
|
+
for i in range(1000):
|
|
156
|
+
out = pytorch_stdplus.log_stdplus(x) # .exp()
|
|
157
|
+
t2 = time.time()
|
|
158
|
+
print(t2 - t1)
|
|
159
|
+
|
|
160
|
+
# y**2 - 2*ln(y) - 1 = x**2
|
|
161
|
+
# torch.expm(out*2) - 2*out - 1
|
|
162
|
+
|
|
163
|
+
# import slashed_normal
|
|
164
|
+
# yGT = slashed_normal.log_stdplus(x)
|
|
165
|
+
|
|
166
|
+
d_GT = torch.autograd.grad(yGT.sum(), x)[0]
|
|
167
|
+
out_d = torch.autograd.grad(out.sum(), x)[0]
|
|
168
|
+
|
|
169
|
+
plt.plot(x.cpu().detach(), out.cpu().detach())
|
|
170
|
+
plt.plot(x.cpu().detach(), out_d.cpu().detach())
|
|
171
|
+
plt.show()
|
|
172
|
+
print((yGT - out).abs().max())
|
|
173
|
+
print((d_GT - out_d).abs().max())
|
stdplus_ext/stdplus.py
ADDED
|
@@ -0,0 +1,246 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import math
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
def handles_scalar_tensor(func):
|
|
6
|
+
def wrapper(x):
|
|
7
|
+
if torch.is_tensor(x) and x.dim() == 0:
|
|
8
|
+
x = x.unsqueeze(-1)
|
|
9
|
+
return func(x).squeeze(-1)
|
|
10
|
+
else:
|
|
11
|
+
return func(x)
|
|
12
|
+
|
|
13
|
+
return wrapper
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
@torch.jit.script
|
|
17
|
+
def _log_stdplus(x, eps: float, exponential: bool = False):
|
|
18
|
+
|
|
19
|
+
# iters = int(math.ceil(math.log(- math.log(eps)/math.log(2))/math.log(2)-1))
|
|
20
|
+
iters = 4
|
|
21
|
+
|
|
22
|
+
# Pre-compute constants once
|
|
23
|
+
b1 = 1.0 / math.sqrt(2)
|
|
24
|
+
b2 = 0.25
|
|
25
|
+
b3 = b1 / 90.0
|
|
26
|
+
a1 = 5.0 * b1 / 6.0
|
|
27
|
+
a2 = 17.0 / 180.0
|
|
28
|
+
|
|
29
|
+
# Compute all intermediate values in one go
|
|
30
|
+
x_sq = x * x
|
|
31
|
+
|
|
32
|
+
# Small values optimization (-0.04 < x < 0.04)
|
|
33
|
+
small_vals_mask = (x > -0.04) & (x < 0.04)
|
|
34
|
+
|
|
35
|
+
# Initialize result tensor
|
|
36
|
+
result = torch.empty_like(x)
|
|
37
|
+
|
|
38
|
+
# Handle small values first
|
|
39
|
+
if small_vals_mask.any():
|
|
40
|
+
# Horner's method for polynomial evaluation
|
|
41
|
+
fx = x[small_vals_mask] * (
|
|
42
|
+
b1 + x[small_vals_mask] * (b2 + x[small_vals_mask] * b3)
|
|
43
|
+
)
|
|
44
|
+
gx = 1.0 + x[small_vals_mask] * (a1 + x[small_vals_mask] * a2)
|
|
45
|
+
small_result = fx / gx
|
|
46
|
+
|
|
47
|
+
result[small_vals_mask] = small_result
|
|
48
|
+
|
|
49
|
+
# Handle non-small values
|
|
50
|
+
non_small_mask = ~small_vals_mask
|
|
51
|
+
if non_small_mask.any():
|
|
52
|
+
# Compute initial values for non-small x
|
|
53
|
+
x_non_small = x[non_small_mask]
|
|
54
|
+
x_sq_non_small = x_sq[non_small_mask]
|
|
55
|
+
|
|
56
|
+
# Optimize sqrt computation
|
|
57
|
+
# xsqr_p_4 = torch.sqrt(x_sq_non_small + 4.0)
|
|
58
|
+
xsqr_p_4 = torch.hypot(x_non_small, torch.ones_like(x_non_small) * 2.0)
|
|
59
|
+
|
|
60
|
+
tmp = torch.where(
|
|
61
|
+
x_non_small < 0, 4.0 / (xsqr_p_4 - x_non_small), x_non_small + xsqr_p_4
|
|
62
|
+
)
|
|
63
|
+
|
|
64
|
+
# logy = 2.0 * torch.log((x_non_small + xsqr_p_4) * 0.5)
|
|
65
|
+
logy = 2.0 * torch.log(tmp * 0.5)
|
|
66
|
+
|
|
67
|
+
xSqr_p_1 = x_sq_non_small + 1.0
|
|
68
|
+
# print(logy)
|
|
69
|
+
|
|
70
|
+
# Optimization: pre-compute is_positive
|
|
71
|
+
is_positive = x_non_small >= 0
|
|
72
|
+
|
|
73
|
+
# Main iteration loop
|
|
74
|
+
for _ in range(iters):
|
|
75
|
+
abs_logy = torch.abs(logy)
|
|
76
|
+
exp_nabslogy = torch.exp(-abs_logy)
|
|
77
|
+
|
|
78
|
+
# Vectorized condition handling
|
|
79
|
+
exp_r_a = torch.where(is_positive, torch.ones_like(logy), exp_nabslogy)
|
|
80
|
+
exp_a = torch.where(is_positive, exp_nabslogy, torch.ones_like(logy))
|
|
81
|
+
|
|
82
|
+
# Combined computation
|
|
83
|
+
numerator = (logy - 1) * exp_r_a + xSqr_p_1 * exp_a
|
|
84
|
+
denominator = exp_r_a - exp_a
|
|
85
|
+
|
|
86
|
+
# Update in-place
|
|
87
|
+
logy = numerator / denominator
|
|
88
|
+
|
|
89
|
+
# Final computations
|
|
90
|
+
logy = logy * 0.5
|
|
91
|
+
|
|
92
|
+
result[non_small_mask] = logy
|
|
93
|
+
|
|
94
|
+
if exponential:
|
|
95
|
+
result = torch.exp(result)
|
|
96
|
+
|
|
97
|
+
return result
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
@torch.jit.script
|
|
101
|
+
def _d_log_stdplus(x, eps: float, exponential: bool = False):
|
|
102
|
+
"""
|
|
103
|
+
Optimized PyTorch implementation of log_stdplus function.
|
|
104
|
+
|
|
105
|
+
Args:
|
|
106
|
+
x (torch.Tensor): Input tensor
|
|
107
|
+
iters (int): Number of iterations for the approximation
|
|
108
|
+
exponential (bool): If True, returns exp(result) for stdplus computation
|
|
109
|
+
|
|
110
|
+
Returns:
|
|
111
|
+
torch.Tensor: Computed result
|
|
112
|
+
"""
|
|
113
|
+
iters = int(math.ceil(math.log(-math.log(eps) / math.log(2)) / math.log(2) - 1))
|
|
114
|
+
|
|
115
|
+
# Pre-compute constants once
|
|
116
|
+
b1 = 1.0 / math.sqrt(2)
|
|
117
|
+
b2 = 0.25
|
|
118
|
+
b3 = b1 / 90.0
|
|
119
|
+
a1 = 5.0 * b1 / 6.0
|
|
120
|
+
a2 = 17.0 / 180.0
|
|
121
|
+
|
|
122
|
+
# Compute all intermediate values in one go
|
|
123
|
+
x_sq = x * x
|
|
124
|
+
|
|
125
|
+
# Small values optimization (-0.04 < x < 0.04)
|
|
126
|
+
small_vals_mask = (x > -0.04) & (x < 0.04)
|
|
127
|
+
|
|
128
|
+
# Initialize result tensor
|
|
129
|
+
result = torch.empty_like(x)
|
|
130
|
+
|
|
131
|
+
# Handle small values first
|
|
132
|
+
if small_vals_mask.any():
|
|
133
|
+
# Horner's method for polynomial evaluation
|
|
134
|
+
fx = x[small_vals_mask] * (
|
|
135
|
+
b1 + x[small_vals_mask] * (b2 + x[small_vals_mask] * b3)
|
|
136
|
+
)
|
|
137
|
+
gx = 1.0 + x[small_vals_mask] * (a1 + x[small_vals_mask] * a2)
|
|
138
|
+
d_fx = b1 + x[small_vals_mask] * (2.0 * b2 + x[small_vals_mask] * (3.0 * b3))
|
|
139
|
+
d_gx = a1 + x[small_vals_mask] * (2.0 * a2)
|
|
140
|
+
|
|
141
|
+
small_result = (d_fx * gx - fx * d_gx) / gx / gx
|
|
142
|
+
|
|
143
|
+
if exponential:
|
|
144
|
+
small_result = torch.exp(fx / gx) * small_result
|
|
145
|
+
|
|
146
|
+
result[small_vals_mask] = small_result
|
|
147
|
+
|
|
148
|
+
# Handle non-small values
|
|
149
|
+
non_small_mask = ~small_vals_mask
|
|
150
|
+
if non_small_mask.any():
|
|
151
|
+
# Compute initial values for non-small x
|
|
152
|
+
x_non_small = x[non_small_mask]
|
|
153
|
+
x_sq_non_small = x_sq[non_small_mask]
|
|
154
|
+
|
|
155
|
+
# Optimize sqrt computation
|
|
156
|
+
# xsqr_p_4 = torch.sqrt(x_sq_non_small + 4.0)
|
|
157
|
+
# logy = 2.0 * torch.log((x_non_small + xsqr_p_4) * 0.5)
|
|
158
|
+
xsqr_p_4 = torch.hypot(x_non_small, torch.ones_like(x_non_small) * 2.0)
|
|
159
|
+
tmp = torch.where(
|
|
160
|
+
x_non_small < 0, 4.0 / (xsqr_p_4 - x_non_small), x_non_small + xsqr_p_4
|
|
161
|
+
)
|
|
162
|
+
|
|
163
|
+
# logy = 2.0 * torch.log((x_non_small + xsqr_p_4) * 0.5)
|
|
164
|
+
logy = 2.0 * torch.log(tmp * 0.5)
|
|
165
|
+
xSqr_p_1 = x_sq_non_small + 1.0
|
|
166
|
+
|
|
167
|
+
# Optimization: pre-compute is_positive
|
|
168
|
+
is_positive = x_non_small >= 0
|
|
169
|
+
|
|
170
|
+
# Main iteration loop
|
|
171
|
+
for _ in range(iters):
|
|
172
|
+
abs_logy = torch.abs(logy)
|
|
173
|
+
exp_nabslogy = torch.exp(-abs_logy)
|
|
174
|
+
|
|
175
|
+
# Vectorized condition handling
|
|
176
|
+
exp_r_a = torch.where(is_positive, torch.ones_like(logy), exp_nabslogy)
|
|
177
|
+
exp_a = torch.where(is_positive, exp_nabslogy, torch.ones_like(logy))
|
|
178
|
+
|
|
179
|
+
# Combined computation
|
|
180
|
+
numerator = (logy - 1) * exp_r_a + xSqr_p_1 * exp_a
|
|
181
|
+
denominator = exp_r_a - exp_a
|
|
182
|
+
|
|
183
|
+
# Update in-place
|
|
184
|
+
logy = numerator / denominator
|
|
185
|
+
|
|
186
|
+
# Final computations
|
|
187
|
+
logy = logy * 0.5
|
|
188
|
+
|
|
189
|
+
d_logy = x_non_small / torch.expm1(logy * 2)
|
|
190
|
+
if exponential:
|
|
191
|
+
d_logy = torch.exp(logy) * d_logy
|
|
192
|
+
|
|
193
|
+
result[non_small_mask] = d_logy
|
|
194
|
+
|
|
195
|
+
return result
|
|
196
|
+
|
|
197
|
+
|
|
198
|
+
from torch.utils._triton import has_triton as _has_triton
|
|
199
|
+
|
|
200
|
+
if _has_triton:
|
|
201
|
+
from .log_stdplus_triton import log_stdplus_triton as log_stdplus_triton
|
|
202
|
+
from .stdplus_triton import stdplus_triton as stdplus_triton
|
|
203
|
+
|
|
204
|
+
|
|
205
|
+
class _LOG_STDPLUS(torch.autograd.Function):
|
|
206
|
+
|
|
207
|
+
@staticmethod
|
|
208
|
+
def forward(ctx, x):
|
|
209
|
+
ctx.save_for_backward(x)
|
|
210
|
+
eps = torch.finfo(x.dtype).eps
|
|
211
|
+
return _log_stdplus(x, eps, exponential=False)
|
|
212
|
+
|
|
213
|
+
@staticmethod
|
|
214
|
+
def backward(ctx, grad_output):
|
|
215
|
+
(x,) = ctx.saved_tensors
|
|
216
|
+
eps = torch.finfo(x.dtype).eps
|
|
217
|
+
derivative = _d_log_stdplus(x, eps, exponential=False)
|
|
218
|
+
|
|
219
|
+
return grad_output * derivative
|
|
220
|
+
|
|
221
|
+
|
|
222
|
+
@handles_scalar_tensor
|
|
223
|
+
def stdplus(x):
|
|
224
|
+
if x.is_cuda and _has_triton:
|
|
225
|
+
return stdplus_triton(x)
|
|
226
|
+
|
|
227
|
+
return _LOG_STDPLUS.apply(x).exp()
|
|
228
|
+
|
|
229
|
+
|
|
230
|
+
@handles_scalar_tensor
|
|
231
|
+
def log_stdplus(x):
|
|
232
|
+
if x.is_cuda and _has_triton:
|
|
233
|
+
return log_stdplus_triton(x)
|
|
234
|
+
return _LOG_STDPLUS.apply(x)
|
|
235
|
+
|
|
236
|
+
|
|
237
|
+
@handles_scalar_tensor
|
|
238
|
+
def inv_log_stdplus(y):
|
|
239
|
+
|
|
240
|
+
y = y * 2
|
|
241
|
+
return torch.sgn(y) * (torch.expm1(y) - y).sqrt()
|
|
242
|
+
|
|
243
|
+
|
|
244
|
+
@handles_scalar_tensor
|
|
245
|
+
def inv_stdplus(y):
|
|
246
|
+
return inv_log_stdplus(y.log())
|
|
@@ -0,0 +1,205 @@
|
|
|
1
|
+
import math
|
|
2
|
+
|
|
3
|
+
import torch
|
|
4
|
+
from torch.utils._triton import has_triton
|
|
5
|
+
|
|
6
|
+
import triton
|
|
7
|
+
from triton import language as tl
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
@triton.autotune(
|
|
11
|
+
configs=[
|
|
12
|
+
triton.Config(kwargs={"BLOCK_SIZE": 64}, num_warps=8, num_stages=4),
|
|
13
|
+
triton.Config(kwargs={"BLOCK_SIZE": 64}, num_warps=16, num_stages=4),
|
|
14
|
+
triton.Config(kwargs={"BLOCK_SIZE": 64}, num_warps=32, num_stages=4),
|
|
15
|
+
triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=8, num_stages=4),
|
|
16
|
+
triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=16, num_stages=4),
|
|
17
|
+
triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=32, num_stages=4),
|
|
18
|
+
triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=8, num_stages=4),
|
|
19
|
+
triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=16, num_stages=4),
|
|
20
|
+
triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=32, num_stages=4),
|
|
21
|
+
],
|
|
22
|
+
key=["n_iterations"],
|
|
23
|
+
)
|
|
24
|
+
@triton.jit
|
|
25
|
+
def stdplus_kernel(
|
|
26
|
+
in_ptr0,
|
|
27
|
+
out_ptr,
|
|
28
|
+
out_d_ptr,
|
|
29
|
+
n_elements,
|
|
30
|
+
n_iterations,
|
|
31
|
+
BLOCK_SIZE: "tl.constexpr",
|
|
32
|
+
):
|
|
33
|
+
pid = tl.program_id(axis=0)
|
|
34
|
+
block_start = pid * BLOCK_SIZE
|
|
35
|
+
offsets = block_start + tl.arange(0, BLOCK_SIZE)
|
|
36
|
+
|
|
37
|
+
mask = offsets < n_elements
|
|
38
|
+
|
|
39
|
+
x = tl.load(in_ptr0 + offsets, mask=mask)
|
|
40
|
+
|
|
41
|
+
small_mask = (x < 0.04) & (x > -0.04)
|
|
42
|
+
|
|
43
|
+
# firstly computes log_stdplus
|
|
44
|
+
|
|
45
|
+
# b1 = 1.0/tl.math.sqrt(2.0)
|
|
46
|
+
# b2 = 1.0/4
|
|
47
|
+
# b3 = 1.0/(90*tl.math.sqrt(2.0))
|
|
48
|
+
|
|
49
|
+
# a1 = 5.0/(6*tl.math.sqrt(2.0))
|
|
50
|
+
|
|
51
|
+
b1 = 0.7071067811865476 # M_SQRT1_2
|
|
52
|
+
b2 = 0.25 # 1.0/4
|
|
53
|
+
b3 = 0.007856742013183984 # M_SQRT1_2/90.0
|
|
54
|
+
a1 = 0.5892556509887896 # 5.0*M_SQRT1_2/6.0
|
|
55
|
+
a2 = 0.09444444444444444 # 17.0/180
|
|
56
|
+
|
|
57
|
+
# result = tl.zeros_like(x)
|
|
58
|
+
# result_d = tl.zeros_like(x)
|
|
59
|
+
|
|
60
|
+
# logy = x1*(b1+ x1*(b2+ x1*b3)) \
|
|
61
|
+
# /(1+x1*(a1+ x1*(a2)))
|
|
62
|
+
x_small = tl.where(small_mask, x, 0.0)
|
|
63
|
+
fx = x_small * (b1 + x_small * (b2 + x_small * b3))
|
|
64
|
+
d_fx = b1 + x_small * (2.0 * b2 + 3.0 * x_small * b3)
|
|
65
|
+
gx = 1.0 + x_small * (a1 + x_small * (a2))
|
|
66
|
+
d_gx = a1 + 2.0 * x_small * (a2)
|
|
67
|
+
|
|
68
|
+
# result = tl.where(small_mask, fx/gx, result)
|
|
69
|
+
|
|
70
|
+
# result_d = tl.where(small_mask, (d_fx*gx - fx*d_gx)/ (gx*gx), result_d)
|
|
71
|
+
result = fx / gx
|
|
72
|
+
|
|
73
|
+
result_d = (d_fx * gx - fx * d_gx) / (gx * gx)
|
|
74
|
+
|
|
75
|
+
# branch2
|
|
76
|
+
regular_mask = ~small_mask
|
|
77
|
+
x = tl.where(regular_mask, x, 1.0)
|
|
78
|
+
|
|
79
|
+
xSqr_p_1 = x * x + 1
|
|
80
|
+
logy = tl.math.log((x + tl.math.sqrt(x * x + 4.0)) / 2.0) * 2.0
|
|
81
|
+
|
|
82
|
+
# for i in range(n_iterations):
|
|
83
|
+
|
|
84
|
+
# a = tl.math.max(logy, 0.0)
|
|
85
|
+
# exp_r_a = tl.math.exp(logy-a)
|
|
86
|
+
# exp_a = tl.math.exp(-a)
|
|
87
|
+
# logy = ((logy-1)*exp_r_a + xSqr_p_1* exp_a)/( exp_r_a - exp_a)
|
|
88
|
+
|
|
89
|
+
is_positive = x >= 0
|
|
90
|
+
|
|
91
|
+
# for i in range(n_iterations):
|
|
92
|
+
for i in range(n_iterations):
|
|
93
|
+
exp_nabslogy = tl.exp(-tl.abs(logy))
|
|
94
|
+
|
|
95
|
+
exp_r_a = tl.where(is_positive, 1.0, exp_nabslogy)
|
|
96
|
+
exp_a = tl.where(is_positive, exp_nabslogy, 1.0)
|
|
97
|
+
|
|
98
|
+
logy = ((logy - 1) * exp_r_a + xSqr_p_1 * exp_a) / (exp_r_a - exp_a)
|
|
99
|
+
|
|
100
|
+
logy = logy / 2
|
|
101
|
+
# expm1_2_logy = torch.expm1(logy*2)
|
|
102
|
+
d_logy = x / (tl.math.exp(logy * 2) - 1)
|
|
103
|
+
|
|
104
|
+
result = tl.where(regular_mask, logy, result)
|
|
105
|
+
result_d = tl.where(regular_mask, d_logy, result_d)
|
|
106
|
+
|
|
107
|
+
## transform
|
|
108
|
+
result = tl.math.exp(result)
|
|
109
|
+
result_d = result_d * result
|
|
110
|
+
|
|
111
|
+
tl.store(out_ptr + offsets, result, mask=mask)
|
|
112
|
+
tl.store(out_d_ptr + offsets, result_d, mask=mask)
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
def _stdplus_triton(x, eps: float):
|
|
116
|
+
|
|
117
|
+
n_iterations = int(
|
|
118
|
+
math.ceil(math.log(-math.log(eps) / math.log(2)) / math.log(2) - 1)
|
|
119
|
+
)
|
|
120
|
+
|
|
121
|
+
output = torch.empty_like(x)
|
|
122
|
+
output_d = torch.empty_like(x)
|
|
123
|
+
|
|
124
|
+
n_elements = output.numel()
|
|
125
|
+
|
|
126
|
+
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
|
|
127
|
+
|
|
128
|
+
stdplus_kernel[grid](x, output, output_d, n_elements, n_iterations)
|
|
129
|
+
|
|
130
|
+
return output, output_d
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
class _STDPLUS_TRITON(torch.autograd.Function):
|
|
134
|
+
|
|
135
|
+
@staticmethod
|
|
136
|
+
def forward(ctx, x):
|
|
137
|
+
eps = torch.finfo(x.dtype).eps
|
|
138
|
+
|
|
139
|
+
logy, d_logy = _stdplus_triton(x, eps)
|
|
140
|
+
|
|
141
|
+
ctx.save_for_backward(d_logy)
|
|
142
|
+
|
|
143
|
+
return logy
|
|
144
|
+
|
|
145
|
+
@staticmethod
|
|
146
|
+
def backward(ctx, grad_output):
|
|
147
|
+
(derivative,) = ctx.saved_tensors
|
|
148
|
+
|
|
149
|
+
return grad_output * derivative
|
|
150
|
+
|
|
151
|
+
|
|
152
|
+
stdplus_triton = _STDPLUS_TRITON.apply
|
|
153
|
+
|
|
154
|
+
|
|
155
|
+
if __name__ == "__main__":
|
|
156
|
+
|
|
157
|
+
import matplotlib.pyplot as plt
|
|
158
|
+
import time
|
|
159
|
+
import torch.nn.functional as F
|
|
160
|
+
import pytorch_stdplus
|
|
161
|
+
|
|
162
|
+
# x = torch.linspace(-0.08,0.08,65535, device="cuda")
|
|
163
|
+
# x = torch.linspace(-1.00,1.00,65535, device="cuda")
|
|
164
|
+
x = torch.linspace(-5.00, 5.00, 6553500, device="cuda")
|
|
165
|
+
|
|
166
|
+
x.requires_grad_()
|
|
167
|
+
|
|
168
|
+
out = stdplus_triton(x)
|
|
169
|
+
|
|
170
|
+
print("gelu")
|
|
171
|
+
t1 = time.time()
|
|
172
|
+
for i in range(10000):
|
|
173
|
+
out1 = F.gelu(x)
|
|
174
|
+
# out = torch.log(torch.exp(x)+1)
|
|
175
|
+
t2 = time.time()
|
|
176
|
+
print("gelu", t2 - t1)
|
|
177
|
+
|
|
178
|
+
t1 = time.time()
|
|
179
|
+
for i in range(10000):
|
|
180
|
+
out2 = stdplus_triton(x) # .exp()
|
|
181
|
+
t2 = time.time()
|
|
182
|
+
print("triton", t2 - t1)
|
|
183
|
+
|
|
184
|
+
out3 = pytorch_stdplus.stdplus(x) # .exp()
|
|
185
|
+
out3 = pytorch_stdplus.stdplus(x) # .exp()
|
|
186
|
+
t1 = time.time()
|
|
187
|
+
for i in range(10000):
|
|
188
|
+
out3 = pytorch_stdplus.stdplus(x) # .exp()
|
|
189
|
+
t2 = time.time()
|
|
190
|
+
print("cuda", t2 - t1)
|
|
191
|
+
|
|
192
|
+
print((out3**2 - 2 * torch.log(out3) - 1 - x**2).abs().max())
|
|
193
|
+
|
|
194
|
+
print((out3 - out2).abs().max())
|
|
195
|
+
# import slashed_normal
|
|
196
|
+
# yGT = slashed_normal.log_stdplus(x)
|
|
197
|
+
|
|
198
|
+
out_d_2 = torch.autograd.grad(out2.sum(), x)[0]
|
|
199
|
+
out_d_3 = torch.autograd.grad(out3.sum(), x)[0]
|
|
200
|
+
print((out_d_2 - out_d_3).abs().max())
|
|
201
|
+
|
|
202
|
+
plt.plot(x.cpu().detach(), out2.cpu().detach())
|
|
203
|
+
plt.plot(x.cpu().detach(), out_d_2.cpu().detach())
|
|
204
|
+
plt.show()
|
|
205
|
+
# print((yGT-out).abs().max())
|
|
@@ -0,0 +1,42 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: stdplus_ext
|
|
3
|
+
Version: 0.0.1
|
|
4
|
+
Summary: The stdplus activation function
|
|
5
|
+
Author: Yujia Yan
|
|
6
|
+
License: MIT
|
|
7
|
+
Requires-Python: >=3.6
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Dynamic: license-file
|
|
11
|
+
|
|
12
|
+
# stdplus
|
|
13
|
+
|
|
14
|
+
`stdplus(x)` is a smooth activation function for computing **Gaussian standard deviations** in variational models. It is designed as a numerically stable alternative to `softplus` / `exp`, and is introduced in *Chapter 6* of Yujia Yan’s PhD dissertation (University of Rochester, 2025).
|
|
15
|
+
It resembles a **1-centered softplus** (i.e., equals 1 at `x = 0`) while providing a clean analytic derivative and a stable evaluation recipe.
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
## Definition
|
|
19
|
+
|
|
20
|
+
stdplus function is defined as a function y = f(x) such that
|
|
21
|
+
$$y^2- 2 \log(y) - 1 = x^2$$
|
|
22
|
+
|
|
23
|
+
This is proposed in Chapter 6 of Yujia Yan's dissertation: Slashed Normal Parameterization for Approximate Normal Posterior Distributions
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
## Why stdplus?
|
|
27
|
+
|
|
28
|
+
In the **Slashed Normal / KL-amplitude** parameterization, Gaussian parameters are expressed as a complex “KL amplitude” `ψ = a + bi`, and the standard deviation is computed with:
|
|
29
|
+
|
|
30
|
+
- `μ = sqrt(2) * a`
|
|
31
|
+
- `σ = stdplus(sqrt(2) * b)`
|
|
32
|
+
|
|
33
|
+
This yields an exact closed-form KL term equal to `||ψ||_2^2` (enabling direct, interpretable rate control).
|
|
34
|
+
|
|
35
|
+
## Reference
|
|
36
|
+
|
|
37
|
+
- Yujia Yan, *Structured Analysis and Generation in Music, Audio, and Beyond*, PhD Dissertation, University of Rochester, 2025. (See Chapter 6, Sec. 6.6 “More on stdplus Function”.)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
## Install
|
|
41
|
+
pip3 install -e .
|
|
42
|
+
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
stdplus_ext/__init__.py,sha256=zoFOJTXy_FJ2uS-l3edTTpiYnR1HQmtYOsvkqfzFwqY,145
|
|
2
|
+
stdplus_ext/log_stdplus_triton.py,sha256=o0LnAAjOobN0KSl3sTs3UiRUUe5dmnPTMxuPWqzsLIs,4660
|
|
3
|
+
stdplus_ext/stdplus.py,sha256=IubHvqmtqj8ke1UVBlrLa89W2RPzl7dyRCZeWXTAfFE,7019
|
|
4
|
+
stdplus_ext/stdplus_triton.py,sha256=gXc_Fx7yD7B_k_KIIm2bOVnbmQEYUwEJcKScUs6M3Yc,5620
|
|
5
|
+
stdplus_ext-0.0.1.dist-info/licenses/LICENSE,sha256=sWiaHDoxM4R3BORdvarD7InLZKH0xYG4e_ggJVgw27s,1066
|
|
6
|
+
stdplus_ext-0.0.1.dist-info/METADATA,sha256=bY2TAUJ_-bdtkfwhDvnkaJJXeVHDqJC1NRbN4JXYEK4,1532
|
|
7
|
+
stdplus_ext-0.0.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
|
|
8
|
+
stdplus_ext-0.0.1.dist-info/top_level.txt,sha256=RuYSH0DbYHpESQBQME0_98sGE9xkMOhi3Zv6ZGtYLZE,12
|
|
9
|
+
stdplus_ext-0.0.1.dist-info/RECORD,,
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Yujia Yan
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
stdplus_ext
|