I don’t have a strong opinion here, but I’d like to convey my thought behind this problem.
On one hand using arith / math tensor ops are handy in not having to generalize code (but you still need to check the type, shape and element type every time). On the other hand, without formal semantics for tensors (dynamic shapes, unknown rank, type cast/broadcast/reduction semantics etc.) you get “implementation defined behavior” and downstream implementations diverge from upstream and themselves.
It’s not because it works for some cases that it’s safe, as @mehdi_amini alludes to above. The main decision is between a safe arith on tensors and no tensor support, not between what we have today and no tensor support. So the result of this thread should not be to not do nothing.
We started adding arithmetic and maths operations to linalg to fix it in the other way, which I believed was the consensus until the Triton case. So I want to make sure we’re not missing anything.
IIUC, Triton uses pointers and offsets as tensor representations, so the operations themselves can be thought of as on the “memory that those pointers refer to”, which is already odd in the usual MLIR progression, but should not harm the tensor story. Now, if there is a tensor usage and that gets reduced to loop on pointers, then this is the clear case for linalg → loops+arith.
Also, the Triton transform example above actually matches on mlir::Operation, not mlir::arith::SelectOp, so you still wouldn’t need to generalize the function signatures (not that hard to do either as @ftynse suggests). The main problem there is that you have to pick the information from different APIs if they’re ops in different dialects, but that’s also not too convoluted (and you can have facade helpers to extract information).
Either way, you’ll have to support shape analysis, dynamic types, etc. whether you use tensors, memrefs or pointers and offsets, it’s just a matter of where the logic lies.
Whatever the result of this, we really should have formal semantics for shapes and casts, or downstream dialects will start to diverge more and more and it will be harder for us to fix this as time goes on.