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

[XLA GPU] Simplify tiling prologue/epilogue calling logic by initializing them to no-op lambdas

Browse Source 
PiperOrigin-RevId: 259412438
This commit is contained in:
George Karpenkov 2019-07-22 14:54:46 -07:00 committed by TensorFlower Gardener
parent a9429e942a
commit 6b0dba99c3
2 changed files with 6 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
View File

@ -3191,20 +3191,9 @@ LaunchDimensions IrEmitterUnnested::EmitKernel(
} }
}; };
const BlockPrologueGenerator& block_prologue_generator = kernel_generator.GetBlockPrologueGenerator()(unnested_hlo, kernel_info);
kernel_generator.GetBlockPrologueGenerator();
if (block_prologue_generator) {
block_prologue_generator(unnested_hlo, kernel_info);
}
EmitBlock(std::move(emit_one_tile), kernel_info, &ksl, index_ty); EmitBlock(std::move(emit_one_tile), kernel_info, &ksl, index_ty);
kernel_generator.GetBlockEpilogueGenerator()(unnested_hlo, kernel_info);
const BlockEpilogueGenerator& block_epilogue_generator =
kernel_generator.GetBlockEpilogueGenerator();
if (block_epilogue_generator) {
block_epilogue_generator(unnested_hlo, kernel_info);
}
return launch_dimensions; return launch_dimensions;
} }

6
tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h
View File

@ -121,8 +121,10 @@ class IrEmitterUnnested : public IrEmitter {
public: public:
explicit KernelCodeGenerator( explicit KernelCodeGenerator(
TileElementGenerator tile_element_generator, TileElementGenerator tile_element_generator,
BlockPrologueGenerator block_prologue_generator = {}, BlockPrologueGenerator block_prologue_generator =
BlockEpilogueGenerator block_epilogue_generator = {}) [](HloInstruction*, KernelCodegenInfo*) {},
BlockEpilogueGenerator block_epilogue_generator =
[](HloInstruction*, KernelCodegenInfo*) {})
: tile_element_generator_(std::move(tile_element_generator)), : tile_element_generator_(std::move(tile_element_generator)),
block_prologue_generator_(std::move(block_prologue_generator)), block_prologue_generator_(std::move(block_prologue_generator)),
block_epilogue_generator_(std::move(block_epilogue_generator)) {} block_epilogue_generator_(std::move(block_epilogue_generator)) {}