Fix handling of rank-1 tensors in tosa.reduce_sum

[ramiro050]

Summary

The conversion of tosa.reduce_sum to linalg creates a
linalg.generic op that produces a tensor of rank `input_rank -
1`. This tensor is then expanded back into a tensor of rank
input_rank. In the case where the tensor being expanded is rank-0,
the reassociation map used must be empty. However, the current
implementation indexes and modifies the reassociation map independent
of the rank of the tensor being expanded, resulting in out-of-bounds
indexing when the tensor being expanded is rank-0. This commit adds a
guard to the reassociation map indexing.

[jpienaar]

Good catch

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

787

I'd suggest CHECK-DAG as these two are effecitively constants.

[ramiro050]

Change CHECK to CHECK-DAG for constants in test

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

static LogicalResult reduceMatchAndRewriteHelper(Operation *op, uint64_t axis,
   uint64_t expandInputRank =
       linalgOp.getResults()[0].getType().cast<ShapedType>().getRank();
   reassociationMap.resize(expandInputRank);
 
   for (uint64_t i = 0; i < expandInputRank; i++) {
     int32_t dimToPush = i > axis ? i + 1 : i;
     reassociationMap[i].push_back(rewriter.getAffineDimExpr(dimToPush));
   }
+
+  if (expandInputRank != 0) {
     int32_t expandedDim = axis < expandInputRank ? axis : expandInputRank - 1;
     reassociationMap[expandedDim].push_back(
         rewriter.getAffineDimExpr(expandedDim + 1));
+  }
 
   rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
       op, resultTy, linalgOp.getResults()[0], reassociationMap);
   return success();
 }
 
 static bool findIntermediateShape(ArrayRef<int64_t> lhsShape,
                                   ArrayRef<int64_t> rhsShape,

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

func.func @reduce_float_dyn(%arg0: tensor<?x5x4xf32>) -> () {
   // CHECK:   linalg.yield %[[RES]] : f32
   // CHECK: tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1, 2]] : tensor<?x4xf32> into tensor<?x1x4xf32>
   %0 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<?x5x4xf32>) -> tensor<?x1x4xf32>
   return
 }
 
 // -----
 
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> ()>
+
+// CHECK-LABEL: @reduce_float_dyn_rank_1
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<?xf32>
+func.func @reduce_float_dyn_rank_1(%arg0: tensor<?xf32>) -> () {
+  // CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<f32>
+  // CHECK-DAG: %[[CST0:.+]] = arith.constant 0.0
+  // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST0]]{{.*}}outs(%[[INIT]]
+  // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction"]} ins(%[[ARG0]] : tensor<?xf32>) outs(%[[FILL]] : tensor<f32>)
+  // CHECK: ^bb0(%[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32)
+  // CHECK:   %[[RES:.+]] = arith.addf %[[ARG1]], %[[ARG2]] : f32
+  // CHECK:   linalg.yield %[[RES]] : f32
+  // CHECK: tensor.expand_shape %[[GENERIC]] {{\[}}] : tensor<f32> into tensor<1xf32>
+  %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<?xf32>) -> tensor<1xf32>
+  return
+}
+
+// -----
+
 // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
 
 // CHECK-LABEL: @reduce_float_dyn_nonzero_batch
 // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
 func.func @reduce_float_dyn_nonzero_batch(%arg0: tensor<5x?x4xf32>) -> () {
   // CHECK: %[[C1:.+]] = arith.constant 1
   // CHECK: %[[DYN:.+]] = tensor.dim %[[ARG0]], %[[C1]]