ista-daslab-optimizers-cuda 1.0.0__tar.gz → 1.1.0__tar.gz

{ista_daslab_optimizers_cuda-1.0.0 → ista_daslab_optimizers_cuda-1.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ista_daslab_optimizers_cuda
-Version: 1.0.0
+Version: 1.1.0
 Summary: CUDA kernels for ISTA-DASLab-Optimizers project developed in the Distributed Algorithms and Systems group (DASLab) @ Institute of Science and Technology Austria (ISTA)
 Author-email: Ionut-Vlad Modoranu <ionut-vlad.modoranu@ist.ac.at>
 Maintainer-email: Ionut-Vlad Modoranu <ionut-vlad.modoranu@ist.ac.at>
@@ -29,7 +29,6 @@ License: MIT License
 Project-URL: Repository, https://github.com/IST-DASLab/ISTA-DASLab-Optimizers-CUDA
 Keywords: adaptive optimization,deep learning,low memory optimization
 Classifier: Programming Language :: Python :: 3.8
-Classifier: License :: OSI Approved :: Apache Software License
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
@@ -38,3 +37,13 @@ Requires-Dist: torchaudio
 Requires-Dist: torchvision
 Requires-Dist: numpy
 Dynamic: license-file
+
+# Core dependency of ISTA DAS Lab Optimization Package containing CUDA kernels
+This project contains CUDA kernels designed for [ISTA-DASLab-Optimizers](https://github.com/IST-DASLab/ISTA-DASLab-Optimizers) as a 
+dependency.
+
+# Versions summary:
+- **1.1.0** @ February 5th, 2026:
+  - added kernels for the Sparse M-FAC Pruner
+- **1.0.0** @ February 5th, 2026:
+  - created this repository to decouple the CUDA kernels from the mai **ISTA-DASLab-Optimizers** project

ista_daslab_optimizers_cuda-1.1.0/README.md

@@ -0,0 +1,9 @@
+# Core dependency of ISTA DAS Lab Optimization Package containing CUDA kernels
+This project contains CUDA kernels designed for [ISTA-DASLab-Optimizers](https://github.com/IST-DASLab/ISTA-DASLab-Optimizers) as a 
+dependency.
+
+# Versions summary:
+- **1.1.0** @ February 5th, 2026:
+  - added kernels for the Sparse M-FAC Pruner
+- **1.0.0** @ February 5th, 2026:
+  - created this repository to decouple the CUDA kernels from the mai **ISTA-DASLab-Optimizers** project

{ista_daslab_optimizers_cuda-1.0.0 → ista_daslab_optimizers_cuda-1.1.0}/ista_daslab_optimizers_cuda.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ista_daslab_optimizers_cuda
-Version: 1.0.0
+Version: 1.1.0
 Summary: CUDA kernels for ISTA-DASLab-Optimizers project developed in the Distributed Algorithms and Systems group (DASLab) @ Institute of Science and Technology Austria (ISTA)
 Author-email: Ionut-Vlad Modoranu <ionut-vlad.modoranu@ist.ac.at>
 Maintainer-email: Ionut-Vlad Modoranu <ionut-vlad.modoranu@ist.ac.at>
@@ -29,7 +29,6 @@ License: MIT License
 Project-URL: Repository, https://github.com/IST-DASLab/ISTA-DASLab-Optimizers-CUDA
 Keywords: adaptive optimization,deep learning,low memory optimization
 Classifier: Programming Language :: Python :: 3.8
-Classifier: License :: OSI Approved :: Apache Software License
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
@@ -38,3 +37,13 @@ Requires-Dist: torchaudio
 Requires-Dist: torchvision
 Requires-Dist: numpy
 Dynamic: license-file
+
+# Core dependency of ISTA DAS Lab Optimization Package containing CUDA kernels
+This project contains CUDA kernels designed for [ISTA-DASLab-Optimizers](https://github.com/IST-DASLab/ISTA-DASLab-Optimizers) as a 
+dependency.
+
+# Versions summary:
+- **1.1.0** @ February 5th, 2026:
+  - added kernels for the Sparse M-FAC Pruner
+- **1.0.0** @ February 5th, 2026:
+  - created this repository to decouple the CUDA kernels from the mai **ISTA-DASLab-Optimizers** project

{ista_daslab_optimizers_cuda-1.0.0 → ista_daslab_optimizers_cuda-1.1.0}/ista_daslab_optimizers_cuda.egg-info/SOURCES.txt

@@ -1,5 +1,6 @@
 LICENSE
 MANIFEST.in
+README.md
 pyproject.toml
 setup.py
 ./kernels/dense_mfac/dense_mfac.cpp
@@ -11,6 +12,10 @@ setup.py
 ./kernels/sparse_mfac/sparse_mfac.cpp
 ./kernels/sparse_mfac/sparse_mfac_LCG_kernel.cu
 ./kernels/sparse_mfac/sparse_mfac_SP_kernel.cu
+./kernels/sparse_mfac_pruner/mfac_pruner_cpp.cpp
+./kernels/sparse_mfac_pruner/mfac_pruner_dense.cu
+./kernels/sparse_mfac_pruner/mfac_pruner_initial.cu
+./kernels/sparse_mfac_pruner/mfac_pruner_sparse.cu
 ./kernels/tools/tools.cpp
 ./kernels/tools/tools_kernel.cu
 ista_daslab_optimizers_cuda.egg-info/PKG-INFO
@@ -28,5 +33,9 @@ kernels/micro_adam/micro_adam_update.cu
 kernels/sparse_mfac/sparse_mfac.cpp
 kernels/sparse_mfac/sparse_mfac_LCG_kernel.cu
 kernels/sparse_mfac/sparse_mfac_SP_kernel.cu
+kernels/sparse_mfac_pruner/mfac_pruner_cpp.cpp
+kernels/sparse_mfac_pruner/mfac_pruner_dense.cu
+kernels/sparse_mfac_pruner/mfac_pruner_initial.cu
+kernels/sparse_mfac_pruner/mfac_pruner_sparse.cu
 kernels/tools/tools.cpp
 kernels/tools/tools_kernel.cu

{ista_daslab_optimizers_cuda-1.0.0 → ista_daslab_optimizers_cuda-1.1.0}/ista_daslab_optimizers_cuda.egg-info/top_level.txt

@@ -1,4 +1,5 @@
 ista_daslab_cuda_dense_mfac
 ista_daslab_cuda_micro_adam
 ista_daslab_cuda_sparse_mfac
+ista_daslab_cuda_sparse_mfac_pruner
 ista_daslab_cuda_tools

ista_daslab_optimizers_cuda-1.1.0/kernels/sparse_mfac_pruner/mfac_pruner_cpp.cpp

@@ -0,0 +1,150 @@
+#include <torch/extension.h>
+#include <c10/cuda/CUDAGuard.h>
+#include "../utils.h"
+//#include "parallel_reduce.h"
+
+__global__ void compute_row_initial_kernel (float *global_V, float *global_g, float *global_q, float *global_out, int row_start, int row_end, int m, float damp, int N, int B, int nbits, int use_kahan, int grad_const, int do_init, int do_debug);
+void compute_row_initial_cuda (TT V, TT g, TT q, TT out, int row_start, int row_end, int m, float damp, int N, int B, int nblocks, int nthreads, int nbits, int use_kahan, int grad_const, int do_init, int do_debug);
+void compute_row_initial      (TT V, TT g, TT q, TT out, int row_start, int row_end, int m, float damp, int N, int B, int nblocks, int nthreads, int nbits, int use_kahan, int grad_const, int do_init, int do_debug)
+{
+    assert((nbits == 32) || (nbits == 64));
+    assert((use_kahan == 0) || (use_kahan == 1));
+    assert((grad_const == 0) || (grad_const == 512));
+    assert((0 <= row_start) && (row_start < row_end) && (row_end <= m));
+	CHECK_INPUT(V);
+	CHECK_INPUT(g);
+	CHECK_INPUT(q);
+	CHECK_INPUT(out);
+
+	const at::cuda::OptionalCUDAGuard device_guard(device_of(V));
+    compute_row_initial_cuda(V, g, q, out, row_start, row_end, m, damp, N, B, nblocks, nthreads, nbits, use_kahan, grad_const, do_init, do_debug);
+}
+
+__global__ void compute_row_dense_kernel(float *global_V, float *global_g, float *global_q, float *global_out, int row_start, int row_end, int m, float damp, int N, int B, int grad_const);
+void compute_row_dense_cuda (TT V, TT g, TT q, TT out, int row_start, int row_end, int m, float damp, int N, int B, int nblocks, int nthreads, int grad_const);
+void compute_row_dense      (TT V, TT g, TT q, TT out, int row_start, int row_end, int m, float damp, int N, int B, int nblocks, int nthreads, int grad_const)
+{
+    //assert((grad_const == 0) || (grad_const == 512));
+    assert((0 <= row_start) && (row_start < row_end) && (row_end <= m));
+	CHECK_INPUT(V);
+	CHECK_INPUT(g);
+	CHECK_INPUT(q);
+	CHECK_INPUT(out);
+
+	const at::cuda::OptionalCUDAGuard device_guard(device_of(V));
+    compute_row_dense_cuda(V, g, q, out, row_start, row_end, m, damp, N, B, nblocks, nthreads, grad_const);
+}
+
+__global__ void compute_row_sparse_kernel(float *global_V, int16 *global_gi, float *global_gv, float *global_q, float *global_out, int row_start, int row_end, int m, float damp, int density, int N, int B, int do_init);
+void compute_row_sparse_cuda (TT V, TT gi, TT gv, TT q, TT out, int row_start, int row_end, int m, float damp, int density, int N, int B, int nblocks, int nthreads, int do_init);
+void compute_row_sparse      (TT V, TT gi, TT gv, TT q, TT out, int row_start, int row_end, int m, float damp, int density, int N, int B, int nblocks, int nthreads, int do_init)
+{
+    assert((0 <= row_start) && (row_start < row_end) && (row_end <= m));
+	CHECK_INPUT(V);
+	CHECK_INPUT(gi);
+	CHECK_INPUT(gv);
+	CHECK_INPUT(q);
+	CHECK_INPUT(out);
+
+	const at::cuda::OptionalCUDAGuard device_guard(device_of(V));
+    compute_row_sparse_cuda(V, gi, gv, q, out, row_start, row_end, m, damp, density, N, B, nblocks, nthreads, do_init);
+}
+/*
+void pipeline_copy_compute(TT Vcpu, TT Vgpu0, TT Vgpu1,
+                           TT Qcpu, TT Qgpu0, TT Qgpu1,
+                           TT Vtmp, TT grad, TT gi, TT gv,
+                           int start_copy_cpu, int end_copy_cpu,
+                           int start_copy_gpu, int end_copy_gpu,
+                           int start_comp_gpu, int end_comp_gpu,
+                           int half_copy, int half_compute,
+                           int m, int N, int B, float damp,
+                           int grad_const, int kernel_call_count, int topk_type,
+                           int nblocks, int nthreads)
+{
+    float *pVcpu = (float*) Vcpu.data_ptr();
+    float *pVgpu0 = (float*) Vgpu0.data_ptr();
+    float *pVgpu1 = (float*) Vgpu1.data_ptr();
+    float *pQcpu = (float*) Qcpu.data_ptr();
+    float *pQgpu0 = (float*) Qgpu0.data_ptr();
+    float *pQgpu1 = (float*) Qgpu1.data_ptr();
+    float *pVtmp = (float*) Vtmp.data_ptr();
+    float *pgrad = (float*) grad.data_ptr();
+    int16 *pgi = (int16*) gi.data_ptr();
+    float *pgv = (float*) gv.data_ptr();
+
+    dim3 blocks(nblocks, 1, 1);
+    dim3 threads(nthreads, 1, 1);
+
+    cudaStream_t stream_copy_V, stream_copy_Q, stream_compute;
+    cudaStreamCreate(&stream_copy_V);
+    cudaStreamCreate(&stream_copy_Q);
+    cudaStreamCreate(&stream_compute);
+
+    /// START SECTION COPY
+    int NB = N * B;
+    int rows_copy = end_copy_cpu - start_copy_cpu;
+    int sizeVcpu = rows_copy * N * B;
+    int sizeQcpu = rows_copy * N;
+    int offsetVcpu = start_copy_cpu * NB;
+    int offsetQcpu = start_copy_cpu * N;
+
+    float *copyVgpu = (half_copy == 0) ? pVgpu0 : pVgpu1;
+    float *copyQgpu = (half_copy == 0) ? pQgpu0 : pQgpu1;
+    /// END SECTION COPY
+
+    /// START SECTION COMPUTE
+    int density = gi.sizes()[1];
+    long shmem_initial = 4 * B * sizeof(float);
+    long shmem_sparse = (2 * B + 2 * density) * sizeof(float) + density * sizeof(int16);
+
+    float *compVgpu = (half_compute == 0) ? pVgpu0 : pVgpu1;
+    float *compQgpu = (half_compute == 0) ? pQgpu0 : pQgpu1;
+
+    // kernel_call_count=-1 means kernel_call_count=None
+    int do_init = (kernel_call_count == -1) ? 0 : static_cast<int>(kernel_call_count == 1);
+    /// END SECTION COMPUTE
+
+    /// COPY V
+    cudaMemcpyAsync(
+        copyVgpu,               // device pointer
+        pVcpu + offsetVcpu,     // host pointer
+        sizeVcpu,               // size
+        cudaMemcpyHostToDevice, // direction
+        stream_copy_V           // stream
+    );
+
+    /// COPY Q
+    cudaMemcpyAsync(
+        copyQgpu,               // device pointer
+        pQcpu + offsetQcpu,     // host pointer
+        sizeQcpu,               // size
+        cudaMemcpyHostToDevice, // direction
+        stream_copy_Q           // stream
+    );
+
+    /// COMPUTE
+    if(topk_type == 0) { // global topk
+        if(shmem_initial > 48 * 1024) {
+            cudaFuncSetAttribute(compute_row_initial_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_initial);
+        }
+        compute_row_initial_kernel<<<blocks, threads, shmem_initial, stream_compute>>>(compVgpu, pgrad, compQgpu, pVtmp, start_comp_gpu, end_comp_gpu, m, damp, N, B, 32, 0, grad_const, do_init, 0);
+    } else { // row topk
+        if(shmem_sparse > 48 * 1024){
+            cudaFuncSetAttribute(shmem_sparse, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_sparse);
+        }
+        compute_row_sparse_kernel<<<blocks, threads, shmem_sparse, stream_compute>>>(compVgpu, pgi, pgv, compQgpu, pVtmp, start_comp_gpu, end_comp_gpu, m, damp, density, N, B, do_init);
+    }
+
+	GPU_ERROR_CHECK(cudaGetLastError());
+	GPU_ERROR_CHECK(cudaPeekAtLastError());
+ 	GPU_ERROR_CHECK(cudaDeviceSynchronize());
+}
+*/
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+	m.def("compute_row_initial", &compute_row_initial, "Computes one row of matrix V used for pruning");
+	m.def("compute_row_dense", &compute_row_dense, "Computes one row of matrix V used for pruning using dense gradients");
+	m.def("compute_row_sparse", &compute_row_sparse, "Computes one row of matrix V used for pruning using sparse gradients");
+	//m.def("pipeline_copy_compute", &pipeline_copy_compute, "CPU-GPU transfer and GPU computation using streams");
+}
+
+

ista_daslab_optimizers_cuda-1.1.0/kernels/sparse_mfac_pruner/mfac_pruner_dense.cu

@@ -0,0 +1,132 @@
+#include "../utils.h"
+#include <float.h>
+#include <limits>
+// #include "parallel_reduce.h"
+
+__device__ inline void generic_parallel_reduce(float *mem, int N, const long THREADS, const long Tid) {
+	/*
+	    Compute parallel reduce on a shared memory array of `n` elements using `T` threads, even when n > T.
+	*/
+
+	// perform addition to the first T elements (when n > T)
+	for(int i = Tid + THREADS; i < N; i += THREADS) {
+	    mem[Tid] += mem[i];
+	}
+	__syncthreads();
+	for(int stride = (THREADS >> 1); stride > 0; stride >>= 1) {
+	    if(Tid < stride && Tid + stride < N) {
+	        mem[Tid] += mem[Tid + stride];
+	    }
+	    __syncthreads();
+	}
+}
+
+__global__ void
+compute_row_dense_kernel(float *global_V, float *global_g, float *global_q, float *global_out, int row_start, int row_end, int m, float damp, int N, int B, int grad_const)
+{
+    const long Bid = blockIdx.x; // block id
+	const long THREADS = blockDim.x; // number of threads
+	const long Tid = threadIdx.x; // thread id
+
+    extern __shared__ float mem[];
+    float *V = mem; // size B, stores one row of V, e.g. V[i, Bid, :]
+    float *g = mem + B; // size B, stores one row of g, e.g. g[Bid, :]
+    float *prods = mem + 2 * B; // size B, stores products V*g before summing up.
+    float *Vout = mem + 3 * B; // size B, accumulates dot * V
+
+    // predefined constants to avoid computing the same quantities multiple times
+    long N_B = N * B;
+    long Bid_B = Bid * B;
+
+    int i, j, j_global;
+    float dot, q, delta;
+    long V_start;
+
+    // g = global_g[Bid, :]
+    copy_global_to_shmem(global_g, g, Bid_B, Bid_B + B, THREADS, Tid);
+    __syncthreads();
+
+//     copy_global_to_shmem(global_out, Vout, Bid_B, Bid_B + B, THREADS, Tid); // Vout = out[Bid, :]
+    if(row_end < m) { // we call the kernel to compute rows of V
+//         V_start = 0 * N_B + Bid_B;
+//         copy_global_to_shmem<T>(global_V, Vout, V_start, V_start + B, THREADS, Tid);
+        for(i = Tid; i < B; i += THREADS) {
+            if(should_skip(g[i], grad_const)) {
+                Vout[i] = static_cast<float>(0);
+            } else {
+                Vout[i] = static_cast<float>(damp) * static_cast<float>(g[i]);
+            }
+        }
+    } else if(row_end == m) { // we call the kernel to compute the final update that prunes the model
+        for(i = Tid; i < B; i += THREADS) {
+            Vout[i] = static_cast<float>(0);
+        }
+    }
+    __syncthreads();
+
+    for(j = row_start; j < row_end; ++j) {
+        // V = global_V[j, Bid, :]
+        V_start = j * N_B + Bid_B;
+        copy_global_to_shmem(global_V, V, V_start, V_start + B, THREADS, Tid);
+        __syncthreads();
+
+        // (1) compute dot products
+        for(i = Tid; i < B; i += THREADS) {
+            if(should_skip(g[i], grad_const)) {
+                prods[i] = static_cast<float>(0);
+            } else {
+                prods[i] = V[i] * g[i];
+            }
+        }
+
+        __syncthreads();
+
+        generic_parallel_reduce(prods, B, THREADS, Tid);
+        dot = prods[0];
+
+        // read q from global memory: q = global_q[j, Bid]
+        if(Tid == 0) {
+            prods[0] = global_q[j * N + Bid];
+        }
+        __syncthreads();
+        q = prods[0];
+        delta = dot / q;
+
+        for(i = Tid; i < B; i += THREADS) {
+            Vout[i] -= delta * V[i];
+        }
+    } // end for j < row
+
+    // out[Bid, :] = Vout
+    for(j_global = Bid_B + Tid, j = Tid;
+        j_global < Bid_B + B;
+        j_global += THREADS, j += THREADS)
+    {
+        global_out[j_global] = Vout[j];
+    }
+
+    // TODO: compute q here, based on Vout: q[row, Bid] = m + dot_product(Vout, g)
+}
+
+void
+compute_row_dense_cuda (TT V, TT g, TT q, TT out, int row_start, int row_end, int m, float damp, int N, int B, int nblocks, int nthreads, int grad_const)
+{
+    dim3 blocks(nblocks, 1, 1);
+    dim3 threads(nthreads, 1, 1);
+    long sh_mem_size_bytes = 4 * B * sizeof(float);
+
+    if(sh_mem_size_bytes > 48 * 1024) {
+        cudaFuncSetAttribute(compute_row_dense_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, sh_mem_size_bytes);
+    }
+
+    float* fpV = (float*) V.data_ptr();
+    float* fpg = (float*) g.data_ptr();
+    float* fpq = (float*) q.data_ptr();
+    float* fpout = (float*) out.data_ptr();
+
+    compute_row_dense_kernel<<<blocks, threads, sh_mem_size_bytes>>>(fpV, fpg, fpq, fpout, row_start, row_end, m, damp, N, B, grad_const);
+
+	GPU_ERROR_CHECK(cudaGetLastError());
+	GPU_ERROR_CHECK(cudaPeekAtLastError());
+// 	GPU_ERROR_CHECK(cudaDeviceSynchronize());
+}

ista_daslab_optimizers_cuda-1.1.0/kernels/sparse_mfac_pruner/mfac_pruner_initial.cu

@@ -0,0 +1,244 @@
+#include "../utils.h"
+// #include "parallel_reduce.h"
+
+__device__ inline void generic_parallel_reduce(float *mem, int N, const long THREADS, const long Tid) {
+	/*
+	    Compute parallel reduce on a shared memory array of `n` elements using `T` threads, even when n > T.
+	*/
+
+	// perform addition to the first T elements (when n > T)
+	for(int i = Tid + THREADS; i < N; i += THREADS) {
+	    mem[Tid] += mem[i];
+	}
+	__syncthreads();
+	for(int stride = (THREADS >> 1); stride > 0; stride >>= 1) {
+	    if(Tid < stride && Tid + stride < N) {
+	        mem[Tid] += mem[Tid + stride];
+	    }
+	    __syncthreads();
+	}
+}
+
+__device__ inline void kahan_parallel_reduce(float* mem, int N, int THREADS, int Tid) {
+   // initially, we sum everything in interval [THREADS, N-1] to [0, THREADS-1]
+   // the already existing values in mem[Tid] serve as initial values, so sum=mem[Tid]
+   double sum = static_cast<double>(mem[Tid]);
+   double c = static_cast<double>(0);
+   double y, t;
+
+   // the following for-loop implements mem[Tid] += mem[i] using Kahan summation
+   // for the values at indices i > THREADS
+   for(int i = Tid + THREADS; i < N; i += THREADS) {
+       y = static_cast<double>(mem[i]) - c;
+       t = sum + y;
+       c = (t - sum) - y;
+       sum = t;
+   }
+   mem[Tid] = static_cast<float>(sum);
+   __syncthreads();
+
+   // the following for-loop implements mem[Tid] += mem[Tid + stride] using
+   // Kahan summation and parallel reduce in logarithmic time
+   for(int stride = (THREADS >> 1); stride > 0; stride >>= 1) {
+       if(Tid < stride && Tid + stride < N) {
+           y = static_cast<double>(mem[Tid + stride]) - c; // mem[Tid+stride] is the value to be summed up
+           t = sum + y; // mem[Tid] stores the sum
+           c = (t - sum) - y;
+           sum = t; // update sum
+           mem[Tid] = static_cast<float>(sum);
+       }
+       __syncthreads();
+   }
+}
+
+__global__ void compute_row_initial_kernel(float *global_V, float *global_g, float *global_q, float *global_out, int row_start, int row_end, int m, float damp, int N, int B, int nbits, int use_kahan, int grad_const, int do_init, int do_debug)
+{
+    const long Bid = blockIdx.x; // block id
+	const long THREADS = blockDim.x; // number of threads
+	const long Tid = threadIdx.x; // thread id
+
+    if(do_debug) {
+        printf("[Bid=%ld][Tid=%ld][rows=%d-%d] THREADS=%ld, damp=%f\n", Bid, Tid, row_start, row_end, THREADS, damp);
+    }
+
+    extern __shared__ unsigned char shmem[];
+    float *mem = reinterpret_cast<float*>(shmem);
+    float *V = mem; // size B, stores one row of V, e.g. V[i, Bid, :]
+    float *g = mem + B; // size B, stores one row of g, e.g. g[Bid, :]
+    float *prods = mem + 2 * B; // size B, stores products V*g before summing up.
+    float *Vout = mem + 3 * B; // size B, accumulates dot * V
+    double *comps = 0;
+    if(use_kahan) { //
+        comps = reinterpret_cast<double*>(shmem); // Kahan compensations for each component in Vout
+        if(nbits == 32) {
+            comps += 2 * B;
+        } else if(nbits == 64) {
+            comps += 4 * B;
+        }
+    }
+
+    // predefined constants to avoid computing the same quantities multiple times
+    long N_B = N * B;
+    long Bid_B = Bid * B;
+
+    long g_start = Bid_B;
+    long V_start;
+
+    int i, j, j_global;
+    float dot, q, delta;
+
+    // g = global_g[Bid, :]
+    copy_global_to_shmem(global_g, g, g_start, g_start + B, THREADS, Tid);
+    __syncthreads();
+
+    if(do_debug) {
+        for(i = Tid; i < B; i += THREADS) {
+            printf("[Bid=%ld][Tid=%ld][rows=%d-%d][step-1] g[%d]=%lf\n", Bid, Tid, row_start, row_end, i, g[i]);
+        }
+    }
+
+    if(row_end < m) { // we call the kernel to compute rows of V
+//         V_start = 0 * N_B + Bid_B;
+//         copy_global_to_shmem(global_V, Vout, V_start, V_start + B, THREADS, Tid);
+        for(i = Tid; i < B; i += THREADS) {
+            if(do_init) {
+                if(should_skip(g[i], grad_const)) {
+                    Vout[i] = static_cast<float>(0);
+                } else {
+                    Vout[i] = static_cast<float>(damp) * static_cast<float>(g[i]);
+                }
+            } else {
+                Vout[i] = static_cast<float>(0);
+            }
+        }
+
+    } else if(row_end == m) { // we call the kernel to compute the final update that prunes the model
+        for(i = Tid; i < B; i += THREADS) {
+            Vout[i] = static_cast<float>(0);
+        }
+    }
+    __syncthreads();
+
+    double y, t; // for Kahan
+
+    // initialize compensations to zero
+    if(use_kahan) {
+        for(i = Tid; i < B; i += THREADS) {
+            comps[i] = static_cast<float>(0);
+        }
+    }
+
+    for(j = row_start; j < row_end; ++j) {
+        // V = global_V[j, Bid, :]
+        V_start = j * N_B + Bid_B;
+        copy_global_to_shmem(global_V, V, V_start, V_start + B, THREADS, Tid);
+        __syncthreads();
+
+        if(do_debug) {
+            for(i = Tid; i < B; i += THREADS) {
+                printf("[Bid=%ld][Tid=%ld][rows=%d-%d][step-2] v[%d, %ld, %d]=%lf\n", Bid, Tid, row_start, row_end, j, Bid, i, V[i]);
+            }
+        }
+
+        // (1) compute dot products
+        for(i = Tid; i < B; i += THREADS) {
+            if(should_skip(g[i], grad_const)) {
+                prods[i] = static_cast<float>(0);
+            } else {
+                prods[i] = V[i] * g[i];
+            }
+        }
+
+        if(do_debug) {
+            for(i = Tid; i < B; i += THREADS) {
+                printf("[Bid=%ld][Tid=%ld][rows=%d-%d][step-3] prods[%d]=%lf (pre-reduce)\n", Bid, Tid, row_start, row_end, i, prods[i]);
+            }
+        }
+        __syncthreads();
+        if(use_kahan) {
+            kahan_parallel_reduce(prods, B, THREADS, Tid);
+        } else {
+            generic_parallel_reduce(prods, B, THREADS, Tid);
+        }
+        dot = prods[0];
+
+        if(do_debug) {
+            for(i = Tid; i < B; i += THREADS) {
+                printf("[Bid=%ld][Tid=%ld][rows=%d-%d][step-4] prods[%d]=%lf (post-reduce)\n", Bid, Tid, row_start, row_end, i, prods[i]);
+            }
+            printf("[Bid=%ld][Tid=%ld][rows=%d-%d][step-5] dot=%lf\n", Bid, Tid, row_start, row_end, dot);
+        }
+
+        // read q from global memory: q = global_q[j, Bid]
+        if(Tid == 0) {
+            prods[0] = static_cast<float>(global_q[j * N + Bid]);
+        }
+        __syncthreads();
+        q = prods[0];
+        delta = dot / q;
+
+        if(do_debug) {
+            printf("[Bid=%ld][Tid=%ld][rows=%d-%d][step-6] q=%lf, delta=%lf\n", Bid, Tid, row_start, row_end, q, delta);
+        }
+
+        if(use_kahan) {
+            for(i = Tid; i < B; i += THREADS) {
+                y = static_cast<double>(-delta) * static_cast<double>(V[i]) - comps[i];
+                t = static_cast<double>(Vout[i]) + y;
+                comps[i] = (t - static_cast<double>(Vout[i])) - y;
+                Vout[i] = t;
+            }
+        } else {
+            for(i = Tid; i < B; i += THREADS) {
+                Vout[i] -= delta * V[i];
+                if(do_debug) {
+                    printf("[Bid=%ld][Tid=%ld][rows=%d-%d][step-7] delta*V[%d, %ld, %d]=%lf\n", Bid, Tid, row_start, row_end, j, Bid, i, delta * V[i]);
+                }
+            }
+        }
+    } // end for j < row
+
+    for(j_global = Bid_B + Tid, j = Tid;
+        j_global < Bid_B + B;
+        j_global += THREADS, j += THREADS)
+    {
+        global_out[j_global] = static_cast<float>(Vout[j]);
+    }
+    if(do_debug) {
+        for(i = Tid; i < B; i += THREADS) {
+            printf("[Bid=%ld][Tid=%ld][rows=%d-%d][step-8] vout[%d]=%lf[OUT]\n", Bid, Tid, row_start, row_end, i, Vout[i]);
+        }
+    }
+
+    // TODO: compute q here, based on Vout: q[row, Bid] = m + dot_product(Vout, g)
+}
+
+void
+compute_row_initial_cuda (TT V, TT g, TT q, TT out, int row_start, int row_end, int m, float damp, int N, int B, int nblocks, int nthreads, int nbits, int use_kahan, int grad_const, int do_init, int do_debug) {
+    assert(nbits == 32);
+    dim3 blocks(nblocks, 1, 1);
+    dim3 threads(nthreads, 1, 1);
+    long sh_mem_size_bytes = 4 * B * ((nbits == 32) ? sizeof(float) : sizeof(double));
+    if(use_kahan){
+        sh_mem_size_bytes += B * sizeof(double); // add shared memory space for the Kahan compensations
+    }
+
+//     printf("row=%d, N=%d, B=%d, blocks=%d, threads=%d, sh_mem_size_bytes=%ld\n", row, N, B, nblocks, nthreads, sh_mem_size_bytes);
+
+    if(sh_mem_size_bytes > 48 * 1024) {
+        //// if we want to allocate more than 48KB, then we have to call this method
+        cudaFuncSetAttribute(compute_row_initial_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, sh_mem_size_bytes);
+    }
+
+    float* fpV = (float*) V.data_ptr();
+    float* fpg = (float*) g.data_ptr();
+    float* fpq = (float*) q.data_ptr();
+    float* fpout = (float*) out.data_ptr();
+
+    compute_row_initial_kernel<<<blocks, threads, sh_mem_size_bytes>>>(fpV, fpg, fpq, fpout, row_start, row_end, m, damp, N, B, nbits, use_kahan, grad_const, do_init, do_debug);
+
+
+	GPU_ERROR_CHECK(cudaGetLastError());
+	GPU_ERROR_CHECK(cudaPeekAtLastError());
+// 	GPU_ERROR_CHECK(cudaDeviceSynchronize());
+}

ista_daslab_optimizers_cuda-1.1.0/kernels/sparse_mfac_pruner/mfac_pruner_sparse.cu

@@ -0,0 +1,156 @@
+#include "../utils.h"
+// #include "parallel_reduce.h"
+
+__device__ inline void generic_parallel_reduce(float *mem, int N, const long THREADS, const long Tid) {
+	/*
+	    Compute parallel reduce on a shared memory array of `n` elements using `T` threads, even when n > T.
+	*/
+
+	// perform addition to the first T elements (when n > T)
+	for(int i = Tid + THREADS; i < N; i += THREADS) {
+	    mem[Tid] += mem[i];
+	}
+	__syncthreads();
+	for(int stride = (THREADS >> 1); stride > 0; stride >>= 1) {
+	    if(Tid < stride && Tid + stride < N) {
+	        mem[Tid] += mem[Tid + stride];
+	    }
+	    __syncthreads();
+	}
+}
+
+__device__ inline void unpack(float packed, int &i, float &f) { // ct stands for constant
+    /*
+        *i = static_cast<int>(packed);
+        *f = packed - (*i);
+        if((*i) < 0) {
+            (*i) = -(*i);
+        }
+    */
+    i = static_cast<int>(packed);
+    f = packed - i;
+    if(i < 0) {
+        i = -i;
+    }
+}
+
+__global__ void
+compute_row_sparse_kernel(float *global_V, int16 *global_gi, float *global_gv, float *global_q, float *global_out, int row_start, int row_end, int m, float damp, int density, int N, int B, int do_init)
+{
+    const long Bid = blockIdx.x; // block id
+	const long THREADS = blockDim.x; // number of threads
+	const long Tid = threadIdx.x; // thread id
+
+    extern __shared__ unsigned char shmem[];
+
+//     sh_mem_size_bytes += B * sizeof(float); // for a row of V
+//     sh_mem_size_bytes += B * sizeof(float); // for a row of Vout
+//     sh_mem_size_bytes += density * sizeof(float); // for prods
+//     sh_mem_size_bytes += density * sizeof(float); // for gv
+//     sh_mem_size_bytes += density * sizeof(int16); // for gi
+    float *V = (float*) shmem;
+    float *Vout = V + B;
+    float *prods = V + 2 * B;
+    float *gv = V + 2 * B + density;
+    int16 *gi = (int16*)(shmem + (2 * B + 2 * density) * sizeof(float));
+
+    // predefined constants to avoid computing the same quantities multiple times
+    long N_B = N * B;
+    long Bid_B = Bid * B;
+    long Bid_density = Bid * density;
+
+    long V_start;
+
+    int i, j, j_global;
+    float dot, q, delta;
+
+    copy_global_to_shmem(global_gv, gv, Bid_density, Bid_density + density, THREADS, Tid);
+    __syncthreads();
+
+    copy_global_to_shmem(global_gi, gi, Bid_density, Bid_density + density, THREADS, Tid);
+    __syncthreads();
+
+//     for(i = Tid; i < density; i += THREADS) {
+//         printf("[Bid=%ld][Tid=%ld][i=%d] gi=%d, gv=%.8f\n", Bid, Tid, i, gi[i], gv[i]);
+//     }
+
+    // initialize Vout with zeros in the first place
+    for(i = Tid; i < B; i += THREADS) {
+        Vout[i] = static_cast<float>(0);
+    }
+    __syncthreads();
+
+    if(do_init) {
+        // initialize with damp * grad
+        for(i = Tid; i < density; i += THREADS) {
+            Vout[gi[i]] = damp * gv[i];
+        }
+    }
+    __syncthreads();
+
+    for(j = row_start; j < row_end; ++j) {
+        // V = global_V[j, Bid, :]
+        V_start = j * N_B + Bid_B;
+        copy_global_to_shmem(global_V, V, V_start, V_start + B, THREADS, Tid);
+        __syncthreads();
+
+        // (1) compute dot products
+        for(i = Tid; i < density; i += THREADS) {
+            prods[i] = V[gi[i]] * gv[i];
+        }
+        __syncthreads();
+
+        generic_parallel_reduce(prods, density, THREADS, Tid);
+        dot = prods[0];
+
+        // read q from global memory: q = global_q[j, Bid]
+        if(Tid == 0) {
+            prods[0] = static_cast<float>(global_q[j * N + Bid]);
+        }
+        __syncthreads();
+        q = prods[0];
+        delta = dot / q;
+
+        for(i = Tid; i < B; i += THREADS) {
+            Vout[i] -= delta * V[i];
+        }
+    } // end for j < row
+
+    for(j_global = Bid_B + Tid, j = Tid;
+        j_global < Bid_B + B;
+        j_global += THREADS, j += THREADS)
+    {
+        global_out[j_global] = Vout[j];
+    }
+    // TODO: compute q here, based on Vout: q[row, Bid] = m + dot_product(Vout, g)
+}
+
+void compute_row_sparse_cuda (TT V, TT gi, TT gv, TT q, TT out, int row_start, int row_end, int m, float damp, int density, int N, int B, int nblocks, int nthreads, int do_init)
+{
+    dim3 blocks(nblocks, 1, 1);
+    dim3 threads(nthreads, 1, 1);
+
+    long sh_mem_size_bytes = 0;
+    sh_mem_size_bytes += B * sizeof(float); // for a row of V
+    sh_mem_size_bytes += B * sizeof(float); // for a row of Vout
+    sh_mem_size_bytes += density * sizeof(float); // for prods
+    sh_mem_size_bytes += density * sizeof(float); // for gv
+    sh_mem_size_bytes += density * sizeof(int16); // for gi
+
+    if(sh_mem_size_bytes > 48 * 1024) {
+        //// if we want to allocate more than 48KB, then we have to call this method
+        cudaFuncSetAttribute(compute_row_sparse_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, sh_mem_size_bytes);
+    }
+
+    float* pV = (float*) V.data_ptr();
+    int16* pgi = (int16*) gi.data_ptr();
+    float* pgv = (float*) gv.data_ptr();
+    float* pq = (float*) q.data_ptr();
+    float* pout = (float*) out.data_ptr();
+
+    compute_row_sparse_kernel<<<blocks, threads, sh_mem_size_bytes>>>(pV, pgi, pgv, pq, pout, row_start, row_end, m, damp, density, N, B, do_init);
+
+	GPU_ERROR_CHECK(cudaGetLastError());
+	GPU_ERROR_CHECK(cudaPeekAtLastError());
+ 	// GPU_ERROR_CHECK(cudaDeviceSynchronize());
+}

{ista_daslab_optimizers_cuda-1.0.0 → ista_daslab_optimizers_cuda-1.1.0}/kernels/utils.h

@@ -18,6 +18,7 @@
 #include <limits> // for epsilon
 
 using namespace std;
+using TT = torch::Tensor;
 
 typedef __nv_bfloat16 bfloat16;
 typedef __nv_bfloat162 bfloat162;
@@ -73,6 +74,7 @@ __device__ inline long log_threads(long T) {
 	if(T == 256) return 8;
 	if(T == 512) return 9;
 	if(T == 1024) return 10;
+	return 1; // default to avoid compilation errors
 }
 
 inline LL get_threads(LL max_threads) {
@@ -86,6 +88,19 @@ inline LL get_threads(LL max_threads) {
     return threads;
 }
 
+__device__ inline void dynamically_assign_float(void *out, int out_index, float value, int out_bits) {
+    /*
+        This function assigns out[out_index] = value.
+        If nbits=16, then it means out is bfloat16 and we need to convert value to bfloat16.
+        If nbits=32, then it means out is float and no conversion is needed
+    */
+    if(out_bits == 16) {
+        ((bfloat16*) out)[out_index] = __float2bfloat16(value);
+    } else {
+        ((float*) out)[out_index] = value;
+    }
+}
+
 __device__ inline void dynamically_assign(void *out, void *inp, int out_index, int inp_index, int out_bits, int inp_bits) {
     /*
         This function assigns out[out_index] = inp[inp_index] based on the types and performs the conversions when needed:
@@ -108,6 +123,34 @@ __device__ inline void dynamically_assign(void *out, void *inp, int out_index, i
     }
 }
 
+__device__ inline bool should_skip(float x, int ct) { // ct stands for constant
+    if(ct == 0) {
+        return false;
+    }
+    int x_int = static_cast<int>(x);
+    if(((ct - 10) <= x_int) && (x_int <= (ct + 10))) {
+        return true;
+    }
+    return false;
+}
+
+template<typename T>
+__device__ inline void copy_global_to_shmem(T *global,
+                                            T *shmem,
+                                            long global_start,
+                                            long global_end,
+                                            const long THREADS,
+                                            const long Tid) {
+    long j_global, j_shmem; // used in for-loops to read from global memory to shared memory
+
+    for(j_global = global_start + Tid, j_shmem = Tid;
+        j_global < global_end;
+        j_global += THREADS, j_shmem += THREADS)
+    {
+        shmem[j_shmem] = global[j_global];
+    }
+}
+
 #define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
 #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
 #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
@@ -116,7 +159,10 @@ __device__ inline void dynamically_assign(void *out, void *inp, int out_index, i
 #define FLOAT_EPS std::numeric_limits<float>::epsilon()
 #define DOUBLE_EPS std::numeric_limits<double>::epsilon()
 #define GPU_ERROR_CHECK(ans) { gpuAssert((ans), __FILE__, __LINE__); }
+#define IS_BF16(x) torch::ScalarType::BFloat16 == x.scalar_type()
+#define IS_FLOAT(x) torch::ScalarType::Float == x.scalar_type()
 #define ASSERT_BF16(x) { assert(torch::ScalarType::BFloat16 == x.scalar_type()); }
+#define ASSERT_FLOAT(x) { assert(IS_FLOAT(x)); }
 #define ASSERT_FLOAT_16_OR_32(x) { assert(torch::ScalarType::BFloat16 == x.scalar_type() || torch::ScalarType::Float == x.scalar_type()); }
 
 #define COPY_DIRECTION_k2d 0

{ista_daslab_optimizers_cuda-1.0.0 → ista_daslab_optimizers_cuda-1.1.0}/pyproject.toml

@@ -1,42 +1,41 @@
-[build-system]
-requires = ["setuptools", "wheel", "torch"]
-build-backend = "setuptools.build_meta"
-
-[project]
-name='ista_daslab_optimizers_cuda'
-version='1.0.0'
-dependencies = [
-    "torch", # >=2.3.1",
-    "torchaudio", # >=2.3.1",
-    "torchvision", #>=0.18.1",
-    "numpy", # >=1.24.1",
-#    "wandb",#>=0.17.1",
-#    "gpustat",#>=1.1.1",
-#    "timm", # >=1.0.3",
-#    "einops", # >=0.7.0",
-#    "psutil", # >=5.9.8",
-#    "fast-hadamard-transform",
-#    "fast-hadamard-transform @ git+https://github.com/Dao-AILab/fast-hadamard-transform.git",
-]
-requires-python = '>= 3.8'
-authors = [
-    {name = "Ionut-Vlad Modoranu", email = "ionut-vlad.modoranu@ist.ac.at"}
-]
-maintainers = [
-    {name = "Ionut-Vlad Modoranu", email = "ionut-vlad.modoranu@ist.ac.at"},
-]
-description = 'CUDA kernels for ISTA-DASLab-Optimizers project developed in the Distributed Algorithms and Systems group (DASLab) @ Institute of Science and Technology Austria (ISTA)'
-readme = "README.md"
-license = {file = "LICENSE"}
-keywords = [
-    "adaptive optimization",
-    "deep learning",
-    "low memory optimization",
-]
-classifiers = [
-    "Programming Language :: Python :: 3.8",
-    "License :: OSI Approved :: Apache Software License",
-]
-
-[project.urls]
-Repository = 'https://github.com/IST-DASLab/ISTA-DASLab-Optimizers-CUDA'
+[build-system]
+requires = ["setuptools", "wheel", "torch"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name='ista_daslab_optimizers_cuda'
+version='1.1.0'
+dependencies = [
+    "torch", # >=2.3.1",
+    "torchaudio", # >=2.3.1",
+    "torchvision", #>=0.18.1",
+    "numpy", # >=1.24.1",
+#    "wandb",#>=0.17.1",
+#    "gpustat",#>=1.1.1",
+#    "timm", # >=1.0.3",
+#    "einops", # >=0.7.0",
+#    "psutil", # >=5.9.8",
+#    "fast-hadamard-transform",
+#    "fast-hadamard-transform @ git+https://github.com/Dao-AILab/fast-hadamard-transform.git",
+]
+requires-python = '>= 3.8'
+authors = [
+    {name = "Ionut-Vlad Modoranu", email = "ionut-vlad.modoranu@ist.ac.at"}
+]
+maintainers = [
+    {name = "Ionut-Vlad Modoranu", email = "ionut-vlad.modoranu@ist.ac.at"},
+]
+description = 'CUDA kernels for ISTA-DASLab-Optimizers project developed in the Distributed Algorithms and Systems group (DASLab) @ Institute of Science and Technology Austria (ISTA)'
+readme = "README.md"
+license = {file = "LICENSE"}
+keywords = [
+    "adaptive optimization",
+    "deep learning",
+    "low memory optimization",
+]
+classifiers = [
+    "Programming Language :: Python :: 3.8",
+]
+
+[project.urls]
+Repository = 'https://github.com/IST-DASLab/ISTA-DASLab-Optimizers-CUDA'

{ista_daslab_optimizers_cuda-1.0.0 → ista_daslab_optimizers_cuda-1.1.0}/setup.py

@@ -51,6 +51,15 @@ setup(
                 './kernels/micro_adam/micro_adam_asymm_block_quant_inv.cu',
             ],
         ),
+        get_cuda_extension(
+            name=f'ista_daslab_cuda_sparse_mfac_pruner',
+            sources=[
+                './kernels/sparse_mfac_pruner/mfac_pruner_cpp.cpp',
+                './kernels/sparse_mfac_pruner/mfac_pruner_dense.cu',
+                './kernels/sparse_mfac_pruner/mfac_pruner_initial.cu',
+                './kernels/sparse_mfac_pruner/mfac_pruner_sparse.cu',
+            ],
+        ),
     ],
     cmdclass={'build_ext': BuildExtension.with_options(verbose=True)},
 )