Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

WIP Add in-place variant of transfer_to_logical_device #220

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

sogartar
Copy link
Contributor

The purpose of this is to insert IREE ABI affinity attributes for function arguments.
iree.flow.transfer is not sufficient for cases of in-place updates of function arguments.

An alternative approach is to set the attribute of a Torch FX graph function argument. Then during/after conversion to MLIR to set the corresponding attribute on the MLIR function.

@sogartar
Copy link
Contributor Author

The output IR in the test will look like this

  func.func @main(%arg0: !torch.tensor<[9,8],f32> {iree.abi.affinity = #hal.device.promise<@__device_1>}) -> !torch.vtensor<[9,8],f32> attributes {torch.assume_strict_symbolic_shapes} {
    %0 = torch.copy.to_vtensor %arg0 : !torch.vtensor<[9,8],f32>
    %1 = torch_c.to_builtin_tensor %0 : !torch.vtensor<[9,8],f32> -> tensor<9x8xf32>
    %2 = torch_c.from_builtin_tensor %1 : tensor<9x8xf32> -> !torch.vtensor<[9,8],f32>
    %int1 = torch.constant.int 1
    %int1_0 = torch.constant.int 1
    %3 = torch.aten.add.Scalar %2, %int1, %int1_0 : !torch.vtensor<[9,8],f32>, !torch.int, !torch.int -> !torch.vtensor<[9,8],f32>
    torch.overwrite.tensor.contents %3 overwrites %arg0 : !torch.vtensor<[9,8],f32>, !torch.tensor<[9,8],f32>
    return %3 : !torch.vtensor<[9,8],f32>
  }

@rsuderman
Copy link
Contributor

We should try to include execution based testing. Having a simple script that performs the transfer, does some simple math, then transfers back running via iree would tell us more about functionality. E.g. IR is good, execution is better.

The purpose of this is to insert IREE ABI affinity attributes for
function arguments.
iree.flow.transfer is not sufficient for cases of in-place updates of
function arguments.

An alternative approach is to set the attribute of a Torch FX graph
function argument. Then during/after conversion to MLIR to set the
corresponding attribute on the MLIR function.

Signed-off-by: Boian Petkantchin <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants