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();
if (block_prologue_generator) {
block_prologue_generator(unnested_hlo, kernel_info);
}
kernel_generator.GetBlockPrologueGenerator()(unnested_hlo, kernel_info);
EmitBlock(std::move(emit_one_tile), kernel_info, &ksl, index_ty);
const BlockEpilogueGenerator& block_epilogue_generator =
kernel_generator.GetBlockEpilogueGenerator();
if (block_epilogue_generator) {
block_epilogue_generator(unnested_hlo, kernel_info);
}
kernel_generator.GetBlockEpilogueGenerator()(unnested_hlo, kernel_info);
return launch_dimensions;
}

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

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