Apply the trait `NoMemoryEffect` to most `ReadOnly` ops

[zjgarvey]

In my understanding, the trait ReadOnly implies that the memory for input values does not get modified by the operation.

This is similar to NoMemoryEffect, with a few exceptions. Ops like prim.RaiseException and prim.Print are ReadOnly, but have effects like terminating the program and printing, respectively.

There may be other ReadOnly ops with side effects. As a hypothetical example, there might be a comparison op that checks if two tensors are the same type, prints some debug info, and also returns a boolean value. In the event that such an op is discovered, I've added a keyword argument to the emit_op function to explicitly label an operation as has_memory_effects=True to avoid adding the NoMemoryEffect trait to the tablegen definition.

My primary motivation for this change is to remove the need to have one-off RemoveUnused patterns for ops that we want to remove when dead, but adding this trait also has the very appealing effect of dramatically simplifying torch-IR for some models. Specifically for ONNX models, it is not uncommon to have patterns that consistently call "onnx.Shape" on the same exact tensor, only to extract a different element from the shape. That redundant IR doesn't get CSE'd until converting to a core mlir dialect like linalg/tensor/arith.

[zjgarvey]

I'm wondering if it is more appropriate to label ops more carefully with specific MemoryEffects? Perhaps NoMemoryEffect would imply that it also doesn't allocate memory for a new output tensor, or write to the new output tensor?

Please give some feedback if you have some opinions on what would be the best way forward on this.

[IanWood1]

I'm wondering if it is more appropriate to label ops more carefully with specific MemoryEffects? Perhaps NoMemoryEffect would imply that it also doesn't allocate memory for a new output tensor, or write to the new output tensor?

Please give some feedback if you have some opinions on what would be the best way forward on this.

I looked into this a bit ago and one thing I found is that there are several ops that can throw exceptions that aren't as obvious as prim.RaiseException (e.g. torch.aminmax). But I'm not sure if torch-mlir's goal is to respect these errors so it may not be a problem.

[kuhar]

I'm wondering if it is more appropriate to label ops more carefully with specific MemoryEffects? Perhaps NoMemoryEffect would imply that it also doesn't allocate memory for a new output tensor, or write to the new output tensor?

Tensors are not considered memory. See https://mlir.llvm.org/docs/Rationale/SideEffectsAndSpeculation/.

[zjgarvey]

I skimmed it and it looks fine to me. I'd probably put all of these standard properties in a new def and give them a name instead of repeating everywhere. This should make it easier to update if you realize you need to tweak it later on. For example, see how Pure is defined here: https://github.com/llvm/llvm-project/blob/132de3a71f581dcb008a124d52c83ccca8158d98/mlir/include/mlir/Interfaces/SideEffectInterfaces.td#L144C1-L146C60

Edit:

Actually, I don't think you need ReadOnly on most/any of these. If they only operate on tensors, there's nor memory accessed. Printing, however, would be considered a side effect.

Edit 2:

I don't know where ReadOnly comes from, I don't see it in MLIR. Is this a torch thing?

Yeah, ReadOnly is a torch-mlir specific trait. I believe ValueTensor is roughly equivalent to builtin mlir Tensor, but NonValueTensor is considered a mutable resource. At some point, we attempt to maximize value semantics through a pass, basically attempting to convert all NonValueTensor types to ValueTensorTypes by replacing ops which do not have value semantics with those that do, or in the case of isViewLikeOp, simply converting the NonValueTensor types to ValueTensor types. ReadOnly is only really used in torch-mlir to determine if list operands would get mutated by an operation in mlir::torch::Torch::potentiallyMutatesListOperands. My reasoning in this change is: if ReadOnly ops aren't mutating operands and don't have some additional behaviour like terminating the program or printing something, then it should be able to get dce'd or cse'd, even if they are acting on NonValueTensor types.

[ibsidorenko]

Hi, @zjgarvey ! Thank you for this PR. Adding of trait NoMemoryEffect - this is something I’ve been dreaming about for a long time. 👍 As I see this PR was approved, but still was not merged. Are there any blockers to merging it?

[zjgarvey]

Hi, @zjgarvey ! Thank you for this PR. Adding of trait NoMemoryEffect - this is something I’ve been dreaming about for a long time. 👍 As I see this PR was approved, but still was not merged. Are there any blockers to merging it?

@ibsidorenko Sorry for losing track of this comment. I'll clean this up and try to land this week.