llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Create scatter op with unique indicies (#3853) 

For the op `index_put_`, if accumulate == false, the behavior is
undefined if the indicies aren't unique
(https://pytorch.org/docs/stable/generated/torch.Tensor.index_put_.html).
So, when converting `AtenIndexPutHackedTwinOp` to a TMTensor scatter op,
mark the indices as unique if when `accumulate == false`.

This should have no functional effect (unless users are relying on UB)
and assuming unique indices has the benefit of unlocking better
optimizations in further compiler stages.

Signed-off-by: Ian Wood <ianwood2024@u.northwestern.edu>
pull/3804/head
Ian Wood 2024-11-05 12:48:34 -08:00 committed by GitHub
parent b75d0e3f8b
commit e88faf08ff
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
lib/Conversion/TorchToTMTensor/TorchToTMTensor.cpp
View File

@ -932,9 +932,12 @@ public:
// 2.) `values` is mapped to `updates` in scatter op. // 2.) `values` is mapped to `updates` in scatter op.
// 3.) `input` is mapped to `original` in scatter op. // 3.) `input` is mapped to `original` in scatter op.
bool invalidInputTypeFound = false; bool invalidInputTypeFound = false;
// If accumulate == false, the behavior is undefined if the indicies aren't
// unique.
bool uniqueIndices = !accumulate;
Value scatterOp = createTMTensorScatterOp( Value scatterOp = createTMTensorScatterOp(
rewriter, loc, values, indices, input, indicesMap, rewriter, loc, values, indices, input, indicesMap,
/*uniqueIndices=*/false, /*uniqueIndices=*/uniqueIndices,
[&](OpBuilder &b, Location loc, Value valuesElement, [&](OpBuilder &b, Location loc, Value valuesElement,
Value inputElement) { Value inputElement) {
Value yieldValue = valuesElement; Value yieldValue = valuesElement;