[mlir][tosa] make Select operator broadcastable in the pass

[tatwaichong]

Summary

Making Select broadcastable can let this op easier to use.

Change-Id: I4a4bec4f7cbe532e954a5b4fe53136676ab4300c

[rsuderman]

mlir/lib/Dialect/Tosa/Transforms/TosaMakeBroadcastable.cpp

252

Include an notifyMatchFailure for the return. Its useful for debugging.

254

What happens if pred broadcasts both times? They may not broadcast in the same way.

265

Ditto notifyMatchRewriteFailure.

271

Same

277

Seeing this nested ternary makes me concerned the broadcasted predicate may have some issues. It assumes that if both newPred1 and newPred2 are new, they could still be the same broadcasting.

[tatwaichong]

• replace failure with notifyMatchFailure
• address review comment
• add more mlir tests
• add comment in code

[tatwaichong]

mlir/lib/Dialect/Tosa/Transforms/TosaMakeBroadcastable.cpp

254

This pattern assumes the rank of output is known. In broadcasting, the reshaped rank of inputs must match the rank of output that must be equal to the maximum of the input ranks. And [[ https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Tosa/Transforms/TosaMakeBroadcastable.cpp#L124 | reshapeLowerToHigher ]] checks if the reshaped rank of input arguments match with the rank of output.

So both newPred1 and newPred2 are still the same broadcasting if they are new at the same round of the pass.
e.g. given an example that both newPred1 and newPred2 are new.
%3 = "tosa.select"(%0, %1, %2) : (tensor<1xi1>, tensor<1x32x32x8xf32>, tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32>

When a predicate gets changed, it cannot have the highest rank. On the other hand, a lower-rank predicate gets broadcast success in reshapeLowerToHigher only when it aligns with an input having the highest rank (=output rank).

Further, in the example below.
%3 = "tosa.select"(%0, %1, %2) : (tensor<1xi1>, tensor<32x8xf32>, tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32>
after calling reshapeLowerToHigher(pred=<1xi1>, input1=<32x8xf32>, output=<1x32x32x8xf32>) will return failure, and won’t get a new predicate.

To my understanding, Patterns rewriting pass will repeatedly try to do pattern matching and stop when it meets convergence. The broadcasting can be done in multiple rounds of the pass. Let's give an example that the three input tensors are a, b, and c where rank(a) < rank(b) < rank(c). Hence there are six permutations of the set. They are (a, b, c), (a, c, b), (b, a, c), (b, c, a), (c, a, b), and (c, b, a). The rank of output shape is c as this is the maximum.

Let me take the case (a, b, c) where (pred, input1, input2) = (32x8, 32x32x8, 1x32x32x8) as an example,
a) The first round of the pass
Fail on broadcast(perd, input1) as the rank result doesn’t match the output.
Success on broadcast(perd, input2), then we get "tosa.reshape"(%arg0) {new_shape = [1, 1, 32, 8]} : (tensor<32x8xi1>) -> tensor<1x1x32x8xi1>

b) Second round of the pass
Success on broadcast(perd`, input1), then we get "tosa.reshape"(%arg1) {new_shape = [1, 32, 32, 8]} : (tensor<32x32x8xf32>) -> tensor<1x32x32x8xf32>

c) Nothing changes within further rounds because the code cannot be rewritten more as it's already correct.

I added more mlir tests to show its coverage.

277

I more or less replied with my thoughts above on your question.

About this nested ternary, I'm also not satisfied with this. Do you come up with better writing?

[rsuderman]

mlir/lib/Dialect/Tosa/Transforms/TosaMakeBroadcastable.cpp

254

I see, it should be possible by tweaking reshapeLowerToHigher's behavior with returning values and guarantee the broadcasting happens all at once. See comment below.

312

Rather than doing this predicates, set outInput1 and outInput2 to the input values in reshapeLowerToHigher. This avoids all of the predicate selection as the values will be updated at each point.

While you are there, make it return a bool instead of an integer.

[tatwaichong]

overload reshapeLowerToHigher to support 3 inputs so that the broadcasting happens all at once.
rebase.

[tatwaichong]

mlir/lib/Dialect/Tosa/Transforms/TosaMakeBroadcastable.cpp

254

Yes, it can. I create an overloading reshapeLowerToHigher to support 3 inputs (the function logic should be able to extend to support more operands). Sort the tensors by rank to make broadcasting easier as the tensor in the last position of the sorted list has the highest rank.

[rsuderman]

mlir/lib/Dialect/Tosa/Transforms/TosaMakeBroadcastable.cpp

147

We shouldn't need another function. The change I wanted was just to invoke the broadcast twice. Specifically:

rebroadcast(input1, input2, result1, result2)
rebroadcast(result1, input3, result1, result3)

Then check if inputN == resultN for each input.

The idea is we can apply each broadcasting behavior and know the result is legal as they compound. Previously you had extra mutations / checks and no interaction between the first broadcast and the second. By guaranteeing they chain in this way it minimizes the number of cases.

[tatwaichong]

mlir/lib/Dialect/Tosa/Transforms/TosaMakeBroadcastable.cpp

147

I see. I'm not sure if I fully catch your message, Let me double-check my understanding.

If we want to guarantee the whole broadcasting for all inputs happens at once in one round of the pass.
Given the highest rank tensor is input3, and the others have a lower rank. In this case, we still need to do an extra reshapeLowerToHigher, is this right?

matchAndRewrite () {
...
  reshapeLowerToHigher(input1, input2, result1, result2)
  reshapeLowerToHigher(result1, input3, result1, result3)

  if (result2.rank != result3.rank)
    reshapeLowerToHigher(result2, result3, result2, result3)
  }

[tatwaichong]

roll back to use the existing pair-wise broadcasting function. Apply each broadcasting to 3 different pairs of inputs. By chaining them this way as a compound the broadcasting happens all at once.
update mlir tests.

[rsuderman]

mlir/lib/Dialect/Tosa/Transforms/TosaMakeBroadcastable.cpp

84

You don't need separate input1 and outInput1. Just mutate the referenced value passed in. This avoids doing the swapping in the lowering below.

260

See comment above - if you just mutate input1 and input2 via references you can more easily track the three inputs and avoid the extra reassigning.

[tatwaichong]

As suggested, mutate the referenced inputs passed in to avoid doing the extra reassigning and stateful variabled in the lowering.

mlir/lib/Dialect/Tosa/Transforms/TosaMakeBroadcastable.cpp

 
 /// Common code to create the reshape op where necessary to make the rank of the
 /// operations equal. Returns the updated input1 and input2 for the original
 /// input. The caller is expected to use these to rewrite the original operator
 /// with the RESHAPE now in the graph.
 static LogicalResult reshapeLowerToHigher(PatternRewriter &rewriter,
                                           Location loc,
                                           RankedTensorType outputType,
                                           Value input1, Value input2,
                                           Value &outInput1, Value &outInput2) {
   auto input1Ty = input1.getType().dyn_cast<RankedTensorType>();
   auto input2Ty = input2.getType().dyn_cast<RankedTensorType>();
 
-  if (!input1Ty || !input2Ty)
-    return failure();
+  if (!input1Ty || !input2Ty) {
+    return rewriter.notifyMatchFailure(loc, "input not a ranked tensor");
+  }
 
   int64_t input1Rank = input1Ty.getRank();
   int64_t input2Rank = input2Ty.getRank();
 
-  Value higherTensorValue, lowerTensorValue;
-  // Cannot rewrite as its already correct.
   if (input1Rank == input2Rank)
-    return failure();
+    return rewriter.notifyMatchFailure(loc,
+                                       "cannot rewrite as its already correct");
 
+  Value higherTensorValue, lowerTensorValue;
   if (input1Rank > input2Rank) {
     higherTensorValue = input1;
     lowerTensorValue = input2;
   } else {
     higherTensorValue = input2;
     lowerTensorValue = input1;
   }
 
   ArrayRef<int64_t> higherRankShape =
       higherTensorValue.getType().cast<RankedTensorType>().getShape();
-  (void)higherRankShape;
   ArrayRef<int64_t> lowerRankShape =
       lowerTensorValue.getType().cast<RankedTensorType>().getShape();
 
   SmallVector<int64_t, 4> reshapeOutputShape;
 
   if (computeReshapeOutput(higherRankShape, lowerRankShape, reshapeOutputShape)
           .failed())
-    return failure();
+    return rewriter.notifyMatchFailure(loc, "fail to compute a reshape type");
 
   auto reshapeInputType = lowerTensorValue.getType().cast<RankedTensorType>();
   auto reshapeOutputType = RankedTensorType::get(
       ArrayRef<int64_t>(reshapeOutputShape), reshapeInputType.getElementType());
 
   // Verify the rank agrees with the output type if the output type is ranked.
   if (outputType) {
     if (outputType.getShape().size() != reshapeOutputShape.size() ||
         outputType.getShape().size() != higherRankShape.size())
-      return failure();
+      return rewriter.notifyMatchFailure(
+          loc, "the reshaped type doesn't agrees with the ranked output type");
   }
 
   auto reshapeLower = rewriter.create<tosa::ReshapeOp>(
       loc, reshapeOutputType, lowerTensorValue,
       rewriter.getDenseI64ArrayAttr(reshapeOutputShape));
 
   if (input1Rank > input2Rank) {
     outInput1 = higherTensorValue;
     outInput2 = reshapeLower.getResult();
   } else {
     outInput1 = reshapeLower.getResult();
     outInput2 = higherTensorValue;
   }
 
   return success();
 }
 
 namespace {
-template <typename OpTy> struct ConvertTosaOp : public OpRewritePattern<OpTy> {
+template <typename OpTy>
+struct ConvertTosaOp : public OpRewritePattern<OpTy> {
   using OpRewritePattern<OpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(OpTy tosaBinaryOp,
                                 PatternRewriter &rewriter) const override {
 
     Value input1 = tosaBinaryOp.getInput1();
     Value input2 = tosaBinaryOp.getInput2();
     Value output = tosaBinaryOp.getResult();
if (reshapeLowerToHigher(rewriter, tosaBinaryOp.getLoc(), outputType,
       return failure();
 
     rewriter.replaceOpWithNewOp<tosa::ArithmeticRightShiftOp>(
         tosaBinaryOp, outputType, outInput1, outInput2, round);
 
     return success();
   }
 };
+
+template <>
+struct ConvertTosaOp<tosa::SelectOp> : public OpRewritePattern<tosa::SelectOp> {
+  using OpRewritePattern<tosa::SelectOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::SelectOp tosaOp,
+                                PatternRewriter &rewriter) const override {
+
+    Value input1 = tosaOp.getPred();
+    Value input2 = tosaOp.getOnTrue();
+    Value input3 = tosaOp.getOnFalse();
+    Value output = tosaOp.getResult();
+
+    auto outputType = output.getType().dyn_cast<RankedTensorType>();
+    if (!outputType)
+      return rewriter.notifyMatchFailure(tosaOp, "output not a ranked tensor");
+
+    Value result1 = input1;
+    Value result2 = input2;
+    Value result3 = input3;
+
+    // Apply broadcasting to each pair of inputs separately, and chain them as
+    // compound as below so that the broadcasting happens all at once.
+    bool reshaped1 = reshapeLowerToHigher(rewriter, tosaOp.getLoc(), outputType,
+                                          input1, input2, result1, result2)
+                         .succeeded();
+
+    bool reshaped2 = reshapeLowerToHigher(rewriter, tosaOp.getLoc(), outputType,
+                                          result1, input3, result1, result3)
+                         .succeeded();
+
+    bool reshaped3 = reshapeLowerToHigher(rewriter, tosaOp.getLoc(), outputType,
+                                          result2, input3, result2, result3)
+                         .succeeded();
+
+    if (!reshaped1 && !reshaped2 && !reshaped3)
+      return rewriter.notifyMatchFailure(
+          tosaOp,
+          "cannot rewrite as the rank of all operands is already aligned");
+
+    int32_t result1Rank = result1.getType().cast<RankedTensorType>().getRank();
+    int32_t result2Rank = result2.getType().cast<RankedTensorType>().getRank();
+    int32_t result3Rank = result3.getType().cast<RankedTensorType>().getRank();
+
+    if ((result1Rank != result2Rank) || (result2Rank != result3Rank))
+      return rewriter.notifyMatchFailure(
+          tosaOp, "not all ranks are aligned with each other");
+
+    rewriter.replaceOpWithNewOp<tosa::SelectOp>(tosaOp, outputType, result1,
+                                                result2, result3);
+
+    return success();
+  }
+};
 } // namespace
 
 namespace {
 /// Pass that enables broadcast by making all input arrays have the same
 /// number of dimensions. Insert RESHAPE operations to lower rank operand
 struct TosaMakeBroadcastable
     : public tosa::impl::TosaMakeBroadcastableBase<TosaMakeBroadcastable> {
 public:
   void runOnOperation() override {
     auto func = getOperation();
     RewritePatternSet patterns(func.getContext());
     MLIRContext *ctx = func.getContext();
     // Add the generated patterns to the list.
     patterns.add<ConvertTosaOp<tosa::BitwiseAndOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::BitwiseOrOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::BitwiseXorOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::AddOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::SubOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::MulOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::DivOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::MaximumOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::MinimumOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::EqualOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::GreaterOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::GreaterEqualOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::LogicalLeftShiftOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::ArithmeticRightShiftOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::LogicalRightShiftOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::LogicalAndOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::LogicalOrOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::LogicalXorOp>>(ctx);
+    patterns.add<ConvertTosaOp<tosa::SelectOp>>(ctx);
     patterns.add<ConvertTosaOp<tosa::PowOp>>(ctx);
     (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::tosa::createTosaMakeBroadcastablePass() {
   return std::make_unique<TosaMakeBroadcastable>();
 }

mlir/test/Dialect/Tosa/broadcast.mlir

 // -----
 // CHECK-LABEL: broadcast_scalar
 func.func @test_broadcast_scalar(%arg0: tensor<i32>, %arg1: tensor<17x16x15x14xi32>) -> tensor<17x16x15x14xi32> {
   // CHECK-DAG: %[[VAR0:.*]] = "tosa.reshape"(%arg0) {new_shape = array<i64: 1, 1, 1, 1>}
   // CHECK: %[[VAR1:.*]] = "tosa.add"(%[[VAR0]], %arg1)
   %0 = "tosa.add"(%arg0, %arg1) : (tensor<i32>, tensor<17x16x15x14xi32>) -> tensor<17x16x15x14xi32>
   return %0 : tensor<17x16x15x14xi32>
 }
+
+// -----
+// CHECK-LABEL: broadcast_select_both_input
+func.func @test_broadcast_select_both_input(%arg0: tensor<1x16x16xi1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<1x16x16xf32> {
+  // CHECK-DAG: %[[VAL_0:.*]] = "tosa.reshape"(%arg1) {new_shape = array<i64: 1, 1, 1>}
+  // CHECK-DAG: %[[VAL_1:.*]] = "tosa.reshape"(%arg2) {new_shape = array<i64: 1, 1, 1>}
+  // CHECK: %[[VAL_2:.*]] = "tosa.select"(%arg0, %[[VAL_0]], %[[VAL_1]])
+  %0 = "tosa.select"(%arg0, %arg1, %arg2) : (tensor<1x16x16xi1>, tensor<f32>, tensor<f32>) -> tensor<1x16x16xf32>
+  return %0 : tensor<1x16x16xf32>
+}
+
+// -----
+// CHECK-LABEL: broadcast_select_one_input
+func.func @test_broadcast_select_one_input(%arg0: tensor<17x16x15x14xi1>, %arg1: tensor<17x16x15x14xf32>, %arg2: tensor<f32>) -> tensor<17x16x15x14xf32> {
+  // CHECK-DAG: %[[VAL_0:.*]] = "tosa.reshape"(%arg2) {new_shape = array<i64: 1, 1, 1, 1>}
+  // CHECK: %[[VAL_1:.*]] = "tosa.select"(%arg0, %arg1, %[[VAL_0]])
+  %0 = "tosa.select"(%arg0, %arg1, %arg2) : (tensor<17x16x15x14xi1>, tensor<17x16x15x14xf32>, tensor<f32>) -> tensor<17x16x15x14xf32>
+  return %0 : tensor<17x16x15x14xf32>
+}
+
+// -----
+// CHECK-LABEL: broadcast_select_predicate
+func.func @test_broadcast_select_predicate(%arg0: tensor<i1>, %arg1: tensor<1x32x32x8xf32>, %arg2: tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32> {
+  // CHECK-DAG: %[[VAL_0:.*]] = "tosa.reshape"(%arg0) {new_shape = array<i64: 1, 1, 1, 1>}
+  // CHECK: %[[VAL_1:.*]] = "tosa.select"(%[[VAL_0]], %arg1, %arg2)
+  %0 = "tosa.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<1x32x32x8xf32>, tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32>
+  return %0 : tensor<1x32x32x8xf32>
+}
+
+// -----
+// CHECK-LABEL: broadcast_select_abc
+func.func @test_broadcast_select_abc(%arg0: tensor<i1>, %arg1: tensor<32x8xf32>, %arg2: tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32> {
+  // CHECK-DAG: %[[VAL_0:.*]] = "tosa.reshape"(%arg0) {new_shape = array<i64: 1, 1, 1, 1>}
+  // CHECK-DAG: %[[VAL_1:.*]] = "tosa.reshape"(%arg1) {new_shape = array<i64: 1, 1, 32, 8>}
+  // CHECK: %[[VAL_2:.*]] = "tosa.select"(%[[VAL_0]], %[[VAL_1]], %arg2)
+  %0 = "tosa.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<32x8xf32>, tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32>
+  return %0 : tensor<1x32x32x8xf32>
+}
+
+// -----
+// CHECK-LABEL: broadcast_select_acb
+func.func @test_broadcast_select_acb(%arg0: tensor<i1>, %arg1: tensor<1x32x32x8xf32>, %arg2: tensor<32x8xf32>) -> tensor<1x32x32x8xf32> {
+  // CHECK-DAG: %[[VAL_0:.*]] = "tosa.reshape"(%arg0) {new_shape = array<i64: 1, 1, 1, 1>}
+  // CHECK-DAG: %[[VAL_1:.*]] = "tosa.reshape"(%arg2) {new_shape = array<i64: 1, 1, 32, 8>}
+  // CHECK: %[[VAL_2:.*]] = "tosa.select"(%[[VAL_0]], %arg1, %[[VAL_1]])
+  %0 = "tosa.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<1x32x32x8xf32>, tensor<32x8xf32>) -> tensor<1x32x32x8xf32>
+  return %0 : tensor<1x32x32x8xf32>
+}
+
+// -----
+// CHECK-LABEL: broadcast_select_bac
+func.func @test_broadcast_select_bac(%arg0: tensor<32x8xi1>, %arg1: tensor<f32>, %arg2: tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32> {
+  // CHECK-DAG: %[[VAL_0:.*]] = "tosa.reshape"(%arg0) {new_shape = array<i64: 1, 1, 32, 8>}
+  // CHECK-DAG: %[[VAL_1:.*]] = "tosa.reshape"(%arg1) {new_shape = array<i64: 1, 1, 1, 1>}
+  // CHECK: %[[VAL_2:.*]] = "tosa.select"(%[[VAL_0]], %[[VAL_1]], %arg2)
+  %0 = "tosa.select"(%arg0, %arg1, %arg2) : (tensor<32x8xi1>, tensor<f32>, tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32>
+  return %0 : tensor<1x32x32x8xf32>
+}
+
+// -----
+// CHECK-LABEL: broadcast_select_bca
+func.func @test_broadcast_select_bca(%arg0: tensor<32x8xi1>, %arg1: tensor<1x32x32x8xf32>, %arg2: tensor<i1>) -> tensor<1x32x32x8xf32> {
+  // CHECK-DAG: %[[VAL_0:.*]] = "tosa.reshape"(%arg0) {new_shape = array<i64: 1, 1, 32, 8>}
+  // CHECK-DAG: %[[VAL_1:.*]] = "tosa.reshape"(%arg2) {new_shape = array<i64: 1, 1, 1, 1>}
+  // CHECK: %[[VAL_2:.*]] = "tosa.select"(%[[VAL_0]], %arg1, %[[VAL_1]])
+  %0 = "tosa.select"(%arg0, %arg1, %arg2) : (tensor<32x8xi1>, tensor<1x32x32x8xf32>, tensor<i1>) -> tensor<1x32x32x8xf32>
+  return %0 : tensor<1x32x32x8xf32>
+}
+
+// -----
+// CHECK-LABEL: broadcast_select_cab
+func.func @test_broadcast_select_cab(%arg0: tensor<1x32x32x8xi1>, %arg1: tensor<f32>, %arg2: tensor<32x8xf32>) -> tensor<1x32x32x8xf32> {
+  // CHECK-DAG: %[[VAL_0:.*]] = "tosa.reshape"(%arg1) {new_shape = array<i64: 1, 1, 1, 1>}
+  // CHECK-DAG: %[[VAL_1:.*]] = "tosa.reshape"(%arg2) {new_shape = array<i64: 1, 1, 32, 8>}
+  // CHECK: %[[VAL_2:.*]] = "tosa.select"(%arg0, %[[VAL_0]], %[[VAL_1]])
+  %0 = "tosa.select"(%arg0, %arg1, %arg2) : (tensor<1x32x32x8xi1>, tensor<f32>, tensor<32x8xf32>) -> tensor<1x32x32x8xf32>
+  return %0 : tensor<1x32x32x8xf32>
+}
+
+// -----
+// CHECK-LABEL: broadcast_select_cba
+func.func @test_broadcast_select_cba(%arg0: tensor<1x32x32x8xi1>, %arg1: tensor<32x8xf32>, %arg2: tensor<i1>) -> tensor<1x32x32x8xf32> {
+  // CHECK-DAG: %[[VAL_0:.*]] = "tosa.reshape"(%arg1) {new_shape = array<i64: 1, 1, 32, 8>}
+  // CHECK-DAG: %[[VAL_1:.*]] = "tosa.reshape"(%arg2) {new_shape = array<i64: 1, 1, 1, 1>}
+  // CHECK: %[[VAL_2:.*]] = "tosa.select"(%arg0, %[[VAL_0]], %[[VAL_1]])
+  %0 = "tosa.select"(%arg0, %arg1, %arg2) : (tensor<1x32x32x8xi1>, tensor<32x8xf32>, tensor<i1>) -> tensor<1x32x32x8xf32>
+  return %0 : tensor<1x32x32x8xf32>
+}