[MLIR][TOSA] Tosa elementwise broadcasting

[rsuderman]

Summary

Added support for broadcasting size-1 dimensions for TOSA elementwise
operations.

[rriddle]

There is a typo in the commit description: elemtnwise

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

155

The way the error message is worded makes me think that this should be an assert.

166

Cache end iterators to avoid recomputing during the loop.

https://llvm.org/docs/CodingStandards.html#don-t-evaluate-end-every-time-through-a-loop

211

nit: Spell out auto here. Type is the same number of characters as auto.

218

nit: Drop trivial braces here.

229

Same comment here on braces, and caching the number of results. Although in this case, seems better to do something like:

indexingMaps.append(operation->getNumResults(), rewriter.getMultiDimIdentityMap(nloops));

[rsuderman]

rriddle@ comments

[silvas]

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

17

add tests

[rsuderman]

ssilva@ comments for additional broadcast tests.

[rsuderman]

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

17

I revamped the previous failing broadcasting test. Updated to include additional broadcasting tests.

[silvas]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

161

add assertion message.

169

Tosa op verifiers should enforce this already. Can be deleted.

204

Now it supports broadcast :)

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

59–60

Capture with filecheck. Example of "non-local" capture requires "$": https://github.com/llvm/llvm-project/blob/main/mlir/test/Conversion/SCFToGPU/parallel_loop.mlir#L18

[rsuderman]

Removed unneeded checks. Updated broadcasting comments.

[rsuderman]

Updated to reference affine maps through variable capture

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

elementwiseMatchAndRewriteHelper(Operation *operation,
                                  PatternRewriter &rewriter) {
   auto loc = operation->getLoc();
   auto results = operation->getResults();
   auto t0 = operation->getOperand(0).getType().template dyn_cast<ShapedType>();
   if (!t0)
     return rewriter.notifyMatchFailure(operation,
                                        "All results must be a shaped type");
 
-  // For now require no broadcasting. Consider making it support broadcasting
-  // operations.
-  Type uniqueInTy = operation->getOperand(0).getType();
-  bool allInputTypesEqual =
-      llvm::all_of(operation->getOperandTypes(),
-                   [&](Type operandTy) { return operandTy == uniqueInTy; });
-  if (!allInputTypesEqual)
-    return rewriter.notifyMatchFailure(operation,
-                                       "All operands must have the same type");
-  bool resultAndInputShapeEqual =
-      llvm::all_of(operation->getResultTypes(), [&](Type resultTy) {
-        return resultTy.cast<ShapedType>().getShape() == t0.getShape();
-      });
+  if (operation->getNumResults() != 1) {
+    return rewriter.notifyMatchFailure(
+        operation, "TOSA ops should only produce one result");
+  }
 
-  if (!resultAndInputShapeEqual)
+  // Dimension sizes should all be broadcastable to the output shape.
+  auto resultTy = operation->getResultTypes().front().cast<ShapedType>();
+  auto shapeComparison = [&](Type operandTy) {
+    auto shapedTy = operandTy.cast<ShapedType>();
+    if (shapedTy.getRank() != resultTy.getRank())
+      return false;
+    for (int i = 0; i < shapedTy.getRank(); i++) {
+      if (resultTy.getDimSize(i) != shapedTy.getDimSize(i) &&
+          shapedTy.getDimSize(i) != 1)
+        return false;
+    }
+    return true;
+  };
+
+  bool allInputShapesEqual =
+      llvm::all_of(operation->getOperandTypes(), shapeComparison);
+  if (!allInputShapesEqual)
     return rewriter.notifyMatchFailure(
-        operation, "All results must have the same shape as the input");
+        operation, "All operands must broadcast compatible shapes");
 
   // Construct the indexing maps needed for linalg.generic ops.
   SmallVector<Type> bodyArgTypes;
 
   for (Value in : operation->getOperands())
     bodyArgTypes.emplace_back(getElementTypeOrSelf(in.getType()));
 
   SmallVector<Type> opResultTypes;
initTensors.push_back(rewriter.create<linalg::InitTensorOp>(
         loc, ArrayRef<Value>({}), resultType.getShape(),
         resultType.getElementType()));
     opResultTypes.push_back(result.getType());
   }
 
   auto bodyResultTypes = llvm::to_vector<4>(llvm::map_range(
       initTensors, [](Value v) { return getElementTypeOrSelf(v); }));
 
-  // Supports only non-broadcasted operation. Shoudl consider update indexing
+  // Supports only non-broadcasted operation. Should consider updating the index
   // map to be multidimensional.
   unsigned nloops = t0.getRank();
