diff --git a/paddle/fluid/operators/math/gru_compute.cu b/paddle/fluid/operators/math/gru_compute.cu index b564f990b4920a3a01b6ce0dd53e8f5e5d0464aa..cf3d57b0630dab2c06028b1afb8c7adaf0fe12cf 100644 --- a/paddle/fluid/operators/math/gru_compute.cu +++ b/paddle/fluid/operators/math/gru_compute.cu @@ -31,23 +31,41 @@ struct GRUUnitFunctor { dim3 grid; if (batch_size == 1) { if (context.GetComputeCapability() >= 70) { - constexpr int tiled_size = 16; - int frame_blocks = (frame_size * 2 + tiled_size - 1) / tiled_size; - threads = dim3(tiled_size, 1); - grid = dim3(frame_blocks, 1); - detail::KeFastCollectiveGruGate< - T, tiled_size><<>>( - value.gate_value, value.prev_out_value, value.gate_weight, - value.reset_output_value, frame_size, active_gate); - - frame_blocks = (frame_size + tiled_size - 1) / tiled_size; - grid = dim3(frame_blocks, 1); - detail::KeFastCollectiveGruOut< - T, tiled_size><<>>( - value.state_weight, value.prev_out_value, value.output_value, - value.gate_value, value.reset_output_value, frame_size, active_node, - origin_mode); - + if (frame_size < 16) { + constexpr int tiled_size = 8; + int frame_blocks = (frame_size * 2 + tiled_size - 1) / tiled_size; + threads = dim3(tiled_size, 1); + grid = dim3(frame_blocks, 1); + detail::KeFastCollectiveGruGate< + T, tiled_size><<>>( + value.gate_value, value.prev_out_value, value.gate_weight, + value.reset_output_value, frame_size, active_gate); + + frame_blocks = (frame_size + tiled_size - 1) / tiled_size; + grid = dim3(frame_blocks, 1); + detail::KeFastCollectiveGruOut< + T, tiled_size><<>>( + value.state_weight, value.prev_out_value, value.output_value, + value.gate_value, value.reset_output_value, frame_size, + active_node, origin_mode); + } else { + constexpr int tiled_size = 16; + int frame_blocks = (frame_size * 2 + tiled_size - 1) / tiled_size; + threads = dim3(tiled_size, 1); + grid = dim3(frame_blocks, 1); + detail::KeFastCollectiveGruGate< + T, tiled_size><<>>( + value.gate_value, value.prev_out_value, value.gate_weight, + value.reset_output_value, frame_size, active_gate); + + frame_blocks = (frame_size + tiled_size - 1) / tiled_size; + grid = dim3(frame_blocks, 1); + detail::KeFastCollectiveGruOut< + T, tiled_size><<>>( + value.state_weight, value.prev_out_value, value.output_value, + value.gate_value, value.reset_output_value, frame_size, + active_node, origin_mode); + } return; } else { int frame_per_block = frame_size <= 1024 ? frame_size : 1024;