[mlir][arith] Fold redundant broadcasts

[awarzynski]

Summary

This revision folds addi(vector.broadcast(x), vector.broadcast(y)) as
vector.broadcast(addi(x, y)).

[awarzynski]

The idea for this canonicalization came up in https://github.com/openxla/iree/issues/12991. I wasn't entirely sure whether it belongs in the Airth or the Vector dialect, probably a bit of both.

I'm also not sure about the CMake dependencies (it's a bit tricky without a good architectual overview), so I'd appreciate suggestions how to improve this. For a bit of context, here's the build error that I get without the CMake changes:

Undefined symbols for architecture arm64:
  "mlir::detail::TypeIDResolver<mlir::vector::BroadcastOp, void>::id", referenced from:
      llvm::DefaultDoCastIfPossible<mlir::vector::BroadcastOp, mlir::Operation*, llvm::CastInfo<mlir::vector::BroadcastOp, mlir::Operation*, void>>::doCastIfPossible(mlir::Operation*) in libMLIRArithDialect.a(ArithOps.cpp.o)
  "mlir::vector::BroadcastOp::getODSResults(unsigned int)", referenced from:
      (anonymous namespace)::AddIVectorBroadcast::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&) const in libMLIRArithDialect.a(ArithOps.cpp.o)
  "mlir::vector::BroadcastOp::getODSOperands(unsigned int)", referenced from:
      (anonymous namespace)::AddIVectorBroadcast::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&) const in libMLIRArithDialect.a(ArithOps.cpp.o)
  "mlir::vector::BroadcastOp::build(mlir::OpBuilder&, mlir::OperationState&, mlir::TypeRange, mlir::ValueRange, llvm::ArrayRef<mlir::NamedAttribute>)", referenced from:
      mlir::vector::BroadcastOp mlir::OpBuilder::create<mlir::vector::BroadcastOp, llvm::SmallVector<mlir::Type, 4u>&, llvm::SmallVector<mlir::Value, 4u>&, llvm::SmallVector<mlir::NamedAttribute, 4u>&>(mlir::Location, llvm::SmallVector<mlir::Type, 4u>&, llvm::SmallVector<mlir::Value, 4u>&, llvm::SmallVector<mlir::NamedAttribute, 4u>&) in libMLIRArithDialect.a(ArithOps.cpp.o)

Thanks for taking a look!

[kuhar]

I'm not sure if this belongs to general arith canonicalizations; the vector dialect depends on artih, so I'd expect this transform to be archored at vector.broadcast.
Could we have a more general pattern that fold any elementwise op that supports both scalars and vectors when the vector operand(s) are either broadcasted or splat constants?

I'm also not sure about the CMake dependencies (it's a bit tricky without a good architectual overview), so I'd appreciate suggestions how to improve this.

I typically check with bazel in cases like this -- it's more picky than CMake when it comes to implicit/circular dependencies, which makes it tell when some are missing.

[awarzynski]

@kuhar , thanks for taking a look and apologies for the delay getting back to this. I will abandon this in favor of a more general approach: https://reviews.llvm.org/D152812. Cheers for the feedback!

mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td

 //===- ArithPatterns.td - Arith dialect patterns -*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef ARITH_PATTERNS
 #define ARITH_PATTERNS
 
 include "mlir/IR/PatternBase.td"
 include "mlir/Dialect/Arith/IR/ArithOps.td"
+include "mlir/Dialect/Vector/IR/VectorOps.td"
 
 // Create zero attribute of type matching the argument's type.
 def GetZeroAttr : NativeCodeCall<"$_builder.getZeroAttr($0.getType())">;
 
 // Add two integer attributes and create a new one with the result.
 def AddIntAttrs : NativeCodeCall<"addIntegerAttrs($_builder, $0, $1, $2)">;
 
 // Subtract two integer attributes and createa a new one with the result.
 def SubIntAttrs : NativeCodeCall<"subIntegerAttrs($_builder, $0, $1, $2)">;
 
 //===----------------------------------------------------------------------===//
 // AddIOp
 //===----------------------------------------------------------------------===//
 
 // addi is commutative and will be canonicalized to have its constants appear
 // as the second operand.
 
 // addi(addi(x, c0), c1) -> addi(x, c0 + c1)
 def AddIAddConstant :
     Pat<(Arith_AddIOp:$res
           (Arith_AddIOp $x, (ConstantLikeMatcher APIntAttr:$c0)),
           (ConstantLikeMatcher APIntAttr:$c1)),
         (Arith_AddIOp $x, (Arith_ConstantOp (AddIntAttrs $res, $c0, $c1)))>;
 
