netcl 0.3.0__tar.gz → 0.3.2__tar.gz
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.
- {netcl-0.3.0/netcl.egg-info → netcl-0.3.2}/PKG-INFO +2 -1
- {netcl-0.3.0 → netcl-0.3.2}/README.md +1 -1
- {netcl-0.3.0 → netcl-0.3.2}/README_PACKAGE.md +1 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/autograd/ops.py +63 -7
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/conv2d.py +20 -0
- netcl-0.3.2/netcl/ops/conv2d_pointwise.py +520 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/depthwise_conv2d.py +236 -11
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/elementwise.py +44 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/reduction.py +39 -2
- {netcl-0.3.0 → netcl-0.3.2}/netcl/optim/clip.py +71 -18
- {netcl-0.3.0 → netcl-0.3.2/netcl.egg-info}/PKG-INFO +2 -1
- {netcl-0.3.0 → netcl-0.3.2}/netcl.egg-info/SOURCES.txt +1 -0
- {netcl-0.3.0 → netcl-0.3.2}/pyproject.toml +1 -1
- {netcl-0.3.0 → netcl-0.3.2}/LICENSE +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/MANIFEST.in +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/amp.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/autograd/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/autograd/compiler.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/autograd/debug.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/autograd/engine.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/autograd/forward_compiler.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/autograd/graph.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/autograd/shape_ops.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/autograd/training_compiler.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/cluster/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/cluster/dbscan.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/cluster/gmm.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/cluster/kmeans.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/cluster/metrics.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/cluster/ops.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/cluster/spectral.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/autotune.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/backend/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/backend/base.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/backend/cpu.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/backend/opencl.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/capabilities.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/device.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/device_profile.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/kernel_selector.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/kernels/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/kernels/primitives.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/memory.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/parameter.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/shape.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/core/tensor.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/data/augment.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/data/augment_gpu.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/data/dataloader.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/data/filters.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/data/gpu_batch.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/data/shared_batch.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/distributed/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/distributed/collectives.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/distributed/data_parallel.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/distributed/device_manager.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/distributed/trainer.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/evo/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/evo/benchmarks.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/evo/engine.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/evo/fitness.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/evo/genome.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/evo/ops.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/evo/population.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/evo/random.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/evo/strategies.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/io/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/io/checkpoint.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/io/interop.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/io/onnx_export.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/io/serialization.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/batchnorm.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/containers.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/decorators.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/factory.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/functional.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/groupnorm.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/init.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/layernorm.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/layers.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/loss.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/margin.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/modern.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/modules.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/padding.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/pooling.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/residual.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/resnet.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/simple.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/nn/transformer.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/attention.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/bmm.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/broadcast.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/conv2d_cpu.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/conv2d_dilated.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/conv2d_optimized.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/conv2d_planner.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/conv_transpose2d.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/elementwise_optimized.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/fused_ops.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/im2col.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/indexing.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/jit_fusion.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/lm_ops.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/matmul.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/matmul_optimized.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/norm_rope.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/permute.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/softmax.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/softmax_fp16.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/softmax_fused.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/tensor_ops.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/transpose.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/upsample.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ops/winograd_fused.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/optim/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/optim/_utils.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/optim/adam.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/optim/adamw.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/optim/amp.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/optim/ema.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/optim/lr_plateau.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/optim/lr_scheduler.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/optim/momentum.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/optim/rmsprop.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/optim/sgd.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/profiling/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/profiling/autotuner.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/profiling/timing.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/runtime/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/runtime/cache.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/runtime/capture.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/runtime/graph.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/runtime/graph_fusion.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/runtime/perf.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/runtime/scheduler.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ssl/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ssl/deep_cluster.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ssl/evaluate.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ssl/heads.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ssl/losses.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ssl/methods.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/ssl/views.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/text/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/text/bpe.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/text/chat.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/trainer/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/trainer/trainer.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/utils/__init__.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/utils/data.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl/utils/progress.py +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl.egg-info/dependency_links.txt +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl.egg-info/requires.txt +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/netcl.egg-info/top_level.txt +0 -0
- {netcl-0.3.0 → netcl-0.3.2}/setup.cfg +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.1
|
|
2
2
|
Name: netcl
|
|
3
|
-
Version: 0.3.
|
|
3
|
+
Version: 0.3.2
|
|
4
4
|
Summary: PyOpenCL-based deep learning playground with autograd, kernels, and high-level APIs.
|
|
5
5
|
Author: netcl contributors
|
|
6
6
|
License: MIT
|
|
@@ -148,3 +148,4 @@ model2 = load_model("model_artifact")
|
|
|
148
148
|
## Notes
|
|
149
149
|
- Mixed precision is experimental; it is disabled on the CPU backend.
|
|
150
150
|
- Conv2d algorithms can be tuned via env flags like `NETCL_CONV_AUTOTUNE=1`.
|
|
151
|
+
- 1x1 convolutions (stride 1, no padding) are timed once per shape class against the general conv path and run as batched GEMMs where that is faster (`NETCL_CONV_POINTWISE=auto|1|0`); the decision is cached per device in the autotune cache. Depthwise convolutions with filters up to 7x7 and stride 1/2 use local-memory tiled kernels (`NETCL_DWCONV_ALGO=local|naive|tile`). `NETCL_AUTOTUNE=0` disables all measurements and keeps the previous conv paths.
|
|
@@ -49,4 +49,4 @@ trainer.fit(loader, epochs=1, loss_fn=functional.cross_entropy)
|
|
|
49
49
|
|
|
50
50
|
## Docs
|
|
51
51
|
- `README_PACKAGE.md` for the package summary used on PyPI.
|
|
52
|
-
- `wiki/
|
|
52
|
+
- `web-wiki/content/` for the documentation (API reference, concepts, tutorials); serve it with `pip install -r web-wiki/requirements.txt && python web-wiki/app.py`.
|
|
@@ -134,3 +134,4 @@ model2 = load_model("model_artifact")
|
|
|
134
134
|
## Notes
|
|
135
135
|
- Mixed precision is experimental; it is disabled on the CPU backend.
|
|
136
136
|
- Conv2d algorithms can be tuned via env flags like `NETCL_CONV_AUTOTUNE=1`.
|
|
137
|
+
- 1x1 convolutions (stride 1, no padding) are timed once per shape class against the general conv path and run as batched GEMMs where that is faster (`NETCL_CONV_POINTWISE=auto|1|0`); the decision is cached per device in the autotune cache. Depthwise convolutions with filters up to 7x7 and stride 1/2 use local-memory tiled kernels (`NETCL_DWCONV_ALGO=local|naive|tile`). `NETCL_AUTOTUNE=0` disables all measurements and keeps the previous conv paths.
|
|
@@ -138,6 +138,47 @@ def _strides_for_dims(dims):
|
|
|
138
138
|
return strides_for_dims(dims)
|
|
139
139
|
|
|
140
140
|
|
|
141
|
+
_REDUCE_TO_SHAPE_WG_SRC = r"""
|
|
142
|
+
// One work-group (256) per output element o of the target shape t0..t3; it sums the gradient over the
|
|
143
|
+
// broadcast axes (extents r0..r3, 1 for kept axes). The reduced index runs fastest along the last axis, so
|
|
144
|
+
// neighbouring work-items read neighbouring memory.
|
|
145
|
+
__kernel __attribute__((reqd_work_group_size(256, 1, 1)))
|
|
146
|
+
void reduce_to_shape_wg(__global const float* grad, __global float* out,
|
|
147
|
+
const int t0, const int t1, const int t2, const int t3,
|
|
148
|
+
const int r0, const int r1, const int r2, const int r3,
|
|
149
|
+
const int s0, const int s1, const int s2, const int s3,
|
|
150
|
+
const int R)
|
|
151
|
+
{
|
|
152
|
+
const int o = get_group_id(0);
|
|
153
|
+
const int lid = get_local_id(0);
|
|
154
|
+
int tmp = o;
|
|
155
|
+
const int c3 = tmp % t3; tmp /= t3;
|
|
156
|
+
const int c2 = tmp % t2; tmp /= t2;
|
|
157
|
+
const int c1 = tmp % t1; tmp /= t1;
|
|
158
|
+
const int c0 = tmp;
|
|
159
|
+
const size_t base = (size_t)(r0 == 1 ? c0 : 0) * s0 + (size_t)(r1 == 1 ? c1 : 0) * s1
|
|
160
|
+
+ (size_t)(r2 == 1 ? c2 : 0) * s2 + (size_t)(r3 == 1 ? c3 : 0) * s3;
|
|
161
|
+
float acc = 0.0f;
|
|
162
|
+
for (int r = lid; r < R; r += 256) {
|
|
163
|
+
int q = r;
|
|
164
|
+
const int d3 = q % r3; q /= r3;
|
|
165
|
+
const int d2 = q % r2; q /= r2;
|
|
166
|
+
const int d1 = q % r1; q /= r1;
|
|
167
|
+
const int d0 = q;
|
|
168
|
+
acc += grad[base + (size_t)d0 * s0 + (size_t)d1 * s1 + (size_t)d2 * s2 + (size_t)d3 * s3];
|
|
169
|
+
}
|
|
170
|
+
__local float red[256];
|
|
171
|
+
red[lid] = acc;
|
|
172
|
+
barrier(CLK_LOCAL_MEM_FENCE);
|
|
173
|
+
for (int s = 128; s > 0; s >>= 1) {
|
|
174
|
+
if (lid < s) red[lid] += red[lid + s];
|
|
175
|
+
barrier(CLK_LOCAL_MEM_FENCE);
|
|
176
|
+
}
|
|
177
|
+
if (lid == 0) out[o] = red[0];
|
|
178
|
+
}
|
|
179
|
+
"""
|
|
180
|
+
|
|
181
|
+
|
|
141
182
|
def _reduce_to_shape(grad: Tensor, target_shape: Tuple[int, ...]) -> Tensor:
|
|
142
183
|
target_shape = tuple(target_shape)
|
|
143
184
|
if grad.shape == target_shape:
|
|
@@ -214,6 +255,22 @@ def _reduce_to_shape(grad: Tensor, target_shape: Tuple[int, ...]) -> Tensor:
|
|
|
214
255
|
_REDUCE_TO_SHAPE_CACHE[cache_key] = cl.Program(ctx, src).build().reduce_to_shape
|
|
215
256
|
|
|
216
257
|
total_out = int(np.prod(target_shape))
|
|
258
|
+
reduced = 1
|
|
259
|
+
for gd, td in zip(gdims, tdims):
|
|
260
|
+
if td == 1:
|
|
261
|
+
reduced *= gd
|
|
262
|
+
if reduced >= 256:
|
|
263
|
+
# many elements per output (e.g. a broadcast scalar or bias): one work-group per output element with a
|
|
264
|
+
# parallel tree reduction instead of one thread summing everything serially
|
|
265
|
+
wg_key = (ctx.int_ptr, "reduce_to_shape_wg", grad.dtype)
|
|
266
|
+
kern = _REDUCE_TO_SHAPE_CACHE.get(wg_key)
|
|
267
|
+
if kern is None:
|
|
268
|
+
kern = _REDUCE_TO_SHAPE_CACHE[wg_key] = cl.Program(ctx, _REDUCE_TO_SHAPE_WG_SRC).build().reduce_to_shape_wg
|
|
269
|
+
rd = [gd if td == 1 else 1 for gd, td in zip(gdims, tdims)]
|
|
270
|
+
kern(grad.queue, (total_out * 256,), (256,), grad.buffer, out.buffer,
|
|
271
|
+
*[np.int32(v) for v in tdims], *[np.int32(v) for v in rd], *[np.int32(v) for v in gstrides],
|
|
272
|
+
np.int32(reduced))
|
|
273
|
+
return out
|
|
217
274
|
gsize = (int(np.ceil(total_out / 256.0)) * 256,)
|
|
218
275
|
mf = cl.mem_flags
|
|
219
276
|
gdims_buf = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=np.int32(gdims))
|
|
@@ -649,15 +706,14 @@ def hard_sigmoid(x: Node, tape=None) -> Node:
|
|
|
649
706
|
def hard_swish(x: Node, tape=None) -> Node:
|
|
650
707
|
def grad_fn(grad_out: Tensor):
|
|
651
708
|
# forward is x * hard_sigmoid(x) with hard_sigmoid(x) = clip(0.2x + 0.5, 0, 1)
|
|
652
|
-
# d/dx = hard_sigmoid(x) + 0.2 * x * [-2.5 < x < 2.5]
|
|
653
|
-
|
|
654
|
-
|
|
655
|
-
|
|
656
|
-
|
|
657
|
-
grad = elementwise_binary(grad_out, deriv, expression="MUL(v0, v1)")
|
|
709
|
+
# d/dx = hard_sigmoid(x) + 0.2 * x * [-2.5 < x < 2.5]; one fused kernel (v0 = grad_out, v1 = x)
|
|
710
|
+
grad = elementwise_binary(
|
|
711
|
+
grad_out, x.value,
|
|
712
|
+
expression="v0 * (fmax(0.0f, fmin(1.0f, 0.2f * v1 + 0.5f)) + ((v1 > -2.5f && v1 < 2.5f) ? 0.2f * v1 : 0.0f))",
|
|
713
|
+
)
|
|
658
714
|
return [grad]
|
|
659
715
|
|
|
660
|
-
return apply_op(lambda t: hard_swish_op(t), grad_fn, x, tape=tape)
|
|
716
|
+
return apply_op(lambda t: hard_swish_op(t), grad_fn, x, tape=tape, op_name="hard_swish")
|
|
661
717
|
|
|
662
718
|
|
|
663
719
|
def clamp(x: Node, min_val: float, max_val: float, tape=None) -> Node:
|
|
@@ -850,6 +850,17 @@ def conv2d(
|
|
|
850
850
|
from netcl.ops.conv2d_dilated import conv2d_dilated
|
|
851
851
|
|
|
852
852
|
return conv2d_dilated(x, w, bias=bias, stride=stride, pad=pad, dilation=dilation, out=out, pool=pool)
|
|
853
|
+
# 1x1 / stride 1 convolutions are GEMMs: measured choice between the dedicated pointwise kernels and the
|
|
854
|
+
# path below (see netcl.ops.conv2d_pointwise). An explicitly requested algorithm is always honoured.
|
|
855
|
+
if not use_im2col and algo in (None, "auto", "optimized") and strategy is None:
|
|
856
|
+
from netcl.ops.conv2d_pointwise import pointwise_eligible, dispatch_forward
|
|
857
|
+
|
|
858
|
+
if pointwise_eligible(x, w, stride=stride, pad=pad, dilation=dilation, groups=groups):
|
|
859
|
+
return dispatch_forward(
|
|
860
|
+
x, w, bias, out, pool, _ENV_CONV_FUSE_RELU if fuse_relu is None else bool(fuse_relu),
|
|
861
|
+
legacy=lambda: conv2d(x, w, bias=bias, out=out, pool=pool, stride=stride, pad=pad,
|
|
862
|
+
use_im2col=use_im2col, algo=algo, strategy=strategy, fuse_relu=fuse_relu),
|
|
863
|
+
)
|
|
853
864
|
if backend == "cpu":
|
|
854
865
|
return conv2d_cpu(x, w, bias=bias, stride=stride, pad=pad)
|
|
855
866
|
if cl is None or np is None:
|
|
@@ -1538,6 +1549,15 @@ def conv2d_backward(
|
|
|
1538
1549
|
x, w, grad_out, bias=bias, stride=stride, pad=pad, dilation=dilation,
|
|
1539
1550
|
pool=pool, need_grad_x=need_grad_x,
|
|
1540
1551
|
)
|
|
1552
|
+
if not use_im2col and algo in (None, "auto", "optimized"):
|
|
1553
|
+
from netcl.ops.conv2d_pointwise import pointwise_eligible, dispatch_backward
|
|
1554
|
+
|
|
1555
|
+
if pointwise_eligible(x, w, stride=stride, pad=pad, dilation=dilation, groups=groups):
|
|
1556
|
+
return dispatch_backward(
|
|
1557
|
+
x, w, grad_out, bias, pool, need_grad_x,
|
|
1558
|
+
legacy=lambda: conv2d_backward(x, w, grad_out, bias=bias, pool=pool, stride=stride, pad=pad,
|
|
1559
|
+
use_im2col=use_im2col, algo=algo, need_grad_x=need_grad_x),
|
|
1560
|
+
)
|
|
1541
1561
|
if backend == "cpu":
|
|
1542
1562
|
return conv2d_backward_cpu(x, w, grad_out, stride=stride, pad=pad, bias=bias)
|
|
1543
1563
|
if cl is None or np is None:
|
|
@@ -0,0 +1,520 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Pointwise (1x1, stride 1, no padding, no groups/dilation) convolution as batched GEMM.
|
|
3
|
+
|
|
4
|
+
A 1x1 convolution on NCHW data is a matrix product per sample,
|
|
5
|
+
|
|
6
|
+
out[n] (F x P) = W (F x C) @ x[n] (C x P), P = H * W,
|
|
7
|
+
|
|
8
|
+
so the dense conv kernels (tuned for 3x3 filters with many channels at small
|
|
9
|
+
resolution) are a poor fit for it: MobileNet-style networks run their 1x1
|
|
10
|
+
layers with few channels at high resolution, where those kernels reach only a
|
|
11
|
+
few percent of the device's throughput. This module computes the three passes
|
|
12
|
+
directly as register-tiled GEMMs on the NCHW buffers, without im2col or
|
|
13
|
+
layout copies:
|
|
14
|
+
|
|
15
|
+
forward out[n] = W @ x[n] (+ bias, optional ReLU)
|
|
16
|
+
grad_x gx[n] = W^T @ go[n]
|
|
17
|
+
grad_w gw = sum_n go[n] @ x[n]^T (split-K over N * P, then reduced)
|
|
18
|
+
grad_b gb = sum_{n,p} go[n, :, p]
|
|
19
|
+
|
|
20
|
+
Selection is measured, not guessed: for every new shape the pointwise kernels
|
|
21
|
+
(several tile configurations) and the existing ("legacy") conv path are timed
|
|
22
|
+
once with the real buffers via netcl.core.autotune; the winner is cached per
|
|
23
|
+
device (memory + JSON file). Shapes where the legacy path is faster keep using
|
|
24
|
+
it, so the change can only make a given shape faster.
|
|
25
|
+
|
|
26
|
+
Environment:
|
|
27
|
+
NETCL_CONV_POINTWISE=auto (default) measured choice per shape
|
|
28
|
+
NETCL_CONV_POINTWISE=1 always use the pointwise kernels when eligible
|
|
29
|
+
NETCL_CONV_POINTWISE=0 never use them (previous behaviour)
|
|
30
|
+
With NETCL_AUTOTUNE=0 no measurement happens and the legacy path is used,
|
|
31
|
+
i.e. behaviour is exactly that of earlier releases.
|
|
32
|
+
"""
|
|
33
|
+
from __future__ import annotations
|
|
34
|
+
|
|
35
|
+
import os
|
|
36
|
+
import threading
|
|
37
|
+
import time
|
|
38
|
+
from contextlib import contextmanager
|
|
39
|
+
from typing import Optional, Tuple
|
|
40
|
+
|
|
41
|
+
from netcl.core.memory import BufferPool
|
|
42
|
+
from netcl.core.tensor import Tensor
|
|
43
|
+
|
|
44
|
+
try:
|
|
45
|
+
import numpy as np
|
|
46
|
+
except ImportError: # pragma: no cover
|
|
47
|
+
np = None
|
|
48
|
+
|
|
49
|
+
try:
|
|
50
|
+
import pyopencl as cl
|
|
51
|
+
except ImportError: # pragma: no cover
|
|
52
|
+
cl = None
|
|
53
|
+
|
|
54
|
+
__all__ = [
|
|
55
|
+
"pointwise_eligible",
|
|
56
|
+
"conv2d_pointwise",
|
|
57
|
+
"conv2d_pointwise_backward",
|
|
58
|
+
"GEMM_CONFIGS",
|
|
59
|
+
"GW_CONFIGS",
|
|
60
|
+
]
|
|
61
|
+
|
|
62
|
+
_MODE = os.environ.get("NETCL_CONV_POINTWISE", "auto").strip().lower()
|
|
63
|
+
|
|
64
|
+
# (TM, TN, TK, RM, RN): output tile M x N per work-group, K step, register block per thread
|
|
65
|
+
GEMM_CONFIGS: Tuple[Tuple[int, int, int, int, int], ...] = (
|
|
66
|
+
(32, 64, 16, 4, 4),
|
|
67
|
+
(64, 64, 16, 4, 4),
|
|
68
|
+
(32, 128, 16, 4, 8),
|
|
69
|
+
(64, 128, 16, 8, 8),
|
|
70
|
+
)
|
|
71
|
+
# (TF, TC, TK, RF, RC): grad_w tile F x C per work-group, reduction step, register block per thread
|
|
72
|
+
GW_CONFIGS: Tuple[Tuple[int, int, int, int, int], ...] = (
|
|
73
|
+
(32, 32, 32, 4, 4),
|
|
74
|
+
(64, 32, 32, 4, 4),
|
|
75
|
+
(64, 64, 16, 4, 4),
|
|
76
|
+
)
|
|
77
|
+
|
|
78
|
+
_KERNEL_CACHE: dict = {}
|
|
79
|
+
_state = threading.local()
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
# --------------------------------------------------------------------------- eligibility / bypass
|
|
83
|
+
|
|
84
|
+
@contextmanager
|
|
85
|
+
def _bypass():
|
|
86
|
+
"""While active, conv2d / conv2d_backward take their previous (legacy) path for pointwise shapes."""
|
|
87
|
+
prev = getattr(_state, "bypass", False)
|
|
88
|
+
_state.bypass = True
|
|
89
|
+
try:
|
|
90
|
+
yield
|
|
91
|
+
finally:
|
|
92
|
+
_state.bypass = prev
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
def pointwise_eligible(x: Tensor, w: Tensor, stride: int = 1, pad: int = 0, dilation: int = 1, groups: int = 1) -> bool:
|
|
96
|
+
"""True when the convolution is a plain 1x1 GEMM this module can compute (OpenCL, float32)."""
|
|
97
|
+
if cl is None or np is None or _MODE in ("0", "off", "false", "no"):
|
|
98
|
+
return False
|
|
99
|
+
if getattr(_state, "bypass", False):
|
|
100
|
+
return False
|
|
101
|
+
if getattr(x, "backend", "cl") == "cpu" or getattr(w, "backend", "cl") == "cpu":
|
|
102
|
+
return False
|
|
103
|
+
if groups != 1 or dilation != 1 or stride != 1 or pad != 0:
|
|
104
|
+
return False
|
|
105
|
+
if len(x.shape) != 4 or len(w.shape) != 4 or w.shape[2] != 1 or w.shape[3] != 1:
|
|
106
|
+
return False
|
|
107
|
+
if x.dtype not in ("float32", "float") or w.dtype != x.dtype or w.shape[1] != x.shape[1]:
|
|
108
|
+
return False
|
|
109
|
+
return True
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
def _forced() -> bool:
|
|
113
|
+
return _MODE in ("1", "on", "true", "yes", "force")
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
# --------------------------------------------------------------------------- kernels
|
|
117
|
+
|
|
118
|
+
_GEMM_SRC = r"""
|
|
119
|
+
// C[b] (M x P) = op(A) (M x K) @ B[b] (K x P) (+ bias[m], optional ReLU)
|
|
120
|
+
// op(A)[m][k] = TRANS_A ? A[k * M + m] : A[m * K + k]
|
|
121
|
+
__kernel __attribute__((reqd_work_group_size(TN / RN, TM / RM, 1)))
|
|
122
|
+
void pw_gemm(__global const float* restrict A,
|
|
123
|
+
__global const float* restrict B,
|
|
124
|
+
__global float* restrict Cout,
|
|
125
|
+
__global const float* restrict bias,
|
|
126
|
+
const int M, const int K, const int P,
|
|
127
|
+
const int has_bias, const int relu)
|
|
128
|
+
{
|
|
129
|
+
const int tx = get_local_id(0);
|
|
130
|
+
const int ty = get_local_id(1);
|
|
131
|
+
const int bn = get_group_id(0) * TN;
|
|
132
|
+
const int bm = get_group_id(1) * TM;
|
|
133
|
+
const int b = get_group_id(2);
|
|
134
|
+
const int NTX = TN / RN, NTY = TM / RM, NT = NTX * NTY;
|
|
135
|
+
const int tid = ty * NTX + tx;
|
|
136
|
+
__global const float* Bb = B + (size_t)b * K * P;
|
|
137
|
+
__global float* Cb = Cout + (size_t)b * M * P;
|
|
138
|
+
|
|
139
|
+
__local float As[TK][TM];
|
|
140
|
+
__local float Bs[TK][TN];
|
|
141
|
+
float acc[RM][RN];
|
|
142
|
+
for (int i = 0; i < RM; ++i)
|
|
143
|
+
for (int j = 0; j < RN; ++j)
|
|
144
|
+
acc[i][j] = 0.0f;
|
|
145
|
+
|
|
146
|
+
for (int k0 = 0; k0 < K; k0 += TK) {
|
|
147
|
+
for (int i = tid; i < TK * TM; i += NT) {
|
|
148
|
+
#if TRANS_A
|
|
149
|
+
const int kk = i / TM, mm = i % TM; // consecutive threads -> consecutive m (contiguous)
|
|
150
|
+
const int m = bm + mm, k = k0 + kk;
|
|
151
|
+
As[kk][mm] = (m < M && k < K) ? A[(size_t)k * M + m] : 0.0f;
|
|
152
|
+
#else
|
|
153
|
+
const int mm = i / TK, kk = i % TK; // consecutive threads -> consecutive k (contiguous)
|
|
154
|
+
const int m = bm + mm, k = k0 + kk;
|
|
155
|
+
As[kk][mm] = (m < M && k < K) ? A[(size_t)m * K + k] : 0.0f;
|
|
156
|
+
#endif
|
|
157
|
+
}
|
|
158
|
+
for (int i = tid; i < TK * TN; i += NT) {
|
|
159
|
+
const int kk = i / TN, nn = i % TN;
|
|
160
|
+
const int k = k0 + kk, n = bn + nn;
|
|
161
|
+
Bs[kk][nn] = (k < K && n < P) ? Bb[(size_t)k * P + n] : 0.0f;
|
|
162
|
+
}
|
|
163
|
+
barrier(CLK_LOCAL_MEM_FENCE);
|
|
164
|
+
for (int kk = 0; kk < TK; ++kk) {
|
|
165
|
+
float a[RM], bv[RN];
|
|
166
|
+
for (int i = 0; i < RM; ++i) a[i] = As[kk][ty + i * NTY];
|
|
167
|
+
for (int j = 0; j < RN; ++j) bv[j] = Bs[kk][tx + j * NTX];
|
|
168
|
+
for (int i = 0; i < RM; ++i)
|
|
169
|
+
for (int j = 0; j < RN; ++j)
|
|
170
|
+
acc[i][j] = mad(a[i], bv[j], acc[i][j]);
|
|
171
|
+
}
|
|
172
|
+
barrier(CLK_LOCAL_MEM_FENCE);
|
|
173
|
+
}
|
|
174
|
+
for (int i = 0; i < RM; ++i) {
|
|
175
|
+
const int m = bm + ty + i * NTY;
|
|
176
|
+
if (m >= M) continue;
|
|
177
|
+
const float bb = has_bias ? bias[m] : 0.0f;
|
|
178
|
+
for (int j = 0; j < RN; ++j) {
|
|
179
|
+
const int n = bn + tx + j * NTX;
|
|
180
|
+
if (n < P) {
|
|
181
|
+
float v = acc[i][j] + bb;
|
|
182
|
+
if (relu) v = fmax(v, 0.0f);
|
|
183
|
+
Cb[(size_t)m * P + n] = v;
|
|
184
|
+
}
|
|
185
|
+
}
|
|
186
|
+
}
|
|
187
|
+
}
|
|
188
|
+
"""
|
|
189
|
+
|
|
190
|
+
_GW_SRC = r"""
|
|
191
|
+
// partial[s] (F x C) = sum_{p in chunk} go[n][f][p] * x[n][c][p], s = n * chunks + chunk
|
|
192
|
+
__kernel __attribute__((reqd_work_group_size(TC / RC, TF / RF, 1)))
|
|
193
|
+
void pw_grad_w_partial(__global const float* restrict go,
|
|
194
|
+
__global const float* restrict x,
|
|
195
|
+
__global float* restrict partial,
|
|
196
|
+
const int F, const int C, const int P,
|
|
197
|
+
const int chunks, const int L)
|
|
198
|
+
{
|
|
199
|
+
const int tx = get_local_id(0);
|
|
200
|
+
const int ty = get_local_id(1);
|
|
201
|
+
const int bc = get_group_id(0) * TC;
|
|
202
|
+
const int bf = get_group_id(1) * TF;
|
|
203
|
+
const int s = get_group_id(2);
|
|
204
|
+
const int n = s / chunks;
|
|
205
|
+
const int p_begin = (s % chunks) * L;
|
|
206
|
+
const int p_end = min(P, p_begin + L);
|
|
207
|
+
const int NTX = TC / RC, NTY = TF / RF, NT = NTX * NTY;
|
|
208
|
+
const int tid = ty * NTX + tx;
|
|
209
|
+
__global const float* gon = go + (size_t)n * F * P;
|
|
210
|
+
__global const float* xn = x + (size_t)n * C * P;
|
|
211
|
+
|
|
212
|
+
__local float Gs[TK][TF + 1];
|
|
213
|
+
__local float Xs[TK][TC + 1];
|
|
214
|
+
float acc[RF][RC];
|
|
215
|
+
for (int i = 0; i < RF; ++i)
|
|
216
|
+
for (int j = 0; j < RC; ++j)
|
|
217
|
+
acc[i][j] = 0.0f;
|
|
218
|
+
|
|
219
|
+
for (int p0 = p_begin; p0 < p_end; p0 += TK) {
|
|
220
|
+
for (int i = tid; i < TK * TF; i += NT) {
|
|
221
|
+
const int ff = i / TK, kk = i % TK; // consecutive threads -> consecutive p (contiguous)
|
|
222
|
+
const int f = bf + ff, p = p0 + kk;
|
|
223
|
+
Gs[kk][ff] = (f < F && p < p_end) ? gon[(size_t)f * P + p] : 0.0f;
|
|
224
|
+
}
|
|
225
|
+
for (int i = tid; i < TK * TC; i += NT) {
|
|
226
|
+
const int cc = i / TK, kk = i % TK;
|
|
227
|
+
const int c = bc + cc, p = p0 + kk;
|
|
228
|
+
Xs[kk][cc] = (c < C && p < p_end) ? xn[(size_t)c * P + p] : 0.0f;
|
|
229
|
+
}
|
|
230
|
+
barrier(CLK_LOCAL_MEM_FENCE);
|
|
231
|
+
for (int kk = 0; kk < TK; ++kk) {
|
|
232
|
+
float g[RF], xv[RC];
|
|
233
|
+
for (int i = 0; i < RF; ++i) g[i] = Gs[kk][ty + i * NTY];
|
|
234
|
+
for (int j = 0; j < RC; ++j) xv[j] = Xs[kk][tx + j * NTX];
|
|
235
|
+
for (int i = 0; i < RF; ++i)
|
|
236
|
+
for (int j = 0; j < RC; ++j)
|
|
237
|
+
acc[i][j] = mad(g[i], xv[j], acc[i][j]);
|
|
238
|
+
}
|
|
239
|
+
barrier(CLK_LOCAL_MEM_FENCE);
|
|
240
|
+
}
|
|
241
|
+
__global float* out = partial + (size_t)s * F * C;
|
|
242
|
+
for (int i = 0; i < RF; ++i) {
|
|
243
|
+
const int f = bf + ty + i * NTY;
|
|
244
|
+
if (f >= F) continue;
|
|
245
|
+
for (int j = 0; j < RC; ++j) {
|
|
246
|
+
const int c = bc + tx + j * NTX;
|
|
247
|
+
if (c < C) out[(size_t)f * C + c] = acc[i][j];
|
|
248
|
+
}
|
|
249
|
+
}
|
|
250
|
+
}
|
|
251
|
+
"""
|
|
252
|
+
|
|
253
|
+
_REDUCE_SRC = r"""
|
|
254
|
+
// out[i] = sum_s partial[s * n + i]
|
|
255
|
+
__kernel void pw_reduce_partials(__global const float* restrict partial, __global float* restrict out,
|
|
256
|
+
const int n, const int S)
|
|
257
|
+
{
|
|
258
|
+
const int i = get_global_id(0);
|
|
259
|
+
if (i >= n) return;
|
|
260
|
+
float acc = 0.0f;
|
|
261
|
+
for (int s = 0; s < S; ++s) acc += partial[(size_t)s * n + i];
|
|
262
|
+
out[i] = acc;
|
|
263
|
+
}
|
|
264
|
+
|
|
265
|
+
// gb[f] = sum_{n, p} go[n][f][p]; one work-group of 256 per f
|
|
266
|
+
__kernel __attribute__((reqd_work_group_size(256, 1, 1)))
|
|
267
|
+
void pw_grad_bias(__global const float* restrict go, __global float* restrict gb,
|
|
268
|
+
const int N, const int F, const int P)
|
|
269
|
+
{
|
|
270
|
+
const int f = get_group_id(0);
|
|
271
|
+
const int t = get_local_id(0);
|
|
272
|
+
__local float red[256];
|
|
273
|
+
float acc = 0.0f;
|
|
274
|
+
for (int n = 0; n < N; ++n) {
|
|
275
|
+
__global const float* row = go + ((size_t)n * F + f) * P;
|
|
276
|
+
for (int p = t; p < P; p += 256) acc += row[p];
|
|
277
|
+
}
|
|
278
|
+
red[t] = acc;
|
|
279
|
+
barrier(CLK_LOCAL_MEM_FENCE);
|
|
280
|
+
for (int s = 128; s > 0; s >>= 1) {
|
|
281
|
+
if (t < s) red[t] += red[t + s];
|
|
282
|
+
barrier(CLK_LOCAL_MEM_FENCE);
|
|
283
|
+
}
|
|
284
|
+
if (t == 0) gb[f] = red[0];
|
|
285
|
+
}
|
|
286
|
+
"""
|
|
287
|
+
|
|
288
|
+
|
|
289
|
+
def _kernel(ctx: "cl.Context", key: tuple, src: str, defines: dict, name: str) -> "cl.Kernel":
|
|
290
|
+
"""Build each (source, defines) once per context and keep the kernel object (retrieving a kernel from a
|
|
291
|
+
program on every launch is slow in pyopencl)."""
|
|
292
|
+
k = (ctx.int_ptr, name) + key
|
|
293
|
+
kern = _KERNEL_CACHE.get(k)
|
|
294
|
+
if kern is None:
|
|
295
|
+
opts = [f"-D{d}={v}" for d, v in defines.items()]
|
|
296
|
+
kern = _KERNEL_CACHE[k] = getattr(cl.Program(ctx, src).build(options=opts), name)
|
|
297
|
+
return kern
|
|
298
|
+
|
|
299
|
+
|
|
300
|
+
def _gemm_kernel(ctx, cfg, trans_a: bool):
|
|
301
|
+
TM, TN, TK, RM, RN = cfg
|
|
302
|
+
return _kernel(ctx, (cfg, trans_a), _GEMM_SRC, dict(TM=TM, TN=TN, TK=TK, RM=RM, RN=RN, TRANS_A=int(trans_a)), "pw_gemm")
|
|
303
|
+
|
|
304
|
+
|
|
305
|
+
def _gw_kernel(ctx, cfg):
|
|
306
|
+
TF, TC, TK, RF, RC = cfg
|
|
307
|
+
return _kernel(ctx, (cfg,), _GW_SRC, dict(TF=TF, TC=TC, TK=TK, RF=RF, RC=RC), "pw_grad_w_partial")
|
|
308
|
+
|
|
309
|
+
|
|
310
|
+
def _aux_kernel(ctx, name: str):
|
|
311
|
+
return _kernel(ctx, (), _REDUCE_SRC, {}, name)
|
|
312
|
+
|
|
313
|
+
|
|
314
|
+
def _ceil(a: int, b: int) -> int:
|
|
315
|
+
return -(-a // b)
|
|
316
|
+
|
|
317
|
+
|
|
318
|
+
def _bucket(v: int) -> int:
|
|
319
|
+
"""Next power of two: tuning keys group similar batch sizes / resolutions, so inputs of varying size
|
|
320
|
+
(inference on arbitrary images, a smaller last batch) reuse one measurement instead of re-measuring and
|
|
321
|
+
growing the cache with every new shape."""
|
|
322
|
+
return 1 << max(0, int(v) - 1).bit_length()
|
|
323
|
+
|
|
324
|
+
|
|
325
|
+
def _check_config(cfg, dev_max_wg: int) -> bool:
|
|
326
|
+
t1, t2, _, r1, r2 = cfg
|
|
327
|
+
return t1 % r1 == 0 and t2 % r2 == 0 and (t1 // r1) * (t2 // r2) <= dev_max_wg
|
|
328
|
+
|
|
329
|
+
|
|
330
|
+
# --------------------------------------------------------------------------- launches
|
|
331
|
+
|
|
332
|
+
def _gemm(q, A_buf, B: Tensor, out: Tensor, M: int, K: int, P: int, batch: int, cfg, trans_a: bool,
|
|
333
|
+
bias: Optional[Tensor] = None, relu: bool = False) -> None:
|
|
334
|
+
TM, TN, TK, RM, RN = cfg
|
|
335
|
+
kernel = _gemm_kernel(q.context, cfg, trans_a)
|
|
336
|
+
lsize = (TN // RN, TM // RM, 1)
|
|
337
|
+
gsize = (_ceil(P, TN) * lsize[0], _ceil(M, TM) * lsize[1], batch)
|
|
338
|
+
kernel(q, gsize, lsize, A_buf, B.buffer, out.buffer, bias.buffer if bias is not None else A_buf,
|
|
339
|
+
np.int32(M), np.int32(K), np.int32(P), np.int32(bias is not None), np.int32(bool(relu)))
|
|
340
|
+
|
|
341
|
+
|
|
342
|
+
def _gw_split(F: int, C: int, N: int, P: int, cfg) -> Tuple[int, int]:
|
|
343
|
+
"""Chunk length L along P and chunks per sample: enough work-groups to fill the device, without
|
|
344
|
+
growing the partial buffer beyond what the reduction needs."""
|
|
345
|
+
TF, TC, TK, _, _ = cfg
|
|
346
|
+
tiles = _ceil(F, TF) * _ceil(C, TC)
|
|
347
|
+
target_groups = 1024
|
|
348
|
+
splits = max(N, _ceil(target_groups, tiles))
|
|
349
|
+
L = max(TK * 8, _ceil(_ceil(N * P, splits), TK) * TK)
|
|
350
|
+
L = min(L, _ceil(P, TK) * TK)
|
|
351
|
+
return L, _ceil(P, L)
|
|
352
|
+
|
|
353
|
+
|
|
354
|
+
def _grad_w(q, x: Tensor, go: Tensor, gw: Tensor, N: int, C: int, F: int, P: int, cfg, pool) -> None:
|
|
355
|
+
L, chunks = _gw_split(F, C, N, P, cfg)
|
|
356
|
+
S = N * chunks
|
|
357
|
+
TF, TC, TK, RF, RC = cfg
|
|
358
|
+
partial = Tensor.from_shape(q, (S, F, C), dtype="float32", pool=pool)
|
|
359
|
+
lsize = (TC // RC, TF // RF, 1)
|
|
360
|
+
gsize = (_ceil(C, TC) * lsize[0], _ceil(F, TF) * lsize[1], S)
|
|
361
|
+
_gw_kernel(q.context, cfg)(q, gsize, lsize, go.buffer, x.buffer, partial.buffer,
|
|
362
|
+
np.int32(F), np.int32(C), np.int32(P), np.int32(chunks), np.int32(L))
|
|
363
|
+
n = F * C
|
|
364
|
+
_aux_kernel(q.context, "pw_reduce_partials")(q, (_ceil(n, 256) * 256,), (256,), partial.buffer, gw.buffer,
|
|
365
|
+
np.int32(n), np.int32(S))
|
|
366
|
+
|
|
367
|
+
|
|
368
|
+
def _grad_bias(q, go: Tensor, gb: Tensor, N: int, F: int, P: int) -> None:
|
|
369
|
+
_aux_kernel(q.context, "pw_grad_bias")(q, (F * 256,), (256,), go.buffer, gb.buffer,
|
|
370
|
+
np.int32(N), np.int32(F), np.int32(P))
|
|
371
|
+
|
|
372
|
+
|
|
373
|
+
# --------------------------------------------------------------------------- tuning helpers
|
|
374
|
+
|
|
375
|
+
def _timed(q, fn) -> float:
|
|
376
|
+
q.finish()
|
|
377
|
+
t0 = time.perf_counter()
|
|
378
|
+
fn()
|
|
379
|
+
q.finish()
|
|
380
|
+
return time.perf_counter() - t0
|
|
381
|
+
|
|
382
|
+
|
|
383
|
+
def _valid_configs(q, configs):
|
|
384
|
+
max_wg = int(getattr(q.device, "max_work_group_size", 256))
|
|
385
|
+
return [c for c in configs if _check_config(c, max_wg)]
|
|
386
|
+
|
|
387
|
+
|
|
388
|
+
def _tuned_gemm_config(q, M: int, K: int, P: int, batch: int, trans_a: bool, A_buf, B: Tensor, out: Tensor):
|
|
389
|
+
from netcl.core.autotune import autotune
|
|
390
|
+
|
|
391
|
+
cands = _valid_configs(q, GEMM_CONFIGS)
|
|
392
|
+
return tuple(autotune(
|
|
393
|
+
q, "conv2d_pointwise_gemm", (M, K, _bucket(P), _bucket(batch), int(trans_a)), [list(c) for c in cands],
|
|
394
|
+
bench=lambda c: _timed(q, lambda: _gemm(q, A_buf, B, out, M, K, P, batch, tuple(c), trans_a)),
|
|
395
|
+
default=list(cands[0]), reps=3,
|
|
396
|
+
))
|
|
397
|
+
|
|
398
|
+
|
|
399
|
+
def _tuned_gw_config(q, x: Tensor, go: Tensor, gw: Tensor, N: int, C: int, F: int, P: int, pool):
|
|
400
|
+
from netcl.core.autotune import autotune
|
|
401
|
+
|
|
402
|
+
cands = _valid_configs(q, GW_CONFIGS)
|
|
403
|
+
return tuple(autotune(
|
|
404
|
+
q, "conv2d_pointwise_grad_w", (F, C, _bucket(N), _bucket(P)), [list(c) for c in cands],
|
|
405
|
+
bench=lambda c: _timed(q, lambda: _grad_w(q, x, go, gw, N, C, F, P, tuple(c), pool)),
|
|
406
|
+
default=list(cands[0]), reps=3,
|
|
407
|
+
))
|
|
408
|
+
|
|
409
|
+
|
|
410
|
+
# --------------------------------------------------------------------------- public ops
|
|
411
|
+
|
|
412
|
+
def conv2d_pointwise(
|
|
413
|
+
x: Tensor,
|
|
414
|
+
w: Tensor,
|
|
415
|
+
bias: Optional[Tensor] = None,
|
|
416
|
+
out: Optional[Tensor] = None,
|
|
417
|
+
pool: Optional[BufferPool] = None,
|
|
418
|
+
fuse_relu: bool = False,
|
|
419
|
+
config: Optional[Tuple[int, int, int, int, int]] = None,
|
|
420
|
+
) -> Tensor:
|
|
421
|
+
"""1x1 convolution x (N, C, H, W) * w (F, C, 1, 1) -> (N, F, H, W) as one batched GEMM launch.
|
|
422
|
+
`config` forces a tile configuration (GEMM_CONFIGS entry); default: autotuned per shape."""
|
|
423
|
+
N, C, H, W = x.shape
|
|
424
|
+
F = w.shape[0]
|
|
425
|
+
P = H * W
|
|
426
|
+
q = x.queue
|
|
427
|
+
if out is None:
|
|
428
|
+
out = Tensor.from_shape(q, (N, F, H, W), dtype=x.dtype, pool=pool)
|
|
429
|
+
if N == 0 or F == 0 or P == 0:
|
|
430
|
+
return out
|
|
431
|
+
if C == 0:
|
|
432
|
+
raise ValueError("conv2d_pointwise: zero input channels")
|
|
433
|
+
if config is None:
|
|
434
|
+
config = _tuned_gemm_config(q, F, C, P, N, False, w.buffer, x, out)
|
|
435
|
+
_gemm(q, w.buffer, x, out, F, C, P, N, config, False, bias=bias, relu=bool(fuse_relu))
|
|
436
|
+
return out
|
|
437
|
+
|
|
438
|
+
|
|
439
|
+
def conv2d_pointwise_backward(
|
|
440
|
+
x: Tensor,
|
|
441
|
+
w: Tensor,
|
|
442
|
+
grad_out: Tensor,
|
|
443
|
+
bias: Optional[Tensor] = None,
|
|
444
|
+
pool: Optional[BufferPool] = None,
|
|
445
|
+
need_grad_x: bool = True,
|
|
446
|
+
gemm_config: Optional[Tuple[int, int, int, int, int]] = None,
|
|
447
|
+
gw_config: Optional[Tuple[int, int, int, int, int]] = None,
|
|
448
|
+
) -> Tuple[Optional[Tensor], Tensor, Optional[Tensor]]:
|
|
449
|
+
"""Gradients of conv2d_pointwise: (grad_x or None, grad_w (F, C, 1, 1), grad_b or None)."""
|
|
450
|
+
N, C, H, W = x.shape
|
|
451
|
+
F = w.shape[0]
|
|
452
|
+
P = H * W
|
|
453
|
+
q = x.queue
|
|
454
|
+
grad_w = Tensor.from_shape(q, (F, C, 1, 1), dtype=x.dtype, pool=pool)
|
|
455
|
+
grad_b = Tensor.from_shape(q, (F,), dtype=x.dtype, pool=pool) if bias is not None else None
|
|
456
|
+
grad_x = Tensor.from_shape(q, (N, C, H, W), dtype=x.dtype, pool=pool) if need_grad_x else None
|
|
457
|
+
if N * P == 0:
|
|
458
|
+
for t in (grad_w, grad_b):
|
|
459
|
+
if t is not None and t.shape[0] > 0:
|
|
460
|
+
cl.enqueue_fill_buffer(q, t.buffer, np.float32(0), 0, int(np.prod(t.shape)) * 4)
|
|
461
|
+
return grad_x, grad_w, grad_b
|
|
462
|
+
if need_grad_x:
|
|
463
|
+
cfg = gemm_config or _tuned_gemm_config(q, C, F, P, N, True, w.buffer, grad_out, grad_x)
|
|
464
|
+
_gemm(q, w.buffer, grad_out, grad_x, C, F, P, N, cfg, True)
|
|
465
|
+
cfg = gw_config or _tuned_gw_config(q, x, grad_out, grad_w, N, C, F, P, pool)
|
|
466
|
+
_grad_w(q, x, grad_out, grad_w, N, C, F, P, cfg, pool)
|
|
467
|
+
if grad_b is not None:
|
|
468
|
+
_grad_bias(q, grad_out, grad_b, N, F, P)
|
|
469
|
+
return grad_x, grad_w, grad_b
|
|
470
|
+
|
|
471
|
+
|
|
472
|
+
# --------------------------------------------------------------------------- measured dispatch
|
|
473
|
+
|
|
474
|
+
def _shape_key(x: Tensor, w: Tensor, bias, extra) -> tuple:
|
|
475
|
+
N, C, H, W = (int(d) for d in x.shape)
|
|
476
|
+
return (_bucket(N), C, _bucket(H * W), int(w.shape[0]), int(bias is not None)) + tuple(extra)
|
|
477
|
+
|
|
478
|
+
|
|
479
|
+
def dispatch_forward(x: Tensor, w: Tensor, bias: Optional[Tensor], out: Optional[Tensor], pool, fuse_relu,
|
|
480
|
+
legacy) -> Tensor:
|
|
481
|
+
"""Called by conv2d for eligible shapes. `legacy()` runs the previous conv path (pointwise bypassed)."""
|
|
482
|
+
if not _forced():
|
|
483
|
+
from netcl.core.autotune import autotune
|
|
484
|
+
|
|
485
|
+
q = x.queue
|
|
486
|
+
scratch = Tensor.from_shape(q, (x.shape[0], w.shape[0], x.shape[2], x.shape[3]), dtype=x.dtype, pool=pool)
|
|
487
|
+
|
|
488
|
+
def bench(cand):
|
|
489
|
+
if cand == "legacy":
|
|
490
|
+
with _bypass():
|
|
491
|
+
return _timed(q, legacy)
|
|
492
|
+
return _timed(q, lambda: conv2d_pointwise(x, w, bias, out=scratch, pool=pool, fuse_relu=fuse_relu))
|
|
493
|
+
|
|
494
|
+
choice = autotune(q, "conv2d_pointwise_fwd_select", _shape_key(x, w, bias, (int(bool(fuse_relu)),)),
|
|
495
|
+
["legacy", "pointwise"], bench=bench, default="legacy", reps=3)
|
|
496
|
+
if choice == "legacy":
|
|
497
|
+
with _bypass():
|
|
498
|
+
return legacy()
|
|
499
|
+
return conv2d_pointwise(x, w, bias, out=out, pool=pool, fuse_relu=fuse_relu)
|
|
500
|
+
|
|
501
|
+
|
|
502
|
+
def dispatch_backward(x: Tensor, w: Tensor, grad_out: Tensor, bias: Optional[Tensor], pool, need_grad_x, legacy):
|
|
503
|
+
"""Called by conv2d_backward for eligible shapes; same measured choice as dispatch_forward."""
|
|
504
|
+
if not _forced():
|
|
505
|
+
from netcl.core.autotune import autotune
|
|
506
|
+
|
|
507
|
+
q = x.queue
|
|
508
|
+
|
|
509
|
+
def bench(cand):
|
|
510
|
+
if cand == "legacy":
|
|
511
|
+
with _bypass():
|
|
512
|
+
return _timed(q, legacy)
|
|
513
|
+
return _timed(q, lambda: conv2d_pointwise_backward(x, w, grad_out, bias, pool=pool, need_grad_x=need_grad_x))
|
|
514
|
+
|
|
515
|
+
choice = autotune(q, "conv2d_pointwise_bwd_select", _shape_key(x, w, bias, (int(bool(need_grad_x)),)),
|
|
516
|
+
["legacy", "pointwise"], bench=bench, default="legacy", reps=3)
|
|
517
|
+
if choice == "legacy":
|
|
518
|
+
with _bypass():
|
|
519
|
+
return legacy()
|
|
520
|
+
return conv2d_pointwise_backward(x, w, grad_out, bias, pool=pool, need_grad_x=need_grad_x)
|