NVIDIA
Adding Warp implementation of SDF
Modulus Pull Request
Description
Checklist
- [x] I am familiar with the Contributing Guidelines.
- [x] New or existing tests cover these changes.
- [x] The documentation is up to date with these changes.
- [ ] The CHANGELOG.md is up to date with these changes.
- [ ] An issue is linked to this pull request.
Dependencies
Probably not needed for this PR but I have found the following structure to be extremely helpful for keeping track of warp kernels and functions. The bellow code is to update the electric field in a FDTD solver. You can construct the function with electric_field_update = ElectricFeildUpdate(), and then call with electric_field_update(...). If you have this kind of structure all the warp kernels and function can be contained within. If you have kinda meta programming and want to change things inside the warp kernel then this can work as well. For example, suppose we wanted to make it work with 3D and 2D. You could implement the 2D kernels in the class and then when you make the function/operator electric_field_update = ElectricFieldUpdate(dim=3). Anyways, maybe there is a common structure we can use for this going forward.
class ElectricFieldUpdate(Operator):
"""
Electric field update operator
"""
@wp.func
def _sample_electric_property(
solid_id: wp.array4d(dtype=wp.uint8),
material_property: wp.array(dtype=wp.float32),
i: wp.int32,
j: wp.int32,
k: wp.int32,
):
# Get material property
prop_0_0_1 = material_property[wp.int32(solid_id[0, i - 1, j - 1, k])]
prop_0_1_0 = material_property[wp.int32(solid_id[0, i - 1, j, k - 1])]
prop_0_1_1 = material_property[wp.int32(solid_id[0, i - 1, j, k])]
prop_1_0_0 = material_property[wp.int32(solid_id[0, i, j - 1, k - 1])]
prop_1_0_1 = material_property[wp.int32(solid_id[0, i, j - 1, k])]
prop_1_1_0 = material_property[wp.int32(solid_id[0, i, j, k - 1])]
prop_1_1_1 = material_property[wp.int32(solid_id[0, i, j, k])]
# Get average property
prop_x = (prop_1_1_1 + prop_1_1_0 + prop_1_0_1 + prop_1_0_0) / 4.0
prop_y = (prop_1_1_1 + prop_1_1_0 + prop_0_1_1 + prop_0_1_0) / 4.0
prop_z = (prop_1_1_1 + prop_1_0_1 + prop_0_1_1 + prop_0_0_1) / 4.0
return wp.vec3(prop_x, prop_y, prop_z)
@wp.kernel
def _update_electric_field(
electric_field: wp.array4d(dtype=wp.float32),
magnetic_field: wp.array4d(dtype=wp.float32),
impressed_current: wp.array4d(dtype=wp.float32),
solid_id: wp.array4d(dtype=wp.uint8),
eps_mapping: wp.array(dtype=wp.float32),
sigma_e_mapping: wp.array(dtype=wp.float32),
spacing: wp.vec3f,
dt: wp.float32,
nr_ghost_cells: wp.int32,
):
# get index
i, j, k = wp.tid()
# Skip ghost cells
i += nr_ghost_cells
j += nr_ghost_cells
k += nr_ghost_cells
# get properties
eps = ElectricFieldUpdate._sample_electric_property(
solid_id, eps_mapping, i, j, k
)
sigma_e = ElectricFieldUpdate._sample_electric_property(
solid_id, sigma_e_mapping, i, j, k
)
# Get coefficients
_denom = 2.0 * eps + sigma_e * dt
c_ee = wp.cw_div(2.0 * eps - sigma_e * dt, _denom)
c_eh = (2.0 * dt) / wp.cw_mul(spacing, _denom)
c_ej = (-2.0 * dt) / _denom
# Get curl of magnetic field
curl_h_x = (magnetic_field[2, i, j, k] - magnetic_field[2, i, j - 1, k]) - (
magnetic_field[1, i, j, k] - magnetic_field[1, i, j, k - 1]
)
curl_h_y = (magnetic_field[0, i, j, k] - magnetic_field[0, i, j, k - 1]) - (
magnetic_field[2, i, j, k] - magnetic_field[2, i - 1, j, k]
)
curl_h_z = (magnetic_field[1, i, j, k] - magnetic_field[1, i - 1, j, k]) - (
magnetic_field[0, i, j, k] - magnetic_field[0, i, j - 1, k]
)
curl_h = wp.vec3(curl_h_x, curl_h_y, curl_h_z)
# compute new electric field
e = wp.vec3f(
electric_field[0, i, j, k],
electric_field[1, i, j, k],
electric_field[2, i, j, k],
)
cur = wp.vec3f(
impressed_current[0, i, j, k],
impressed_current[1, i, j, k],
impressed_current[2, i, j, k],
)
new_e = wp.cw_mul(c_ee, e) + wp.cw_mul(c_eh, curl_h) + wp.cw_mul(c_ej, cur)
# Set electric field
electric_field[0, i, j, k] = new_e[0]
electric_field[1, i, j, k] = new_e[1]
electric_field[2, i, j, k] = new_e[2]
def __call__(
self,
electric_field: wp.array4d(dtype=wp.float32),
magnetic_field: wp.array4d(dtype=wp.float32),
impressed_current: wp.array4d(dtype=wp.float32),
solid_id: wp.array4d(dtype=wp.uint8),
eps_mapping: wp.array(dtype=wp.float32),
sigma_e_mapping: wp.array(dtype=wp.float32),
spacing: Union[float, tuple[float, float, float]],
dt: float,
nr_ghost_cells: int = 1,
):
# Launch kernel
wp.launch(
self._update_electric_field,
inputs=[
electric_field,
magnetic_field,
impressed_current,
solid_id,
eps_mapping,
sigma_e_mapping,
spacing,
dt,
nr_ghost_cells,
],
dim=[x - 2 * nr_ghost_cells for x in solid_id.shape[1:]],
)
return electric_field
```
@mnabian added a couple of comments which I feel will have to be addressed. Let's discuss this further
@mnabian fixed the issues I pointed out and also used this for Sym's tessellation module here: https://github.com/NVIDIA/modulus-sym/pull/159