Move LinalgQuantized* passes to GlobalOptimization (iree-org#18287)

These passes are a better fit for GlobalOptimization given they have to
do with Linalg ops. Additionally this turns on the passes for all
frontends now that other frontends are starting to generate these ops
too.

compiler/plugins/input/TOSA/InputConversion/Passes.cpp

@@ -65,12 +65,6 @@ void buildTOSAInputConversionPassPipeline(OpPassManager &passManager) {
       iree_compiler::createStripSignednessPass());
   passManager.addNestedPass<func::FuncOp>(mlir::createCanonicalizerPass());
 
-  passManager.addNestedPass<func::FuncOp>(
-      InputConversion::createLinalgQuantizedMatmulToMatmulPass());
-  passManager.addNestedPass<func::FuncOp>(
-      InputConversion::createLinalgQuantizedConvToConvPass());
-  passManager.addNestedPass<func::FuncOp>(mlir::createCanonicalizerPass());
-
   //----------------------------------------------------------------------------
   // Entry dialect cleanup
   //----------------------------------------------------------------------------

compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel

@@ -61,6 +61,8 @@ iree_compiler_cc_library(
         "OptimizeNumerics.cpp",
         "Passes.cpp",
         "PropagateLinalgTranspose.cpp",
+        "QuantizedConvToConv.cpp",
+        "QuantizedMatmulToMatmul.cpp",
         "RaiseSpecialOps.cpp",
         "RemoveZeroExtentTensors.cpp",
         "SimplifyPackUnpack.cpp",

compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt

@@ -57,6 +57,8 @@ iree_cc_library(
     "OptimizeNumerics.cpp"
     "Passes.cpp"
     "PropagateLinalgTranspose.cpp"
+    "QuantizedConvToConv.cpp"
+    "QuantizedMatmulToMatmul.cpp"
     "RaiseSpecialOps.cpp"
     "RemoveZeroExtentTensors.cpp"
     "SimplifyPackUnpack.cpp"

compiler/src/iree/compiler/GlobalOptimization/Passes.cpp

@@ -91,6 +91,9 @@ void buildGlobalOptimizationPassPipeline(
 
   // Preprocessing passes to get the program into a canonical state.
   FunctionLikeNest(mainPassManager)
+      .addPass(createLinalgQuantizedConvToConvPass)
+      .addPass(createLinalgQuantizedMatmulToMatmulPass)
+      .addPass(IREE::Flow::createCanonicalizerPass)
       .addPass(createRemoveZeroExtentTensorsPass)
       .addPass(createDetachElementwiseFromNamedOpsPass)
       .addPass(mlir::createLinalgNamedOpConversionPass)

compiler/src/iree/compiler/GlobalOptimization/Passes.td

@@ -114,6 +114,16 @@ def PropagateLinalgTransposePass :
   ];
 }
 
+def LinalgQuantizedConvToConvPass
+    : InterfacePass<"iree-global-opt-quantized-conv-to-conv", "mlir::FunctionOpInterface"> {
+  let summary = "lower quantized_conv to conv";
+}
+
+def LinalgQuantizedMatmulToMatmulPass
+    : InterfacePass<"iree-global-opt-quantized-matmul-to-matmul", "mlir::FunctionOpInterface"> {
+  let summary = "lower quantized_matmul to matmul";
+}
+
 def RaiseSpecialOpsPass :
     Pass<"iree-global-opt-raise-special-ops", ""> {
   let summary = "Raises special ops like softmax to the high level linalg.ext representation.";

compiler/src/iree/compiler/InputConversion/Common/QuantizedConvToConv.cpp → compiler/src/iree/compiler/GlobalOptimization/QuantizedConvToConv.cpp

@@ -4,8 +4,8 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include "iree/compiler/InputConversion/Common/Passes.h"
-#include "iree/compiler/InputConversion/Common/Utils.h"
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "iree/compiler/GlobalOptimization/Utils.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
@@ -21,10 +21,10 @@
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
-namespace mlir::iree_compiler::InputConversion {
+namespace mlir::iree_compiler::GlobalOptimization {
 
 #define GEN_PASS_DEF_LINALGQUANTIZEDCONVTOCONVPASS
-#include "iree/compiler/InputConversion/Common/Passes.h.inc"
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
 
 namespace {
 
@@ -357,4 +357,4 @@ class LinalgQuantizedConvToConvPass final
 };
 
 } // namespace
-} // namespace mlir::iree_compiler::InputConversion
+} // namespace mlir::iree_compiler::GlobalOptimization

compiler/src/iree/compiler/InputConversion/Common/QuantizedMatmulToMatmul.cpp → compiler/src/iree/compiler/GlobalOptimization/QuantizedMatmulToMatmul.cpp

@@ -4,8 +4,8 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include "iree/compiler/InputConversion/Common/Passes.h"
-#include "iree/compiler/InputConversion/Common/Utils.h"
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "iree/compiler/GlobalOptimization/Utils.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
@@ -18,10 +18,10 @@
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
-namespace mlir::iree_compiler::InputConversion {
+namespace mlir::iree_compiler::GlobalOptimization {
 
 #define GEN_PASS_DEF_LINALGQUANTIZEDMATMULTOMATMULPASS
-#include "iree/compiler/InputConversion/Common/Passes.h.inc"
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
 
 namespace {
 
@@ -193,4 +193,4 @@ class LinalgQuantizedMatmulToMatmulPass final
 };
 
 } // namespace
-} // namespace mlir::iree_compiler::InputConversion
+} // namespace mlir::iree_compiler::GlobalOptimization

compiler/src/iree/compiler/GlobalOptimization/Utils.cpp

@@ -122,4 +122,66 @@ wrapConsecutiveOpsInDispatchRegion(RewriterBase &rewriter,
                                                         regionOp);
 }
 
