experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Perform SCCP early in the TF to HLO pipeline.

Browse Source 
SCCP propagates constants into callee functions and has limited support for propagation across functions of control flow ops like tf.While. Performing it early in the pipeline should benefit shape inference and legalization.

Specifically, it should at least be moved prior to conversion of region based control flow to functional control flow. This is to support the following example:

```
func @foo () {
  %const = constant dense<2> : tensor<i32>
  %while = tf.WhileRegion(...) ({
               // cond
               ...
             }, {
               // body
               call @bar(%const)
               ...
             })
}
func @bar (%arg0 : tensor<i32>) {
  %0 = tf.SomeOp(%arg0)
}
```

Say we have a call op (with a constant argument) nested inside a control flow op (if/case/while). SCCP would handle it easily in region based control flow representation. In functional control flow representation, the constant argument of call op would be explicitly captured by the outer functional control flow op. SCCP would not be able to propagate constant across the branches of functional control flow ops, thus preventing the constant propagation to the nested callee.

PiperOrigin-RevId: 347642576
Change-Id: I82c519361d3912e1076d7feb76cb5da463fba940
This commit is contained in:
Prakalp Srivastava 2020-12-15 10:28:36 -08:00 committed by TensorFlower Gardener
parent 8b853090bc
commit 6aec6478d3
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.cc
View File

@ -282,6 +282,9 @@ void CreateConvertMlirToXlaHloPipeline(
pm.addPass(mlir::TF::CreateTFFunctionalControlFlowToRegions());
pm.addNestedPass<mlir::FuncOp>(mlir::TF::CreateDropWhileShapeInvariantPass());
pm.addNestedPass<mlir::FuncOp>(mlir::createCanonicalizerPass());
// The SCCP pass performs constant propagation across the IR, which, for
// example, propagates constant arguments into callee functions.
pm.addPass(mlir::createSCCPPass());
// Guarantee all functions have one use, which enables shape inference.
pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
// Run shape inference pass before tensorlist decomposition to get buffer
@ -303,9 +306,6 @@ void CreateConvertMlirToXlaHloPipeline(
// with a tuple argument which break the assumption of resource lifting
// inside PromoteResourcesToArgs.
pm.addPass(mlir::mhlo::createLegalizeTFControlFlowPass());
// The SCCP pass performs constant propagation across the IR, which, for
// example, propagates constant arguments into callee functions.
pm.addPass(mlir::createSCCPPass());
pm.addNestedPass<mlir::FuncOp>(mlir::mhlo::createLegalizeTFPass(
/*allow_partial_conversion=*/true, /*legalize_chlo=*/true,