Enable ConstEval for CPU data tiling path. (#14792)

- Moves SetEncoding pass to GlobalOptimization stage, and renames `iree-flow-enable-data-tiling` to `iree-opt-enable-data-tiling` - Introduces MaterializeHomogeneousEncodings pass to GlobalOptimization phase. It materializes encodings to physicacl ops when there is a single executable. - Make CPU materialization passes take optional executable attributes. They look into the attributes if they are passed. Otherwise, they look into existing attributes. Fixes #11360
iree-org · Sep 1, 2023 · 73b04d3 · 73b04d3
1 parent fe8f79e
commit 73b04d3
Show file tree

Hide file tree

Showing 26 changed files with 377 additions and 88 deletions.
diff --git a/build_tools/benchmarks/comparisons/setup_mobile.sh b/build_tools/benchmarks/comparisons/setup_mobile.sh
@@ -121,7 +121,7 @@ for i in $(ls ${ROOT_DIR}/models/tflite/); do
     --iree-input-type=tosa \
     --iree-hal-target-backends=llvm-cpu \
     --iree-llvmcpu-target-triple=aarch64-none-linux-android29 \
-    --iree-flow-enable-data-tiling \
+    --iree-opt-data-tiling \
     --iree-llvmcpu-target-cpu-features=+dotprod \
     "--iree-flow-enable-fuse-padding-into-linalg-consumer-ops" \
     "--iree-llvmcpu-enable-pad-consumer-fusion" \
@@ -136,7 +136,7 @@ for i in $(ls ${ROOT_DIR}/models/tflite/); do
     --iree-input-type=tosa \
     --iree-hal-target-backends=llvm-cpu \
     --iree-llvmcpu-target-triple=aarch64-none-linux-android29 \
-    --iree-flow-enable-data-tiling \
+    --iree-opt-data-tiling \
     --iree-llvmcpu-target-cpu-features=+dotprod \
     "--iree-flow-enable-fuse-padding-into-linalg-consumer-ops" \
     "--iree-llvmcpu-enable-pad-consumer-fusion" \

diff --git a/build_tools/python/benchmark_suites/iree/armv8_a_benchmarks.py b/build_tools/python/benchmark_suites/iree/armv8_a_benchmarks.py
@@ -43,7 +43,7 @@ class Android_ARMv8_A_Benchmarks(object):
         tags=["experimental-flags", "mmt4d"],
         compile_targets=[ARMV8_A_CPU_TARGET],
         extra_flags=[
-            "--iree-flow-enable-data-tiling",
+            "--iree-opt-data-tiling",
             "--iree-flow-enable-fuse-padding-into-linalg-consumer-ops",
             "--iree-llvmcpu-enable-pad-consumer-fusion",
         ],
@@ -53,7 +53,7 @@ class Android_ARMv8_A_Benchmarks(object):
         tags=["experimental-flags", "mmt4d", "dotprod"],
         compile_targets=[ARMV8_A_CPU_TARGET],
         extra_flags=[
-            "--iree-flow-enable-data-tiling",
+            "--iree-opt-data-tiling",
             "--iree-llvmcpu-target-cpu-features=+dotprod",
             "--iree-flow-enable-fuse-padding-into-linalg-consumer-ops",
             "--iree-llvmcpu-enable-pad-consumer-fusion",

diff --git a/build_tools/python/benchmark_suites/iree/x86_64_benchmarks.py b/build_tools/python/benchmark_suites/iree/x86_64_benchmarks.py
@@ -33,7 +33,7 @@ class Linux_x86_64_Benchmarks(object):
         tags=["experimental-flags", "data-tiling", "ukernel"],
         compile_targets=[CASCADELAKE_CPU_TARGET],
         extra_flags=[
-            "--iree-flow-enable-data-tiling",
+            "--iree-opt-data-tiling",
             "--iree-llvmcpu-enable-microkernels",
         ],
     )

diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/CPUMaterializeEncodingPass.cpp b/compiler/src/iree/compiler/Codegen/Common/CPU/CPUMaterializeEncodingPass.cpp
@@ -144,20 +144,33 @@ static MatmulTileParams chooseMatmulTileParams(EncodingUser user,
 
 struct CPUMaterializeEncodingPass
     : public CPUMaterializeEncodingBase<CPUMaterializeEncodingPass> {
+  CPUMaterializeEncodingPass() : targetAttr(nullptr) {}
+  explicit CPUMaterializeEncodingPass(IREE::HAL::ExecutableTargetAttr attr)
+      : targetAttr(attr) {}
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<arith::ArithDialect, IREE::LinalgExt::IREELinalgExtDialect,
                     IREE::Codegen::IREECodegenDialect>();
   }
   void runOnOperation() override;
+
+private:
+  IREE::HAL::ExecutableTargetAttr targetAttr;
 };
 
 struct CPUMaterializeUpperBoundTileSizePass
     : public CPUMaterializeUpperBoundTileSizeBase<
           CPUMaterializeUpperBoundTileSizePass> {
+  CPUMaterializeUpperBoundTileSizePass() = default;
+  explicit CPUMaterializeUpperBoundTileSizePass(
+      ArrayRef<IREE::HAL::ExecutableTargetAttr> attrs)
+      : targetAttrs(attrs) {}
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<arith::ArithDialect>();
   }
   void runOnOperation() override;
+
+private:
+  SmallVector<IREE::HAL::ExecutableTargetAttr, 4> targetAttrs;
 };
 
 FailureOr<MaterializeEncodingInfo>
