If you could translate your computations to MLIR then you have the ability to lift your instructions to various dialects. For ex, you can lift to the 'metal' d

If you could translate your computations to MLIR then you have the ability to lift your instructions to various dialects. For ex, you can lift to the 'metal' dialect or 'nvgpu' dialect. There's also an 'affine' dialect, which implements a small set of polyhedral transformations and works on memrefs like most compilers do. However, it seems that the 'affine' dialect is not as actively used these days. Most optimizations are now performed on tensors, but depending on your use case, it may still make sense.