[mlir][linalg] Do not check if added dimension are static in linalg.broadcast.

[olegshyshkov]

Summary

Added dimensions can be both static and dinamic. Mapped dimension should be the same in the input and the init.

[olegshyshkov]

Fix tests.

[nicolasvasilache]

what was the rationale behind this check in the first place ?

Also, looking at the test you remove, the attribute looks unnatural to me:

%bcast = linalg.broadcast
    ins(%input:tensor<4x16xf32>)
    outs(%init:tensor<4x?x16xf32>)
    dimensions = [0, 2]
func.return %bcast : tensor<4x?x16xf32>

I would have expected to have

%bcast = linalg.broadcast
    ins(%input:tensor<4x16xf32>)
    outs(%init:tensor<4x?x16xf32>)
    dimensions = [1]
func.return %bcast : tensor<4x?x16xf32>

I.e. the dimension 1 of the result is added by the broadcast.
As opposed to the dimensions 0, and 2 of the output were already present in the input.

Thoughts?

[pifon2a]

So, you would prefer having the attribute that specifies what dimensions are added, instead of what input dimensions are mapped to? These two representations are indeed equivalent and can be converted back-and-forth.

In D138291#3936862, @nicolasvasilache wrote:

what was the rationale behind this check in the first place ?

Also, looking at the test you remove, the attribute looks unnatural to me:

%bcast = linalg.broadcast
    ins(%input:tensor<4x16xf32>)
    outs(%init:tensor<4x?x16xf32>)
    dimensions = [0, 2]
func.return %bcast : tensor<4x?x16xf32>

I would have expected to have

%bcast = linalg.broadcast
    ins(%input:tensor<4x16xf32>)
    outs(%init:tensor<4x?x16xf32>)
    dimensions = [1]
func.return %bcast : tensor<4x?x16xf32>

I.e. the dimension 1 of the result is added by the broadcast.
As opposed to the dimensions 0, and 2 of the output were already present in the input.

Thoughts?

[nicolasvasilache]

In D138291#3936883, @pifon2a wrote:

So, you would prefer having the attribute that specifies what dimensions are added, instead of what input dimensions are mapped to? These two representations are indeed equivalent and can be converted back-and-forth.

I just noted that it felt a bit unnatural to me and I am wondering if there is another interpretation than what I applied, i.e. "broadcast %input into a new tensor along result dimension 1" ?
I guess we could use an explanation in plain words as a judge of whether the syntax is straightforward or a bit convoluted?

[pifon2a]

In D138291#3936888, @nicolasvasilache wrote:

In D138291#3936883, @pifon2a wrote:

So, you would prefer having the attribute that specifies what dimensions are added, instead of what input dimensions are mapped to? These two representations are indeed equivalent and can be converted back-and-forth.

I just noted that it felt a bit unnatural to me and I am wondering if there is another interpretation than what I applied, i.e. "broadcast %input into a new tensor along result dimension 1" ?
I guess we could use an explanation in plain words as a judge of whether the syntax is straightforward or a bit convoluted?

I think it's just two different representation. We just borrowed the current one from mHLO, but the bcast there can do transpose, reshape and, possibly, double espresso shot. With no transpose present in the semantics of the op, we can actually switch from `dimensions = [0, 2]" to "dimensions = [1]". It will make it more obvious how many and what dimensions are added. The current representation make it clear how the input dims are mapped to the output dims, but it's less important, i think.

Anyway, I suggest we do that in a follow-up patch and submit this one for now.

[nicolasvasilache]

In D138291#3936925, @pifon2a wrote:

In D138291#3936888, @nicolasvasilache wrote:

In D138291#3936883, @pifon2a wrote:

So, you would prefer having the attribute that specifies what dimensions are added, instead of what input dimensions are mapped to? These two representations are indeed equivalent and can be converted back-and-forth.

I just noted that it felt a bit unnatural to me and I am wondering if there is another interpretation than what I applied, i.e. "broadcast %input into a new tensor along result dimension 1" ?
I guess we could use an explanation in plain words as a judge of whether the syntax is straightforward or a bit convoluted?

I think it's just two different representation. We just borrowed the current one from mHLO, but the bcast there can do transpose, reshape and, possibly, double espresso shot. With no transpose present in the semantics of the op, we can actually switch from `dimensions = [0, 2]" to "dimensions = [1]". It will make it more obvious how many and what dimensions are added. The current representation make it clear how the input dims are mapped to the output dims, but it's less important, i think.

Anyway, I suggest we do that in a follow-up patch and submit this one for now.

Yes, the mhlo "heritage" is visible in the semantics, great let's do as a it followup.

Thanks!

[olegshyshkov]

Thank you! I'll prepare the follow-up patch.

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

if (dim < 0 || dim >= initRank)
       return emitOpError() << "dimension " << idx
                            << " is out of range. expected range: [0, "
                            << initRank - 1 << "], got: " << dim;
 
     reverseDimMap[dim] = idx;
   }
 
   for (const auto &[idx, inputDimIdx] : llvm::enumerate(reverseDimMap)) {
-    if (inputDimIdx == kUnmappedDim) {
-      // This dimensions is being added. Should be statically known.
-      if (ShapedType::isDynamic(initShape[idx]))
-        return emitOpError()
-               << "init dim " << idx
-               << " can't be dynamic, because it's not matched to input";
-    } else {
+    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];
     }

mlir/test/Dialect/Linalg/invalid.mlir

%bcast = linalg.broadcast
       ins(%input:tensor<4x16xf32>)
       outs(%init:tensor<5x8x16xf32>)
       dimensions = [0, 2]
   func.return %bcast : tensor<5x8x16xf32>
 }
 
 // -----
 
-func.func @broadcast_added_dynamic_mismatch(
-    %input: tensor<4x16xf32>, %init: tensor<4x?x16xf32>)
-    -> tensor<4x?x16xf32> {
-  // expected-error @+1 {{'linalg.broadcast' op init dim 1 can't be dynamic, because it's not matched to input}}
-  %bcast = linalg.broadcast
-      ins(%input:tensor<4x16xf32>)
-      outs(%init:tensor<4x?x16xf32>)
-      dimensions = [0, 2]
-  func.return %bcast : tensor<4x?x16xf32>
-}
-
-// -----
-
 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]
   func.return %bcast : tensor<4x?x16xf32>
 }