Making `iree-util-strip-and-splat-constants` use byte_patterns. (#14627)
Had to remove the path someone added for arith.constant as that can't
take byte patterns (and it wasn't the intent of the pass). If someone
wants that they can hoist into globals first before running this pass
with `--iree-util-outline-constants`.
A common usage of this would be to take a giant mlir file and run it
through iree-opt with
`--iree-util-outline-constants --iree-util-strip-and-splat-constants` to
outline any arith.constant values into globals and then strip them.
Note that the final VMFB will contain the same number of bytes as the
original unstripped file.
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel
index dc3fb79..91d5562 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel
@@ -67,7 +67,6 @@
"RemoveZeroExtentTensors.cpp",
"SetEncoding.cpp",
"SplitReduction.cpp",
- "StripAndSplatConstantVariables.cpp",
"StripSignedness.cpp",
"TensorPadToTensorInsertSlice.cpp",
"VerifyInputLegality.cpp",
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
index 491afe1..296d646 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -66,7 +66,6 @@
"RemoveZeroExtentTensors.cpp"
"SetEncoding.cpp"
"SplitReduction.cpp"
- "StripAndSplatConstantVariables.cpp"
"StripSignedness.cpp"
"TensorPadToTensorInsertSlice.cpp"
"VerifyInputLegality.cpp"
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h
index 4969d84..f8c9b01 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h
@@ -254,14 +254,6 @@
// Simplification and Development Tools
//===----------------------------------------------------------------------===//
-// Strips constant util.globals and replaces them with splats.
-// This destructively removes data (often model weights and other parameters)
-// and is intended for use as a development tool.
-// TODO(scotttodd): pass pipeline with this and other development passes to
-// generate test cases / models suitable for check-in
-std::unique_ptr<OperationPass<mlir::ModuleOp>>
-createStripAndSplatConstantVariablesPass();
-
/// Creates a pass to dump a graph for dispatches
std::unique_ptr<Pass>
createDumpDispatchGraphPass(raw_ostream &os = llvm::errs());
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td
index dbefc54..6953102 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td
@@ -109,7 +109,7 @@
def CollapseDimensions :
InterfacePass<"iree-flow-collapse-dimensions", "mlir::FunctionOpInterface"> {
let summary = "Collapse dimensions of Linalg Ops on tensor ops.";
- let constructor = "mlir::iree_compiler::IREE::Flow::createCollapseDimensionsPass()";
+ let constructor = "mlir::iree_compiler::IREE::Flow::createCollapseDimensionsPass()";
}
def DispatchWithTransformDialect :
@@ -317,12 +317,6 @@
let constructor = "mlir::iree_compiler::IREE::Flow::createStripSignednessPass()";
}
-def StripAndSplatConstantVariables :
- Pass<"iree-flow-strip-and-splat-constant-variables", "mlir::ModuleOp"> {
- let summary = "Strips constant util.globals and replaces them with splats.";
- let constructor = "mlir::iree_compiler::IREE::Flow::createStripAndSplatConstantVariablesPass()";
-}
-
def VerifyInputLegality: Pass<"iree-verify-input-legality", ""> {
let summary = "Checks the legality of the IR at the start of IREE flow transformation pipeline.";
let constructor = "mlir::iree_compiler::IREE::Flow::createVerifyInputLegalityPass()";
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/StripAndSplatConstantVariables.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/StripAndSplatConstantVariables.cpp
deleted file mode 100644
index 381475a..0000000
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/StripAndSplatConstantVariables.cpp
+++ /dev/null
@@ -1,107 +0,0 @@
-// Copyright 2020 The IREE Authors
-//
-// Licensed 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 <utility>
-
-#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
-#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
-#include "iree/compiler/Dialect/Flow/Transforms/PassDetail.h"
-#include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
-#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
-#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/IR/Attributes.h"
-#include "mlir/IR/Builders.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/Pass/Pass.h"
-
-namespace mlir {
-namespace iree_compiler {
-namespace IREE {
-namespace Flow {
-
-class StripAndSplatConstantVariablesPass
- : public StripAndSplatConstantVariablesBase<
- StripAndSplatConstantVariablesPass> {
-public:
- StripAndSplatConstantVariablesPass() = default;
-
- void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<IREE::Flow::FlowDialect, IREE::Util::UtilDialect>();
- }
-
- void runOnOperation() override {
- auto moduleOp = getOperation();
- auto builder = OpBuilder::atBlockBegin(moduleOp.getBody());
-
- // Use a heuristic to space out splat values in hopes of avoiding NaN and
- // INF values at runtime:
- // floats: 1/1, 1/2, 1/3, ...
- // ints: 1, 2, 3, 4, ...
- // TODO(scotttodd): flags to control numbers used (all 0, all 1, increasing)
- int replaceIndex = 1;
-
- auto getSplatAttr = [&](TensorType tensorType) {
- auto elementType = tensorType.getElementType();
- TypedAttr newAttr;
- if (llvm::isa<FloatType>(elementType)) {
- newAttr = DenseElementsAttr::get(
- tensorType, FloatAttr::get(elementType, 1.0 / replaceIndex));
- } else {
- assert(elementType.isa<IntegerType>());
- newAttr = DenseElementsAttr::get(
- tensorType, IntegerAttr::get(elementType, replaceIndex));
- }
-
- replaceIndex++;
- return newAttr;
- };
-
- moduleOp.walk([&](Operation *op) {
- if (auto globalOp = dyn_cast<Util::GlobalOp>(op)) {
- // Only strip constant variables.
- if (globalOp.getIsMutable())
- return;
-
- // Only strip tensor type constants (to replace with dense<>).
- if (!llvm::isa<TensorType>(globalOp.getType()))
- return;
-
- auto tensorType = llvm::cast<TensorType>(globalOp.getType());
- TypedAttr newValue = getSplatAttr(tensorType);
-
- builder.setInsertionPoint(globalOp);
- auto newOp = builder.create<IREE::Util::GlobalOp>(
- globalOp.getLoc(), globalOp.getSymName(), globalOp.getIsMutable(),
- globalOp.getType(), newValue);
- newOp.setVisibility(globalOp.getVisibility());
- newOp->setAttr("noinline", UnitAttr::get(builder.getContext()));
- globalOp.erase();
- } else if (auto cstOp = dyn_cast<arith::ConstantOp>(op)) {
- if (!llvm::isa<TensorType>(cstOp.getType()))
- return;
-
- auto tensorType = llvm::cast<TensorType>(cstOp.getType());
- TypedAttr newValue = getSplatAttr(tensorType);
- builder.setInsertionPoint(cstOp);
- auto newOp =
- builder.create<arith::ConstantOp>(cstOp.getLoc(), newValue);
- cstOp->replaceAllUsesWith(newOp);
- cstOp.erase();
- }
- });
- }
-};
-
-std::unique_ptr<OperationPass<mlir::ModuleOp>>
-createStripAndSplatConstantVariablesPass() {
- return std::make_unique<StripAndSplatConstantVariablesPass>();
-}
-
-} // namespace Flow
-} // namespace IREE
-} // namespace iree_compiler
-} // namespace mlir
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel
index e41e9eb..e11e6c7 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel
@@ -52,7 +52,6 @@
"raise_special_ops.mlir",
"remove_zero_extent_tensors.mlir",
"set_encoding.mlir",
- "strip_and_splat_constant_variables.mlir",
"strip_signedness.mlir",
"tensor_pad_to_tensor_insert_slice.mlir",
"transform_dispatch_region_formation.mlir",
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
index 23a255d..4ae1f62 100644
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
@@ -50,7 +50,6 @@
"raise_special_ops.mlir"
"remove_zero_extent_tensors.mlir"
"set_encoding.mlir"
- "strip_and_splat_constant_variables.mlir"
"strip_signedness.mlir"
"tensor_pad_to_tensor_insert_slice.mlir"
"transform_dispatch_region_formation.mlir"
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/strip_and_splat_constant_variables.mlir b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/strip_and_splat_constant_variables.mlir
deleted file mode 100644
index 93515a9..0000000
--- a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/strip_and_splat_constant_variables.mlir
+++ /dev/null
@@ -1,23 +0,0 @@
-// RUN: iree-opt --split-input-file --iree-flow-strip-and-splat-constant-variables %s | FileCheck %s
-
-func.func @fn() -> () {
- // CHECK: util.global private @float_0 {noinline} = dense<1.000000e+00> : tensor<3xf32>
- util.global private @float_0 = dense<"0x012345670123456701234567"> : tensor<3xf32>
- // CHECK: util.global private @float_1 {noinline} = dense<5.000000e-01> : tensor<3xf32>
- util.global private @float_1 = dense<"0x89ABCDEF89ABCDEF89ABCDEF"> : tensor<3xf32>
- // CHECK: util.global private @int_1 {noinline} = dense<3> : tensor<3xi32>
- util.global private @int_1 = dense<"0x89ABCDEF89ABCDEF89ABCDEF"> : tensor<3xi32>
- return
-}
-
-// -----
-
-func.func @fn() -> (tensor<3xf32>, tensor<3xf32>, tensor<3xi32>) {
- // CHECK: arith.constant dense<1.000000e+00> : tensor<3xf32>
- %0 = arith.constant dense<"0x012345670123456701234567"> : tensor<3xf32>
- // CHECK: arith.constant dense<5.000000e-01> : tensor<3xf32>
- %1 = arith.constant dense<"0x89ABCDEF89ABCDEF89ABCDEF"> : tensor<3xf32>
- // CHECK: arith.constant dense<3> : tensor<3xi32>
- %2 = arith.constant dense<"0x89ABCDEF89ABCDEF89ABCDEF"> : tensor<3xi32>
- return %0, %1, %2 : tensor<3xf32>, tensor<3xf32>, tensor<3xi32>
-}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
index dda0c82..6387723 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
@@ -27,7 +27,6 @@
"LayoutSlices.cpp",
"MaterializeBuiltins.cpp",
"MaterializeCopyOnWrite.cpp",
- "OutlineConstants.cpp",
"PackAllocations.cpp",
"PackConstants.cpp",
"PackDispatchOperands.cpp",
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
index 2bb7f3f..9fb6314 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
@@ -29,7 +29,6 @@
"LayoutSlices.cpp"
"MaterializeBuiltins.cpp"
"MaterializeCopyOnWrite.cpp"
- "OutlineConstants.cpp"
"PackAllocations.cpp"
"PackConstants.cpp"
"PackDispatchOperands.cpp"
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
index 83a57a3..cbb2413 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
@@ -68,7 +68,7 @@
// Turn all constant ops into global variables and fix up the IR.
// As many locations change and constants are deduplicated we'll end up with
// a lot of extraneous IR (mostly global loads) and clean those up here.
- passManager.addPass(IREE::Stream::createOutlineConstantsPass());
+ passManager.addPass(IREE::Util::createOutlineConstantsPass());
// Perform cleanup after constant simplification as more canonicalizers may be
// able to kick in.
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.h b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.h
index 91aef3b..324937f 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.h
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.h
@@ -92,12 +92,6 @@
void registerStreamTransformPassPipelines();
//===----------------------------------------------------------------------===//
-// Optimizations
-//===----------------------------------------------------------------------===//
-
-std::unique_ptr<OperationPass<mlir::ModuleOp>> createOutlineConstantsPass();
-
-//===----------------------------------------------------------------------===//
// Conversion
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
index ed72ab0..c11e395 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
@@ -10,18 +10,6 @@
include "mlir/Pass/PassBase.td"
//===----------------------------------------------------------------------===//
-// Optimizations
-//===----------------------------------------------------------------------===//
-
-def OutlineConstants :
- Pass<"iree-stream-outline-constants", "mlir::ModuleOp"> {
- let summary = "Outlines tensor constants into util.globals at the module level.";
- let constructor = [{
- mlir::iree_compiler::IREE::Stream::createOutlineConstantsPass()
- }];
-}
-
-//===----------------------------------------------------------------------===//
// Conversion
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
index 713f6ac..d0a3103 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
@@ -34,7 +34,6 @@
"layout_slices.mlir",
"materialize_builtins.mlir",
"materialize_copy_on_write.mlir",
- "outline_constants.mlir",
"pack_allocations.mlir",
"pack_constants.mlir",
"pack_dispatch_operands.mlir",
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
index 3b3448c..2a1d56a 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
@@ -32,7 +32,6 @@
"layout_slices.mlir"
"materialize_builtins.mlir"
"materialize_copy_on_write.mlir"
- "outline_constants.mlir"
"pack_allocations.mlir"
"pack_constants.mlir"
"pack_dispatch_operands.mlir"
diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Util/Transforms/BUILD.bazel
index 524113a..9efa5ce 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Transforms/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/BUILD.bazel
@@ -26,11 +26,13 @@
"HoistIntoGlobals.cpp",
"IPO.cpp",
"ImportResources.cpp",
+ "OutlineConstants.cpp",
"PassDetail.h",
"Passes.cpp",
"Patterns.cpp",
"PropagateSubranges.cpp",
"SimplifyGlobalAccesses.cpp",
+ "StripAndSplatConstants.cpp",
"StripDebugOps.cpp",
"TestConversion.cpp",
"TestFloatRangeAnalysis.cpp",
diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Util/Transforms/CMakeLists.txt
index 8ae9513..0b8ef3a 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/CMakeLists.txt
@@ -29,11 +29,13 @@
"HoistIntoGlobals.cpp"
"IPO.cpp"
"ImportResources.cpp"
+ "OutlineConstants.cpp"
"PassDetail.h"
"Passes.cpp"
"Patterns.cpp"
"PropagateSubranges.cpp"
"SimplifyGlobalAccesses.cpp"
+ "StripAndSplatConstants.cpp"
"StripDebugOps.cpp"
"TestConversion.cpp"
"TestFloatRangeAnalysis.cpp"
diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp b/compiler/src/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp
index c685e56..9736055 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp
@@ -28,8 +28,8 @@
// Maps an original value in the program to the symbol name of a global.
using HoistedValueMap = llvm::DenseMap<Value, GlobalOp>;
-// expressions into globals. It is not expected that such a greedy algorithm
-// is great, but it is simple. Naive use of this algorithm very likely
+// Hoist expressions into globals. It is not expected that such a greedy
+// algorithm is great, but it is simple. Naive use of this algorithm very likely
// favors programs that consume more memory at runtime than is strictly
// necessary. Either this algorithm can be made smarter or a follow-on pass
// can sink globals into the program where it is profitable to reduce
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/OutlineConstants.cpp b/compiler/src/iree/compiler/Dialect/Util/Transforms/OutlineConstants.cpp
similarity index 93%
rename from compiler/src/iree/compiler/Dialect/Stream/Transforms/OutlineConstants.cpp
rename to compiler/src/iree/compiler/Dialect/Util/Transforms/OutlineConstants.cpp
index ae2ceee..f0d31e3 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/OutlineConstants.cpp
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/OutlineConstants.cpp
@@ -6,12 +6,10 @@
#include <utility>
-#include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
-#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
-#include "iree/compiler/Dialect/Stream/Transforms/PassDetail.h"
-#include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+#include "iree/compiler/Dialect/Util/Transforms/PassDetail.h"
+#include "iree/compiler/Dialect/Util/Transforms/Passes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Attributes.h"
@@ -23,7 +21,7 @@
namespace mlir {
namespace iree_compiler {
namespace IREE {
-namespace Stream {
+namespace Util {
// Returns true if |value| is worth outlining (large, etc).
static bool isOutlinableValue(Attribute value) {
@@ -129,7 +127,7 @@
return std::make_unique<OutlineConstantsPass>();
}
-} // namespace Stream
+} // namespace Util
} // namespace IREE
} // namespace iree_compiler
} // namespace mlir
diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.h b/compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.h
index abf1c19..51f6877 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.h
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.h
@@ -25,8 +25,11 @@
std::unique_ptr<OperationPass<mlir::ModuleOp>> createFuseGlobalsPass();
std::unique_ptr<OperationPass<mlir::ModuleOp>> createHoistIntoGlobalsPass();
std::unique_ptr<OperationPass<mlir::ModuleOp>> createIPOPass();
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createOutlineConstantsPass();
std::unique_ptr<OperationPass<mlir::ModuleOp>> createPropagateSubrangesPass();
std::unique_ptr<OperationPass<void>> createSimplifyGlobalAccessesPass();
+std::unique_ptr<OperationPass<mlir::ModuleOp>>
+createStripAndSplatConstantsPass();
std::unique_ptr<OperationPass<void>> createStripDebugOpsPass();
// Resource Management.
diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.td
index 873b1ab..8794c56 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.td
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.td
@@ -54,6 +54,14 @@
}];
}
+def OutlineConstants :
+ Pass<"iree-util-outline-constants", "mlir::ModuleOp"> {
+ let summary = "Outlines tensor constants into util.globals at the module level.";
+ let constructor = [{
+ mlir::iree_compiler::IREE::Util::createOutlineConstantsPass()
+ }];
+}
+
def PropagateSubranges : Pass<"iree-util-propagate-subranges", "mlir::ModuleOp"> {
let summary = "Propagates resource subranges across the program.";
let constructor = [{
@@ -61,6 +69,12 @@
}];
}
+def StripAndSplatConstants :
+ Pass<"iree-util-strip-and-splat-constants", "mlir::ModuleOp"> {
+ let summary = "Strips constant util.global ops and replaces them with splats.";
+ let constructor = "mlir::iree_compiler::IREE::Util::createStripAndSplatConstantsPass()";
+}
+
def StripDebugOps : Pass<"iree-util-strip-debug-ops", ""> {
let summary = "Strips debug ops, like assertions.";
let constructor = [{
@@ -112,7 +126,7 @@
derive from *ElementsAttr. Given the uniquing/inline behavior, this exacts
very large runtime and memory overhead costs.
- This is a temporary pass to convert a majority of the legacy
+ This is a temporary pass to convert a majority of the legacy
DenseElementsAttr attributes to DenseResourceElementsAttr. Ideally this
is done at the source (frontend), but this pass is provided to aid
transition and testing by doing a manual conversion with iree-opt.
diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/StripAndSplatConstants.cpp b/compiler/src/iree/compiler/Dialect/Util/Transforms/StripAndSplatConstants.cpp
new file mode 100644
index 0000000..47e11c0
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/StripAndSplatConstants.cpp
@@ -0,0 +1,64 @@
+// Copyright 2020 The IREE Authors
+//
+// Licensed 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 <utility>
+
+#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
+#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+#include "iree/compiler/Dialect/Util/Transforms/PassDetail.h"
+#include "iree/compiler/Dialect/Util/Transforms/Passes.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace IREE {
+namespace Util {
+
+class StripAndSplatConstantsPass
+ : public StripAndSplatConstantsBase<StripAndSplatConstantsPass> {
+public:
+ StripAndSplatConstantsPass() = default;
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<IREE::Util::UtilDialect>();
+ }
+
+ void runOnOperation() override {
+ auto moduleOp = getOperation();
+
+ // Give each splatted value a module-unique byte value so that it's easier
+ // to track back to where it came from in the final output.
+ int replaceIndex = 1;
+ auto getSplatAttr = [&](ShapedType type) {
+ return IREE::Util::BytePatternAttr::get(moduleOp.getContext(), type,
+ replaceIndex++);
+ };
+
+ moduleOp.walk([&](Operation *op) {
+ if (auto globalOp = dyn_cast<Util::GlobalOp>(op)) {
+ if (auto initialValue = globalOp.getInitialValueAttr()) {
+ if (auto shapedType = dyn_cast<ShapedType>(initialValue.getType())) {
+ globalOp.setInitialValueAttr(getSplatAttr(shapedType));
+ }
+ }
+ }
+ });
+ }
+};
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>>
+createStripAndSplatConstantsPass() {
+ return std::make_unique<StripAndSplatConstantsPass>();
+}
+
+} // namespace Util
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Util/Transforms/test/BUILD.bazel
index c473c64..90ac39c 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Transforms/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/test/BUILD.bazel
@@ -29,10 +29,12 @@
"hoist_into_globals_linalg.mlir",
"import_resources.mlir",
"ipo.mlir",
+ "outline_constants.mlir",
"promote_bf16_to_f32.mlir",
"promote_f16_to_f32.mlir",
"propagate_subranges.mlir",
"simplify_global_accesses.mlir",
+ "strip_and_splat_constants.mlir",
"strip_debug_ops.mlir",
"test_float_range_analysis.mlir",
"test_float_range_analysis_linalg.mlir",
diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Util/Transforms/test/CMakeLists.txt
index dfc917f..e7f92f5 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Transforms/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/test/CMakeLists.txt
@@ -27,10 +27,12 @@
"hoist_into_globals_linalg.mlir"
"import_resources.mlir"
"ipo.mlir"
+ "outline_constants.mlir"
"promote_bf16_to_f32.mlir"
"promote_f16_to_f32.mlir"
"propagate_subranges.mlir"
"simplify_global_accesses.mlir"
+ "strip_and_splat_constants.mlir"
"strip_debug_ops.mlir"
"test_float_range_analysis.mlir"
"test_float_range_analysis_linalg.mlir"
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/outline_constants.mlir b/compiler/src/iree/compiler/Dialect/Util/Transforms/test/outline_constants.mlir
similarity index 86%
rename from compiler/src/iree/compiler/Dialect/Stream/Transforms/test/outline_constants.mlir
rename to compiler/src/iree/compiler/Dialect/Util/Transforms/test/outline_constants.mlir
index 40e0954..46f14be 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/outline_constants.mlir
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/test/outline_constants.mlir
@@ -1,4 +1,4 @@
-// RUN: iree-opt --split-input-file --iree-stream-outline-constants %s | FileCheck %s
+// RUN: iree-opt --split-input-file --iree-util-outline-constants %s | FileCheck %s
// CHECK-LABEL: @scalarConstant
func.func @scalarConstant() {
@@ -9,8 +9,8 @@
// -----
-// CHECK-LABEL: @sparseConstant
-func.func @sparseConstant() {
+// CHECK-LABEL: @splatConstant
+func.func @splatConstant() {
// CHECK: = arith.constant dense<1.200000e+00> : tensor<512x128xf32>
%cst = arith.constant dense<1.2> : tensor<512x128xf32>
return
diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/test/strip_and_splat_constants.mlir b/compiler/src/iree/compiler/Dialect/Util/Transforms/test/strip_and_splat_constants.mlir
new file mode 100644
index 0000000..1990522
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/test/strip_and_splat_constants.mlir
@@ -0,0 +1,11 @@
+// RUN: iree-opt --split-input-file --iree-util-strip-and-splat-constants %s | FileCheck %s
+
+func.func @fn() -> () {
+ // CHECK: util.global private @float_0 = #util.byte_pattern<1> : tensor<3xf32>
+ util.global private @float_0 = dense<"0x012345670123456701234567"> : tensor<3xf32>
+ // CHECK: util.global private @float_1 = #util.byte_pattern<2> : tensor<3xf32>
+ util.global private @float_1 = dense<"0x89ABCDEF89ABCDEF89ABCDEF"> : tensor<3xf32>
+ // CHECK: util.global private @int_1 = #util.byte_pattern<3> : tensor<3xi32>
+ util.global private @int_1 = dense<"0x89ABCDEF89ABCDEF89ABCDEF"> : tensor<3xi32>
+ return
+}