+Value sumReduceDimensionSubset(ImplicitLocOpBuilder &rewriter, Value val,
+                               Type accETy, ArrayRef<bool> is_reduction) {
+  auto context = val.getContext();
+  RankedTensorType ty = llvm::cast<RankedTensorType>(val.getType());
+
+  llvm::SmallVector<int64_t> staticSizes;
+  SmallVector<Value> dynSizes;
+  for (int i = 0, s = is_reduction.size(); i < s; i++) {
+    if (is_reduction[i])
+      continue;
+
+    staticSizes.push_back(ty.getDimSize(i));
+    if (ty.isDynamicDim(i)) {
+      dynSizes.push_back(rewriter.create<tensor::DimOp>(val, i));
+    }
+  }
+
+  // Create a zero-filled accumulator.
+  Value initAcc =
+      rewriter.create<tensor::EmptyOp>(staticSizes, accETy, dynSizes);
+  Value zeroInt = rewriter.create<arith::ConstantIntOp>(0, accETy).getResult();
+  Value zeroAcc =
+      rewriter.create<linalg::FillOp>(zeroInt, initAcc).getResult(0);
+
+  SmallVector<AffineExpr> filterExprs(ty.getRank());
+  SmallVector<AffineExpr> outputExprs;
+  SmallVector<utils::IteratorType> iterators;
+
+  for (int i = 0, s = ty.getRank(); i < s; i++) {
+    if (!is_reduction[i]) {
+      auto expr = rewriter.getAffineDimExpr(iterators.size());
+      iterators.push_back(utils::IteratorType::parallel);
+
+      outputExprs.push_back(expr);
+      filterExprs[i] = expr;
+    }
+  }
+
+  for (int i = 0, s = filterExprs.size(); i < s; i++) {
+    if (!filterExprs[i]) {
+      auto expr = mlir::getAffineDimExpr(iterators.size(), context);
+      iterators.push_back(utils::IteratorType::reduction);
+      filterExprs[i] = expr;
+    }
+  }
+
+  SmallVector<AffineMap> affineMaps{
+      AffineMap::get(ty.getRank(), 0, filterExprs, context),
+      AffineMap::get(ty.getRank(), 0, outputExprs, context)};
+
+  return rewriter
+      .create<linalg::GenericOp>(
+          zeroAcc.getType(), ValueRange{val}, ValueRange{zeroAcc}, affineMaps,
+          iterators,
+          [=](OpBuilder &b, Location loc, ValueRange args) {
+            Value ext = b.create<arith::ExtSIOp>(loc, accETy, args[0]);
+            Value sum = b.create<arith::AddIOp>(loc, ext, args[1]);
+            b.create<linalg::YieldOp>(loc, sum);
+          })
+      .getResult(0);
+}
+
 } // namespace mlir::iree_compiler::GlobalOptimization

compiler/src/iree/compiler/GlobalOptimization/Utils.h

@@ -55,6 +55,10 @@ FailureOr<IREE::Flow::DispatchRegionOp>
 wrapConsecutiveOpsInDispatchRegion(RewriterBase &rewriter,
                                    SmallVector<Operation *> ops);
 
+// Reduce the input value along the reduction dimensions.
+Value sumReduceDimensionSubset(ImplicitLocOpBuilder &rewriter, Value val,
+                               Type accETy, ArrayRef<bool> is_reduction);
+
 } // namespace mlir::iree_compiler::GlobalOptimization
 
 #endif // IREE_COMPILER_GLOBALOPTIMIZATION_UTILS_H_

compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel

