[mlir][linalg] Change linalg.broadcast `dimensions` attribute to represent added dimensions.

[olegshyshkov]

Summary

Original RFC defined dimensions as a map from input to init, but a discussion in D138291 concluded that it's more natural for dimensions to represent added dims. Also this way is more consistent with linalg.reduce.

[pifon2a]

Thank you, Oleg!

[olegshyshkov]

Fix tiling test.

[nicolasvasilache]

Nice, thanks @olegshyshkov !

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

 //===----------------------------------------------------------------------===//
 
 def BroadcastOp : LinalgStructuredBase_Op<"broadcast", [
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
     SameVariadicOperandSize,
     SingleBlockImplicitTerminator<"YieldOp">]> {
   let summary = "Static broadcast operator";
   let description = [{
-    Broadcast the input into the given shape by adding dimensions.
-
-    Each index in `dimensions` attribute maps input dimension into the
-    corresponding target dimension. The length of the `dimensions` list should
-    match the `input` rank and dimensions should be in sorted order. There is no
-    ambiguity at compile-time about shape information.
+    Broadcast the input into the given shape by adding `dimensions`.
 
     Example:
     ```
       %bcast = linalg.broadcast
           ins(%input:tensor<16xf32>)
           inits(%init:tensor<16x64xf32>)
-          dimensions = [0]
+          dimensions = [1]
     ```
   }];
 
   let arguments = (ins
     // Input arg
     TensorOrMemref:$input,
     // Output arg
     TensorOrMemref:$init,

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

void BroadcastOp::print(OpAsmPrinter &p) {
   printDenseI64ArrayAttr(p, getDimensionsAttrName(), getDimensions());
   p.printOptionalAttrDict((*this)->getAttrs(), {getDimensionsAttrName()});
   p.decreaseIndent();
 }
 
 LogicalResult BroadcastOp::verify() {
   ArrayRef<int64_t> dimensionsRef = getDimensions();
 
-  if (!llvm::is_sorted(dimensionsRef))
-    return emitOpError() << "dimensions should be in sorted order, implicit "
-                            "transpose is not supported";
-
   auto inputType = getInput().getType();
   auto initType = getInit().getType();
 
   int64_t inputRank = inputType.getRank();
   int64_t initRank = initType.getRank();
 
   auto inputShape = inputType.getShape();
   auto initShape = initType.getShape();
 
-  if ((size_t)inputRank != dimensionsRef.size())
-    return emitOpError()
-           << "input rank does match the number of dimensions. expected: "
-           << inputRank << ", got: " << dimensionsRef.size();
-
-  // Mapping from init dims to input dims.
-  const int64_t kUnmappedDim = -1;
-  SmallVector<int64_t> reverseDimMap(initRank, kUnmappedDim);
+  if ((size_t)inputRank + dimensionsRef.size() != (size_t)initRank)
+    return emitOpError() << "input rank plus added dimensions does not "
+                            "match init rank. input rank: "
+                         << inputRank
+                         << ", dimensions size: " << dimensionsRef.size()
+                         << ", init rank: " << initRank;
 
   for (const auto &[idx, dim] : llvm::enumerate(dimensionsRef)) {
     if (dim < 0 || dim >= initRank)
       return emitOpError() << "dimension " << idx
                            << " is out of range. expected range: [0, "
                            << initRank - 1 << "], got: " << dim;
+  }
 
-    reverseDimMap[dim] = idx;
+  // Mapping from input dims to init dims.
+  SmallVector<int64_t> dimMap;
+  for (auto dim : llvm::seq<int64_t>(0, initRank)) {
+    if (!llvm::is_contained(dimensionsRef, dim))
+      dimMap.push_back(dim);
   }
 
-  for (const auto &[idx, inputDimIdx] : llvm::enumerate(reverseDimMap)) {
-    if (inputDimIdx != kUnmappedDim) {
-      // This dimensions is mapped from the input. Init and input dims should
-      // match.
-      if (inputShape[inputDimIdx] != initShape[idx])
-        return emitOpError()
-               << "input dim " << inputDimIdx << " should match init dim "
-               << idx << ". input: " << inputShape[inputDimIdx]
-               << ", init: " << initShape[idx];
-    }
+  for (const auto &[inputDimIdx, initDimIdx] : llvm::enumerate(dimMap)) {
+    if (inputShape[inputDimIdx] != initShape[initDimIdx])
+      return emitOpError() << "input dim " << inputDimIdx
+                           << " should match init dim " << initDimIdx
+                           << ". input: " << inputShape[inputDimIdx]
+                           << ", init: " << initShape[initDimIdx];
   }
 
   return success();
 }
 
 SmallVector<utils::IteratorType> BroadcastOp::getIteratorTypesArray() {
   int64_t rank = getInit().getType().getRank();
   return SmallVector<utils::IteratorType>(rank, utils::IteratorType::parallel);
 }
 
 ArrayAttr BroadcastOp::getIndexingMaps() {
   Builder builder(getContext());
   int64_t rank = getInit().getType().getRank();
   return builder.getAffineMapArrayAttr(
-      {builder.getMultiDimIdentityMap(rank).getSubMap(
-           llvm::to_vector_of<unsigned>(getDimensions())),
+      {builder.getMultiDimIdentityMap(rank).dropResults(getDimensions()),
        builder.getMultiDimIdentityMap(rank)});
 }
 
 void BroadcastOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   getGenericEffectsImpl(effects, getOperation()->getResults(),
                         getDpsInputOperands(), getDpsInitOperands());

mlir/test/Dialect/Linalg/invalid.mlir

%transpose = linalg.transpose
       ins(%input:tensor<16x32xf32>)
       outs(%init:tensor<32x64x16xf32>)
       permutation = [1, 0, 2]
   func.return %transpose : tensor<32x64x16xf32>
 }
 
 // -----
 
-func.func @broadcast_unsorted_dims(
-    %input: tensor<4x16xf32>, %init: tensor<4x8x16xf32>)
-    -> tensor<4x8x16xf32> {
-  // expected-error @+1 {{'linalg.broadcast' op dimensions should be in sorted order}}
-  %bcast = linalg.broadcast
-      ins(%input:tensor<4x16xf32>)
-      outs(%init:tensor<4x8x16xf32>)
-      dimensions = [1, 0]
-  func.return %bcast : tensor<4x8x16xf32>
-}
-
-// -----
-
 func.func @broadcast_input_dims_rank_mismatch(
     %input: tensor<4x16xf32>, %init: tensor<4x8x16xf32>)
     -> tensor<4x8x16xf32> {
-  // expected-error @+1 {{'linalg.broadcast' op input rank does match the number of dimensions. expected: 2, got: 1}}
+  // expected-error @+1 {{'linalg.broadcast' op input rank plus added dimensions does not match init rank. }}
   %bcast = linalg.broadcast
       ins(%input:tensor<4x16xf32>)
       outs(%init:tensor<4x8x16xf32>)
-      dimensions = [0]
+      dimensions = [1, 2]
   func.return %bcast : tensor<4x8x16xf32>
 }
 
 // -----
 
 func.func @broadcast_unsorted_dims(
     %input: tensor<4x16xf32>, %init: tensor<4x8x16xf32>)
     -> tensor<4x8x16xf32> {
-  // expected-error @+1 {{'linalg.broadcast' op dimension 1 is out of range. expected range: [0, 2], got: 5}}
+  // expected-error @+1 {{'linalg.broadcast' op dimension 0 is out of range. expected range: [0, 2], got: 5}}
   %bcast = linalg.broadcast
       ins(%input:tensor<4x16xf32>)
       outs(%init:tensor<4x8x16xf32>)
-      dimensions = [0, 5]
+      dimensions = [5]
   func.return %bcast : tensor<4x8x16xf32>
 }
 
 // -----
 
 func.func @broadcast_mapped_dim_mismatch(
     %input: tensor<4x16xf32>, %init: tensor<5x8x16xf32>)
     -> tensor<5x8x16xf32> {
   // expected-error @+1 {{'linalg.broadcast' op input dim 0 should match init dim 0. input: 4, init: 5}}
   %bcast = linalg.broadcast
       ins(%input:tensor<4x16xf32>)
       outs(%init:tensor<5x8x16xf32>)
-      dimensions = [0, 2]
+      dimensions = [1]
   func.return %bcast : tensor<5x8x16xf32>
 }
 
 // -----
 
 func.func @broadcast_size_1_extension_not_supported(
     %input: tensor<1x16xf32>, %init: tensor<4x?x16xf32>)
     -> tensor<4x?x16xf32> {
   // expected-error @+1 {{'linalg.broadcast' op input dim 0 should match init dim 0. input: 1, init: 4}}
   %bcast = linalg.broadcast
       ins(%input:tensor<1x16xf32>)
       outs(%init:tensor<4x?x16xf32>)
-      dimensions = [0, 2]
+      dimensions = [1]
   func.return %bcast : tensor<4x?x16xf32>
 }

mlir/test/Dialect/Linalg/one-shot-bufferize.mlir

 
 // CHECK-LABEL: func @broadcast
 // CHECK-SAME:  %[[ARG0:.*]]: memref<8x32xf32
 func.func @broadcast(%input: tensor<8x32xf32>,
                      %init: tensor<8x16x32xf32>) -> tensor<8x16x32xf32> {
   %bcast = linalg.broadcast
       ins(%input:tensor<8x32xf32>)
       outs(%init:tensor<8x16x32xf32>)
-      dimensions = [0, 2]
+      dimensions = [1]
   func.return %bcast : tensor<8x16x32xf32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // AllocTensorOp elimination would produce SSA violations for the example below.
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Linalg/roundtrip.mlir

 
 // -----
 
 func.func @broadcast_static_sizes(%input: tensor<8x32xf32>,
                             %init: tensor<8x16x32xf32>) -> tensor<8x16x32xf32> {
   %bcast = linalg.broadcast
       ins(%input:tensor<8x32xf32>)
       outs(%init:tensor<8x16x32xf32>)
-      dimensions = [0, 2]
+      dimensions = [1]
   func.return %bcast : tensor<8x16x32xf32>
 }
 // CHECK-LABEL: func @broadcast_static_sizes
 //      CHECK:    linalg.broadcast
 // CHECK-NEXT:    ins
 // CHECK-NEXT:    outs
 // CHECK-NEXT:    dimensions
 
 // -----
 
 func.func @broadcast_with_dynamic_sizes(
               %input: tensor<8x?xf32>, %init: tensor<8x16x?xf32>)
               -> tensor<8x16x?xf32> {
   %bcast = linalg.broadcast
       ins(%input:tensor<8x?xf32>)
       outs(%init:tensor<8x16x?xf32>)
-      dimensions = [0, 2]
+      dimensions = [1]
   func.return %bcast : tensor<8x16x?xf32>
 }
 // CHECK-LABEL: func @broadcast_with_dynamic_sizes
 //      CHECK:    linalg.broadcast
 // CHECK-NEXT:    ins
 // CHECK-NEXT:    outs
 // CHECK-NEXT:    dimensions
 
 // -----
 
 func.func @broadcast_memref(%input: memref<8x32xf32>,
                             %init: memref<8x16x32xf32>) {
   linalg.broadcast
       ins(%input:memref<8x32xf32>)
       outs(%init:memref<8x16x32xf32>)
-      dimensions = [0, 2]
+      dimensions = [1]
   func.return
 }
 
 // CHECK-LABEL: func @broadcast_memref
 //      CHECK:    linalg.broadcast
 // CHECK-NEXT:    ins
 // CHECK-NEXT:    outs
 // CHECK-NEXT:    dimensions

mlir/test/Interfaces/TilingInterface/lower-to-loops-using-interface.mlir

 
 // -----
 
 func.func @broadcast(%input: memref<8x32xf32>,
                      %init: memref<8x16x32xf32>) {
   linalg.broadcast
       ins(%input:memref<8x32xf32>)
       outs(%init:memref<8x16x32xf32>)
-      dimensions = [0, 2]
+      dimensions = [1]
   func.return
 }
 // CHECK-LABEL: func.func @broadcast(
 // CHECK-SAME:    %[[IN:[a-zA-Z0-9]+]]: memref<8x32xf32>,
 // CHECK-SAME:    %[[OUT:[a-zA-Z0-9]+]]: memref<8x16x32xf32>
 
 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index