Repository: rapidsai/cucim (language: python)

Issue:
Improve performance of Euclidean distance transform
This PR gives a further 2-4x speedup to the 2D and 3D distance transform functions by creating dedicated elementwise CUDA kernels for the pre- and post-processing steps. These include, for example, packing and upacking the 3D coordinates into a single integer. This resolves various "TODO" comments that were in the source, but I don't think we have an issue open for this.

Aside from the performance benefit there is also a substantial reduction in memory usage for the default case of `return_distance=True, return_indices=False` as in this case the explicit arrays of indices never need to be created.

# Benchmarks

## Implementation from release 22.08.01

|      shape      | return dist. | return ind. | SciPy  | ITK | cuCIM  | accel. (CuPy vs SciPy) | accel. (CuPy vs ITK) |
|-----------------|--------------|-------------|--------|-----|--------|------------------------------|--------------------------- |
|256, 256       | False            | True           | 0.00379 | ---- | 0.00081 | 4.66                         | ---- |
|256, 256 | True | False | 0.00428 | 0.00383 | 0.00102 | 4.21 | 3.76 |
|256, 256 | True | True | 0.00429 | ---- | 0.00107 | 4.02 | ---- |
|512, 512 | False | True | 0.01724 | ---- | 0.00091 | 18.97 | ---- |
|512, 512 | True | False | 0.01866 | 0.00788 | 0.00104 | 17.94 | 7.58 |
|512, 512 | True | True | 0.01876 | ---- | 0.00107 | 17.48 | ---- |
|2048, 2048 | False | True | 0.35053 | ---- | 0.00256 | 137.03 | ---- |
|2048, 2048 | True | False | 0.39468 | 0.08220 | 0.00288 | 137.02 | 28.54 |
|2048, 2048 | True | True | 0.39356 | ---- | 0.00296 | 133.04 | ---- |
|4096, 4096 | False | True | 1.57430 | ---- | 0.00736 | 213.85 | ---- |
|4096, 4096 | True | False | 1.79513 | 0.54968 | 0.00858 | 209.34 | 64.10 |
|4096, 4096 | True | True | 1.81570 | ---- | 0.00889 | 204.25 | ---- |
|32, 32, 32 | False | True | 0.00319 | ---- | 0.00069 | 4.62 | ---- |
|32, 32, 32 | True | False | 0.00356 | 0.00481 | 0.00096 | 3.71 | 5.01 |
|32, 32, 32 | True | True | 0.00360 | ---- | 0.00103 | 3.49 | ---- |
|128, 128, 128 | False | True | 0.27901 | ---- | 0.00131 | 213.20 | ---- |
|128, 128, 128 | True | False | 0.31013 | 0.04504 | 0.00156 | 199.32 | 28.94 |
|128, 128, 128 | True | True | 0.31289 | ---- | 0.00163 | 192.07 | ---- |
|256, 256, 256 | False | True | 3.02435 | ---- | 0.00696 | 434.25 | ---- |
|256, 256, 256 | True | False | 3.34077 | 0.33528 | 0.00855 | 390.57 | 39.20 |
|256, 256, 256 | True | True | 3.32920 | ---- | 0.00915 | 363.89 | ---- |
|384, 384, 384 | False | True | 10.68973 | ---- | 0.02284 | 468.07 | ---- |
|384, 384, 384 | True | False | 11.73340 | 1.16172 | 0.02809 | 417.65 | 41.35 |
|384, 384, 384 | True | True | 11.68002 | ---- | 0.03007 | 388.40 | ---- |

## Proposed Implementation

|       shape      | return dist. | return ind.|  SciPy  | ITK | cuCIM  | accel. (CuPy vs SciPy) | accel. (CuPy vs ITK) |
| -----------------|--------------|------------|---------|-----|--------|------------------------------|--------------------------- |
| 256, 256 | False | True | 0.00385 | N/A | 0.00021 | 18.49 | N/A |
| 256, 256 | True | False | 0.00433 | 0.00370 | 0.00021 | 20.84 | 17.83 |
| 256, 256 | True | True | 0.00436 | N/A | 0.00022 | 19.55 | N/A |
| 512, 512 | False | True | 0.01724 | N/A | 0.00030 | 58.44 | N/A |
| 512, 512 | True | False | 0.01893 | 0.00790 | 0.00029 | 66.42 | 27.73 |
| 512, 512 | True | True | 0.01867 | N/A | 0.00030 | 62.65 | N/A | |
| 2048, 2048 | False | True | 0.33893 | N/A | 0.00154 | 220.25 | N/A |
| 2048, 2048 | True | False | 0.40912 | 0.08219 | 0.00154 | 266.11 | 53.46 |
| 2048, 2048 | True | True | 0.39056 | N/A | 0.00162 | 241.76 | N/A |
| 4096, 4096 | False | True | 1.58069 | N/A | 0.00418 | 377.92 | N/A |
| 4096, 4096 | True | False | 1.81860 | 0.55152 | 0.00418 | 435.36 | 132.03 |
| 4096, 4096 | True | True | 1.80671 | N/A | 0.00448 | 402.85 | N/A |
| 32, 32, 32 | False | True | 0.00319 | N/A | 0.00017 | 19.05 | N/A |
| 32, 32, 32 | True | False | 0.00359 | 0.00514 | 0.00011 | 33.25 | 47.68 |
| 32, 32, 32 | True | True | 0.00355 | N/A | 0.00019 | 18.37 | N/A |
| 128, 128, 128 | False | True | 0.27217 | N/A | 0.00059 | 458.70 | N/A |
| 128, 128, 128 | True | False | 0.31409 | 0.03384 | 0.00050 | 633.12 | 68.21 |
| 128, 128, 128 | True | True | 0.31207 | N/A | 0.00066 | 474.26 | N/A |
| 256, 256, 256 | False | True | 3.01718 | N/A | 0.00304 | 993.26 | N/A |
| 256, 256, 256 | True | False | 3.32660 | 0.28979 | 0.00235 | 1418.07 | 123.53 |
| 256, 256, 256 | True | True | 3.31874 | N/A | 0.00352 | 942.49 | N/A |
| 384, 384, 384 | False | True | 10.67328 | N/A | 0.01000 | 1067.34 | N/A |
| 384, 384, 384 | True | False | 11.69219 | 1.11274 | 0.00771 | 1517.46 | 144.42 |
| 384, 384, 384 | True | True | 11.76073 | N/A | 0.01161 | 1013.03 | N/A |



Discussion / hints:
@gpucibot merge

Fix patch:
diff --git a/python/cucim/src/cucim/core/operations/morphology/_distance_transform.py b/python/cucim/src/cucim/core/operations/morphology/_distance_transform.py
index 70523c5bf..7991fb9f8 100644
--- a/python/cucim/src/cucim/core/operations/morphology/_distance_transform.py
+++ b/python/cucim/src/cucim/core/operations/morphology/_distance_transform.py
@@ -3,9 +3,8 @@
 from ._pba_2d import _pba_2d
 from ._pba_3d import _pba_3d
 
