fix: add reassociation indices for expand_shape op

I noticed a bug when compiling the Pytorch MLP that if the reassociation map for the expand shape op that is used to canonicalize the linalg.broadcast didn't include reassociation indices.

I guess it was invoking default op builders that don't require the map, and maybe other canonicalization patterns were taking precedence? Anyway, added a failing test.

PiperOrigin-RevId: 798365675

Author: asraa

lib/Transforms/DropUnitDims/BUILD

@@ -12,6 +12,7 @@ cc_library(
     hdrs = ["DropUnitDims.h"],
     deps = [
         ":pass_inc_gen",
+        "@heir//lib/Utils:TensorUtils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithUtils",
         "@llvm-project//mlir:IR",

lib/Transforms/DropUnitDims/DropUnitDims.cpp

@@ -4,6 +4,7 @@
 #include <cstdint>
 #include <utility>
 
+#include "lib/Utils/TensorUtils.h"
 #include "llvm/include/llvm/ADT/STLExtras.h"              // from @llvm-project
 #include "llvm/include/llvm/ADT/SmallVector.h"            // from @llvm-project
 #include "llvm/include/llvm/ADT/SmallVectorExtras.h"      // from @llvm-project
@@ -36,40 +37,6 @@ namespace heir {
 
 namespace {
 
-/// The following functions are copied from
-/// llvm-project/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp, where they
-/// are in an anonymous namespace.
-
-/// Returns reassociation indices for collapsing/expanding a
-/// tensor of rank `rank` at positions in `positions`.
-static SmallVector<ReassociationIndices> getReassociationForReshapeAtDim(
-    int64_t rank, ArrayRef<int64_t> positions) {
-  SmallVector<ReassociationIndices> reassociation;
-  reassociation.reserve(rank - positions.size());
-
-  llvm::DenseMap<int64_t, bool> positionsMap;
-  for (int64_t pos : positions) {
-    positionsMap[pos] = true;
-  }
-  auto isUnitDim = [&](int64_t dim) { return positionsMap.contains(dim); };
-
-  ReassociationIndices reassociationGroup;
-  unsigned dim = 0;
-  while (dim < rank && isUnitDim(dim)) reassociationGroup.push_back(dim++);
-  while (dim < rank) {
-    assert(!isUnitDim(dim) && "expected non unit-extent");
-    reassociationGroup.push_back(dim);
-    ++dim;
-    // Fold all following dimensions that are unit-extent.
-    while (dim < rank && isUnitDim(dim)) {
-      reassociationGroup.push_back(dim++);
-    }
-    reassociation.push_back(reassociationGroup);
-    reassociationGroup.clear();
-  }
-  return reassociation;
-}
-
 /// Collapse the given `value` so that the type matches the type of
 /// `origOutput`.
 static Value collapseValue(RewriterBase& rewriter, Location loc, Value operand,

lib/Transforms/LinalgCanonicalizations/BUILD

@@ -14,6 +14,7 @@ cc_library(
     ],
     deps = [
         ":pass_inc_gen",
+        "@heir//lib/Utils:TensorUtils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:DialectUtils",

lib/Transforms/LinalgCanonicalizations/LinalgCanonicalizations.cpp

@@ -3,6 +3,7 @@
 #include <cstdint>
 #include <utility>
 
+#include "lib/Utils/TensorUtils.h"
 #include "llvm/include/llvm/ADT/STLExtras.h"             // from @llvm-project
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
 #include "mlir/include/mlir/Dialect/Linalg/IR/Linalg.h"  // from @llvm-project
@@ -272,8 +273,13 @@ struct BroadcastToExpandShape
         broadcastOp.getInit().getType(), broadcastOp.getInput().getType());
     if (res != SliceVerificationResult::Success) return failure();
 
+    SmallVector<ReassociationIndices> expandingMap =
+        getReassociationForReshapeAtDim(
+            broadcastOp.getInit().getType().getRank(),
+            broadcastOp.getDimensions());
     rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
-        broadcastOp, broadcastOp.getInit().getType(), broadcastOp.getInput());
+        broadcastOp, broadcastOp.getInit().getType(), broadcastOp.getInput(),
+        expandingMap);
     return success();
   }
 };

lib/Utils/TensorUtils.cpp

@@ -1,8 +1,10 @@
 #include "lib/Utils/TensorUtils.h"
 
+#include <cassert>
 #include <cstdint>
 
 #include "llvm/include/llvm/ADT/SmallVectorExtras.h"  // from @llvm-project
+#include "mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Utils/StaticValueUtils.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypes.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/OpDefinition.h"  // from @llvm-project
@@ -48,5 +50,33 @@ SmallVector<int64_t> getIndicesFromRowMajorShape(int64_t flattenedIndex,
   return indices;
 }
 
+SmallVector<ReassociationIndices> getReassociationForReshapeAtDim(
+    int64_t rank, ArrayRef<int64_t> positions) {
+  SmallVector<ReassociationIndices> reassociation;
+  reassociation.reserve(rank - positions.size());
+
+  llvm::DenseMap<int64_t, bool> positionsMap;
+  for (int64_t pos : positions) {
+    positionsMap[pos] = true;
+  }
+  auto isUnitDim = [&](int64_t dim) { return positionsMap.contains(dim); };
+
+  ReassociationIndices reassociationGroup;
+  unsigned dim = 0;
+  while (dim < rank && isUnitDim(dim)) reassociationGroup.push_back(dim++);
+  while (dim < rank) {
+    assert(!isUnitDim(dim) && "expected non unit-extent");
+    reassociationGroup.push_back(dim);
+    ++dim;
+    // Fold all following dimensions that are unit-extent.
+    while (dim < rank && isUnitDim(dim)) {
+      reassociationGroup.push_back(dim++);
+    }
+    reassociation.push_back(reassociationGroup);
+    reassociationGroup.clear();
+  }
+  return reassociation;
+}
+
 }  // namespace heir
 }  // namespace mlir

