[MLIR][Shape] Support >2 args in `shape.broadcast` folder #126808

mtsokol · 2025-02-11T22:08:31Z

Hi!

As the title says, this PR adds support for >2 arguments in shape.broadcast folder by sequentially calling getBroadcastedShape.

llvmbot · 2025-02-11T22:09:03Z

@llvm/pr-subscribers-mlir

Author: Mateusz Sokół (mtsokol)

Changes

Hi!

As the title says, this PR adds support for >2 arguments in shape.broadcast folder by sequentially calling getBroadcastedShape.

Full diff: https://github.com/llvm/llvm-project/pull/126808.diff

3 Files Affected:

(modified) mlir/lib/Dialect/Shape/IR/Shape.cpp (+21-13)
(modified) mlir/lib/Dialect/Traits.cpp (+1-1)
(modified) mlir/test/Dialect/Shape/canonicalize.mlir (+13)

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 65efc88e9c403..daa33ea865a5c 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -649,24 +649,32 @@ OpFoldResult BroadcastOp::fold(FoldAdaptor adaptor) {
     return getShapes().front();
   }
 
-  // TODO: Support folding with more than 2 input shapes
-  if (getShapes().size() > 2)
+  if (!adaptor.getShapes().front())
     return nullptr;
 
-  if (!adaptor.getShapes()[0] || !adaptor.getShapes()[1])
-    return nullptr;
-  auto lhsShape = llvm::to_vector<6>(
-      llvm::cast<DenseIntElementsAttr>(adaptor.getShapes()[0])
-          .getValues<int64_t>());
-  auto rhsShape = llvm::to_vector<6>(
-      llvm::cast<DenseIntElementsAttr>(adaptor.getShapes()[1])
+  auto firstShape = llvm::to_vector<6>(
+      llvm::cast<DenseIntElementsAttr>(adaptor.getShapes().front())
           .getValues<int64_t>());
+
   SmallVector<int64_t, 6> resultShape;
+  resultShape.clear();
+  std::copy(firstShape.begin(), firstShape.end(), std::back_inserter(resultShape));
 
-  // If the shapes are not compatible, we can't fold it.
-  // TODO: Fold to an "error".
-  if (!OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
-    return nullptr;
+  for (auto next : adaptor.getShapes().drop_front()) {
+    if (!next)
+      return nullptr;
+    auto nextShape = llvm::to_vector<6>(
+        llvm::cast<DenseIntElementsAttr>(next).getValues<int64_t>());
+
+    SmallVector<int64_t, 6> tmpShape;
+    // If the shapes are not compatible, we can't fold it.
+    // TODO: Fold to an "error".
+    if (!OpTrait::util::getBroadcastedShape(resultShape, nextShape, tmpShape))
+      return nullptr;
+
+    resultShape.clear();
+    std::copy(tmpShape.begin(), tmpShape.end(), std::back_inserter(resultShape));
+  }
 
   Builder builder(getContext());
   return builder.getIndexTensorAttr(resultShape);
diff --git a/mlir/lib/Dialect/Traits.cpp b/mlir/lib/Dialect/Traits.cpp
index a7aa25eae2644..6e62a33037eb8 100644
--- a/mlir/lib/Dialect/Traits.cpp
+++ b/mlir/lib/Dialect/Traits.cpp
@@ -84,7 +84,7 @@ bool OpTrait::util::getBroadcastedShape(ArrayRef<int64_t> shape1,
     if (ShapedType::isDynamic(*i1) || ShapedType::isDynamic(*i2)) {
       // One or both dimensions is unknown. Follow TensorFlow behavior:
       // - If either dimension is greater than 1, we assume that the program is
-      //   correct, and the other dimension will be broadcast to match it.
+      //   correct, and the other dimension will be broadcasted to match it.
       // - If either dimension is 1, the other dimension is the output.
       if (*i1 > 1) {
         *iR = *i1;
diff --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index cf439c9c1b854..9ed4837a2fe7e 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -86,6 +86,19 @@ func.func @broadcast() -> !shape.shape {
 
 // -----
 
+// Variadic case including extent tensors.
+// CHECK-LABEL: @broadcast_variadic
+func.func @broadcast_variadic() -> !shape.shape {
+  // CHECK: shape.const_shape [7, 2, 10] : !shape.shape
+  %0 = shape.const_shape [2, 1] : tensor<2xindex>
+  %1 = shape.const_shape [7, 2, 1] : tensor<3xindex>
+  %2 = shape.const_shape [1, 10] : tensor<2xindex>
+  %3 = shape.broadcast %0, %1, %2 : tensor<2xindex>, tensor<3xindex>, tensor<2xindex> -> !shape.shape
+  return %3 : !shape.shape
+}
+
+// -----
+
 // Rhs is a scalar.
 // CHECK-LABEL: func @f
 func.func @f(%arg0 : !shape.shape) -> !shape.shape {

llvmbot · 2025-02-11T22:09:04Z

@llvm/pr-subscribers-mlir-shape

Author: Mateusz Sokół (mtsokol)

Changes

Hi!

As the title says, this PR adds support for >2 arguments in shape.broadcast folder by sequentially calling getBroadcastedShape.

Full diff: https://github.com/llvm/llvm-project/pull/126808.diff

3 Files Affected:

(modified) mlir/lib/Dialect/Shape/IR/Shape.cpp (+21-13)
(modified) mlir/lib/Dialect/Traits.cpp (+1-1)
(modified) mlir/test/Dialect/Shape/canonicalize.mlir (+13)

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 65efc88e9c403..daa33ea865a5c 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -649,24 +649,32 @@ OpFoldResult BroadcastOp::fold(FoldAdaptor adaptor) {
     return getShapes().front();
   }
 
-  // TODO: Support folding with more than 2 input shapes
-  if (getShapes().size() > 2)
+  if (!adaptor.getShapes().front())
     return nullptr;
 
-  if (!adaptor.getShapes()[0] || !adaptor.getShapes()[1])
-    return nullptr;
-  auto lhsShape = llvm::to_vector<6>(
-      llvm::cast<DenseIntElementsAttr>(adaptor.getShapes()[0])
-          .getValues<int64_t>());
-  auto rhsShape = llvm::to_vector<6>(
-      llvm::cast<DenseIntElementsAttr>(adaptor.getShapes()[1])
+  auto firstShape = llvm::to_vector<6>(
+      llvm::cast<DenseIntElementsAttr>(adaptor.getShapes().front())
           .getValues<int64_t>());
+
   SmallVector<int64_t, 6> resultShape;
+  resultShape.clear();
+  std::copy(firstShape.begin(), firstShape.end(), std::back_inserter(resultShape));
 
-  // If the shapes are not compatible, we can't fold it.
-  // TODO: Fold to an "error".
-  if (!OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
-    return nullptr;
+  for (auto next : adaptor.getShapes().drop_front()) {
+    if (!next)
+      return nullptr;
+    auto nextShape = llvm::to_vector<6>(
+        llvm::cast<DenseIntElementsAttr>(next).getValues<int64_t>());
+
+    SmallVector<int64_t, 6> tmpShape;
+    // If the shapes are not compatible, we can't fold it.
+    // TODO: Fold to an "error".
+    if (!OpTrait::util::getBroadcastedShape(resultShape, nextShape, tmpShape))
+      return nullptr;
+
+    resultShape.clear();
+    std::copy(tmpShape.begin(), tmpShape.end(), std::back_inserter(resultShape));
+  }
 
   Builder builder(getContext());
   return builder.getIndexTensorAttr(resultShape);
diff --git a/mlir/lib/Dialect/Traits.cpp b/mlir/lib/Dialect/Traits.cpp
index a7aa25eae2644..6e62a33037eb8 100644
--- a/mlir/lib/Dialect/Traits.cpp
+++ b/mlir/lib/Dialect/Traits.cpp
@@ -84,7 +84,7 @@ bool OpTrait::util::getBroadcastedShape(ArrayRef<int64_t> shape1,
     if (ShapedType::isDynamic(*i1) || ShapedType::isDynamic(*i2)) {
       // One or both dimensions is unknown. Follow TensorFlow behavior:
       // - If either dimension is greater than 1, we assume that the program is
-      //   correct, and the other dimension will be broadcast to match it.
+      //   correct, and the other dimension will be broadcasted to match it.
       // - If either dimension is 1, the other dimension is the output.
       if (*i1 > 1) {
         *iR = *i1;
diff --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index cf439c9c1b854..9ed4837a2fe7e 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -86,6 +86,19 @@ func.func @broadcast() -> !shape.shape {
 
 // -----
 
+// Variadic case including extent tensors.
+// CHECK-LABEL: @broadcast_variadic
+func.func @broadcast_variadic() -> !shape.shape {
+  // CHECK: shape.const_shape [7, 2, 10] : !shape.shape
+  %0 = shape.const_shape [2, 1] : tensor<2xindex>
+  %1 = shape.const_shape [7, 2, 1] : tensor<3xindex>
+  %2 = shape.const_shape [1, 10] : tensor<2xindex>
+  %3 = shape.broadcast %0, %1, %2 : tensor<2xindex>, tensor<3xindex>, tensor<2xindex> -> !shape.shape
+  return %3 : !shape.shape
+}
+
+// -----
+
 // Rhs is a scalar.
 // CHECK-LABEL: func @f
 func.func @f(%arg0 : !shape.shape) -> !shape.shape {

mtsokol · 2025-02-11T22:09:55Z

mlir/lib/Dialect/Shape/IR/Shape.cpp

+  for (auto next : adaptor.getShapes().drop_front()) {
+    if (!next)
+      return nullptr;
+    auto nextShape = llvm::to_vector<6>(


In the getBroadcastedShape implementation shape vector size is hardcoded to 6, so I did it similarly here. Does it make sense? Looks like an arbitrary value from the outside.

Yes, semi. If I recall it was either the default elsewhere in an ML framework where this was used or the max rank along set of ML models. But it is a bit arbitrary. Elsewhere folks also use the default of SmallVector. (The latter is probably a little bit more arbitrary, but neither is very fine tuned).

mtsokol · 2025-02-11T22:11:58Z

mlir/lib/Dialect/Shape/IR/Shape.cpp

+    resultShape.clear();
+    std::copy(tmpShape.begin(), tmpShape.end(), std::back_inserter(resultShape));


I followed getBroadcastedShape implementation and I'm not sure if it's the best way to handle Vectors/Shapes here, so a penny for your thoughts!

github-actions · 2025-02-11T22:12:16Z

✅ With the latest revision this PR passed the C/C++ code formatter.

jpienaar

Overall looks good thanks

jpienaar · 2025-02-19T13:54:01Z

mlir/lib/Dialect/Shape/IR/Shape.cpp

+  for (auto next : adaptor.getShapes().drop_front()) {
+    if (!next)
+      return nullptr;
+    auto nextShape = llvm::to_vector<6>(


Yes, semi. If I recall it was either the default elsewhere in an ML framework where this was used or the max rank along set of ML models. But it is a bit arbitrary. Elsewhere folks also use the default of SmallVector. (The latter is probably a little bit more arbitrary, but neither is very fine tuned).

jpienaar · 2025-02-19T13:58:27Z

mlir/lib/Dialect/Shape/IR/Shape.cpp

  SmallVector<int64_t, 6> resultShape;
+  resultShape.clear();
+  std::copy(firstShape.begin(), firstShape.end(),


Why is firstShape needed vs directly initializing resultShape?

I think it isn't needed here - updated!

mlir/test/Dialect/Shape/canonicalize.mlir

jpienaar

LG, modulo checking formatting.

jpienaar · 2025-03-07T17:25:46Z

mlir/lib/Dialect/Shape/IR/Shape.cpp

+      return nullptr;
+
+    resultShape.clear();
+    std::copy(tmpShape.begin(), tmpShape.end(),


Was this what clang-format produced?

@jpienaar Yes, that's correct - it was produced by a clang-format. Here's another place where std::copy is formatted the same way:

llvm-project/clang/include/clang/Lex/MacroInfo.h

Lines 536 to 537 in 74ca579

std::copy(Overrides.begin(), Overrides.end(),

reinterpret_cast<ModuleMacro **>(this + 1));

mtsokol · 2025-03-13T13:05:48Z

@jpienaar Should we merge it?

mtsokol · 2025-04-03T09:39:22Z

I think it's ready, right?

* origin/main: (199 commits) [NFC][AsmPrinter] Refactor AsmPrinter and AArch64AsmPrinter to prepare for jump table partitions on aarch64 (llvm#125993) [HEXAGON] Fix corner cases for hwloops pass (llvm#135439) [flang] Handle volatility in lowering and codegen (llvm#135311) [MLIR][Shape] Support >2 args in `shape.broadcast` folder (llvm#126808) [DirectX] Use scalar arguments for @llvm.dx.dot intrinsics (llvm#134570) Remove the redundant check for "WeakPtr" in isSmartPtrClass to fix the issue 135612. (llvm#135629) [BOLT] Support relative vtable (llvm#135449) [flang] Fix linking to libMLIR (llvm#135483) [AsmPrinter] Link .section_sizes to the correct section (llvm#135583) [ctxprof] Handle instrumenting functions with `musttail` calls (llvm#135121) [SystemZ] Consider VST/VL as SimpleBDXStore/Load (llvm#135623) [libc++][CI] Pin the XCode version. (llvm#135412) [lldb-dap] Fix win32 build. (llvm#135638) [Interp] Mark inline-virtual.cpp as unsupported with ASan (llvm#135402) [libc++] Removes the _LIBCPP_VERBOSE_ABORT_NOT_NOEXCEPT macro. (llvm#135494) [CVP] Add tests for ucmp/scmp with switch (NFC) [mlir][tosa] Align AbsOp example variable names (llvm#135268) [mlir][tosa] Align AddOp examples to spec (llvm#135266) [mlir][tosa] Align RFFT2d and FFT2d operator examples (llvm#135261) [flang][OpenMP][HLFIR] Support vector subscripted array sections for DEPEND (llvm#133892) ...

Hi! As the title says, this PR adds support for >2 arguments in `shape.broadcast` folder by sequentially calling `getBroadcastedShape`.

llvmbot added mlir mlir:shape labels Feb 11, 2025

mtsokol commented Feb 11, 2025

View reviewed changes

mtsokol force-pushed the shape-broadcast-fold-vararg branch from c4a815d to c3cf613 Compare February 11, 2025 22:15

adam-smnk requested review from jpienaar and zero9178 February 19, 2025 12:55

jpienaar reviewed Feb 19, 2025

View reviewed changes

mtsokol added 2 commits February 19, 2025 16:32

[MLIR][Shape] Support >2 args in shape.broadcast folder

47398c2

Initialize resultShape directly

b259199

mtsokol force-pushed the shape-broadcast-fold-vararg branch from c3cf613 to b259199 Compare February 19, 2025 16:57

mtsokol requested a review from jpienaar February 20, 2025 09:34

jpienaar approved these changes Mar 7, 2025

View reviewed changes

makslevental merged commit df84aa8 into llvm:main Apr 14, 2025
8 checks passed

mtsokol deleted the shape-broadcast-fold-vararg branch April 14, 2025 19:11

		resultShape.clear();
		std::copy(tmpShape.begin(), tmpShape.end(), std::back_inserter(resultShape));

	std::copy(Overrides.begin(), Overrides.end(),
	reinterpret_cast<ModuleMacro **>(this + 1));

[MLIR][Shape] Support >2 args in shape.broadcast folder #126808

[MLIR][Shape] Support >2 args in shape.broadcast folder #126808

Uh oh!

Conversation

mtsokol commented Feb 11, 2025

Uh oh!

llvmbot commented Feb 11, 2025

Uh oh!

llvmbot commented Feb 11, 2025

Uh oh!

mtsokol Feb 11, 2025

Choose a reason for hiding this comment

Uh oh!

jpienaar Feb 19, 2025

Choose a reason for hiding this comment

Uh oh!

mtsokol Feb 11, 2025

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Feb 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jpienaar left a comment

Choose a reason for hiding this comment

Uh oh!

jpienaar Feb 19, 2025

Choose a reason for hiding this comment

Uh oh!

jpienaar Feb 19, 2025

Choose a reason for hiding this comment

Uh oh!

mtsokol Feb 19, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jpienaar left a comment

Choose a reason for hiding this comment

Uh oh!

jpienaar Mar 7, 2025

Choose a reason for hiding this comment

Uh oh!

mtsokol Mar 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mtsokol commented Mar 13, 2025

Uh oh!

mtsokol commented Apr 3, 2025

Uh oh!

Uh oh!

Uh oh!

[MLIR][Shape] Support >2 args in `shape.broadcast` folder #126808

[MLIR][Shape] Support >2 args in `shape.broadcast` folder #126808

github-actions bot commented Feb 11, 2025 •

edited

Loading

mtsokol Mar 9, 2025 •

edited

Loading