-  AffineMap commonIndexingMap = rewriter.getMultiDimIdentityMap(nloops);
-  SmallVector<AffineMap, 2> indexingMaps(
-      operation->getNumOperands() + bodyResultTypes.size(), commonIndexingMap);
+  SmallVector<AffineMap, 2> indexingMaps;
+  indexingMaps.reserve(operation->getNumOperands() + bodyResultTypes.size());
+
+  // Input indexing maps may be broadcasted.
+  for (auto types : operation->getOperandTypes()) {
+    auto shape = types.cast<ShapedType>().getShape();
+    SmallVector<AffineExpr, 4> dimExprs;
+    dimExprs.reserve(nloops);
+    for (unsigned i = 0; i < nloops; ++i) {
+      // If the dimension is one we can broadcast the input with a constant
+      // affine expression.
+      if (shape[i] == 1) {
+        dimExprs.push_back(rewriter.getAffineConstantExpr(0));
+      } else {
+        dimExprs.push_back(rewriter.getAffineDimExpr(i));
+      }
+    }
+    indexingMaps.push_back(AffineMap::get(/*dimCount=*/nloops,
+                                          /*symbolCount=*/0, dimExprs,
+                                          rewriter.getContext()));
+  }
+
+  for (unsigned int i = 0; i < operation->getNumResults(); i++) {
+    indexingMaps.push_back(rewriter.getMultiDimIdentityMap(nloops));
+  }
 
   bool didEncounterError = false;
   auto linalgOp = rewriter.create<linalg::GenericOp>(
       loc, opResultTypes, operation->getOperands(), initTensors, indexingMaps,
       getNParallelLoopsAttrs(nloops),
       [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) {
         Value opResult = createLinalgBodyCalculationForElementwiseOp(
             operation, blockArgs.take_front(operation->getNumOperands()),

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

 // RUN: mlir-opt --split-input-file --tosa-to-linalg-on-tensors %s -verify-diagnostics -o -| FileCheck %s
 
 // CHECK: #map = affine_map<() -> ()>
 
 // CHECK-LABEL: @test_abs
 func @test_abs(%arg0: tensor<f32>) -> tensor<f32> {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [] : tensor<f32>
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = []} ins(%arg0 : tensor<f32>) outs([[INIT]] : tensor<f32>) {
   // CHECK: ^bb0(%arg1: f32, %arg2: f32):
   // CHECK:   [[ELEMENT:%.+]] = absf %arg1
   // CHECK:   linalg.yield [[ELEMENT]] : f32
   // CHECK: } -> tensor<f32>
 
   %0 = "tosa.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
 
   // CHECK: return [[GENERIC]]
   return %0 : tensor<f32>
 }
 
 // -----
 
 // CHECK: #map = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL: @test_abs
-func @test_abs(%arg0: tensor<1xf32>) -> tensor<1xf32> {
-  // CHECK: [[INIT:%.+]] = linalg.init_tensor [1] : tensor<1xf32>
-  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<1xf32>) outs([[INIT]] : tensor<1xf32>) {
+func @test_abs(%arg0: tensor<2xf32>) -> tensor<2xf32> {
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [2] : tensor<2xf32>
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xf32>) outs([[INIT]] : tensor<2xf32>) {
   // CHECK: ^bb0(%arg1: f32, %arg2: f32):
   // CHECK:   [[ELEMENT:%.+]] = absf %arg1
   // CHECK:   linalg.yield [[ELEMENT]] : f32
-  // CHECK: } -> tensor<1xf32>
-  %0 = "tosa.abs"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
+  // CHECK: } -> tensor<2xf32>
+  %0 = "tosa.abs"(%arg0) : (tensor<2xf32>) -> tensor<2xf32>
 
   // CHECK: return [[GENERIC]]
-  return %0 : tensor<1xf32>
+  return %0 : tensor<2xf32>
 }
 
 // -----
 
 // CHECK: #map = affine_map<(d0, d1) -> (d0, d1)>
 
 // CHECK-LABEL: @test_abs
-func @test_abs(%arg0: tensor<1x2xf32>) -> tensor<1x2xf32> {
-  // CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 2] : tensor<1x2xf32>
-  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<1x2xf32>) outs([[INIT]] : tensor<1x2xf32>) {
+func @test_abs(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 3] : tensor<2x3xf32>
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<2x3xf32>) outs([[INIT]] : tensor<2x3xf32>) {
   // CHECK: ^bb0(%arg1: f32, %arg2: f32):
   // CHECK:   [[ELEMENT:%.+]] = absf %arg1
   // CHECK:   linalg.yield [[ELEMENT]] : f32
-  // CHECK: } -> tensor<1x2xf32>
-  %0 = "tosa.abs"(%arg0) : (tensor<1x2xf32>) -> tensor<1x2xf32>
+  // CHECK: } -> tensor<2x3xf32>
+  %0 = "tosa.abs"(%arg0) : (tensor<2x3xf32>) -> tensor<2x3xf32>
 
   // CHECK: return [[GENERIC]]
-  return %0 : tensor<1x2xf32>
+  return %0 : tensor<2x3xf32>
 }
 
 // -----
 
+// CHECK-LABEL: @test_add
 func @test_add(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
-  // expected-error @+1 {{failed to legalize operation 'tosa.add'}}
   %0 = "tosa.add"(%arg0, %arg1) : (tensor<1xf32>, tensor<2xf32>) -> tensor<2xf32>
   return %0 : tensor<2xf32>
 }
 
 // -----
 
 func @test_add(%arg0: tensor<1xf32>, %arg1: tensor<f32>) -> tensor<1xf32> {
   // expected-error @+1 {{failed to legalize operation 'tosa.add'}}
   %0 = "tosa.add"(%arg0, %arg1) : (tensor<1xf32>, tensor<f32>) -> tensor<1xf32>