Make it straightforward to support other dimensionalities.

Bug: b/129523279
Change-Id: I29f7a55da26d98141ebea74678776c7d4f16c7c6
Reviewed-on: https://swiftshader-review.googlesource.com/c/SwiftShader/+/30152
Presubmit-Ready: Nicolas Capens <nicolascapens@google.com>
Tested-by: Nicolas Capens <nicolascapens@google.com>
Reviewed-by: Chris Forbes <chrisforbes@google.com>
Kokoro-Presubmit: kokoro <noreply+kokoro@google.com>