Skip to content

Commit eaaa9c2

Browse files
ahagro01ahagro01
authored
[TorchToTosa] Support flattened index_put updates in TOSA lowering (#4561)
Fixes TOSA lowering for index_put/scatter updates when the update tensor is already flattened. Previously, fillK was computed from only the first ND dimensions of fillValues, which assumes updates still follow the indexed tensor shape. Flattened updates break that assumption and can produce invalid reshape/scatter shapes. The fix derives fillK from the total number of update elements with a divisibility check. This supports both shaped and flattened updates while preserving the existing {N, fillK, C} scatter form. A regression test was added for flattened torch.aten.index_put.hacked_twin updates.
1 parent 7db2a11 commit eaaa9c2

3 files changed

Lines changed: 76 additions & 15 deletions

File tree

‎lib/Conversion/TorchToTosa/TosaLegalizeCommon.cpp‎

Lines changed: 35 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -448,7 +448,7 @@ std::optional<Value> convertScatterNdOp(PatternRewriter &rewriter,
448448
auto indicesType = dyn_cast<RankedTensorType>(indicesValue.getType());
449449
auto fillValuesType = dyn_cast<RankedTensorType>(fillValues.getType());
450450

451-
if (!resultType || !paramsType || !indicesType)
451+
if (!resultType || !paramsType || !indicesType || !fillValuesType)
452452
return std::nullopt;
453453

454454
// N: number of batches
@@ -522,7 +522,7 @@ std::optional<Value> convertScatterNdOp(PatternRewriter &rewriter,
522522
// !torch.vtensor<[1,4],si64>
523523
// Detail algorithm visualization:
524524

525-
int N = 1, W = 1, K = 1, fillK = 1, C = 1, ND = 1;
525+
int N = 1, W = 1, K = 1, C = 1, ND = 1;
526526

527527
int paramsRank = paramsType.getShape().size(); // 2
528528
int indicesRank = indicesType.getShape().size(); // 2
@@ -553,11 +553,18 @@ std::optional<Value> convertScatterNdOp(PatternRewriter &rewriter,
553553
// input(chould be scatter) C = product(params.shape[ND:] ND = 2, paramsRank,
554554
// C = 1
555555
for (int i = ND; i < paramsRank; i++) {
556-
C *= paramsType.getShape()[i];
556+
int64_t dim = paramsType.getShape()[i];
557+
if (dim < 0) {
558+
(void)rewriter.notifyMatchFailure(
559+
op, "scatter channel dimensions must be static");
560+
return std::nullopt;
561+
}
562+
C *= dim;
557563
}
558564

559565
// int N = 1, W = 3, K = 4, fillk = 3, C = 1, ND = 2;
560566
SmallVector<int64_t, 3> tosaInputValuesShape({N, K, C}); // {1,4,1}
567+
SmallVector<int64_t, 3> tosaFillValuesShape({N, W, C}); // {1,3,1}
561568
SmallVector<int64_t, 2> tosaIndicesShape({N, W}); // {1,3}
562569
SmallVector<int64_t, 2> indicesMatrixShape({W, ND}); // {3,2}
563570
SmallVector<int64_t, 2> indicesMatrixReducesumShape({W, 1}); // {3,1}
@@ -569,18 +576,26 @@ std::optional<Value> convertScatterNdOp(PatternRewriter &rewriter,
569576
// 2. !torch.vtensor<[],si64>
570577
// reshape(1) tile(3) reshape(1,3) reshape(1,3,1)
571578
// [] -> [0] -> [0,0,0] -> [[0,0,0]] -> [[[0], [0], [0]]]
572-
// reshape to [1] and then tile to same number of indicesValue.shape[0],
573-
// [1,1,1]
574-
if (fillValuesType.getRank() == 0) {
579+
// reshape to [1] and then tile to W * C update values.
580+
if (fillValuesType.getRank() == 0 && C == 0) {
581+
auto emptyFillValues = getZerosLikeTensor(
582+
rewriter, op,
583+
GetTypeFromTensorShape(tosaFillValuesShape,
584+
fillValuesType.getElementType()));
585+
if (!emptyFillValues)
586+
return std::nullopt;
587+
fillValues = *emptyFillValues;
588+
fillValuesType = dyn_cast<RankedTensorType>(fillValues.getType());
589+
} else if (fillValuesType.getRank() == 0) {
575590
// [] -> [0]
576-
SmallVector<int64_t, 1> oneShape({1}); // {3,1}
591+
SmallVector<int64_t, 1> oneShape({1}); // {1}
577592
auto tosaFillValuesOneReshapeOp = tosa::CreateOpAndInfer<tosa::ReshapeOp>(
578593
rewriter, op->getLoc(),
579594
GetTypeFromTensorShape(oneShape, fillValuesType.getElementType()),
580595
fillValues, tosa::getTosaConstShape(rewriter, op->getLoc(), oneShape));
581596

582597
// [0] -> [0,0,0]
583-
SmallVector<int64_t, 1> tileShape({W}); // {3}
598+
SmallVector<int64_t, 1> tileShape({W * C}); // {3}
584599
auto tileOpMultiples =
585600
tosa::getTosaConstShape(rewriter, op->getLoc(), tileShape);
586601
auto tosaFillValuesTileOp = tosa::CreateOpAndInfer<tosa::TileOp>(
@@ -589,7 +604,7 @@ std::optional<Value> convertScatterNdOp(PatternRewriter &rewriter,
589604
tosaFillValuesOneReshapeOp.getResult(), tileOpMultiples);
590605

591606
// [0,0,0] -> [[0,0,0]]
592-
SmallVector<int64_t, 2> newTosaFillValuesShape({N, W}); // {1,3}
607+
SmallVector<int64_t, 2> newTosaFillValuesShape({N, W * C}); // {1,3}
593608
auto newTosaFillValuesReshapeOp = tosa::CreateOpAndInfer<tosa::ReshapeOp>(
594609
rewriter, op->getLoc(),
595610
GetTypeFromTensorShape(newTosaFillValuesShape,
@@ -601,12 +616,18 @@ std::optional<Value> convertScatterNdOp(PatternRewriter &rewriter,
601616
fillValuesType = dyn_cast<RankedTensorType>(fillValues.getType());
602617
}
603618

604-
// fillK: range of each index, total number of fillInput(could be scatter)
605-
// after flattened k = 1*1*3 = 3
606-
for (int i = 0; i < ND; i++) {
607-
fillK *= fillValuesType.getShape()[i];
619+
// TOSA scatter update values are shaped [N, W, C], where W is the
620+
// number of flattened scatter indices.
621+
int64_t fillNumElements = 1;
622+
for (int64_t dim : fillValuesType.getShape()) {
623+
fillNumElements *= dim;
624+
}
625+
if (fillNumElements != W * C) {
626+
(void)rewriter.notifyMatchFailure(
627+
op, "scatter update element count must match flattened indices and "
628+
"channels");
629+
return std::nullopt;
608630
}
609-
SmallVector<int64_t, 3> tosaFillValuesShape({N, fillK, C}); // {1,3,1}
610631

611632
// Reshape/Flatten fillValues to 3d tensor
612633
// [[0,0,0]] -> [[[0], [0], [0]]]

‎projects/pt1/e2e_testing/xfail_sets.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3825,7 +3825,6 @@
38253825
"IndexPutImpl1DFloatAccumulateModule_basic",
38263826
"IndexPutImpl1DIntAccumulateModule_basic",
38273827
"IndexPutImpl2DFloatAccumulateModule_basic",
3828-
"IndexPutImpl2DImplicitModule_basic",
38293828
"IndexPutImpl2DIndexModule_basic",
38303829
"IndexPutImpl2DNoneIndexStaticModule_basic",
38313830
"IndexPutImpl3DFloatAccumulateModule_basic",

‎test/Conversion/TorchToTosa/basic.mlir‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2658,6 +2658,27 @@ func.func @torch.aten.scatter.src$basic(%arg0: !torch.vtensor<[10,8,6],f32>, %ar
26582658

26592659
// -----
26602660

2661+
// CHECK-LABEL: func.func @torch.aten.index_put.hacked_twin$c_gt_1_update(
2662+
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[4,4],f32>,
2663+
// CHECK-SAME: %[[VAL_1:.*]]: !torch.vtensor<[2],si64>,
2664+
// CHECK-SAME: %[[VAL_2:.*]]: !torch.vtensor<[2,4],f32>) -> !torch.vtensor<[4,4],f32> {
2665+
// CHECK: %[[UPDATES_BUILTIN:.*]] = torch_c.to_builtin_tensor %[[VAL_2]] : !torch.vtensor<[2,4],f32> -> tensor<2x4xf32>
2666+
// CHECK: %[[INPUT_BUILTIN:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[4,4],f32> -> tensor<4x4xf32>
2667+
// CHECK: %[[INDEX_BUILTIN:.*]] = torch_c.to_builtin_tensor %[[VAL_1]] : !torch.vtensor<[2],si64> -> tensor<2xi64>
2668+
// CHECK: %[[UPDATES_SHAPE:.*]] = tosa.const_shape {values = dense<[1, 2, 4]> : tensor<3xindex>} : () -> !tosa.shape<3>
2669+
// CHECK: %[[UPDATES:.*]] = tosa.reshape %[[UPDATES_BUILTIN]], %[[UPDATES_SHAPE]] : (tensor<2x4xf32>, !tosa.shape<3>) -> tensor<1x2x4xf32>
2670+
// CHECK: %[[INPUT_SHAPE:.*]] = tosa.const_shape {values = dense<[1, 4, 4]> : tensor<3xindex>} : () -> !tosa.shape<3>
2671+
// CHECK: %[[INPUT:.*]] = tosa.reshape %[[INPUT_BUILTIN]], %[[INPUT_SHAPE]] : (tensor<4x4xf32>, !tosa.shape<3>) -> tensor<1x4x4xf32>
2672+
// CHECK: %{{.*}} = tosa.scatter %[[INPUT]], %{{.*}}, %[[UPDATES]] : (tensor<1x4x4xf32>, tensor<1x2xi32>, tensor<1x2x4xf32>) -> tensor<1x4x4xf32>
2673+
func.func @torch.aten.index_put.hacked_twin$c_gt_1_update(%arg0: !torch.vtensor<[4,4],f32>, %arg1: !torch.vtensor<[2],si64>, %arg2: !torch.vtensor<[2,4],f32>) -> !torch.vtensor<[4,4],f32> {
2674+
%false = torch.constant.bool false
2675+
%0 = torch.prim.ListConstruct %arg1 : (!torch.vtensor<[2],si64>) -> !torch.list<vtensor>
2676+
%1 = torch.aten.index_put.hacked_twin %arg0, %0, %arg2, %false : !torch.vtensor<[4,4],f32>, !torch.list<vtensor>, !torch.vtensor<[2,4],f32>, !torch.bool -> !torch.vtensor<[4,4],f32>
2677+
return %1 : !torch.vtensor<[4,4],f32>
2678+
}
2679+
2680+
// -----
2681+
26612682
// CHECK-LABEL: func.func @torch.aten.slice_scatter$basic(
26622683
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[6,8],f32>,
26632684
// CHECK-SAME: %[[VAL_1:.*]]: !torch.vtensor<[6,1],f32>) -> !torch.vtensor<[6,8],f32> {
@@ -2746,6 +2767,26 @@ func.func @torch.aten.diag_embed$basic(%arg0: !torch.vtensor<[2,3,4],f32>) -> !t
27462767

27472768
// -----
27482769

2770+
// CHECK-LABEL: func.func @torch.aten.index_put_hacked_twin_flattened_updates(
2771+
// CHECK: %[[SCATTER:.*]] = tosa.scatter
2772+
// CHECK-SAME: (tensor<1x6x1xf32>, tensor<1x6xi32>, tensor<1x6x1xf32>) -> tensor<1x6x1xf32>
2773+
// CHECK: %[[RESHAPE:.*]] = tosa.reshape %[[SCATTER]]
2774+
// CHECK-SAME: (tensor<1x6x1xf32>, !tosa.shape<3>) -> tensor<1x2x3xf32>
2775+
// CHECK: torch_c.from_builtin_tensor %[[RESHAPE]] : tensor<1x2x3xf32> -> !torch.vtensor<[1,2,3],f32>
2776+
func.func @torch.aten.index_put_hacked_twin_flattened_updates(
2777+
%arg0: !torch.vtensor<[1,2,3],f32>,
2778+
%arg1: !torch.vtensor<[6],si64>,
2779+
%arg2: !torch.vtensor<[6],si64>,
2780+
%arg3: !torch.vtensor<[6],si64>,
2781+
%arg4: !torch.vtensor<[6],f32>) -> !torch.vtensor<[1,2,3],f32> {
2782+
%indices = torch.prim.ListConstruct %arg1, %arg2, %arg3 : (!torch.vtensor<[6],si64>, !torch.vtensor<[6],si64>, !torch.vtensor<[6],si64>) -> !torch.list<vtensor>
2783+
%false = torch.constant.bool false
2784+
%0 = torch.aten.index_put.hacked_twin %arg0, %indices, %arg4, %false : !torch.vtensor<[1,2,3],f32>, !torch.list<vtensor>, !torch.vtensor<[6],f32>, !torch.bool -> !torch.vtensor<[1,2,3],f32>
2785+
return %0 : !torch.vtensor<[1,2,3],f32>
2786+
}
2787+
2788+
// -----
2789+
27492790
// CHECK-LABEL: func.func @torch.aten.index.Tensor_hacked_twin(
27502791
// CHECK-SAME: %[[VAL_0:.*]]: !torch.vtensor<[2,4,2],si64>,
27512792
// CHECK-SAME: %[[VAL_1:.*]]: !torch.vtensor<[],si64>) -> !torch.vtensor<[4,2],si64> {

0 commit comments

Comments
 (0)