[Mlir-commits] [mlir] [mlir][tosa] Fix Map for Bias Broadcast (PR #89059)

[Jack Frankland]

https://github.com/FranklandJack created https://github.com/llvm/llvm-project/pull/89059

Fixes the affine map in the broadcast of Conv2D and Fully Connected operations. Previously this logic did not handle the special case as specified in the TOSA specification that the bias tensor of rank-1 may have a size of 1 (i.e. a tensor containing a single element). In this case the map should be a constant index of zero.

Updates the lit tests for the tosa-to-linalg-named lowering to include a Conv2D test case with a bias of rank-1 size 1. We do not need a test for Fully Connected since both operators use the same broadcasting logic.

>From 24d999fb390126cdf39d9e8959003837b0d0b1c2 Mon Sep 17 00:00:00 2001
From: Jack Frankland <jack.frankland at arm.com>
Date: Wed, 17 Apr 2024 10:59:43 +0100
Subject: [PATCH] [mlir][tosa] Fix Map for Bias Broadcast

Fixes the affine map in the broadcast of Conv2D and Fully Connected
operations. Previously this logic did not handle the special case as
specified in the TOSA specification that the bias tensor of rank-1 may
have a size of 1 (i.e. a tensor containing a single element). In this
case the map should be a constant index of zero.

Updates the lit tests for the tosa-to-linalg-named lowering to include a
Conv2D test case with a bias of rank-1 size 1. We do not need a test for
Fully Connected since both operators use the same broadcasting logic.

Signed-off-by: Jack Frankland <jack.frankland at arm.com>
---
 .../Conversion/TosaToLinalg/TosaToLinalgNamed.cpp | 15 ++++++++++++---
 .../TosaToLinalg/tosa-to-linalg-named.mlir        | 13 +++++++++++++
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index 8fb8d16486560c..0a38b05e361013 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -101,9 +101,18 @@ static mlir::Value linalgBroadcastAndMaybeExtSI(PatternRewriter &rewriter,
   // The source tensor is broadcast to all the outer dimensions of the
   // result tensor.
   SmallVector<AffineExpr> sourceDims;
-  for (auto dim : llvm::seq<int64_t>(0, sourceRank)) {
-    auto expr = rewriter.getAffineDimExpr(dim + resultRank - sourceRank);
-    sourceDims.push_back(expr);
+  // In the case of a rank one source tensor with a single element TOSA
+  // specifies that the value be broadcast meaning we need an edge case for a
+  // constant map.
+  assert(sourceTy.hasStaticShape() &&
+         "Dynamic broadcasting shapes not supported!");
+  if (1 == sourceRank && 1 == sourceTy.getDimSize(0)) {
+    sourceDims.push_back(rewriter.getAffineConstantExpr(0));
+  } else {
+    for (auto dim : llvm::seq<int64_t>(0, sourceRank)) {
+      auto expr = rewriter.getAffineDimExpr(dim + resultRank - sourceRank);
+      sourceDims.push_back(expr);
+    }
   }
 
   // Creating maps for the input and output of the broacast-like generic op.
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
index b4049000c50dc8..39699ee315e6cb 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -503,6 +503,19 @@ func.func @avg_pool_dyn(%arg0: tensor<?x6x34x62xf32>) -> (tensor<?x5x33x62xf32>)
 
 // -----
 
+// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (0)>
+// CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+
+// CHECK-LABEL: @conv2d_scalar_bias_f32
+func.func @conv2d_scalar_bias_f32(%input: tensor<1x49x42x27xf32>, %weights: tensor<28x3x3x27xf32>, %bias: tensor<1xf32>) -> () {
+  // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x45x40x28xf32>
+  // CHECK: %[[BROADCAST:.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2 : tensor<1xf32>) outs(%[[INIT]] : tensor<1x45x40x28xf32>) {
+  %0 = tosa.conv2d %input, %weights, %bias {pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>, dilation = array<i64: 2, 1>} : (tensor<1x49x42x27xf32>, tensor<28x3x3x27xf32>, tensor<1xf32>) -> tensor<1x45x40x28xf32>
+  return
+}
+
+// -----
+
 // CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
 // CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>