+// addi(vector.broadcast(x), vector.broadcast(y)) -> vector.broadcast(addi(x, y))
+def AddIVectorBroadcast :
+    Pat<(Arith_AddIOp
+         (Vector_BroadcastOp $x), (Vector_BroadcastOp $y)),
+         (Vector_BroadcastOp (Arith_AddIOp $x, $y))>;
+
 // addi(subi(x, c0), c1) -> addi(x, c1 - c0)
 def AddISubConstantRHS :
     Pat<(Arith_AddIOp:$res
           (Arith_SubIOp $x, (ConstantLikeMatcher APIntAttr:$c0)),
           (ConstantLikeMatcher APIntAttr:$c1)),
         (Arith_AddIOp $x, (Arith_ConstantOp (SubIntAttrs $res, $c1, $c0)))>;
 
 // addi(subi(c0, x), c1) -> subi(c0 + c1, x)

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

 //===- ArithOps.cpp - MLIR Arith dialect ops implementation -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include <cassert>
 #include <cstdint>
 #include <utility>
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/CommonFolders.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 
OpFoldResult arith::AddIOp::fold(FoldAdaptor adaptor) {
 
   return constFoldBinaryOp<IntegerAttr>(
       adaptor.getOperands(),
       [](APInt a, const APInt &b) { return std::move(a) + b; });
 }
 
 void arith::AddIOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                 MLIRContext *context) {
-  patterns.add<AddIAddConstant, AddISubConstantRHS, AddISubConstantLHS,
-               AddIMulNegativeOneRhs, AddIMulNegativeOneLhs>(context);
+  patterns
+      .add<AddIAddConstant, AddISubConstantRHS, AddISubConstantLHS,
+           AddIMulNegativeOneRhs, AddIMulNegativeOneLhs, AddIVectorBroadcast>(
+          context);
 }
 
 //===----------------------------------------------------------------------===//
 // AddUIExtendedOp
 //===----------------------------------------------------------------------===//
 
 std::optional<SmallVector<int64_t, 4>>
 arith::AddUIExtendedOp::getShapeForUnroll() {

mlir/test/Dialect/Arith/canonicalize.mlir

 // CHECK-LABEL: @foldShrsi0
 // CHECK-SAME: (%[[ARG:.*]]: i64)
 //       CHECK:   return %[[ARG]] : i64
 func.func @foldShrsi0(%x : i64) -> i64 {
   %c0 = arith.constant 0 : i64
   %r = arith.shrsi %x, %c0 : i64
   return %r : i64
 }
+
+// CHECK-LABEL:   addWithBroadcast
+// CHECK-SAME:    %[[VAL_0:.*]]: index,
+// CHECK-SAME:    %[[VAL_1:.*]]: index) -> vector<1x4xindex> {
+// CHECK:         %[[VAL_2:.*]] = arith.addi %[[VAL_0]], %[[VAL_1]] : index
+// CHECK:         %[[VAL_3:.*]] = vector.broadcast %[[VAL_2]] : index to vector<1x4xindex>
+// CHECK:         return %[[VAL_3]] : vector<1x4xindex>
+func.func @addWithBroadcast( %arg1: index, %arg2: index) -> vector<1x4xindex> {
+  %0 = vector.broadcast %arg1 : index to vector<1x4xindex>
+  %1 = vector.broadcast %arg2 : index to vector<1x4xindex>
+  %2 = arith.addi %0, %1 : vector<1x4xindex>
+  return %2 : vector<1x4xindex>
+}

mlir/unittests/Conversion/PDLToPDLInterp/CMakeLists.txt

 add_mlir_unittest(MLIRPDLToPDLInterpTests
   RootOrderingTest.cpp
 )
 target_link_libraries(MLIRPDLToPDLInterpTests
   PRIVATE
-  MLIRArithDialect
+  MLIRArithTransforms
   MLIRPDLToPDLInterp
 )

mlir/unittests/Dialect/Transform/CMakeLists.txt

 add_mlir_unittest(MLIRTransformDialectTests
   BuildOnlyExtensionTest.cpp
 )
 target_link_libraries(MLIRTransformDialectTests
   PRIVATE
+  MLIRArithTransforms
   MLIRFuncDialect
   MLIRTransformDialect
 )

mlir/unittests/Interfaces/CMakeLists.txt

 add_mlir_unittest(MLIRInterfacesTests
   ControlFlowInterfacesTest.cpp
   DataLayoutInterfacesTest.cpp
   InferIntRangeInterfaceTest.cpp
   InferTypeOpInterfaceTest.cpp
 )
 
 target_link_libraries(MLIRInterfacesTests
   PRIVATE
+  MLIRArithTransforms
   MLIRControlFlowInterfaces
   MLIRDataLayoutInterfaces
   MLIRDLTIDialect
   MLIRFuncDialect
   MLIRInferIntRangeInterface
   MLIRInferTypeOpInterface
   MLIRParser
 )

mlir/unittests/Pass/CMakeLists.txt

 add_mlir_unittest(MLIRPassTests
   AnalysisManagerTest.cpp
   PassManagerTest.cpp
   PassPipelineParserTest.cpp
 )
 target_link_libraries(MLIRPassTests
   PRIVATE
+  MLIRArithTransforms
   MLIRFuncDialect
   MLIRPass)