-# TODO: support sampling distances
-#       support the distances and indices output arguments
-#       support chamfer, chessboard and l1/manhattan distances too?
+# TODO: support the distances and indices output arguments
+#       support chamfer/chessboard and taxicab/manhattan distances too?
 
 
 def distance_transform_edt(image, sampling=None, return_distances=True,
@@ -34,14 +33,17 @@ def distance_transform_edt(image, sampling=None, return_distances=True,
         Whether to calculate the distance transform.
     return_indices : bool, optional
         Whether to calculate the feature transform.
-    distances : float32 cupy.ndarray, optional
+    distances : cupy.ndarray, optional
         An output array to store the calculated distance transform, instead of
-        returning it. `return_distances` must be True. It must be the same
-        shape as `image`.
-    indices : int32 cupy.ndarray, optional
+        returning it. `return_distances` must be ``True``. It must be the same
+        shape as `image`. Should have dtype ``cp.float32`` if
+        `float64_distances` is ``False``, otherwise it should be
+        ``cp.float64``.
+    indices : cupy.ndarray, optional
         An output array to store the calculated feature transform, instead of
-        returning it. `return_indicies` must be True. Its shape must be
-        `(image.ndim,) + image.shape`.
+        returning it. `return_indicies` must be ``True``. Its shape must be
+        ``(image.ndim,) + image.shape``. Its dtype must be a signed or unsigned
+        integer type of at least 16-bits in 2D or 32-bits in 3D.
 
     Other Parameters
     ----------------
@@ -57,14 +59,16 @@ def distance_transform_edt(image, sampling=None, return_distances=True,
 
     Returns
     -------
-    distances : float64 ndarray, optional
+    distances : cupy.ndarray, optional
         The calculated distance transform. Returned only when
-        `return_distances` is True and `distances` is not supplied. It will
-        have the same shape as `image`.
-    indices : int32 ndarray, optional
+        `return_distances` is ``True`` and `distances` is not supplied. It will
+        have the same shape as `image`. Will have dtype `cp.float64` if
+        `float64_distances` is ``True``, otherwise it will have dtype
+        ``cp.float32``.
+    indices : ndarray, optional
         The calculated feature transform. It has an image-shaped array for each
         dimension of the image. See example below. Returned only when
-        `return_indices` is True and `indices` is not supplied.
+        `return_indices` is ``True`` and `indices` is not supplied.
 
     Notes
     -----
@@ -139,24 +143,14 @@ def distance_transform_edt(image, sampling=None, return_distances=True,
             [0, 0, 3, 3, 4]]])
 
     """
-    if distances is not None:
-        raise NotImplementedError(
-            "preallocated distances image is not supported"
-        )
-    if indices is not None:
-        raise NotImplementedError(
-            "preallocated indices image is not supported"
-        )
     scalar_sampling = None
     if sampling is not None:
-        sampling = np.unique(np.atleast_1d(sampling))
-        if len(sampling) == 1:
-            scalar_sampling = float(sampling)
+        unique_sampling = np.unique(np.atleast_1d(sampling))
+        if len(unique_sampling) == 1:
+            # In the isotropic case, can use the kernels without sample scaling
+            # and just adjust the final distance accordingly.
+            scalar_sampling = float(unique_sampling)
             sampling = None
-        else:
-            raise NotImplementedError(
-                "non-uniform values in sampling is not currently supported"
-            )
 
     if image.ndim == 3:
         pba_func = _pba_3d
@@ -171,11 +165,16 @@ def distance_transform_edt(image, sampling=None, return_distances=True,
         sampling=sampling,
         return_distances=return_distances,
         return_indices=return_indices,
-        block_params=block_params
+        block_params=block_params,
+        distances=distances,
+        indices=indices,
     )
 
     if return_distances and scalar_sampling is not None:
-        vals = (vals[0] * scalar_sampling,) + vals[1:]
+        # inplace multiply in case distance != None
+        vals = list(vals)
+        vals[0] *= scalar_sampling
+        vals = tuple(vals)
 
     if len(vals) == 1:
         vals = vals[0]
diff --git a/python/cucim/src/cucim/core/operations/morphology/_pba_2d.py b/python/cucim/src/cucim/core/operations/morphology/_pba_2d.py
index 0a9981c1d..3604c9fd3 100644
--- a/python/cucim/src/cucim/core/operations/morphology/_pba_2d.py
+++ b/python/cucim/src/cucim/core/operations/morphology/_pba_2d.py
@@ -67,8 +67,8 @@ def _init_marker(int_dtype):
 
 
 @cupy.memoize(True)
-def get_pba2d_src(block_size_2d=64, marker=-32768, pixel_int2_t='short2'):
-    make_pixel_func = 'make_' + pixel_int2_t
+def get_pba2d_src(block_size_2d=64, marker=-32768, pixel_int2_t="short2"):
+    make_pixel_func = "make_" + pixel_int2_t
 
     pba2d_code = pba2d_defines_template.format(
         block_size_2d=block_size_2d,
@@ -76,9 +76,9 @@ def get_pba2d_src(block_size_2d=64, marker=-32768, pixel_int2_t='short2'):
         pixel_int2_t=pixel_int2_t,
         make_pixel_func=make_pixel_func
     )
-    kernel_directory = os.path.join(os.path.dirname(__file__), 'cuda')
-    with open(os.path.join(kernel_directory, 'pba_kernels_2d.h'), 'rt') as f:
-        pba2d_kernels = '\n'.join(f.readlines())
+    kernel_directory = os.path.join(os.path.dirname(__file__), "cuda")
+    with open(os.path.join(kernel_directory, "pba_kernels_2d.h"), "rt") as f:
+        pba2d_kernels = "\n".join(f.readlines())
 
     pba2d_code += pba2d_kernels
     return pba2d_code
@@ -88,7 +88,7 @@ def _get_block_size(check_warp_size=False):
     if check_warp_size:
         dev = cupy.cuda.runtime.getDevice()
         device_properties = cupy.cuda.runtime.getDeviceProperties(dev)
-        return int(device_properties['warpSize'])
+        return int(device_properties["warpSize"])
     else:
         return 32
 
@@ -120,18 +120,18 @@ def _get_pack_kernel(int_type, marker=-32768):
         in_params="raw B arr",
         out_params="raw I out",
         operation=code,
-        options=('--std=c++11',),
+        options=("--std=c++11",),
     )
 
 
 def _pack_int2(arr, marker=-32768, int_dtype=cupy.int16):
     if arr.ndim != 2:
         raise ValueError("only 2d arr suppported")
-    int2_dtype = cupy.dtype({'names': ['x', 'y'], 'formats': [int_dtype] * 2})
+    int2_dtype = cupy.dtype({"names": ["x", "y"], "formats": [int_dtype] * 2})
     out = cupy.zeros(arr.shape + (2,), dtype=int_dtype)
     assert out.size == 2 * arr.size
     pack_kernel = _get_pack_kernel(
-        int_type='short' if int_dtype == cupy.int16 else 'int',
+        int_type="short" if int_dtype == cupy.int16 else "int",
         marker=marker
     )
     pack_kernel(arr, out, size=arr.size)
@@ -159,7 +159,7 @@ def _determine_padding(shape, padded_size, block_size):
     return padding_width
 
 
-def _generate_shape(ndim, int_type, var_name='out', raw_var=True):
+def _generate_shape(ndim, int_type, var_name="out", raw_var=True):
     code = ""
     if not raw_var:
         var_name = "_raw_" + var_name
@@ -169,16 +169,16 @@ def _generate_shape(ndim, int_type, var_name='out', raw_var=True):
 
 
 def _generate_indices_ops(ndim, int_type):
-    code = f'{int_type} _i = i;\n'
+    code = f"{int_type} _i = i;\n"
     for j in range(ndim - 1, 0, -1):
-        code += f'{int_type} ind_{j} = _i % shape_{j};\n_i /= shape_{j};\n'
-    code += f'{int_type} ind_0 = _i;'
+        code += f"{int_type} ind_{j} = _i % shape_{j};\n_i /= shape_{j};\n"
+    code += f"{int_type} ind_0 = _i;"
     return code
 
 
 def _get_distance_kernel_code(int_type, dist_int_type, raw_out_var=True):
     code = _generate_shape(
-        ndim=2, int_type=int_type, var_name='dist', raw_var=raw_out_var
+        ndim=2, int_type=int_type, var_name="dist", raw_var=raw_out_var
     )
     code += _generate_indices_ops(ndim=2, int_type=int_type)
     code += f"""
@@ -203,13 +203,86 @@ def _get_distance_kernel(int_type, dist_int_type):
         in_params="raw I y, raw I x",
         out_params="raw F dist",
         operation=operation,
-        options=('--std=c++11',),
+        options=("--std=c++11",),
     )
 
 
+def _get_aniso_distance_kernel_code(int_type, raw_out_var=True):
+    code = _generate_shape(
+        ndim=2, int_type=int_type, var_name="dist", raw_var=raw_out_var
+    )
+    code += _generate_indices_ops(ndim=2, int_type=int_type)
+    code += f"""
+    F tmp;
+    F sq_dist;
+    tmp = static_cast<F>(y[i] - ind_0) * sampling[0];
+    sq_dist = tmp * tmp;
+    tmp = static_cast<F>(x[i] - ind_1) * sampling[1];
+    sq_dist += tmp * tmp;
+    dist[i] = sqrt(sq_dist);
+    """
+    return code
+
+
+@cupy.memoize(for_each_device=True)
+def _get_aniso_distance_kernel(int_type):
+    """Returns kernel computing the Euclidean distance from coordinates."""
+    operation = _get_aniso_distance_kernel_code(int_type, raw_out_var=True)
+    return cupy.ElementwiseKernel(
+        in_params="raw I y, raw I x, raw F sampling",
+        out_params="raw F dist",
+        operation=operation,
+        options=("--std=c++11",),
+    )
+
+
+def _distance_tranform_arg_check(distances_out, indices_out,
+                                 return_distances, return_indices):
+    """Raise a RuntimeError if the arguments are invalid"""
+    error_msgs = []
+    if (not return_distances) and (not return_indices):
+        error_msgs.append(
+            "at least one of return_distances/return_indices must be True")
+    if distances_out and not return_distances:
+        error_msgs.append(
+            "return_distances must be True if distances is supplied"
+        )
+    if indices_out and not return_indices:
+        error_msgs.append("return_indices must be True if indices is supplied")
+    if error_msgs:
+        raise RuntimeError(", ".join(error_msgs))
+
+
+def _check_distances(distances, shape, dtype):
+    if distances.shape != shape:
+        raise RuntimeError("distances array has wrong shape")
+    if distances.dtype != dtype:
+        raise RuntimeError(
+            f"distances array must have dtype: {dtype}")
+
+
+def _check_indices(indices, shape, itemsize):
+    if indices.shape != shape:
+        raise RuntimeError("indices array has wrong shape")
+    if indices.dtype.kind not in 'iu':
+        raise RuntimeError(
+            f"indices array must have an integer dtype"
+        )
+    elif indices.dtype.itemsize < itemsize:
+        raise RuntimeError(
+            f"indices dtype must have itemsize > {itemsize}"
+        )
+
+
 def _pba_2d(arr, sampling=None, return_distances=True, return_indices=False,
             block_params=None, check_warp_size=False, *,
-            float64_distances=False):
+            float64_distances=False, distances=None, indices=None):
+
+    indices_inplace = isinstance(indices, cupy.ndarray)
+    dt_inplace = isinstance(distances, cupy.ndarray)
+    _distance_tranform_arg_check(
+        dt_inplace, indices_inplace, return_distances, return_indices
+    )
 
     # input_arr: a 2D image
     #    For each site at (x, y), the pixel at coordinate (x, y) should contain
