Enable support for LTC Input/Output Mapping #764

henrytwo · 2022-04-18T20:36:05Z

Complements this upstream PyTorch PR: pytorch/pytorch#75828

This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration)

MLIR: 
func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) {
  ...
  return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ...
}

Input/Output Alias Mapping: 
Output: 0 -> Input: 0
Output: 1 -> Input: 1
Output: 2 -> Input: 6
Output: 3 -> Input: 7

The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

cc: @antoniojkim @ke1337

With this op enabled, tensors were being copied, which resulted in incorrect aliasing.

As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext

henrytwo · 2022-04-22T19:07:31Z

@silvasean Keeping a map of aliased output -> input tensors is a feature that we're planning to use internally at Cerebras -- is this something you'd like us to keep in the generic TorchMlirLoweringContext too? Not sure if anyone else would find this useful

silvasean · 2022-04-25T13:52:20Z

@silvasean Keeping a map of aliased output -> input tensors is a feature that we're planning to use internally at Cerebras -- is this something you'd like us to keep in the generic TorchMlirLoweringContext too? Not sure if anyone else would find this useful

Isn't that something that upstream already provides or could be moved there? It definitely seems like a generally useful feature and not specific to MLIR or Cerebras in any way.

henrytwo · 2022-04-25T13:57:28Z

Isn't that something that upstream already provides or could be moved there?

Upstream we're just calling a method in the LoweringContext interface to save the alias, so it's up to the lowering context to figure out what to do with it. Other vendors may want to handle this differently than us (e.g. XLA), so that's why we're implementing it on the torch-mlir side.

silvasean · 2022-04-25T14:00:09Z

Oh got it. Let's keep the functionality in there, as it seems generally useful.

python/torch_mlir/csrc/base_lazy_backend/mlir_lowering_context.h

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

@graph

* Save InputOutputAliases to TorchMlirComputation * Implement GetResultShape for TorchMlirLoweringContext * Use optional return type for GetResultShape * Remove support for aten::detach With this op enabled, tensors were being copied, which resulted in incorrect aliasing. * Add newline before printing I/O alias mapping * Changed printout to use "Input param" as label instead of "Input" * Remote shape inference function for aten::detach * Moved implementation of SetUpAlias to MlirLoweringContext As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext * Use updated PyTorch API * Remove GetResultShape Complements this upstream PyTorch PR: pytorch/pytorch#75828 This PR adds support for mapping input and output tensors which alias each other. (e.g. maps input weight tensor in parameter to the same tensor in output after a training iteration) MLIR: func @graph(%arg0: !torch.vtensor<[1,5],f32>, %arg1: !torch.vtensor<[1],si64>, ..., %arg6: !torch.vtensor<[10,5],f32>, %arg7: !torch.vtensor<[10],f32>, ...) { ... return %arg0, %arg1, %17, %23, ... : !torch.vtensor<[1,5],f32>, !torch.vtensor<[1],si64>, !torch.vtensor<[10,5],f32>, !torch.vtensor<[10],f32>, ... } Input/Output Alias Mapping: Output: 0 -> Input: 0 Output: 1 -> Input: 1 Output: 2 -> Input: 6 Output: 3 -> Input: 7 The aten::detach op has also been disabled in this PR to fix the issue of tensors not aliasing properly due to copying.

* Test conv with dynamic dimensions Signed-off-by: Tung D. Le <[email protected]> * Borrow insertAllocandDealloc from the pooling ops Signed-off-by: Tung D. Le <[email protected]> * Test ResNet50 with unknown dimensions Signed-off-by: Tung D. Le <[email protected]> * Use IndexExpr for conv Signed-off-by: Tung D. Le <[email protected]> * Clean up Signed-off-by: Tung D. Le <[email protected]> * Clean up Signed-off-by: Tung D. Le <[email protected]> * Use IndexExpr for pooling ops Signed-off-by: Tung D. Le <[email protected]> * clang-format Signed-off-by: Tung D. Le <[email protected]>

henrytwo added 6 commits April 14, 2022 14:40

Save InputOutputAliases to TorchMlirComputation

e1be96b

Implement GetResultShape for TorchMlirLoweringContext

1aac26c

Use optional return type for GetResultShape

ff558ab

Remove support for aten::detach

c036f83

With this op enabled, tensors were being copied, which resulted in incorrect aliasing.

Add newline before printing I/O alias mapping

ed7fc53

Changed printout to use "Input param" as label instead of "Input"

95dff18

henrytwo requested a review from silvasean April 18, 2022 20:36

henrytwo self-assigned this Apr 18, 2022

Remote shape inference function for aten::detach

87a1455

silvasean approved these changes Apr 18, 2022

View reviewed changes

Moved implementation of SetUpAlias to MlirLoweringContext

52a0b23

As part of this change, TorchMlirComputation has been moved to the end of mlir_lowering_context.h so that it can access some new structs in TorchMlirLoweringContext

henrytwo force-pushed the henrytu/input_output_mapping branch 2 times, most recently from ac6fca3 to 52a0b23 Compare April 25, 2022 13:44

Use updated PyTorch API

aeddbde

antoniojkim reviewed Apr 25, 2022

View reviewed changes

python/torch_mlir/csrc/base_lazy_backend/mlir_lowering_context.h Outdated Show resolved Hide resolved

Remove GetResultShape

33b489b

henrytwo merged commit 199831d into llvm:torch_mlir_ltc_backend Apr 27, 2022

henrytwo deleted the henrytu/input_output_mapping branch April 27, 2022 17:48

tanyokwok mentioned this pull request Sep 21, 2022

features/bladedisc rebase 20220830 pai-disc/torch-mlir#20

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enable support for LTC Input/Output Mapping #764

Enable support for LTC Input/Output Mapping #764

henrytwo commented Apr 18, 2022

henrytwo commented Apr 22, 2022

silvasean commented Apr 25, 2022

henrytwo commented Apr 25, 2022

silvasean commented Apr 25, 2022

Enable support for LTC Input/Output Mapping #764

Enable support for LTC Input/Output Mapping #764

Conversation

henrytwo commented Apr 18, 2022

henrytwo commented Apr 22, 2022

silvasean commented Apr 25, 2022

henrytwo commented Apr 25, 2022

silvasean commented Apr 25, 2022