Skip to content

Commit

Permalink
Enable ConstEval for CPU data tiling path. (#14792)
Browse files Browse the repository at this point in the history
- 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
  • Loading branch information
hanhanW authored Sep 1, 2023
1 parent fe8f79e commit 73b04d3
Show file tree
Hide file tree
Showing 26 changed files with 377 additions and 88 deletions.
4 changes: 2 additions & 2 deletions build_tools/benchmarks/comparisons/setup_mobile.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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" \
Expand All @@ -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" \
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
],
Expand All @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
],
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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);
Expand All @@ -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
Expand Down
7 changes: 5 additions & 2 deletions compiler/src/iree/compiler/Codegen/Common/CPU/Passes.h
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand All @@ -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.
Expand All @@ -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();

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
7 changes: 0 additions & 7 deletions compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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."),
Expand Down Expand Up @@ -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(),
Expand Down
5 changes: 3 additions & 2 deletions compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ iree_cc_library(
"ConstExpr.cpp"
"OpOracle.cpp"
DEPS
IREELinalgExtDialect
LLVMSupport
MLIRArithDialect
MLIRFuncDialect
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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 {};
Expand Down
37 changes: 36 additions & 1 deletion compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
29 changes: 29 additions & 0 deletions compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading

0 comments on commit 73b04d3

Please sign in to comment.