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

Batch support for MaxUnpooling.

Browse Source 
PiperOrigin-RevId: 272992496
This commit is contained in:
A. Unique TensorFlower 2019-10-04 17:44:13 -07:00 committed by TensorFlower Gardener
parent a59e1600e6
commit 8c650bc738
1 changed files with 64 additions and 50 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

114
tensorflow/lite/delegates/gpu/cl/kernels/max_unpooling.cc
View File

@ -28,65 +28,79 @@ namespace {
std::string GetMaxUnoolingKernelCode( std::string GetMaxUnoolingKernelCode(
const OperationDef& op_def, const CLDevice& device, const OperationDef& op_def, const CLDevice& device,
const std::vector<ElementwiseOperation*>& linked_operations) { const std::vector<ElementwiseOperation*>& linked_operations) {
TensorCodeGenerator src("src_data", "src_size", op_def.src_tensors[0]); TensorCodeGenerator src("src_data",
TensorCodeGenerator src_ind("src_data_indices", "src_size", {"src_size.x", "src_size.y", "src_size.z"},
op_def.src_tensors[0]);
TensorCodeGenerator src_ind("src_data_indices",
{"src_size.x", "src_size.y", "src_size.z"},
op_def.src_tensors[1]); op_def.src_tensors[1]);
TensorCodeGenerator dst("dst_data", "dst_size", op_def.dst_tensors[0]); TensorCodeGenerator dst("dst_data",
{"dst_size.x", "dst_size.y", "dst_size.z"},
op_def.dst_tensors[0]);
const auto address_mode = GetFastestZeroMode(device); const auto address_mode = GetFastestZeroMode(device);
std::string code = GetCommonDefines(op_def.precision); std::string c = GetCommonDefines(op_def.precision);
code += "__kernel void main_function(\n"; c += "__kernel void main_function(\n";
code += src.GetDeclaration(AccessType::READ) + ",\n"; c += src.GetDeclaration(AccessType::READ) + ",\n";
code += src_ind.GetDeclaration(AccessType::READ); c += src_ind.GetDeclaration(AccessType::READ);
code += GetArgsDeclaration(linked_operations); c += GetArgsDeclaration(linked_operations);
code += dst.GetDeclaration(AccessType::WRITE) + ",\n"; c += dst.GetDeclaration(AccessType::WRITE) + ",\n";
code += " int4 src_size, \n"; c += " int4 src_size, \n";
code += " int4 dst_size, \n"; c += " int4 dst_size, \n";
code += " int2 kernel_size, \n"; c += " int2 kernel_size, \n";
code += " int2 padding, \n"; c += " int2 padding, \n";
code += " int2 stride \n"; c += " int2 stride \n";
code += ") {\n"; c += ") {\n";
code += " int X = get_global_id(0);\n"; c += " int X = get_global_id(0);\n";
code += " int Y = get_global_id(1);\n"; c += " int Y = get_global_id(1);\n";
code += " int Z = get_global_id(2);\n"; c += " int Z = get_global_id(2);\n";
code += c += " if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.z) return;\n";
" if (X >= dst_size.x || Y >= dst_size.y || Z >= dst_size.w) return; \n"; if (op_def.batch_support) {
code += " int src_x = (X + padding.x) / stride.x;\n"; c += " int B = get_global_id(0) % dst_size.w;\n";
code += " int src_y = (Y + padding.y) / stride.y;\n"; c += " int X0 = get_global_id(0) / dst_size.w;\n";
code += " " + src.GetAddress("src_adr", "src_x", "src_y", "Z") + "\n"; c += " int src_x0 = (X0 + padding.x) / stride.x;\n";
if (op_def.src_tensors[0].storage_type == TensorStorageType::BUFFER) { c += " int src_x = src_x0 * dst_size.w + B;\n";
code += " bool outside = src_x < 0 || src_y < 0 ||";
code += " src_x >= src_size.x || src_y >= src_size.y;\n";
code += " FLT4 src = (FLT4)(0.0f);\n";
code += " int4 ind = (int4)(0);\n";
code += " if (!outside) {\n";
code += " src = " + src.Read("src_adr", TextureAddressMode::DONT_CARE) +
";\n";
code += " ind = convert_int4(" +
src_ind.Read("src_adr", TextureAddressMode::DONT_CARE) + ");\n";
code += " }\n";
} else { } else {
code += " FLT4 src = " + src.Read("src_adr", address_mode) + ";\n"; c += " int src_x = (X + padding.x) / stride.x;\n";
code += " int4 ind = convert_int4(" +
src_ind.Read("src_adr", address_mode) + ");\n";
} }
code += " int t_x = X - (src_x * stride.x - padding.x);\n"; c += " int src_y = (Y + padding.y) / stride.y;\n";
code += " int t_y = Y - (src_y * stride.y - padding.y);\n"; c += " " + src.GetAddress("src_adr", "src_x", "src_y", "Z") + "\n";
code += " int t_index = t_y * kernel_size.x + t_x;\n"; if (op_def.src_tensors[0].storage_type == TensorStorageType::BUFFER) {
code += " FLT4 result;\n"; c += " bool outside = src_x < 0 || src_y < 0 ||";
c += " src_x >= src_size.x || src_y >= src_size.y;\n";
c += " FLT4 src = (FLT4)(0.0f);\n";
c += " int4 ind = (int4)(0);\n";
c += " if (!outside) {\n";
c += " src = " + src.Read("src_adr", TextureAddressMode::DONT_CARE) +
";\n";
c += " ind = convert_int4(" +
src_ind.Read("src_adr", TextureAddressMode::DONT_CARE) + ");\n";
c += " }\n";
} else {
c += " FLT4 src = " + src.Read("src_adr", address_mode) + ";\n";
c += " int4 ind = convert_int4(" + src_ind.Read("src_adr", address_mode) +
");\n";
}
if (op_def.batch_support) {
c += " int t_x = X0 - (src_x0 * stride.x - padding.x);\n";
} else {
c += " int t_x = X - (src_x * stride.x - padding.x);\n";
}
c += " int t_y = Y - (src_y * stride.y - padding.y);\n";
c += " int t_index = t_y * kernel_size.x + t_x;\n";
c += " FLT4 result;\n";
const std::string channels[] = {".x", ".y", ".z", ".w"}; const std::string channels[] = {".x", ".y", ".z", ".w"};
for (int i = 0; i < 4; ++i) { for (int i = 0; i < 4; ++i) {
const auto& s = channels[i]; const auto& s = channels[i];
code += " result" + s + "= t_index == ind" + s + "? src" + s + ": 0.0f;\n"; c += " result" + s + "= t_index == ind" + s + "? src" + s + ": 0.0f;\n";
} }
const LinkingContext context{"result", "X", "Y", "Z"}; c += PostProcess(linked_operations, {"result", "X", "Y", "Z"});
code += PostProcess(linked_operations, context); c += " " + dst.Write3D("result", "X", "Y", "Z");
code += " " + dst.Write3D("result", "X", "Y", "Z"); c += "}\n";
code += "}\n";
return code; return c;
} }
} // namespace } // namespace
@ -131,8 +145,8 @@ Status MaxUnpooling::BindArguments() {
RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[1]->GetMemoryPtr())); RETURN_IF_ERROR(kernel_.SetMemoryAuto(src_[1]->GetMemoryPtr()));
RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_)); RETURN_IF_ERROR(BindArgs(&kernel_, linked_operations_));
RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting())); RETURN_IF_ERROR(kernel_.SetMemoryAuto(dst_[0]->GetMemoryPtrForWriting()));
RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetSizeWithDepth())); RETURN_IF_ERROR(kernel_.SetBytesAuto(src_[0]->GetWBatchedHDB()));
RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetSizeWithDepth())); RETURN_IF_ERROR(kernel_.SetBytesAuto(dst_[0]->GetWBatchedHDB()));
RETURN_IF_ERROR(kernel_.SetBytesAuto(kernel_size_)); RETURN_IF_ERROR(kernel_.SetBytesAuto(kernel_size_));
RETURN_IF_ERROR(kernel_.SetBytesAuto(padding_)); RETURN_IF_ERROR(kernel_.SetBytesAuto(padding_));
RETURN_IF_ERROR(kernel_.SetBytesAuto(stride_)); RETURN_IF_ERROR(kernel_.SetBytesAuto(stride_));
@ -141,7 +155,7 @@ Status MaxUnpooling::BindArguments() {
} }
int3 MaxUnpooling::GetGridSize() const { int3 MaxUnpooling::GetGridSize() const {
const int grid_x = dst_[0]->Width(); const int grid_x = dst_[0]->Width() * dst_[0]->Batch();
const int grid_y = dst_[0]->Height(); const int grid_y = dst_[0]->Height();
const int grid_z = dst_[0]->Depth(); const int grid_z = dst_[0]->Depth();
return int3(grid_x, grid_y, grid_z); return int3(grid_x, grid_y, grid_z);