o װiU@sddlmZddlZddlmZddlmZddlmZddl m Z ddZ ej ej ejgZejejejejgZejejgZejejejejgZeeeeZd d eeeDZd d eDZd d eDZ d Z!dZ"ej#e!ddd e Ddd eDdd eDdd eDdZ$ej#e"ddd e Ddd eDdd eDdd eDdd eDdd eDdZ%ddZ&ddZ'dd Z(d!d"Z)d,d%d&Z*d'd(Z+ $ d-d*d+Z,dS).)productN)_normalize_axis_index) get_typename)runtime) axis_slicecCs&t|}|dkrtjrd}|Sd}|S)Nfloat16__halfhalf)rris_hip)dtypetypenamer Q/home/ubuntu/.local/lib/python3.10/site-packages/cupyx/scipy/signal/_iir_utils.py _get_typename srcCs.g|]\}}t||t|ur||fqSr )cupy promote_typesr .0xyr r r s rcCsg|]}t|qSr r)rtr r rr scCs g|] \}}t|t|fqSr rrr r rr!s a #include #include #include template __global__ void compute_correction_factors( const int m, const int k, const T* b, U* out) { int idx = blockDim.x * blockIdx.x + threadIdx.x; if(idx >= k) { return; } U* out_start = out + idx * (k + m); U* out_off = out_start + k; for(int i = 0; i < m; i++) { U acc = 0.0; for(int j = 0; j < k; j++) { acc += ((U) b[j]) * out_off[i - j - 1]; } out_off[i] = acc; } } template __global__ void first_pass_iir( const int m, const int k, const int n, const int n_blocks, const int carries_stride, const T* factors, T* out, T* carries) { int orig_idx = blockDim.x * (blockIdx.x % n_blocks) + threadIdx.x; int num_row = blockIdx.x / n_blocks; int idx = 2 * orig_idx + 1; if(idx >= n) { return; } int group_num = idx / m; int group_pos = idx % m; T* out_off = out + num_row * n; T* carries_off = carries + num_row * carries_stride; T* group_start = out_off + m * group_num; T* group_carries = carries_off + k * group_num; int pos = group_pos; int up_bound = pos; int low_bound = pos; int rel_pos; for(int level = 1, iter = 1; level < m; level *=2, iter++) { int sz = min(pow(2.0f, ((float) iter)), ((float) m)); if(level > 1) { int factor = ceil(pos / ((float) sz)); up_bound = sz * factor - 1; low_bound = up_bound - level + 1; } if(level == 1) { pos = low_bound; } if(pos < low_bound) { pos += level / 2; } if(pos + m * group_num >= n) { break; } rel_pos = pos % level; T carry = 0.0; for(int i = 1; i <= min(k, level); i++) { T k_value = group_start[low_bound - i]; const T* k_factors = factors + (m + k) * (i - 1) + k; T factor = k_factors[rel_pos]; carry += k_value * factor; } group_start[pos] += carry; __syncthreads(); } if(pos >= m - k) { if(carries != NULL) { group_carries[pos - (m - k)] = group_start[pos]; } } } template __global__ void correct_carries( const int m, const int k, const int n_blocks, const int carries_stride, const int offset, const T* factors, T* carries) { int idx = threadIdx.x; int pos = idx + (m - k); T* row_carries = carries + carries_stride * blockIdx.x; for(int i = offset; i < n_blocks; i++) { T* this_carries = row_carries + k * (i + (1 - offset)); T* prev_carries = row_carries + k * (i - offset); T carry = 0.0; for(int j = 1; j <= k; j++) { const T* k_factors = factors + (m + k) * (j - 1) + k; T factor = k_factors[pos]; T k_value = prev_carries[k - j]; carry += factor * k_value; } this_carries[idx] += carry; __syncthreads(); } } template __global__ void second_pass_iir( const int m, const int k, const int n, const int carries_stride, const int n_blocks, const int offset, const T* factors, T* carries, T* out) { int idx = blockDim.x * (blockIdx.x % n_blocks) + threadIdx.x; idx += offset * m; int row_num = blockIdx.x / n_blocks; int n_group = idx / m; int pos = idx % m; if(idx >= n) { return; } T* out_off = out + row_num * n; T* carries_off = carries + row_num * carries_stride; const T* prev_carries = carries_off + (n_group - offset) * k; T carry = 0.0; for(int i = 1; i <= k; i++) { const T* k_factors = factors + (m + k) * (i - 1) + k; T factor = k_factors[pos]; T k_value = prev_carries[k - i]; carry += factor * k_value; } out_off[idx] += carry; } a #include #include #include template __global__ void pick_carries( const int m, const int n, const int carries_stride, const int n_blocks, const int offset, T* x, T* carries) { int idx = m * (blockIdx.x % n_blocks) + threadIdx.x + m - 2; int pos = threadIdx.x; int row_num = blockIdx.x / n_blocks; int n_group = idx / m; T* x_off = x + row_num * n; T* carries_off = carries + row_num * carries_stride; T* group_carries = carries_off + (n_group + (1 - offset)) * 2; if(idx >= n) { return; } group_carries[pos] = x_off[idx]; } template __global__ void compute_correction_factors_sos( const int m, const T* f_const, U* all_out) { extern __shared__ __align__(sizeof(T)) thrust::complex bc_d[2]; T* b_c = reinterpret_cast(bc_d); extern __shared__ __align__(sizeof(T)) thrust::complex off_d[4]; U* off_cache = reinterpret_cast(off_d); int idx = threadIdx.x; int num_section = blockIdx.x; const int n_const = 6; const int a_off = 3; const int k = 2; const int off_idx = 1; U* out = all_out + num_section * k * m; U* out_start = out + idx * m; const T* b = f_const + num_section * n_const + a_off + 1; b_c[idx] = b[idx]; __syncthreads(); U* this_cache = off_cache + k * idx; this_cache[off_idx - idx] = 1; this_cache[idx] = 0; for(int i = 0; i < m; i++) { U acc = 0.0; for(int j = 0; j < k; j++) { acc += -((U) b_c[j]) * this_cache[off_idx - j]; } this_cache[0] = this_cache[1]; this_cache[1] = acc; out_start[i] = acc; } } template __global__ void first_pass_iir_sos( const int m, const int n, const int n_blocks, const T* factors, T* out, T* carries) { extern __shared__ unsigned int thread_status[2]; extern __shared__ __align__(sizeof(T)) thrust::complex fc_d[2 * 1024]; T* factor_cache = reinterpret_cast(fc_d); int orig_idx = blockDim.x * (blockIdx.x % n_blocks) + threadIdx.x; int num_row = blockIdx.x / n_blocks; int idx = 2 * orig_idx + 1; const int k = 2; if(idx >= n) { return; } int group_num = idx / m; int group_pos = idx % m; T* out_off = out + num_row * n; T* carries_off = carries + num_row * n_blocks * k; T* group_start = out_off + m * group_num; T* group_carries = carries_off + group_num * k; const T* section_factors = factors; T* section_carries = group_carries; factor_cache[group_pos] = section_factors[group_pos]; factor_cache[group_pos - 1] = section_factors[group_pos - 1]; factor_cache[m + group_pos] = section_factors[m + group_pos]; factor_cache[m + group_pos - 1] = section_factors[m + group_pos - 1]; __syncthreads(); int pos = group_pos; int up_bound = pos; int low_bound = pos; int rel_pos; for(int level = 1, iter = 1; level < m; level *= 2, iter++) { int sz = min(pow(2.0f, ((float) iter)), ((float) m)); if(level > 1) { int factor = ceil(pos / ((float) sz)); up_bound = sz * factor - 1; low_bound = up_bound - level + 1; } if(level == 1) { pos = low_bound; } if(pos < low_bound) { pos += level / 2; } if(pos + m * group_num >= n) { break; } rel_pos = pos % level; T carry = 0.0; for(int i = 1; i <= min(k, level); i++) { T k_value = group_start[low_bound - i]; const T* k_factors = factor_cache + m * (i - 1); T factor = k_factors[rel_pos]; carry += k_value * factor; } group_start[pos] += carry; __syncthreads(); } if(pos >= m - k) { if(carries != NULL) { section_carries[pos - (m - k)] = group_start[pos]; } } } template __global__ void correct_carries_sos( const int m, const int n_blocks, const int carries_stride, const int offset, const T* factors, T* carries) { extern __shared__ __align__(sizeof(T)) thrust::complex fcd3[4]; T* factor_cache = reinterpret_cast(fcd3); int idx = threadIdx.x; const int k = 2; int pos = idx + (m - k); T* row_carries = carries + carries_stride * blockIdx.x; factor_cache[2 * idx] = factors[pos]; factor_cache[2 * idx + 1] = factors[m + pos]; __syncthreads(); for(int i = offset; i < n_blocks; i++) { T* this_carries = row_carries + k * (i + (1 - offset)); T* prev_carries = row_carries + k * (i - offset); T carry = 0.0; for(int j = 1; j <= k; j++) { // const T* k_factors = factors + m * (j - 1); // T factor = k_factors[pos]; T factor = factor_cache[2 * idx + (j - 1)]; T k_value = prev_carries[k - j]; carry += factor * k_value; } this_carries[idx] += carry; __syncthreads(); } } template __global__ void second_pass_iir_sos( const int m, const int n, const int carries_stride, const int n_blocks, const int offset, const T* factors, T* carries, T* out) { extern __shared__ __align__(sizeof(T)) thrust::complex fcd2[2 * 1024]; T* factor_cache = reinterpret_cast(fcd2); extern __shared__ __align__(sizeof(T)) thrust::complex c_d[2]; T* carries_cache = reinterpret_cast(c_d); int idx = blockDim.x * (blockIdx.x % n_blocks) + threadIdx.x; idx += offset * m; int row_num = blockIdx.x / n_blocks; int n_group = idx / m; int pos = idx % m; const int k = 2; T* out_off = out + row_num * n; T* carries_off = carries + row_num * carries_stride; const T* prev_carries = carries_off + (n_group - offset) * k; if(pos < k) { carries_cache[pos] = prev_carries[pos]; } if(idx >= n) { return; } factor_cache[pos] = factors[pos]; factor_cache[pos + m] = factors[pos + m]; __syncthreads(); T carry = 0.0; for(int i = 1; i <= k; i++) { const T* k_factors = factor_cache + m * (i - 1); T factor = k_factors[pos]; T k_value = carries_cache[k - i]; carry += factor * k_value; } out_off[idx] += carry; } template __global__ void fir_sos( const int m, const int n, const int carries_stride, const int n_blocks, const int offset, const T* sos, T* carries, T* out) { extern __shared__ __align__(sizeof(T)) thrust::complex fir_cc[1024 + 2]; T* fir_cache = reinterpret_cast(fir_cc); extern __shared__ __align__(sizeof(T)) thrust::complex fir_b[3]; T* b = reinterpret_cast(fir_b); int idx = blockDim.x * (blockIdx.x % n_blocks) + threadIdx.x; int row_num = blockIdx.x / n_blocks; int n_group = idx / m; int pos = idx % m; const int k = 2; T* out_row = out + row_num * n; T* out_off = out_row + n_group * m; T* carries_off = carries + row_num * carries_stride; T* this_carries = carries_off + k * (n_group + (1 - offset)); T* group_carries = carries_off + (n_group - offset) * k; if(pos <= k) { b[pos] = sos[pos]; } if(pos < k) { if(offset && n_group == 0) { fir_cache[pos] = 0; } else { fir_cache[pos] = group_carries[pos]; } } if(idx >= n) { return; } fir_cache[pos + k] = out_off[pos]; __syncthreads(); T acc = 0.0; for(int i = k; i >= 0; i--) { acc += fir_cache[pos + i] * b[k - i]; } out_off[pos] = acc; } )z -std=c++11cC"g|] \}}d|d|dqS)zcompute_correction_factors<, >r rr r rrcCg|]}d|dqS)zcorrect_carriesdd|D}d|}|r|d|dn|}||}|S)NcSsg|]}t|jqSr )rr )rargr r rrsz$_get_module_func..rt||d}|j}|d|jd}|jjs|}||fS)Nrmoveaxisshapereshapeflags c_contiguouscopyraxisx_shaper r r collapse_2ds r:cCsJt||dd}|j}||jdd|jd}|jjs!|}||fS)Nr/rr0r7r r rcollapse_2d_rests r<cCsb|j}tj||d}tj|dddtj||f|df}ttd||}||fd||||f|S)Nr r/compute_correction_factorsr;)sizereyec_emptyr. IIR_MODULE)ablock_szr k correction corr_kernelr r rr>sr>r/c Cs|dur t|j|j}||}|dur||}|j}|j}t||}|j}||} |dkrBt||\}}|durBt||\}} tj ||dd} |jdkrQdn|jd} | |d|} | | }tj ||d}tj |dddtj ||f|df}tj | | |f|d}t td||}t td| }t td | }t td | }||fd ||||f||f|d f||| | | ||| |f|dur|jdkrt|| d|jf}n|jd kr|| d|jd}|jdkr|}n tj||fdd }|jjs|}| dks|dur8t|du}| |}| d||}|| f|f||| ||||f|| |f|f||| |||||| f |dkrR| |} t| d|} | jjsR| } | S)Nr;Tr r6rr=r/r>first_pass_iirsecond_pass_iircorrect_carriesr?)r8)r result_typer astyper2ndimrr@r:arrayrArBrCr.rD broadcast_tor3 concatenater4r5r6intr1)rrEr8zir rFr9x_ndimrGn_outnum_rowsn_blocks total_blocksrHcarriesrIfirst_pass_kernelsecond_pass_kernelcarry_correction_kernelstarting_groupblocks_to_mergecarries_strider r r apply_iirs                rfcCsD|jd}tj|d|f|d}ttd||}||fd|||f|S)NrrOr=compute_correction_factors_sos)rO)r2rrCr.IIR_SOS_MODULE)sosrFr n_sectionsrHrIr r rrges rgTc Cs|dur t|j|j}||}|dur||}|j}|j} |jd} t|| }d} ||} d} | dkr=t||\}}|durHt||\}} |durTtj ||dd}|jdkr[dn|jd}| |d|}||}t |||}tj ||| f|d}|}d}|durt |}tj ||d| f|d}t td|}t td|}t td |}t td |}t td |}t|du}||}|d|| }|||f| f|| |||||ft| D]}||}|dur ||ddddf}||dddddf<t|| d| ||ddddf<|r!|||f|f|| ||||||f||f|df|| |||||f|dks=|dur|dure||ddddf}||dddddf<||ddddddf<||f| f|||||||f|||f|f|| |||||||f|r|||f| f|| |||||f|durt|| d| ||ddddf<q| dkr||}t|d |}|jjs|}|dur|| }t| dkrt|d |}|jjs|}|dur||fS|S) NrrOr;TrKr=first_pass_iir_sossecond_pass_iir_soscorrect_carries_sosfir_sos pick_carriesr/)rrPr rQr2rRrr:r<rSrgrC empty_liker.rhrVrangerr3r1r4r5r6len) rrir8rWr rF apply_firr[r9rXrjrGrYzi_shaper\r]r^rHr_ all_carrieszi_outr`rarb fir_kernelcarries_kernelrcrdresb section_zir r r apply_iir_sosns             "        "       r|)r/NNrJ)r/NNrJTN)- itertoolsrrcupy._core.internalrcupy._core._scalarrcupy_backends.cuda.apircupyx.scipy.signal._arraytoolsrrrfloat32float64 FLOAT_TYPESint8int16int32int64 INT_TYPES complex64 complex128 COMPLEX_TYPESuint8uint16uint32uint64UNSIGNED_TYPESTYPES TYPE_PAIRS TYPE_NAMESTYPE_PAIR_NAMES IIR_KERNELIIR_SOS_KERNEL RawModulerDrhr.r:r<r>rfrgr|r r r rsr                  S