@@ -249,7 +262,8 @@ void CPUMaterializeEncodingPass::runOnOperation() {
   MLIRContext *context = &getContext();
   auto operation = getOperation();
   RewritePatternSet materializeEncodingPattern(context);
-  auto targetAttr = ExecutableTargetAttr::lookup(operation);
+  if (!targetAttr)
+    targetAttr = ExecutableTargetAttr::lookup(operation);
   auto materializeEncodingFn = getMaterializeEncodingFn(targetAttr);
   if (!materializeEncodingFn) {
     return signalPassFailure();
@@ -283,8 +297,10 @@ void CPUMaterializeEncodingPass::runOnOperation() {
 void CPUMaterializeUpperBoundTileSizePass::runOnOperation() {
   MLIRContext *context = &getContext();
   auto operation = getOperation();
-  auto targetAttrs =
-      IREE::HAL::DeviceTargetAttr::lookupExecutableTargets(operation);
+  if (targetAttrs.empty()) {
+    targetAttrs =
+        IREE::HAL::DeviceTargetAttr::lookupExecutableTargets(operation);
+  }
   RewritePatternSet patterns(context);
   MaterializeEncodingFn materializeEncodingFn =
       getUpperBoundMaterializeEncodingFn(targetAttrs);
@@ -301,13 +317,14 @@ void CPUMaterializeUpperBoundTileSizePass::runOnOperation() {
 }
 
 std::unique_ptr<OperationPass<func::FuncOp>>
-createCPUMaterializeEncodingPass() {
-  return std::make_unique<CPUMaterializeEncodingPass>();
+createCPUMaterializeEncodingPass(IREE::HAL::ExecutableTargetAttr targetAttr) {
+  return std::make_unique<CPUMaterializeEncodingPass>(targetAttr);
 }
 
 std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
-createCPUMaterializeUpperBoundTileSizePass() {
-  return std::make_unique<CPUMaterializeUpperBoundTileSizePass>();
+createCPUMaterializeUpperBoundTileSizePass(
+    ArrayRef<IREE::HAL::ExecutableTargetAttr> targetAttrs) {
+  return std::make_unique<CPUMaterializeUpperBoundTileSizePass>(targetAttrs);
 }
 
 } // namespace iree_compiler

diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/Passes.h b/compiler/src/iree/compiler/Codegen/Common/CPU/Passes.h
@@ -12,6 +12,7 @@
 #ifndef IREE_COMPILER_CODEGEN_COMMON_CPU_PASSES_H_
 #define IREE_COMPILER_CODEGEN_COMMON_CPU_PASSES_H_
 
+#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Pass/Pass.h"
 
@@ -22,7 +23,8 @@ namespace iree_compiler {
 ///   linalg_ext.set_encoding   -> tensor.pack
 ///   linalg_ext.unset_encoding -> tensor.unpack
 ///   linalg.matmul             -> linalg.mmt4d
-std::unique_ptr<OperationPass<func::FuncOp>> createCPUMaterializeEncodingPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createCPUMaterializeEncodingPass(
+    IREE::HAL::ExecutableTargetAttr targetAttr = nullptr);
 
 /// Like createLLVMCPUMaterializeEncodingPass, but specifically for
 /// linalg_ext.upper_bound_tile_size, converting it to constants.
@@ -40,7 +42,8 @@ std::unique_ptr<OperationPass<func::FuncOp>> createCPUMaterializeEncodingPass();
 /// that is the largest tile size that we can use in VMVX, and can be adjusted
 // as needed.
 std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
-createCPUMaterializeUpperBoundTileSizePass();
+createCPUMaterializeUpperBoundTileSizePass(
+    ArrayRef<IREE::HAL::ExecutableTargetAttr> targetAttrs = {});
 
 void registerCodegenCommonCPUPasses();
 

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormScalarDispatches.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/FormScalarDispatches.cpp
@@ -62,7 +62,7 @@ static bool isComputeOperation(Operation *op) {
   }
   if (op->getDialect() == context->getLoadedDialect<tensor::TensorDialect>()) {
     return !isa<tensor::CastOp, tensor::CollapseShapeOp, tensor::EmptyOp,
-                tensor::ExpandShapeOp>(op);
+                tensor::ExpandShapeOp, tensor::PackOp, tensor::UnPackOp>(op);
   }
   return false;
 }

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -79,11 +79,6 @@ static llvm::cl::opt<bool> clDispatchGenerateWorkloadRegion(
     "iree-flow-dispatch-generate-workload-region",
     llvm::cl::desc("Generate the workload region."), llvm::cl::init(true));
 
-static llvm::cl::opt<bool>
-    clEnableDataTiling("iree-flow-enable-data-tiling",
-                       llvm::cl::desc("Enable data tiling path."),
-                       llvm::cl::init(false));
-
 static llvm::cl::opt<bool> clNormalizeInputIndexingMap(
     "iree-flow-normalize-input-indexing-map",
     llvm::cl::desc("Enable normalizing input indexing map to identity."),
@@ -154,8 +149,6 @@ void buildFlowTransformPassPipeline(OpPassManager &passManager,
       // transpose.
       .addPredicatedPass(clNormalizeInputIndexingMap,
                          createInterchangeTransposeGenericOpsPass)
-      // Enable data tiling after all linalg level transformations.
-      .addPredicatedPass(clEnableDataTiling, createSetEncodingPass)
       ////////////////////////////////////////////////////////////////////////
       // Dispatch region formation.
       .addPredicatedPass(!clDispatchTransformFileName.empty(),

diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
@@ -227,8 +227,9 @@ void buildHALTransformPassPipeline(OpPassManager &passManager,
     // Executable translation
     //----------------------------------------------------------------------------
 
-    FunctionLikeNest(passManager)
-        .addPass(createCPUMaterializeUpperBoundTileSizePass);
+    FunctionLikeNest(passManager).addPass([]() {
+      return createCPUMaterializeUpperBoundTileSizePass();
+    });
 
     // Preprocess executables using an external tool. The tool may mutate one or
     // more variants and even insert or remove variants.

diff --git a/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/BUILD.bazel
@@ -25,6 +25,7 @@ iree_compiler_cc_library(
     deps = [
         "//compiler/src/iree/compiler/Dialect/Util/Analysis",
         "//compiler/src/iree/compiler/Dialect/Util/IR",
+        "//llvm-external-projects/iree-dialects:IREELinalgExtDialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:FuncDialect",

diff --git a/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/CMakeLists.txt
@@ -20,6 +20,7 @@ iree_cc_library(
     "ConstExpr.cpp"
     "OpOracle.cpp"
   DEPS
+    IREELinalgExtDialect
     LLVMSupport
     MLIRArithDialect
     MLIRFuncDialect

diff --git a/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.cpp b/compiler/src/iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.cpp
@@ -6,6 +6,7 @@
 
 #include "iree/compiler/Dialect/Util/Analysis/Constant/OpOracle.h"
 
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
 #include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
 #include "iree/compiler/Dialect/Util/IR/UtilOps.h"
 #include "llvm/ADT/SmallPtrSet.h"
@@ -100,6 +101,13 @@ ConstExprOpInfo ConstExprOpInfo::getForOp(Operation *op) {
     return {};
   }
 
+  // Target-dependent ops are not const-expr.
+  // TODO(#14887): Use trait/interface instead.
+  if (isa<IREE::LinalgExt::UpperBoundTileSizeOp,
+          IREE::LinalgExt::SetEncodingOp>(op)) {
+    return {};
+  }
+
   // By default, ops without results are not const-expr.
   if (op->getNumResults() == 0) {
     return {};

diff --git a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
@@ -4,24 +4,59 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//build_tools/bazel:build_defs.oss.bzl", "iree_compiler_cc_library")
+load("//build_tools/bazel:build_defs.oss.bzl", "iree_compiler_cc_library", "iree_gentbl_cc_library")
 
 package(
     default_visibility = ["//visibility:public"],
     features = ["layering_check"],
     licenses = ["notice"],  # Apache 2.0
 )
 
+iree_gentbl_cc_library(
+    name = "PassesIncGen",
+    tbl_outs = [
+        (
+            ["--gen-pass-decls"],
+            "Passes.h.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Passes.td",
+    deps = [
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)
+
+iree_compiler_cc_library(
+    name = "PassHeaders",
+    hdrs = [
+        "PassDetail.h",
+        "Passes.h",
+        "Passes.h.inc",
+    ],
+    deps = [
+        ":PassesIncGen",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
 iree_compiler_cc_library(
     name = "GlobalOptimization",
     srcs = [
+        "MaterializeHomogeneousEncodings.cpp",
         "Passes.cpp",
     ],
     hdrs = [
         "Passes.h",
     ],
     deps = [
+        ":PassHeaders",
+        ":PassesIncGen",
+        "//compiler/src/iree/compiler/Codegen/Common/CPU:CommonCPUPasses",
         "//compiler/src/iree/compiler/Dialect/Flow/Transforms",
+        "//compiler/src/iree/compiler/Dialect/HAL/IR",
+        "//compiler/src/iree/compiler/Dialect/HAL/IR:HALDialect",
         "//compiler/src/iree/compiler/Dialect/Util/Transforms",
         "//compiler/src/iree/compiler/Pipelines:Options",
         "//compiler/src/iree/compiler/Utils",

diff --git a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
@@ -10,22 +10,51 @@
 
 iree_add_all_subdirs()
 
+iree_tablegen_library(
+  NAME
+    PassesIncGen
+  TD_FILE
+    "Passes.td"
+  OUTS
+    --gen-pass-decls Passes.h.inc
+)
+
+iree_cc_library(
+  NAME
+    PassHeaders
+  HDRS
+    "PassDetail.h"
+    "Passes.h"
+    "Passes.h.inc"
+  DEPS
+    ::PassesIncGen
+    MLIRPass
+    MLIRTransforms
+  PUBLIC
+)
+
 iree_cc_library(
   NAME
     GlobalOptimization
   HDRS
     "Passes.h"
   SRCS
+    "MaterializeHomogeneousEncodings.cpp"
     "Passes.cpp"
   DEPS
+    ::PassHeaders
+    ::PassesIncGen
     LLVMSupport
     MLIRFuncDialect
     MLIRIR
     MLIRLinalgTransforms
     MLIRMemRefTransforms
     MLIRPass
     MLIRTransforms
+    iree::compiler::Codegen::Common::CPU::CommonCPUPasses
     iree::compiler::Dialect::Flow::Transforms
+    iree::compiler::Dialect::HAL::IR
+    iree::compiler::Dialect::HAL::IR::HALDialect
     iree::compiler::Dialect::Util::Transforms
     iree::compiler::Pipelines::Options
     iree::compiler::Utils