Using torch.overwrite.tensor.contents to overwrite input argument fails at runtime about iree HOT 9 CLOSED

Probably something that's going to need to be identified/fixed in the frontend - we can't really support memmove in general and need to ensure we aren't generating programs that require it. Not only are there no memmove DMA primitives in hardware (there's no cuMemmove, or vkCmdMoveBuffer, etc) but ensuring all dispatches that we generate have memmove semantics when operating in place is not possible.

from iree.

Chatting with Ben, the repro is not strictly legal (it is returning a value that is in-placed). While we could support that form, it is presently a limitation. Also, I think it is a testing artifact. We should work on this case:

module @module {
  func.func @test_index_copy(%arg0: !torch.tensor<[8192,16,8,128],f32>, %arg1: !torch.vtensor<[4],si64>, %arg2: !torch.vtensor<[4,16,8,128],f32>) {
    %0 = torch.copy.to_vtensor %arg0 : !torch.vtensor<[8192,16,8,128],f32>
    %1 = torch.prim.ListConstruct %arg1 : (!torch.vtensor<[4],si64>) -> !torch.list<optional<vtensor>>
    %false = torch.constant.bool false
    %2 = torch.aten.index_put %0, %1, %arg2, %false : !torch.vtensor<[8192,16,8,128],f32>, !torch.list<optional<vtensor>>, !torch.vtensor<[4,16,8,128],f32>, !torch.bool -> !torch.vtensor<[8192,16,8,128],f32>
    torch.overwrite.tensor.contents %2 overwrites %arg0 : !torch.vtensor<[8192,16,8,128],f32>, !torch.tensor<[8192,16,8,128],f32>
    return
  }
}

from iree.

I don't think this is going to work - we have memcpy semantics and such an operation requires memmove. You must ensure when using in-place operations on in-place calls that you aren't trying to do a memmove.

from iree.

Oh, I might be hitting the same issue on nod-ai/sharktank#22 (comment) (among other things)

from iree.

I only see one of these ops in the program I'm looking at. Trying to find where it's coming from... maybe this: https://github.com/llvm/torch-mlir/blob/ec6d7aa5d28f110aa5b893e16e502e6198988801/python/torch_mlir/extras/fx_importer.py#L1111-L1118 ?

from iree.

Probably - something @stellaraccident may have a pointer to/thoughts about - I'm not sure what the right solution is at that level. We can't do much in IREE as the only safe behavior is to silently insert copies for any externally-provided buffer and that defeats the purpose of the in-place ops. I'm not sure if analysis at the torch level could insert the copies, warn/error if they're required, or what :/

(this is one of the reasons we default to not doing in-place operations - they've got footguns :)

from iree.

I think my eyes aren't calibrated right to see where this is becoming a move-like thing. Probably need some help figuring out how to structure it.

from iree.

I don't know torch stuff (vtensor? what?), but maybe it's whatever to_vtensor is?
index_put ties operand 0 to its result so %0 is %2, assuming to_vtensor is not a forced copy %0 is %arg0, and the final overwrite aliases %2 to %arg0. so %arg0 == %0 == %2.

If to_vtensor must remain a clone and doesn't end up as a flow.tensor.clone then the above will happen. Even if it does become a clone maybe we're dropping it later and that's something we could fix locally but will be trickier in more complex programs to prove. A print-after-all would be useful.

from iree.

notes to tomorrow me:

• remove into from hal.tensor.export
• add a hal.tensor.alias (or whatever) op: %tied = hal.tensor.alias %operand, %storage : tensor<..>, !hal.buffer_view -> %operand
• change torch stuff to use that, drop optimization_barrier and hal.tensor.export
• ensure emplace allocations uses the alias
• error checking if two logical tensors (through tied ops) have an alias defined

from iree.