perf(autograd): optimize grey_dilation with striding

[yaugenst-flex]

The previous implementation of grey_dilation was based on convolution, which was slow for both the forward and backward passes.

This PR replaces it with a high-performance implementation that uses NumPy's sliding_window_view to create sliding window views of the input array. I also wrote a custom VJP that uses the same striding technique to make the backward pass faster too.

I also simplified the implementation of grey_erosion so that grey_dilation is now the only function that does the heavy lifting.

Benchmarks show speedups of 10-100x depending on the array and kernel size.

This should make these ops much more usable in topopt @groberts-flex

Greptile Summary

Significant performance optimization of the grey_dilation morphological operation by replacing convolution-based implementation with NumPy's sliding_window_view for strided array operations.

• Replaced convolution-based implementation with strided array approach in tidy3d/plugins/autograd/functions.py, achieving 10-100x speedup
• Added custom VJP (vector-Jacobian product) for efficient backpropagation using the same striding technique
• Simplified grey_erosion by expressing it through duality with grey_dilation
• Updated morphology test cases to use first-order gradients and include kernel structure testing
• Added comprehensive benchmarks showing performance improvements scaling with array and kernel sizes

[greptile-apps]

LGTM

_{3 files reviewed, no comments
Edit PR Review Bot Settings | Greptile}

[github-actions]

Diff Coverage

Diff: origin/develop...HEAD, staged and unstaged changes

• tidy3d/plugins/autograd/functions.py (98.8%): Missing lines 67

Summary

• Total: 83 lines
• Missing: 1 line
• Coverage: 98%

tidy3d/plugins/autograd/functions.py

  63         The indices for padding along the axis.
  64     """
  65     total_pad = sum(pad_width)
  66     if n == 0:
! 67         return numpy_module.zeros(total_pad, dtype=int)
  68 
  69     idx = numpy_module.arange(-pad_width[0], n + pad_width[1])
  70 
  71     if mode == "constant":

---

[groberts-flex]

in the case of a 1D structuring element or size, does this mean the dilation/erosion gets applied in 2D and then just a single dimension is extracted at the end?

[yaugenst-flex]

The operation is still applied to the full 2D array, but with a 1D-like structuring element. You always need to define a 2D structuring element of shape like (1, size) or (size, 1). The dilation/erosion
operation slides this structuring element across the entire 2D array, effectively performing the operation along rows or columns respectively. I updated the docstring to clarify this, because we don't explicitly handle 1d structuring elements.

[groberts-flex]

possibly the other case to catch is that both the size and the structure are specified

[yaugenst-flex]

yeah good catch. i think this was discussed previously at some point and we went for silently having precedence of structuring element over sizes, but since it came up again i added a check instead to forbid this

[groberts-flex]

checking my understanding - this replaces the convolve call in the previous version which was being used to essentially create this same view but doing so through a bunch of unnecessary computation with an identity kernel?

[yaugenst-flex]

yes, exactly. the previous implementation would create the sliding windows (which we now do directly), multiply each window element-wise with the kernel, and then sum the results. you can see how this scales horrendously with image and kernel sizes 😄

[groberts-flex]

not fully understanding this part - is it possible for values to come out of the operation as the max_val?

[yaugenst-flex]

no, values cannot equal max_val in the output. i added a comment explaining this:

# Note: Values can never exceed maxval in the output since we add structure
# values (capped at maxval) to the input array values.

[groberts-flex]

thanks for this implementation, the speed up looks awesome especially for a function that will be in a lot of robust optimizations!

left some comments/questions, some just for my own understanding!