@@ -220,11 +293,6 @@ def _pba_2d(arr, sampling=None, return_distances=True, return_indices=False,
     #       overhead of querying properties
     block_size = _get_block_size(check_warp_size)
 
-    if sampling is not None:
-        raise NotImplementedError("sampling not yet supported")
-        # if len(sampling) != 2:
-        #     raise ValueError("sampling must be a sequence of two values.")
-
     if block_params is None:
         padded_size = math.ceil(max(arr.shape) / block_size) * block_size
 
@@ -268,7 +336,7 @@ def _pba_2d(arr, sampling=None, return_distances=True, return_indices=False,
     shape_max = max(arr.shape)
     if shape_max <= 32768:
         int_dtype = cupy.int16
-        pixel_int2_type = 'short2'
+        pixel_int2_type = "short2"
     else:
         if shape_max > (1 << 24):
             # limit to coordinate range to 2**24 due to use of __mul24 in
@@ -278,20 +346,20 @@ def _pba_2d(arr, sampling=None, return_distances=True, return_indices=False,
                 f"shape {arr.shape}"
             )
         int_dtype = cupy.int32
-        pixel_int2_type = 'int2'
+        pixel_int2_type = "int2"
 
     marker = _init_marker(int_dtype)
 
     orig_sy, orig_sx = arr.shape
     padding_width = _determine_padding(arr.shape, padded_size, block_size)
     if padding_width is not None:
-        arr = cupy.pad(arr, padding_width, mode='constant', constant_values=1)
+        arr = cupy.pad(arr, padding_width, mode="constant", constant_values=1)
     size = arr.shape[0]
 
     input_arr = _pack_int2(arr, marker=marker, int_dtype=int_dtype)
     output = cupy.zeros_like(input_arr)
 
-    int2_dtype = cupy.dtype({'names': ['x', 'y'], 'formats': [int_dtype] * 2})
+    int2_dtype = cupy.dtype({"names": ["x", "y"], "formats": [int_dtype] * 2})
     margin = cupy.empty((2 * m1 * size,), dtype=int2_dtype)
 
     # phase 1 of PBA. m1 must divide texture size and be <= 64
@@ -302,17 +370,25 @@ def _pba_2d(arr, sampling=None, return_distances=True, return_indices=False,
             pixel_int2_t=pixel_int2_type,
         )
     )
-    kernelFloodDown = pba2d.get_function('kernelFloodDown')
-    kernelFloodUp = pba2d.get_function('kernelFloodUp')
-    kernelPropagateInterband = pba2d.get_function('kernelPropagateInterband')
-    kernelUpdateVertical = pba2d.get_function('kernelUpdateVertical')
-    kernelProximatePoints = pba2d.get_function('kernelProximatePoints')
+    kernelFloodDown = pba2d.get_function("kernelFloodDown")
+    kernelFloodUp = pba2d.get_function("kernelFloodUp")
+    kernelPropagateInterband = pba2d.get_function("kernelPropagateInterband")
+    kernelUpdateVertical = pba2d.get_function("kernelUpdateVertical")
     kernelCreateForwardPointers = pba2d.get_function(
-        'kernelCreateForwardPointers'
+        "kernelCreateForwardPointers"
     )
-    kernelMergeBands = pba2d.get_function('kernelMergeBands')
-    kernelDoubleToSingleList = pba2d.get_function('kernelDoubleToSingleList')
-    kernelColor = pba2d.get_function('kernelColor')
+    kernelDoubleToSingleList = pba2d.get_function("kernelDoubleToSingleList")
+
+    if sampling is None:
+        kernelProximatePoints = pba2d.get_function("kernelProximatePoints")
+        kernelMergeBands = pba2d.get_function("kernelMergeBands")
+        kernelColor = pba2d.get_function("kernelColor")
+    else:
+        kernelProximatePoints = pba2d.get_function(
+            "kernelProximatePointsWithSpacing"
+        )
+        kernelMergeBands = pba2d.get_function("kernelMergeBandsWithSpacing")
+        kernelColor = pba2d.get_function("kernelColorWithSpacing")
 
     block = (block_size, 1, 1)
     grid = (math.ceil(size / block[0]), m1, 1)
@@ -347,10 +423,19 @@ def _pba_2d(arr, sampling=None, return_distances=True, return_indices=False,
     block = (block_size, 1, 1)
     grid = (math.ceil(size / block[0]), m2, 1)
     bandSize2 = size // m2
+    if sampling is None:
+        sampling_args = ()
+    else:
+        # Originally the shape is (y, x) and sampling[1] corresponds to y.
+        # However, kernelUpdateVertical transposed the image, so
+        # we are now working with (x, y) instead. Need sampling ordered
+        # accordingly.
+        sampling = tuple(map(float, sampling))
+        sampling_args = (sampling[0], sampling[1])
     kernelProximatePoints(
         grid,
         block,
-        (output, input_arr, size, bandSize2),
+        (output, input_arr, size, bandSize2) + sampling_args,
     )
     kernelCreateForwardPointers(
         grid,
@@ -364,7 +449,7 @@ def _pba_2d(arr, sampling=None, return_distances=True, return_indices=False,
         kernelMergeBands(
             grid,
             block,
-            (output, input_arr, input_arr, size, size // noBand),
+            (output, input_arr, input_arr, size, size // noBand) + sampling_args,  # noqa
         )
         noBand //= 2
     # Replace the forward link with the X coordinate of the seed to remove
@@ -382,7 +467,7 @@ def _pba_2d(arr, sampling=None, return_distances=True, return_indices=False,
     kernelColor(
         grid,
         block,
-        (input_arr, output, size),
+        (input_arr, output, size) + sampling_args,
     )
 
     output = _unpack_int2(output, make_copy=False, int_dtype=int_dtype)
@@ -393,19 +478,35 @@ def _pba_2d(arr, sampling=None, return_distances=True, return_indices=False,
     vals = ()
     if return_distances:
         dtype_out = cupy.float64 if float64_distances else cupy.float32
-        dist = cupy.zeros(y.shape, dtype=dtype_out)
+        if dt_inplace:
+            _check_distances(distances, y.shape, dtype_out)
+        else:
+            distances = cupy.zeros(y.shape, dtype=dtype_out)
 
-        # make sure maximum possible distance doesn't overflow
+        # make sure maximum possible distance doesn"t overflow
         max_possible_dist = sum((s - 1)**2 for s in y.shape)
-        dist_int_type = 'int' if max_possible_dist < 2**31 else 'ptrdiff_t'
+        dist_int_type = "int" if max_possible_dist < 2**31 else "ptrdiff_t"
 
-        distance_kernel = _get_distance_kernel(
-            int_type=_get_inttype(dist),
-            dist_int_type=dist_int_type,
-        )
-        distance_kernel(y, x, dist, size=dist.size)
-        vals = vals + (dist,)
+        if sampling is None:
+            distance_kernel = _get_distance_kernel(
+                int_type=_get_inttype(distances),
+                dist_int_type=dist_int_type,
+            )
+            distance_kernel(y, x, distances, size=distances.size)
+        else:
+            distance_kernel = _get_aniso_distance_kernel(
+                int_type=_get_inttype(distances),
+            )
+            sampling = cupy.asarray(sampling, dtype=dtype_out)
+            distance_kernel(y, x, sampling, distances, size=distances.size)
+
+        vals = vals + (distances,)
     if return_indices:
-        indices = cupy.stack((y, x), axis=0)
+        if indices_inplace:
+            _check_indices(indices, (arr.ndim,) + arr.shape, x.dtype.itemsize)
+            indices[0, ...] = y
+            indices[1, ...] = x
+        else:
+            indices = cupy.stack((y, x), axis=0)
         vals = vals + (indices,)
     return vals
diff --git a/python/cucim/src/cucim/core/operations/morphology/_pba_3d.py b/python/cucim/src/cucim/core/operations/morphology/_pba_3d.py
index aa96edec0..d5fa6064b 100644
--- a/python/cucim/src/cucim/core/operations/morphology/_pba_3d.py
+++ b/python/cucim/src/cucim/core/operations/morphology/_pba_3d.py
@@ -4,7 +4,9 @@
 import cupy
 import numpy as np
 
-from ._pba_2d import (_get_block_size, _generate_shape, _generate_indices_ops,
+from ._pba_2d import (_check_distances, _check_indices,
+                      _distance_tranform_arg_check, _get_block_size,
+                      _generate_shape, _generate_indices_ops,
                       lcm)
 from cucim.skimage._vendored._ndimage_util import _get_inttype
 
@@ -75,9 +77,9 @@ def get_pba3d_src(block_size_3d=32, marker=-2147483648, max_int=2147483647,
         pba3d_code += pba3d_defines_encode_64bit
     else:
         pba3d_code += pba3d_defines_encode_32bit
-    kernel_directory = os.path.join(os.path.dirname(__file__), 'cuda')
-    with open(os.path.join(kernel_directory, 'pba_kernels_3d.h'), 'rt') as f:
-        pba3d_kernels = '\n'.join(f.readlines())
+    kernel_directory = os.path.join(os.path.dirname(__file__), "cuda")
+    with open(os.path.join(kernel_directory, "pba_kernels_3d.h"), "rt") as f:
+        pba3d_kernels = "\n".join(f.readlines())
     pba3d_code += pba3d_kernels
     return pba3d_code
 
@@ -86,9 +88,9 @@ def get_pba3d_src(block_size_3d=32, marker=-2147483648, max_int=2147483647,
 def _get_encode3d_kernel(size_max, marker=-2147483648):
     """Pack array coordinates into a single integer."""
     if size_max > 1024:
-        int_type = 'ptrdiff_t'  # int64_t
+        int_type = "ptrdiff_t"  # int64_t
     else:
-        int_type = 'int'        # int32_t
+        int_type = "int"        # int32_t
 
     # value must match TOID macro in the C++ code!
     if size_max > 1024:
@@ -115,7 +117,7 @@ def _get_encode3d_kernel(size_max, marker=-2147483648):
         in_params="raw B arr",
         out_params="raw I out",
         operation=code,
-        options=('--std=c++11',),
+        options=("--std=c++11",),
     )
 
 
@@ -128,13 +130,13 @@ def encode3d(arr, marker=-2147483648, bit_depth=32, size_max=1024):
         dtype = np.int64
     else:
         dtype = np.int32
-    image = cupy.zeros(arr.shape, dtype=dtype, order='C')
+    image = cupy.zeros(arr.shape, dtype=dtype, order="C")
     kern = _get_encode3d_kernel(size_max, marker=marker)
     kern(arr, image, size=image.size)
     return image
 
 
-def _get_decode3d_code(size_max, int_type=''):
+def _get_decode3d_code(size_max, int_type=""):
     # bit shifts here must match those used in the encode3d kernel
     if size_max > 1024:
         code = f"""
@@ -155,14 +157,14 @@ def _get_decode3d_code(size_max, int_type=''):
 def _get_decode3d_kernel(size_max):
     """Unpack 3 coordinates encoded as a single integer."""
 
-    # int_type = '' here because x, y, z were already allocated externally
-    code = _get_decode3d_code(size_max, int_type='')
+    # int_type = "" here because x, y, z were already allocated externally
+    code = _get_decode3d_code(size_max, int_type="")
 
     return cupy.ElementwiseKernel(
         in_params="E encoded",
         out_params="I x, I y, I z",
         operation=code,
-        options=('--std=c++11',),
+        options=("--std=c++11",),
     )
 
 
@@ -229,7 +231,7 @@ def _generate_distance_computation(int_type, dist_int_type):
 
 def _get_distance_kernel_code(int_type, dist_int_type, raw_out_var=True):
     code = _generate_shape(
-        ndim=3, int_type=int_type, var_name='dist', raw_var=raw_out_var
+        ndim=3, int_type=int_type, var_name="dist", raw_var=raw_out_var
     )
     code += _generate_indices_ops(ndim=3, int_type=int_type)
     code += _generate_distance_computation(int_type, dist_int_type)
@@ -239,7 +241,7 @@ def _get_distance_kernel_code(int_type, dist_int_type, raw_out_var=True):
 @cupy.memoize(for_each_device=True)
 def _get_distance_kernel(int_type, large_dist=False):
     """Returns kernel computing the Euclidean distance from coordinates."""
-    dist_int_type = 'ptrdiff_t' if large_dist else 'int'
+    dist_int_type = "ptrdiff_t" if large_dist else "int"
     operation = _get_distance_kernel_code(
         int_type, dist_int_type, raw_out_var=True
     )
@@ -247,12 +249,51 @@ def _get_distance_kernel(int_type, large_dist=False):
         in_params="I z, I y, I x",
         out_params="raw F dist",
         operation=operation,
-        options=('--std=c++11',),
+        options=("--std=c++11",),
+    )
+
+
+def _generate_aniso_distance_computation():
+    """
+    Compute euclidean distance from current coordinate (ind_0, ind_1, ind_2) to
+    the coordinates of the nearest point (z, y, x)."""
+    return f"""
+    F tmp = static_cast<F>(z - ind_0) * sampling[0];
+    F sq_dist = tmp * tmp;
+    tmp = static_cast<F>(y - ind_1) * sampling[1];
+    sq_dist += tmp * tmp;
+    tmp = static_cast<F>(x - ind_2) * sampling[2];
+    sq_dist += tmp * tmp;
+    dist[i] = sqrt(static_cast<F>(sq_dist));
+    """
+
+
+def _get_aniso_distance_kernel_code(int_type, raw_out_var=True):
+    code = _generate_shape(
+        ndim=3, int_type=int_type, var_name="dist", raw_var=raw_out_var
+    )
+    code += _generate_indices_ops(ndim=3, int_type=int_type)
+    code += _generate_aniso_distance_computation()
+    return code
+
+
+@cupy.memoize(for_each_device=True)
+def _get_aniso_distance_kernel(int_type):
+    """Returns kernel computing the Euclidean distance from coordinates with
+    axis spacing != 1."""
+    operation = _get_aniso_distance_kernel_code(
+        int_type, raw_out_var=True
+    )
+    return cupy.ElementwiseKernel(
+        in_params="I z, I y, I x, raw F sampling",
+        out_params="raw F dist",
+        operation=operation,
+        options=("--std=c++11",),
     )
 
 
 @cupy.memoize(for_each_device=True)
-def _get_decode_as_distance_kernel(size_max, large_dist=False):
+def _get_decode_as_distance_kernel(size_max, large_dist=False, sampling=None):
     """Fused decode3d and distance computation.
 
     This kernel is for use when `return_distances=True`, but
@@ -261,38 +302,45 @@ def _get_decode_as_distance_kernel(size_max, large_dist=False):
     generating full arrays containing the coordinates since the coordinate
     arrays are not going to be returned.
     """
-    dist_int_type = 'ptrdiff_t' if large_dist else 'int'
-    int_type = 'int'
+    if sampling is None:
+        dist_int_type = "ptrdiff_t" if large_dist else "int"
+    int_type = "int"
 
     # Step 1: decode the (z, y, x) coordinate
     code = _get_decode3d_code(size_max, int_type=int_type)
 
     # Step 2: compute the Euclidean distance based on this (z, y, x).
     code += _generate_shape(
-        ndim=3, int_type=int_type, var_name='dist', raw_var=True
+        ndim=3, int_type=int_type, var_name="dist", raw_var=True
     )
     code += _generate_indices_ops(ndim=3, int_type=int_type)
-    code += _generate_distance_computation(int_type, dist_int_type)
-
+    if sampling is None:
+        code += _generate_distance_computation(int_type, dist_int_type)
+        in_params = "E encoded"
+    else:
+        code += _generate_aniso_distance_computation()
+        in_params = "E encoded, raw F sampling"
     return cupy.ElementwiseKernel(
-        in_params="E encoded",
+        in_params=in_params,
         out_params="raw F dist",
         operation=code,
-        options=('--std=c++11',),
+        options=("--std=c++11",),
     )
 
 
 def _pba_3d(arr, sampling=None, return_distances=True, return_indices=False,
             block_params=None, check_warp_size=False, *,
-            float64_distances=False):
+            float64_distances=False, distances=None, indices=None):
+
+    indices_inplace = isinstance(indices, cupy.ndarray)
+    dt_inplace = isinstance(distances, cupy.ndarray)
+    _distance_tranform_arg_check(
+        dt_inplace, indices_inplace, return_distances, return_indices
+    )
+
     if arr.ndim != 3:
         raise ValueError(f"expected a 3D array, got {arr.ndim}D")
 
-    if sampling is not None:
-        raise NotImplementedError("sampling not yet supported")
-        # if len(sampling) != 3:
-        #     raise ValueError("sampling must be a sequence of three values.")
-
     if block_params is None:
         m1 = 1
         m2 = 1
@@ -319,7 +367,7 @@ def _pba_3d(arr, sampling=None, return_distances=True, return_indices=False,
         arr.shape, block_size, m1, m2, m3, blockx, blocky
     )
     if padding_width is not None:
-        arr = cupy.pad(arr, padding_width, mode='constant', constant_values=1)
+        arr = cupy.pad(arr, padding_width, mode="constant", constant_values=1)
     size = arr.shape[0]
 
     # pba algorithm was implemented to use 32-bit integer to store compressed
@@ -336,9 +384,16 @@ def _pba_3d(arr, sampling=None, return_distances=True, return_indices=False,
         code=get_pba3d_src(block_size_3d=block_size, size_max=size_max)
     )
 
-    kernelFloodZ = pba3d.get_function('kernelFloodZ')
-    kernelMaurerAxis = pba3d.get_function('kernelMaurerAxis')
-    kernelColorAxis = pba3d.get_function('kernelColorAxis')
+    kernelFloodZ = pba3d.get_function("kernelFloodZ")
+    if sampling is None:
+        kernelMaurerAxis = pba3d.get_function("kernelMaurerAxis")
+        kernelColorAxis = pba3d.get_function("kernelColorAxis")
+        sampling_args = ()
+    else:
+        kernelMaurerAxis = pba3d.get_function("kernelMaurerAxisWithSpacing")
+        kernelColorAxis = pba3d.get_function("kernelColorAxisWithSpacing")
+        sampling = tuple(map(float, sampling))
+        sampling_args = (sampling[2], sampling[1], sampling[0])
 
     kernelFloodZ(
         grid,
@@ -352,7 +407,7 @@ def _pba_3d(arr, sampling=None, return_distances=True, return_indices=False,
     kernelMaurerAxis(
         grid,
         block,
-        (pba_images[buffer_idx], pba_images[1 - buffer_idx], size),
+        (pba_images[buffer_idx], pba_images[1 - buffer_idx], size) + sampling_args,  # noqa
     )
 
     block = (block_size, m3, 1)
@@ -360,15 +415,20 @@ def _pba_3d(arr, sampling=None, return_distances=True, return_indices=False,
     kernelColorAxis(
         grid,
         block,
-        (pba_images[1 - buffer_idx], pba_images[buffer_idx], size),
+        (pba_images[1 - buffer_idx], pba_images[buffer_idx], size) + sampling_args,  # noqa
     )
 
+    if sampling is not None:
+        # kernelColorAxis transposes the first two axis, so have to reorder
+        # the sampling_args tuple correspondingly
+        sampling_args = (sampling[1], sampling[2], sampling[0])
+
     block = (blockx, blocky, 1)
     grid = (size // block[0], size // block[1], 1)
     kernelMaurerAxis(
         grid,
         block,
-        (pba_images[buffer_idx], pba_images[1 - buffer_idx], size),
+        (pba_images[buffer_idx], pba_images[1 - buffer_idx], size) + sampling_args,  # noqa
     )
 
     block = (block_size, m3, 1)
@@ -376,14 +436,17 @@ def _pba_3d(arr, sampling=None, return_distances=True, return_indices=False,
     kernelColorAxis(
         grid,
         block,
-        (pba_images[1 - buffer_idx], pba_images[buffer_idx], size),
+        (pba_images[1 - buffer_idx], pba_images[buffer_idx], size) + sampling_args,  # noqa
     )
     output = pba_images[buffer_idx]
 
     if return_distances:
         out_shape = (orig_sz, orig_sy, orig_sx)
         dtype_out = cupy.float64 if float64_distances else cupy.float32
-        dist = cupy.zeros(out_shape, dtype=dtype_out)
+        if dt_inplace:
+            _check_distances(distances, out_shape, dtype_out)
+        else:
+            distances = cupy.zeros(out_shape, dtype=dtype_out)
 
         # make sure maximum possible distance doesn't overflow
         max_possible_dist = sum((s - 1)**2 for s in out_shape)
@@ -394,21 +457,37 @@ def _pba_3d(arr, sampling=None, return_distances=True, return_indices=False,
             kern = _get_decode_as_distance_kernel(
                 size_max=size_max,
                 large_dist=large_dist,
+                sampling=sampling
             )
-            kern(output[:orig_sz, :orig_sy, :orig_sx], dist)
-            return (dist,)
+            if sampling is None:
+                kern(output[:orig_sz, :orig_sy, :orig_sx], distances)
+            else:
+                sampling = cupy.asarray(sampling, dtype=distances.dtype)
+                kern(output[:orig_sz, :orig_sy, :orig_sx], sampling, distances)
+            return (distances,)
 
     if return_indices:
         x, y, z = decode3d(output[:orig_sz, :orig_sy, :orig_sx],
                            size_max=size_max)
     vals = ()
     if return_distances:
-        kern = _get_distance_kernel(
-            int_type=_get_inttype(dist), large_dist=large_dist,
-        )
-        kern(z, y, x, dist)
-        vals = vals + (dist,)
+        if sampling is None:
+            kern = _get_distance_kernel(
+                int_type=_get_inttype(distances), large_dist=large_dist,
+            )
+            kern(z, y, x, distances)
+        else:
+            kern = _get_aniso_distance_kernel(int_type=_get_inttype(distances))
+            sampling = cupy.asarray(sampling, dtype=distances.dtype)
+            kern(z, y, x, sampling, distances)
+        vals = vals + (distances,)
     if return_indices:
-        indices = cupy.stack((z, y, x), axis=0)
+        if indices_inplace:
+            _check_indices(indices, (arr.ndim,) + arr.shape, x.dtype.itemsize)
+            indices[0, ...] = z
+            indices[1, ...] = y
+            indices[2, ...] = x
+        else:
+            indices = cupy.stack((z, y, x), axis=0)
         vals = vals + (indices,)
     return vals
diff --git a/python/cucim/src/cucim/core/operations/morphology/cuda/pba_kernels_2d.h b/python/cucim/src/cucim/core/operations/morphology/cuda/pba_kernels_2d.h
index 61677c682..200821f06 100644
--- a/python/cucim/src/cucim/core/operations/morphology/cuda/pba_kernels_2d.h
+++ b/python/cucim/src/cucim/core/operations/morphology/cuda/pba_kernels_2d.h
@@ -8,6 +8,15 @@
 // Modifications by Gregory Lee (2022) (NVIDIA)
 // - add user-defined pixel_int2_t to enable
 // - replace __mul24 operations with standard multiplication operator
+// - Add variant kernels with support for non-isotropic pixel dimensions (
+//   kernelProximatePointsWithSpacing, kernelMergeBandsWithSpacing
+//   and kernelColorWithSpacing. These kernels differ from the originals in that
+//   they also take sx and sy values indicating the pixel size along the
+//   x and y axes. The kernels are identical except that the `dominate` function
+//   is replaced by `dominate_sp` and the physical spacings are used when
+//   computing distances.
+//
+
 
 
 // START OF DEFINITIONS OVERRIDDEN BY THE PYTHON SCRIPT
@@ -46,6 +55,22 @@ __device__ bool dominate(LL x1, LL y1, LL x2, LL y2, LL x3, LL y3, LL x0)
 }
 #undef LL
 
+// version of dominate, but with per-axis floating-point spacing
+__device__ bool dominate_sp(int _x1, int _y1, int _x2, int _y2, int _x3, int _y3, int _x0, float sx, float sy)
+{
+    float x1 = static_cast<float>(_x1) * sx;
+    float x2 = static_cast<float>(_x2) * sx;
+    float x3 = static_cast<float>(_x3) * sx;
+    float y1 = static_cast<float>(_y1) * sy;
+    float y2 = static_cast<float>(_y2) * sy;
+    float y3 = static_cast<float>(_y3) * sy;
+    float x0_2 = static_cast<float>(_x0 << 1) * sx;
+    float k1 = (y2 - y1);
+    float k2 = (y3 - y2);
+    return (k1 * (y1 + y2) + (x2 - x1) * ((x1 + x2) - x0_2)) * k2 > \
+            (k2 * (y2 + y3) + (x3 - x2) * ((x2 + x3) - x0_2)) * k1;
+}
+
 
 extern "C"{
 
@@ -229,6 +254,43 @@ __global__ void kernelProximatePoints(pixel_int2_t *input, pixel_int2_t *stack,
         stack[TOID(tx, ty + bandSize - 1, size)] = make_pixel(MARKER, lasty);
 }
 
+
+__global__ void kernelProximatePointsWithSpacing(pixel_int2_t *input, pixel_int2_t *stack, int size, int bandSize, double sx, double sy)
+{
+    int tx = blockIdx.x * blockDim.x + threadIdx.x;
+    int ty = blockIdx.y * bandSize;
+    int id = TOID(tx, ty, size);
+    int lasty = -1;
+    pixel_int2_t last1, last2, current;
+
+    last1.y = -1; last2.y = -1;
+
+    for (int i = 0; i < bandSize; i++, id += size) {
+        current = input[id];
+
+        if (current.x != MARKER) {
+            while (last2.y >= 0) {
+                if (!dominate_sp(last1.x, last2.y, last2.x, \
+                    lasty, current.x, current.y, tx, sx, sy))
+                    break;
+
+                lasty = last2.y; last2 = last1;
+
+                if (last1.y >= 0)
+                    last1 = stack[TOID(tx, last1.y, size)];
+            }
+
+            last1 = last2; last2 = make_pixel(current.x, lasty); lasty = current.y;
+
+            stack[id] = last2;
+        }
+    }
+
+    // Store the pointer to the tail at the last pixel of this band
+    if (lasty != ty + bandSize - 1)
+        stack[TOID(tx, ty + bandSize - 1, size)] = make_pixel(MARKER, lasty);
+}
+
 __global__ void kernelCreateForwardPointers(pixel_int2_t *input, pixel_int2_t *output, int size, int bandSize)
 {
     int tx = blockIdx.x * blockDim.x + threadIdx.x;
@@ -378,6 +440,126 @@ __global__ void kernelMergeBands(pixel_int2_t *color, pixel_int2_t *link, pixel_
     }
 }
 
+
+__global__ void kernelMergeBandsWithSpacing(pixel_int2_t *color, pixel_int2_t *link, pixel_int2_t *output, int size, int bandSize, double sx, double sy)
+{
+    int tx = blockIdx.x * blockDim.x + threadIdx.x;
+    int band1 = blockIdx.y * 2;
+    int band2 = band1 + 1;
+    int firsty, lasty;
+    pixel_int2_t last1, last2, current;
+    // last1 and last2: x component store the x coordinate of the site,
+    // y component store the backward pointer
+    // current: y component store the x coordinate of the site,
+    // x component store the forward pointer
+
+    // Get the two last items of the first list
+    lasty = band2 * bandSize - 1;
+    last2 = make_pixel(color[TOID(tx, lasty, size)].x,
+        link[TOID(tx, lasty, size)].y);
+
+    if (last2.x == MARKER) {
+        lasty = last2.y;
+
+        if (lasty >= 0)
+            last2 = make_pixel(color[TOID(tx, lasty, size)].x,
+            link[TOID(tx, lasty, size)].y);
+        else
+            last2 = make_pixel(MARKER, MARKER);
+    }
+
+    if (last2.y >= 0) {
+        // Second item at the top of the stack
+        last1 = make_pixel(color[TOID(tx, last2.y, size)].x,
+            link[TOID(tx, last2.y, size)].y);
+    }
+
+    // Get the first item of the second band
+    firsty = band2 * bandSize;
+    current = make_pixel(link[TOID(tx, firsty, size)].x,
+        color[TOID(tx, firsty, size)].x);
+
+    if (current.y == MARKER) {
+        firsty = current.x;
+
+        if (firsty >= 0)
+            current = make_pixel(link[TOID(tx, firsty, size)].x,
+            color[TOID(tx, firsty, size)].x);
+        else
+            current = make_pixel(MARKER, MARKER);
+    }
+
+    // Count the number of item in the second band that survive so far.
+    // Once it reaches 2, we can stop.
+    int top = 0;
+
+    while (top < 2 && current.y >= 0) {
+        // While there's still something on the left
+        while (last2.y >= 0) {
+
+            if (!dominate_sp(last1.x, last2.y, last2.x, \
+                lasty, current.y, firsty, tx, sx, sy))
+                break;
+
+            lasty = last2.y; last2 = last1;
+            top--;
+
+            if (last1.y >= 0)
+                last1 = make_pixel(color[TOID(tx, last1.y, size)].x,
+                link[TOID(tx, last1.y, size)].y);
+        }
+
+        // Update the current pointer
+        output[TOID(tx, firsty, size)] = make_pixel(current.x, lasty);
+
+        if (lasty >= 0)
+            output[TOID(tx, lasty, size)] = make_pixel(firsty, last2.y);
+
+        last1 = last2; last2 = make_pixel(current.y, lasty); lasty = firsty;
+        firsty = current.x;
+
+        top = max(1, top + 1);
+
+        // Advance the current pointer to the next one
+        if (firsty >= 0)
+            current = make_pixel(link[TOID(tx, firsty, size)].x,
+            color[TOID(tx, firsty, size)].x);
+        else
+            current = make_pixel(MARKER, MARKER);
+    }
+
+    // Update the head and tail pointer.
+    firsty = band1 * bandSize;
+    lasty = band2 * bandSize;
+    current = link[TOID(tx, firsty, size)];
+
+    if (current.y == MARKER && current.x < 0) { // No head?
+        last1 = link[TOID(tx, lasty, size)];
+
+        if (last1.y == MARKER)
+            current.x = last1.x;
+        else
+            current.x = lasty;
+
+        output[TOID(tx, firsty, size)] = current;
+    }
+
+    firsty = band1 * bandSize + bandSize - 1;
+    lasty = band2 * bandSize + bandSize - 1;
+    current = link[TOID(tx, lasty, size)];
+
+    if (current.x == MARKER && current.y < 0) { // No tail?
+        last1 = link[TOID(tx, firsty, size)];
+
+        if (last1.x == MARKER)
+            current.y = last1.y;
+        else
+            current.y = firsty;
+
+        output[TOID(tx, lasty, size)] = current;
+    }
+}
+
 __global__ void kernelDoubleToSingleList(pixel_int2_t *color, pixel_int2_t *link, pixel_int2_t *output, int size)
 {
     int tx = blockIdx.x * blockDim.x + threadIdx.x;
@@ -448,4 +630,69 @@ __global__ void kernelColor(pixel_int2_t *input, pixel_int2_t *output, int size)
         __syncthreads();
     }
 }
+
+
+__global__ void kernelColorWithSpacing(pixel_int2_t *input, pixel_int2_t *output, int size, double sx, double sy)
+{
+    __shared__ pixel_int2_t block[BLOCKSIZE][BLOCKSIZE];
+
+    int col = threadIdx.x;
+    int tid = threadIdx.y;
+    int tx = blockIdx.x * blockDim.x + col;
+    int lasty;
+    double dx, dy, best, dist;
+    pixel_int2_t last1, last2;
+
+    lasty = size - 1;
+
+    last2 = input[TOID(tx, lasty, size)];
+
+    if (last2.x == MARKER) {
+        lasty = last2.y;
+        last2 = input[TOID(tx, lasty, size)];
+    }
+
+    if (last2.y >= 0)
+        last1 = input[TOID(tx, last2.y, size)];
+
+    int y_start, y_end, n_step = size / blockDim.x;
+    for(int step = 0; step < n_step; ++step) {
+        y_start = size - step * blockDim.x - 1;
+        y_end = size - (step + 1) * blockDim.x;
+
+        for (int ty = y_start - tid; ty >= y_end; ty -= blockDim.y) {
+            dx = static_cast<double>(last2.x - tx) * sx;
+            dy = static_cast<double>(lasty - ty) * sy;
+            best = dist = dx * dx + dy * dy;
+
+            while (last2.y >= 0) {
+                dx = static_cast<double>(last1.x - tx) * sx;
+                dy = static_cast<double>(last2.y - ty) * sy;
+                dist = dx * dx + dy * dy;
+
+                if (dist > best)
+                    break;
+
+                best = dist; lasty = last2.y; last2 = last1;
+
+                if (last2.y >= 0)
+                    last1 = input[TOID(tx, last2.y, size)];
+            }
+
+            block[threadIdx.x][ty - y_end] = make_pixel(lasty, last2.x);
+        }
+
+        __syncthreads();
+
+        // note: transposes back to original shape here
+        if(!threadIdx.y) {
+            int id = TOID(y_end + threadIdx.x, blockIdx.x * blockDim.x, size);
+            for(int i = 0; i < blockDim.x; ++i, id+=size) {
+                output[id] = block[i][threadIdx.x];
+            }
+        }
+
+        __syncthreads();
+    }
+}
 } // extern C
diff --git a/python/cucim/src/cucim/core/operations/morphology/cuda/pba_kernels_3d.h b/python/cucim/src/cucim/core/operations/morphology/cuda/pba_kernels_3d.h
index c09f4b51f..d4cdec9d6 100644
--- a/python/cucim/src/cucim/core/operations/morphology/cuda/pba_kernels_3d.h
+++ b/python/cucim/src/cucim/core/operations/morphology/cuda/pba_kernels_3d.h
@@ -6,6 +6,12 @@
 //
 // Modifications by Gregory Lee (2022) (NVIDIA)
 // - allow user-defined ENCODED_INT_TYPE, ENCODE, DECODE
+// - Add variant kernels with support for non-isotropic pixel dimensions (
+//   kernelMaurerAxisWithSpacing, kernelColorAxisWithSpacing. These kernels
+//   differ from the originals in that they also take sx, sy and sz values
+//   indicating the pixel size along the x, y and z axes. The kernels are
+//   identical except that the `dominate` function is replaced by
+//   `dominate_sp` and the physical spacings are used when computing distances.
 
 
 // START OF DEFINITIONS OVERRIDDEN BY THE PYTHON SCRIPT
@@ -65,6 +71,26 @@ __device__ bool dominate(LL x_1, LL y_1, LL z_1, LL x_2, LL y_2, LL z_2, LL x_3,
 }
 #undef LL
 
+__device__ bool dominate_sp(int _x_1, int _y_1, int _z_1, int _x_2, int _y_2, int _z_2, int _x_3, int _y_3, int _z_3, int _x_0, int _z_0, float sx, float sy, float sz)
+{
+    float x_1 = static_cast<float>(_x_1) * sx;
+    float y_1 = static_cast<float>(_y_1) * sy;
+    float z_1 = static_cast<float>(_z_1) * sz;
+    float x_2 = static_cast<float>(_x_2) * sx;
+    float y_2 = static_cast<float>(_y_2) * sy;
+    float z_2 = static_cast<float>(_z_2) * sz;
+    float x_3 = static_cast<float>(_x_3) * sx;
+    float y_3 = static_cast<float>(_y_3) * sy;
+    float z_3 = static_cast<float>(_z_3) * sz;
+    float x_0_2 = static_cast<float>(_x_0 << 1) * sx;
+    float z_0_2 = static_cast<float>(_z_0 << 1) * sz;
+    float k_1 = y_2 - y_1;
+    float k_2 = y_3 - y_2;
+
+    return (((y_1 + y_2) * k_1 + ((x_2 - x_1) * (x_1 + x_2 - (x_0_2)) + (z_2 - z_1) * (z_1 + z_2 - (z_0_2)))) * k_2 > \
+            ((y_2 + y_3) * k_2 + ((x_3 - x_2) * (x_2 + x_3 - (x_0_2)) + (z_3 - z_2) * (z_2 + z_3 - (z_0_2)))) * k_1);
+}
+
 #define TOID(x, y, z, size)    ((((z) * (size)) + (y)) * (size) + (x))
 
 
@@ -162,6 +188,54 @@ __global__ void kernelMaurerAxis(ENCODED_INT_TYPE *input, ENCODED_INT_TYPE *stac
         stack[TOID(tx, ty - 1, tz, size)] = ENCODE(ZERO, lasty, ZERO, ONE, flag);
 }
 
+__global__ void kernelMaurerAxisWithSpacing(ENCODED_INT_TYPE *input, ENCODED_INT_TYPE *stack, int size, double sx, double sy, double sz)
+{
+    int tx = blockIdx.x * blockDim.x + threadIdx.x;
+    int tz = blockIdx.y * blockDim.y + threadIdx.y;
+    int ty = 0;
+
+    int id = TOID(tx, ty, tz, size);
+
+    ENCODED_INT_TYPE lasty = 0;
+    ENCODED_INT_TYPE x1, y1, z1, x2, y2, z2, nx, ny, nz;
+    ENCODED_INT_TYPE p = ENCODE(ZERO,ZERO,ZERO,ONE,ZERO), s1 = ENCODE(ZERO,ZERO,ZERO,ONE,ZERO), s2 = ENCODE(ZERO,ZERO,ZERO,ONE,ZERO);
+    ENCODED_INT_TYPE flag = 0;
+
+    for (ty = 0; ty < size; ++ty, id += size) {
+        p = input[id];
+
+        if (!NOTSITE(p)) {
+
+            while (HASNEXT(s2)) {
+                DECODE(s1, x1, y1, z1);
+                DECODE(s2, x2, y2, z2);
+                DECODE(p, nx, ny, nz);
+
+                if (!dominate_sp(x1, y2, z1, x2, lasty, z2, nx, ty, nz, tx, tz, sx, sy, sz))
+                    break;
+
+                lasty = y2; s2 = s1; y2 = y1;
+
+                if (HASNEXT(s2))
+                    s1 = stack[TOID(tx, y2, tz, size)];
+            }
+
+            DECODE(p, nx, ny, nz);
+            s1 = s2;
+            s2 = ENCODE(nx, lasty, nz, ZERO, flag);
+            y2 = lasty;
+            lasty = ty;
+
+            stack[id] = s2;
+
+            flag = ONE;
+        }
+    }
+
+    if (NOTSITE(p))
+        stack[TOID(tx, ty - 1, tz, size)] = ENCODE(ZERO, lasty, ZERO, ONE, flag);
+}
+
 __global__ void kernelColorAxis(ENCODED_INT_TYPE *input, ENCODED_INT_TYPE *output, int size)
 {
     __shared__ ENCODED_INT_TYPE block[BLOCKSIZE][BLOCKSIZE];
@@ -234,4 +308,80 @@ __global__ void kernelColorAxis(ENCODED_INT_TYPE *input, ENCODED_INT_TYPE *outpu
 }
 
 
+__global__ void kernelColorAxisWithSpacing(ENCODED_INT_TYPE *input, ENCODED_INT_TYPE *output, int size, double sx, double sy, double sz)
+{
+    __shared__ ENCODED_INT_TYPE block[BLOCKSIZE][BLOCKSIZE];
+
+    int col = threadIdx.x;
+    int tid = threadIdx.y;
+    int tx = blockIdx.x * blockDim.x + col;
+    int tz = blockIdx.y;
+
+    ENCODED_INT_TYPE x1, y1, z1, x2, y2, z2;
+    ENCODED_INT_TYPE last1 = ENCODE(ZERO,ZERO,ZERO,ONE,ZERO), last2 = ENCODE(ZERO,ZERO,ZERO,ONE,ZERO), lasty;
+    double dx, dy, dz, best, dist;
+
+    lasty = size - 1;
+
+    last2 = input[TOID(tx, lasty, tz, size)];
+    DECODE(last2, x2, y2, z2);
+
+    if (NOTSITE(last2)) {
+        lasty = y2;
+        if(HASNEXT(last2)) {
+            last2 = input[TOID(tx, lasty, tz, size)];
+            DECODE(last2, x2, y2, z2);
+        }
+    }
+
+    if (HASNEXT(last2)) {
+        last1 = input[TOID(tx, y2, tz, size)];
+        DECODE(last1, x1, y1, z1);
+    }
+
+    int y_start, y_end, n_step = size / blockDim.x;
+    for(int step = 0; step < n_step; ++step) {
+        y_start = size - step * blockDim.x - 1;
+        y_end = size - (step + 1) * blockDim.x;
+
+        for (int ty = y_start - tid; ty >= y_end; ty -= blockDim.y) {
+            dx = static_cast<double>(x2 - tx) * sx;
+            dy = static_cast<double>(lasty - ty) * sy;
+            dz = static_cast<double>(z2 - tz) * sz;
+            best = dx * dx + dy * dy + dz * dz;
+
+            while (HASNEXT(last2)) {
+                dx = static_cast<double>(x1 - tx) * sx;
+                dy = static_cast<double>(y2 - ty) * sy;
+                dz = static_cast<double>(z1 - tz) * sz;
+                dist = dx * dx + dy * dy + dz * dz;
+
+                if(dist > best) break;
+
+                best = dist; lasty = y2; last2 = last1;
+                DECODE(last2, x2, y2, z2);
+
+                if (HASNEXT(last2)) {
+                    last1 = input[TOID(tx, y2, tz, size)];
+                    DECODE(last1, x1, y1, z1);
+                }
+            }
+
+            block[threadIdx.x][ty - y_end] = ENCODE(lasty, x2, z2, NOTSITE(last2), ZERO);
+        }
+
+        __syncthreads();
+
+        if(!threadIdx.y) {
+            int id = TOID(y_end + threadIdx.x, blockIdx.x * blockDim.x, tz, size);
+            for(int i = 0; i < blockDim.x; i++, id+=size) {
+                output[id] = block[i][threadIdx.x];
+            }
+        }
+
+        __syncthreads();
+    }
+}
+
+
 } // extern C


Test patch:
diff --git a/python/cucim/src/cucim/core/operations/morphology/tests/test_distance_transform.py b/python/cucim/src/cucim/core/operations/morphology/tests/test_distance_transform.py
index 73aa69771..0515c9a71 100644
--- a/python/cucim/src/cucim/core/operations/morphology/tests/test_distance_transform.py
+++ b/python/cucim/src/cucim/core/operations/morphology/tests/test_distance_transform.py
@@ -30,11 +30,14 @@ def assert_percentile_equal(arr1, arr2, pct=95):
     [
         ((256, 128), None),
         ((384, 256), (1.5, 1.5)),
+        ((384, 256), (3, 2)),  # integer-valued anisotropic
+        ((384, 256), (2.25, .85)),
         ((14, 32, 50), None),
-        ((50, 32, 24), (2, 2, 2)),
+        ((50, 32, 24), (2., 2., 2.)),
+        ((50, 32, 24), (3, 1, 2)),  # integer-valued anisotropic
     ],
 )
-@pytest.mark.parametrize('density', [5, 50, 95])
+@pytest.mark.parametrize('density', ['single_point', 5, 50, 95])
 @pytest.mark.parametrize('block_params', [None, (1, 1, 1)])
 def test_distance_transform_edt(
     shape, sampling, return_distances, return_indices, density, block_params
@@ -50,19 +53,28 @@ def test_distance_transform_edt(
     )
     kwargs_cucim = copy(kwargs_scipy)
     kwargs_cucim['block_params'] = block_params
-    img = binary_image(shape, pct_true=density)
+    if density == 'single_point':
+        img = cp.ones(shape, dtype=bool)
+        img[tuple(s // 2 for s in shape)] = 0
+    else:
+        img = binary_image(shape, pct_true=density)
+
     out = distance_transform_edt(img, **kwargs_cucim)
     expected = ndi_cpu.distance_transform_edt(cp.asnumpy(img), **kwargs_scipy)
+    if sampling is None:
+        target_pct = 95
+    else:
+        target_pct = 90
     if return_indices and return_distances:
         assert len(out) == 2
-        cp.testing.assert_allclose(out[0], expected[0])
+        cp.testing.assert_allclose(out[0], expected[0], rtol=1e-6)
         # May differ at a small % of coordinates where multiple points were
         # equidistant.
-        assert_percentile_equal(out[1], expected[1], pct=95)
+        assert_percentile_equal(out[1], expected[1], pct=target_pct)
     elif return_distances:
-        cp.testing.assert_allclose(out, expected)
+        cp.testing.assert_allclose(out, expected, rtol=1e-6)
     elif return_indices:
-        assert_percentile_equal(out, expected, pct=95)
+        assert_percentile_equal(out, expected, pct=target_pct)
 
 
 @pytest.mark.parametrize(
@@ -119,36 +131,6 @@ def test_distance_transform_edt_block_params_invalid(block_params):
         distance_transform_edt(img, block_params=block_params)
 
 
-@pytest.mark.parametrize('return_indices', [False, True])
-@pytest.mark.parametrize('return_distances', [False, True])
-@pytest.mark.parametrize(
-    'shape, sampling',
-    [
-        ((384, 256), (1, 3)),
-        ((50, 32, 24), (1, 2, 4)),
-    ]
-)
-@pytest.mark.parametrize('density', [5, 50, 95])
-def test_distance_transform_edt_nonuniform_sampling(
-    shape, sampling, return_distances, return_indices, density
-):
-
-    if not (return_indices or return_distances):
-        return
-
-    kwargs_scipy = dict(
-        sampling=sampling,
-        return_distances=return_distances,
-        return_indices=return_indices,
-    )
-    kwargs_cucim = copy(kwargs_scipy)
-    img = binary_image(shape, pct_true=density)
-    if sampling is not None and len(np.unique(sampling)) != 1:
-        with pytest.raises(NotImplementedError):
-            distance_transform_edt(img, **kwargs_cucim)
-        return
-
-
 @pytest.mark.parametrize('value', [0, 1, 3])
 @pytest.mark.parametrize('ndim', [2, 3])
 def test_distance_transform_edt_uniform_valued(value, ndim):
@@ -172,6 +154,101 @@ def test_distance_transform_edt_2d_aniso(sx, sy):
     cp.testing.assert_allclose(out, expected)
 
 
+@pytest.mark.parametrize('ndim', [2, 3])
+@pytest.mark.parametrize('sampling', [None, 'iso', 'aniso'])
+def test_distance_transform_inplace_distance(ndim, sampling):
+    img = binary_image((32, ) * ndim, pct_true=80)
+    distances = cp.empty(img.shape, dtype=cp.float32)
+    if sampling == 'iso':
+        sampling = (1.5,) * ndim
+    elif sampling == 'aniso':
+        sampling = tuple(range(1, ndim + 1))
+    distance_transform_edt(img, sampling=sampling, distances=distances)
+    expected = ndi_cpu.distance_transform_edt(
+        cp.asnumpy(img), sampling=sampling
+    )
+    cp.testing.assert_allclose(distances, expected)
+
+
+@pytest.mark.parametrize('ndim', [2, 3])
+def test_distance_transform_inplace_distance_errors(ndim):
+    img = binary_image((32, ) * ndim, pct_true=80)
+
+    # for binary input, distances output is float32. Other dtypes raise
+    with pytest.raises(RuntimeError):
+        distances = cp.empty(img.shape, dtype=cp.float64)
+        distance_transform_edt(img, distances=distances)
+    with pytest.raises(RuntimeError):
+        distances = cp.empty(img.shape, dtype=cp.int32)
+        distance_transform_edt(img, distances=distances)
+
+    # wrong shape
+    with pytest.raises(RuntimeError):
+        distances = cp.empty(img.shape + (2,), dtype=cp.float32)
+        distance_transform_edt(img, distances=distances)
+
+    # can't provide indices array when return_indices is False
+    with pytest.raises(RuntimeError):
+        distances = cp.empty(img.shape, dtype=cp.float32)
+        distance_transform_edt(img, distances=distances,
+                               return_distances=False, return_indices=True)
+
+
+@pytest.mark.parametrize('ndim', [2, 3])
+@pytest.mark.parametrize('sampling', [None, 'iso', 'aniso'])
+@pytest.mark.parametrize('dtype', [cp.int16, cp.uint16, cp.uint32, cp.int32,
+                                   cp.uint64, cp.int64])
+@pytest.mark.parametrize('return_distances', [False, True])
+def test_distance_transform_inplace_indices(
+    ndim, sampling, dtype, return_distances
+):
+    img = binary_image((32, ) * ndim, pct_true=80)
+    if ndim == 3 and dtype in [cp.int16, cp.uint16]:
+        pytest.skip(reason="3d case requires at least 32-bit integer output")
+    if sampling == 'iso':
+        sampling = (1.5,) * ndim
+    elif sampling == 'aniso':
+        sampling = tuple(range(1, ndim + 1))
+    common_kwargs = dict(
+        sampling=sampling, return_distances=return_distances,
+        return_indices=True
+    )
+    # verify that in-place and out-of-place results agree
+    indices = cp.empty((ndim,) + img.shape, dtype=dtype)
+    distance_transform_edt(img, indices=indices, **common_kwargs)
+    expected = distance_transform_edt(img, **common_kwargs)
+    if return_distances:
+        cp.testing.assert_array_equal(indices, expected[1])
+    else:
+        cp.testing.assert_array_equal(indices, expected)
+
+
+@pytest.mark.parametrize('ndim', [2, 3])
+def test_distance_transform_inplace_indices_errors(ndim):
+    img = binary_image((32, ) * ndim, pct_true=80)
+    common_kwargs = dict(return_distances=False, return_indices=True)
+
+    # int8 has itemsize too small
+    with pytest.raises(RuntimeError):
+        indices = cp.empty((ndim,) + img.shape, dtype=cp.int8)
+        distance_transform_edt(img, indices=indices, **common_kwargs)
+
+    # float not allowed
+    with pytest.raises(RuntimeError):
+        indices = cp.empty((ndim,) + img.shape, dtype=cp.float64)
+        distance_transform_edt(img, indices=indices, **common_kwargs)
+
+    # wrong shape
+    with pytest.raises(RuntimeError):
+        indices = cp.empty((ndim,), dtype=cp.float32)
+        distance_transform_edt(img, indices=indices, **common_kwargs)
+
+    # can't provide indices array when return_indices is False
+    with pytest.raises(RuntimeError):
+        indices = cp.empty((ndim,) + img.shape, dtype=cp.int32)
+        distance_transform_edt(img, indices=indices, return_indices=False)
+
+
 @pytest.mark.parametrize('sx', list(range(4)))
 @pytest.mark.parametrize('sy', list(range(4)))
 @pytest.mark.parametrize('sz', list(range(4)))
@@ -180,6 +257,7 @@ def test_distance_transform_edt_3d_aniso(sx, sy, sz):
     shape = (16 + sz, 32 + sy, 48 + sx)
     img = binary_image(shape, pct_true=80)
     out = distance_transform_edt(img)
+    print(f"{out.dtype=}")
     expected = ndi_cpu.distance_transform_edt(cp.asnumpy(img))
     cp.testing.assert_allclose(out, expected)