tensorflow/compiler/mlir/tfrt/jit/tf_jitrt_pipeline.cc
// -------------------------------------------------------------------------- //
// Assemble a TF JitRt pipeline to lower from Tensorflow dialects to Linalg on
// buffers via progressive lowering to MHLO and Linalg.
// -------------------------------------------------------------------------- //
Assemble a TF JitRt pipeline to lower from Tensorflow dialects to Linalg on
buffers via progressive lowering to MHLO and Linalg.
Assemble a TF JitRt pipeline to lower from Tensorflow dialects to Linalg on buffers via progressive lowering to MHLO and Linalg .
void CreateTfJitRtPipeline(OpPassManager& pm,
                           const TfJitRtPipelineOptions& options) {
  pm.addNestedPass<FuncOp>(CreateFissionPass());
  pm.addPass(std::make_unique<AddTensorflowProducerVersion>());
  pm.addPass(mlir::TF::CreateTFShapeInferencePass());
  pm.addPass(mlir::createCanonicalizerPass());
  pm.addPass(mlir::mhlo::createLegalizeTFControlFlowPass());
  pm.addNestedPass<FuncOp>(mlir::mhlo::createLegalizeTFPass());
  if (options.legalize_i1_tensors) {
    pm.addPass(CreateJitRtLegalizeI1TypesPass());
  }
  pm.addPass(mlir::createCSEPass());
  pm.addNestedPass<FuncOp>(
      CreateSymbolicShapeOptimizationPass(true));
  pm.addNestedPass<FuncOp>(mlir::mhlo::createMergeAssumingOpsPass());
  pm.addNestedPass<FuncOp>(mlir::mhlo::createBroadcastPropagationPass());
  pm.addPass(mlir::createCSEPass());
  pm.addPass(mlir::createCanonicalizerPass());
  pm.addNestedPass<FuncOp>(CreateSymbolicShapeOptimizationPass());
  pm.addNestedPass<FuncOp>(mlir::mhlo::createGroupReductionDimensionsPass());
  pm.addNestedPass<FuncOp>(mlir::createReshapeSimplifierPass());
  pm.addNestedPass<FuncOp>(mlir::mhlo::createLegalizeControlFlowPass());
  pm.addNestedPass<FuncOp>(mlir::mhlo::createLegalizeHloToLinalgPass());
  pm.addPass(mlir::mhlo::createLegalizeToArithmeticPass());
  pm.addNestedPass<FuncOp>(
      mlir::mhlo::createLegalizeHloShapeOpsToStandardPass());
  pm.addPass(mlir::kernel_gen::transforms::CreateConvertToSignlessPass());
  pm.addNestedPass<FuncOp>(
      mlir::kernel_gen::transforms::CreateShapeSimplification());
  pm.addNestedPass<FuncOp>(mlir::createShapeToShapeLowering());
  pm.addPass(mlir::createConvertShapeToStandardPass());
  pm.addNestedPass<FuncOp>(mlir::createConvertShapeConstraintsPass());
  pm.addPass(mlir::createCSEPass());
  pm.addPass(mlir::memref::createResolveShapedTypeResultDimsPass());
  pm.addNestedPass<FuncOp>(
      mlir::kernel_gen::transforms::CreateLowerIndexCastPass());
  pm.addPass(mlir::createCSEPass());
  pm.addPass(mlir::createCanonicalizerPass());
  AddLinalgTransformations(pm, options);
  pm.addPass(mlir::createInlinerPass());
  pm.addPass(mlir::createCanonicalizerPass());
  AddBufferizationPasses(pm);
  pm.addPass(mlir::createCSEPass());
  pm.addPass(mlir::createCanonicalizerPass());
  pm.addNestedPass<FuncOp>(mlir::bufferization::createBufferDeallocationPass());
  pm.addNestedPass<FuncOp>(CreateLinalgTrivialBufferForwardingPass());
  pm.addNestedPass<FuncOp>(CreateLinalgTrivialCopyRemovalPass());
  if (options.vectorize)
    pm.addNestedPass<FuncOp>(mlir::gml_st::createGmlStToScfPass());
  pm.addPass(mlir::createBufferizationToMemRefPass());
  pm.addPass(mlir::createCSEPass());
  pm.addPass(mlir::createCanonicalizerPass());
  if (options.vectorize && options.codegen_transpose)
    pm.addNestedPass<FuncOp>(CreateLowerVectorTransposePass());
  mlir::VectorTransferToSCFOptions vec_to_scf_options;
  vec_to_scf_options.unroll = true;
  pm.addNestedPass<FuncOp>(
      mlir::createConvertVectorToSCFPass(vec_to_scf_options));
  pm.addNestedPass<FuncOp>(createRewriteVectorMultiReductionPass());
  pm.addNestedPass<FuncOp>(CreateMathApproximationPass({"all"}));
}
void CreateTfJitRtPipeline ( OpPassManager & pm , const TfJitRtPipelineOptions & options ) { pm . addNestedPass < FuncOp > ( CreateFissionPass ( ) ) ; pm . addPass ( std :: make_unique < AddTensorflowProducerVersion > ( ) ) ; pm . addPass ( mlir :: TF :: CreateTFShapeInferencePass ( ) ) ; pm . addPass ( mlir :: createCanonicalizerPass ( ) ) ; pm . addPass ( mlir :: mhlo :: createLegalizeTFControlFlowPass ( ) ) ; pm . addNestedPass < FuncOp > ( mlir :: mhlo :: createLegalizeTFPass ( ) ) ; if ( options . legalize_i1_tensors ) { pm . addPass ( CreateJitRtLegalizeI1TypesPass ( ) ) ; } pm . addPass ( mlir :: createCSEPass ( ) ) ; pm . addNestedPass < FuncOp > ( CreateSymbolicShapeOptimizationPass ( true ) ) ; pm . addNestedPass < FuncOp > ( mlir :: mhlo :: createMergeAssumingOpsPass ( ) ) ; pm . addNestedPass < FuncOp > ( mlir :: mhlo :: createBroadcastPropagationPass ( ) ) ; pm . addPass ( mlir :: createCSEPass ( ) ) ; pm . addPass ( mlir :: createCanonicalizerPass ( ) ) ; pm . addNestedPass < FuncOp > ( CreateSymbolicShapeOptimizationPass ( ) ) ; pm . addNestedPass < FuncOp > ( mlir :: mhlo :: createGroupReductionDimensionsPass ( ) ) ; pm . addNestedPass < FuncOp > ( mlir :: createReshapeSimplifierPass ( ) ) ; pm . addNestedPass < FuncOp > ( mlir :: mhlo :: createLegalizeControlFlowPass ( ) ) ; pm . addNestedPass < FuncOp > ( mlir :: mhlo :: createLegalizeHloToLinalgPass ( ) ) ; pm . addPass ( mlir :: mhlo :: createLegalizeToArithmeticPass ( ) ) ; pm . addNestedPass < FuncOp > ( mlir :: mhlo :: createLegalizeHloShapeOpsToStandardPass ( ) ) ; pm . addPass ( mlir :: kernel_gen :: transforms :: CreateConvertToSignlessPass ( ) ) ; pm . addNestedPass < FuncOp > ( mlir :: kernel_gen :: transforms :: CreateShapeSimplification ( ) ) ; pm . addNestedPass < FuncOp > ( mlir :: createShapeToShapeLowering ( ) ) ; pm . addPass ( mlir :: createConvertShapeToStandardPass ( ) ) ; pm . addNestedPass < FuncOp > ( mlir :: createConvertShapeConstraintsPass ( ) ) ; pm . addPass ( mlir :: createCSEPass ( ) ) ; pm . addPass ( mlir :: memref :: createResolveShapedTypeResultDimsPass ( ) ) ; pm . addNestedPass < FuncOp > ( mlir :: kernel_gen :: transforms :: CreateLowerIndexCastPass ( ) ) ; pm . addPass ( mlir :: createCSEPass ( ) ) ; pm . addPass ( mlir :: createCanonicalizerPass ( ) ) ; AddLinalgTransformations ( pm , options ) ; pm . addPass ( mlir :: createInlinerPass ( ) ) ; pm . addPass ( mlir :: createCanonicalizerPass ( ) ) ; AddBufferizationPasses ( pm ) ; pm . addPass ( mlir :: createCSEPass ( ) ) ; pm . addPass ( mlir :: createCanonicalizerPass ( ) ) ; pm . addNestedPass < FuncOp > ( mlir :: bufferization :: createBufferDeallocationPass ( ) ) ; pm . addNestedPass < FuncOp > ( CreateLinalgTrivialBufferForwardingPass ( ) ) ; pm . addNestedPass < FuncOp > ( CreateLinalgTrivialCopyRemovalPass ( ) ) ; if ( options . vectorize ) pm . addNestedPass < FuncOp > ( mlir :: gml_st :: createGmlStToScfPass ( ) ) ; pm . addPass ( mlir :: createBufferizationToMemRefPass ( ) ) ; pm . addPass ( mlir :: createCSEPass ( ) ) ; pm . addPass ( mlir :: createCanonicalizerPass ( ) ) ; if ( options . vectorize && options . codegen_transpose ) pm . addNestedPass < FuncOp > ( CreateLowerVectorTransposePass ( ) ) ; mlir :: VectorTransferToSCFOptions vec_to_scf_options ; vec_to_scf_options . unroll = true ; pm . addNestedPass < FuncOp > ( mlir :: createConvertVectorToSCFPass ( vec_to_scf_options ) ) ; pm . addNestedPass < FuncOp > ( createRewriteVectorMultiReductionPass ( ) ) ; pm . addNestedPass < FuncOp > ( CreateMathApproximationPass ( { " " } ) ) ; }
Assemble a TF JitRt pipeline to lower from Tensorflow dialects to Linalg on
buffers via progressive lowering to MHLO and Linalg.
Assemble a TF JitRt pipeline to lower from Tensorflow dialects to Linalg on buffers via progressive lowering to MHLO and Linalg .
// Break Tensorflow fused operations into primitive operations before // lowering to HLO. // Run shape inference to propagate potentially specialized input shapes. // Transform TF operation to HLO. // Convert 'i1' tensors into 'i8' tensors. // Remove redundant shape operations left after legalizing to HLO. // Resolve all shape constraints (e.g. broadcast constraints that can be // proved statically and changed to const witness) early to allow more // efficient broadcast operations moving. /*constraints_only=*/ // Move up broadcasting operations to allow for more fusion opportunities. // Add the broadcast propagation pass first, because it can help to avoid // exponential complexity from the EarlyBroadcastInDimOp pattern which is used // in the merge assuming ops pass further down. // After all shape constraints removed and broadcasts moved to the top, try // to resolve broadcasts that can be converted to linalg generic operations. // Group reduction and parallel dimensions of reduction operations and realize // them through equivalent 1D or 2D reductions, if possible. // Also, try to simplify reshape operations. // Transform HLO operations to Linalg and Standard. // Now that all compute operations are converted to standard (as a side effect // of bufferizing to memref dialect) we can remove the remaining references // to unsigned types. // Lower shape dialect to standard to enable linalg canonicalizations (e.g. // use linalg inputs instead of outputs for memref.dim operations). // Fuse Linalg on tensors operations. // Lower index cast on tensors to tensor.generate. // Add linalg passes to perform fusion, tiling, peeling and vectorization. // Inline everything, bufferization doesn't model ownership across calls. // Always run canonicalizer (which does dead code removal) before bufferizing // anything. // Deallocate all temporary buffers. // Do trivial buffer forwarding across linalg.generic operations. // Remove trivial copy operations.