lib/Utils/TensorUtils.h

@@ -3,6 +3,7 @@
 
 #include <cstdint>
 
+#include "mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypes.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/OpDefinition.h"  // from @llvm-project
 #include "mlir/include/mlir/Support/LLVM.h"     // from @llvm-project
@@ -20,6 +21,15 @@ FailureOr<int64_t> getFlattenedIndex(RankedTensorType tensorType,
 SmallVector<int64_t> getIndicesFromRowMajorShape(int64_t flattenedIndex,
                                                  SmallVector<int64_t> shape);
 
+/// The following functions are copied from
+/// llvm-project/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp, where they
+/// are in an anonymous namespace.
+
+/// Returns reassociation indices for collapsing/expanding a
+/// tensor of rank `rank` at positions in `positions`.
+SmallVector<ReassociationIndices> getReassociationForReshapeAtDim(
+    int64_t rank, ArrayRef<int64_t> positions);
+
 }  // namespace heir
 }  // namespace mlir
 

tests/Transforms/linalg_canonicalizations/broadcast_to_expand.mlir

@@ -42,3 +42,17 @@ module {
     func.return %collapsed_1 : tensor<2x3xf32>
   }
 }
+
+// -----
+
+module {
+  // CHECK: func @main
+  // CHECK-SAME: %[[arg0:.*]]: tensor<2x3xf32>
+  // CHECK: %[[v0:.*]] = tensor.expand_shape %[[arg0]] {{\[\[}}0, 1], [2, 3]]
+  // CHECK: return %[[v0]] : tensor<1x2x3x1xf32>
+  func.func @main(%arg1 : tensor<2x3xf32>) -> (tensor<1x2x3x1xf32>) {
+    %2 = tensor.empty() : tensor<1x2x3x1xf32>
+    %broadcasted = linalg.broadcast ins(%arg1 : tensor<2x3xf32>) outs(%2 : tensor<1x2x3x1xf32>) dimensions = [0, 3]
+    func.return %broadcasted : tensor<1x2x3x1xf32>
+  }
+}