rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 20 Packages Projects Releases Wiki Activity

Initial commit

Browse Source
This commit is contained in:
redclawsystems
2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/specialized/rtx-cfd/src/kernels/cuda/diffusion.cu
+274
View File
@@ -0,0 +1,274 @@
/**
* CUDA kernels for diffusion equation
*
* Implements various schemes for the diffusion/heat equation:
* ∂T/∂t = α∇²T
*
* Includes explicit, implicit, and Crank-Nicolson schemes.
*/
extern "C" __global__ void diffusion_explicit_1d(
const float* __restrict__ temp,
float* __restrict__ temp_new,
float alpha,
float dt,
float dx,
int n
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
float r = alpha * dt / (dx * dx); // Diffusion number
// Stability condition: r <= 0.5 for explicit scheme
if (r > 0.5f) {
// Issue warning - this should be handled at higher level
return;
}
// Boundary conditions (Dirichlet for now)
if (i == 0 || i == n - 1) {
temp_new[i] = temp[i]; // Fixed boundary values
return;
}
// Explicit finite difference: T_i^{n+1} = T_i^n + r(T_{i+1}^n - 2T_i^n + T_{i-1}^n)
temp_new[i] = temp[i] + r * (temp[i + 1] - 2.0f * temp[i] + temp[i - 1]);
}
extern "C" __global__ void diffusion_explicit_2d(
const float* __restrict__ temp,
float* __restrict__ temp_new,
float alpha,
float dt,
float dx,
float dy,
int nx,
int ny
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
int j = blockIdx.y * blockDim.y + threadIdx.y;
if (i >= nx || j >= ny) return;
int idx = j * nx + i;
float rx = alpha * dt / (dx * dx);
float ry = alpha * dt / (dy * dy);
// Stability condition: rx + ry <= 0.5 for 2D explicit scheme
if (rx + ry > 0.5f) {
return;
}
// Boundary conditions
if (i == 0 || i == nx - 1 || j == 0 || j == ny - 1) {
temp_new[idx] = temp[idx]; // Fixed boundary values
return;
}
int idx_ip1 = j * nx + (i + 1);
int idx_im1 = j * nx + (i - 1);
int idx_jp1 = (j + 1) * nx + i;
int idx_jm1 = (j - 1) * nx + i;
// 2D explicit finite difference
temp_new[idx] = temp[idx] +
rx * (temp[idx_ip1] - 2.0f * temp[idx] + temp[idx_im1]) +
ry * (temp[idx_jp1] - 2.0f * temp[idx] + temp[idx_jm1]);
}
extern "C" __global__ void diffusion_implicit_1d(
const float* __restrict__ temp,
float* __restrict__ temp_new,
float alpha,
float dt,
float dx,
int n
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
// For implicit scheme, this kernel performs one iteration of iterative solver
// (e.g., Jacobi iteration for the linear system)
float r = alpha * dt / (dx * dx);
// Boundary conditions
if (i == 0 || i == n - 1) {
temp_new[i] = temp[i];
return;
}
// Jacobi iteration for implicit scheme:
// (1 + 2r)T_i^{n+1} - r(T_{i+1}^{n+1} + T_{i-1}^{n+1}) = T_i^n
// T_i^{n+1} = (T_i^n + r(T_{i+1}^{n+1} + T_{i-1}^{n+1})) / (1 + 2r)
temp_new[i] = (temp[i] + r * (temp_new[i + 1] + temp_new[i - 1])) / (1.0f + 2.0f * r);
}
extern "C" __global__ void diffusion_crank_nicolson_1d(
const float* __restrict__ temp,
float* __restrict__ temp_new,
float alpha,
float dt,
float dx,
int n
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
float r = alpha * dt / (dx * dx);
// Boundary conditions
if (i == 0 || i == n - 1) {
temp_new[i] = temp[i];
return;
}
// Crank-Nicolson scheme (θ = 0.5):
// T_i^{n+1} - 0.5*r*(T_{i+1}^{n+1} - 2*T_i^{n+1} + T_{i-1}^{n+1}) =
// T_i^n + 0.5*r*(T_{i+1}^n - 2*T_i^n + T_{i-1}^n)
float rhs = temp[i] + 0.5f * r * (temp[i + 1] - 2.0f * temp[i] + temp[i - 1]);
// Jacobi iteration for Crank-Nicolson
temp_new[i] = (rhs + 0.5f * r * (temp_new[i + 1] + temp_new[i - 1])) / (1.0f + r);
}
// Helper kernel for tridiagonal solver (Thomas algorithm)
extern "C" __global__ void thomas_forward_elimination(
float* __restrict__ a, // Lower diagonal
float* __restrict__ b, // Main diagonal
float* __restrict__ c, // Upper diagonal
float* __restrict__ d, // Right hand side
int n
) {
int i = blockIdx.x * blockDim.x + threadIdx.x + 1; // Start from i=1
if (i >= n) return;
// Forward elimination step
float m = a[i] / b[i - 1];
b[i] = b[i] - m * c[i - 1];
d[i] = d[i] - m * d[i - 1];
}
extern "C" __global__ void thomas_backward_substitution(
const float* __restrict__ b, // Modified main diagonal
const float* __restrict__ c, // Upper diagonal
float* __restrict__ d, // Modified RHS / solution
int n
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n - 1) return; // n-2, n-3, ..., 0
int idx = n - 2 - i; // Reverse order
if (idx < n - 1) {
d[idx] = (d[idx] - c[idx] * d[idx + 1]) / b[idx];
}
}
// Advanced diffusion kernels with anisotropic diffusion
extern "C" __global__ void anisotropic_diffusion_2d(
const float* __restrict__ temp,
float* __restrict__ temp_new,
const float* __restrict__ kappa_x, // Diffusivity in x-direction
const float* __restrict__ kappa_y, // Diffusivity in y-direction
float dt,
float dx,
float dy,
int nx,
int ny
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
int j = blockIdx.y * blockDim.y + threadIdx.y;
if (i >= nx || j >= ny) return;
int idx = j * nx + i;
// Boundary conditions
if (i == 0 || i == nx - 1 || j == 0 || j == ny - 1) {
temp_new[idx] = temp[idx];
return;
}
int idx_ip1 = j * nx + (i + 1);
int idx_im1 = j * nx + (i - 1);
int idx_jp1 = (j + 1) * nx + i;
int idx_jm1 = (j - 1) * nx + i;
// Anisotropic diffusion with spatially varying diffusivity
float kx = kappa_x[idx];
float ky = kappa_y[idx];
float rx = kx * dt / (dx * dx);
float ry = ky * dt / (dy * dy);
// Check stability condition
if (rx + ry > 0.5f) {
temp_new[idx] = temp[idx]; // Skip update if unstable
return;
}
temp_new[idx] = temp[idx] +
rx * (temp[idx_ip1] - 2.0f * temp[idx] + temp[idx_im1]) +
ry * (temp[idx_jp1] - 2.0f * temp[idx] + temp[idx_jm1]);
}
// Kernel for nonlinear diffusion (e.g., for turbulent heat transfer)
extern "C" __global__ void nonlinear_diffusion_2d(
const float* __restrict__ temp,
float* __restrict__ temp_new,
float base_alpha,
float dt,
float dx,
float dy,
int nx,
int ny
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
int j = blockIdx.y * blockDim.y + threadIdx.y;
if (i >= nx || j >= ny) return;
int idx = j * nx + i;
// Boundary conditions
if (i == 0 || i == nx - 1 || j == 0 || j == ny - 1) {
temp_new[idx] = temp[idx];
return;
}
int idx_ip1 = j * nx + (i + 1);
int idx_im1 = j * nx + (i - 1);
int idx_jp1 = (j + 1) * nx + i;
int idx_jm1 = (j - 1) * nx + i;
// Compute local gradient magnitude
float grad_x = (temp[idx_ip1] - temp[idx_im1]) / (2.0f * dx);
float grad_y = (temp[idx_jp1] - temp[idx_jm1]) / (2.0f * dy);
float grad_mag = sqrtf(grad_x * grad_x + grad_y * grad_y);
// Nonlinear diffusivity (example: edge-preserving diffusion)
float alpha = base_alpha / (1.0f + grad_mag * grad_mag);
float rx = alpha * dt / (dx * dx);
float ry = alpha * dt / (dy * dy);
// Check stability
if (rx + ry > 0.5f) {
temp_new[idx] = temp[idx];
return;
}
temp_new[idx] = temp[idx] +
rx * (temp[idx_ip1] - 2.0f * temp[idx] + temp[idx_im1]) +
ry * (temp[idx_jp1] - 2.0f * temp[idx] + temp[idx_jm1]);
}