@@ -28,6 +28,8 @@ iree_lit_test_suite(
             "global_loop_invariant_code_motion.mlir",
             "hoist_into_globals.mlir",
             "infer_numeric_narrowing.mlir",
+            "linalg_quantized_conv_to_conv.mlir",
+            "linalg_quantized_matmul_to_matmul.mlir",
             "optimize_numerics.mlir",
             "propagate_linalg_transpose.mlir",
             "raise_special_ops.mlir",

compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt

@@ -26,6 +26,8 @@ iree_lit_test_suite(
     "global_loop_invariant_code_motion.mlir"
     "hoist_into_globals.mlir"
     "infer_numeric_narrowing.mlir"
+    "linalg_quantized_conv_to_conv.mlir"
+    "linalg_quantized_matmul_to_matmul.mlir"
     "optimize_numerics.mlir"
     "propagate_linalg_transpose.mlir"
     "raise_special_ops.mlir"

compiler/src/iree/compiler/InputConversion/Common/test/linalg_quantized_conv_to_conv.mlir → compiler/src/iree/compiler/GlobalOptimization/test/linalg_quantized_conv_to_conv.mlir

@@ -1,4 +1,4 @@
-// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-linalg-quantized-conv-to-conv))" -split-input-file %s | FileCheck %s
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-global-opt-quantized-conv-to-conv))" -split-input-file %s | FileCheck %s
 
 // CHECK-LABEL: func.func @conv2d_zp
 func.func @conv2d_zps(%arg0: tensor<1x14x16x5xi8>, %arg1: tensor<3x4x5x1024xi8>, %arg2: tensor<1x12x13x1024xi32>) -> tensor<1x12x13x1024xi32> {

compiler/src/iree/compiler/InputConversion/Common/test/linalg_quantized_matmul_to_matmul.mlir → compiler/src/iree/compiler/GlobalOptimization/test/linalg_quantized_matmul_to_matmul.mlir

@@ -1,4 +1,4 @@
-// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-linalg-quantized-matmul-to-matmul))" --split-input-file %s | FileCheck %s
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-global-opt-quantized-matmul-to-matmul))" --split-input-file %s | FileCheck %s
 
 // Tests -iree-linalg-quantized-matmul-to-matmul, converting linalg.quantized_matmul
 // ops to linalg.matmul ops plus additional arithmetic to account for any

compiler/src/iree/compiler/InputConversion/Common/BUILD.bazel

@@ -48,14 +48,10 @@ iree_compiler_cc_library(
         "IREEImportPublic.cpp",
         "ImportMLProgram.cpp",
         "Passes.cpp",
-        "QuantizedConvToConv.cpp",
-        "QuantizedMatmulToMatmul.cpp",
         "SanitizeModuleNames.cpp",
-        "Utils.cpp",
     ],
     hdrs = [
         "Passes.h",
-        "Utils.h",
     ],
     deps = [
         ":PassHeaders",
@@ -96,7 +92,6 @@ iree_compiler_cc_library(
     ],
     hdrs = [
         "Passes.h",
-        "Utils.h",
     ],
     deps = [
         ":PassHeaders",

compiler/src/iree/compiler/InputConversion/Common/CMakeLists.txt

@@ -38,16 +38,12 @@ iree_cc_library(
     Common
   HDRS
     "Passes.h"
-    "Utils.h"
   SRCS
     "ConvertPrimitiveType.cpp"
     "IREEImportPublic.cpp"
     "ImportMLProgram.cpp"
     "Passes.cpp"
-    "QuantizedConvToConv.cpp"
-    "QuantizedMatmulToMatmul.cpp"
     "SanitizeModuleNames.cpp"
-    "Utils.cpp"
   DEPS
     ::PassHeaders
     ::PassesIncGen
@@ -85,7 +81,6 @@ iree_cc_library(
     AutoInputConversionPipeline
   HDRS
     "Passes.h"
-    "Utils.h"
   SRCS
     "AutoInputConversionPipeline.cpp"
   DEPS

compiler/src/iree/compiler/InputConversion/Common/Passes.td

@@ -19,16 +19,6 @@ def ImportMLProgramPass :
   let summary = "Imports MLProgram dialect to IREE Equivalents.";
 }
 
-def LinalgQuantizedConvToConvPass
-    : InterfacePass<"iree-linalg-quantized-conv-to-conv", "mlir::FunctionOpInterface"> {
-  let summary = "lower quantized_conv to conv";
-}
-
-def LinalgQuantizedMatmulToMatmulPass
-    : InterfacePass<"iree-linalg-quantized-matmul-to-matmul", "mlir::FunctionOpInterface"> {
-  let summary = "lower quantized_matmul to matmul";
-}
-
 def SanitizeModuleNamesPass :
     Pass<"iree-sanitize-module-names", "ModuleOp"> {
   let summary = "Sanitizes module names for uniformity across target implementations.";