Function Softmax1x1

tensorflow/lite/delegates/gpu/metal/kernels/softmax.cc:172–214

}

std::vector<ComputeTaskDescriptorPtr> Softmax1x1(int id, ValueId input_id,
                                                 ValueId output_id,
                                                 int channels_count) {
  auto desc = std::make_shared<ComputeTaskDescriptor>();
  desc->id = id;
  desc->is_linkable = false;
  desc->shader_source = GetSoftmax1x1Code();

  desc->input_buffers = {
      {input_id, "device FLT4* const src_buffer"},
  };

  desc->output_buffer = {output_id, "device FLT4* dst_buffer",
                         [input_id](const std::map<ValueId, BHWC>& buffers) {
                           return buffers.find(input_id)->second;
                         }};

  desc->uniform_buffers = {
      {"constant uniforms& params",
       [channels_count](const std::map<ValueId, BHWC>& buffers) {
         const int src_depth = IntegralDivideRoundUp(channels_count, 4);
         struct uniforms {
           int4 size;
           float4 mask;
         };
         uniforms params;
         params.size = {src_depth, IntegralDivideRoundUp(src_depth, 32), 1, 1};
         params.mask = {0.0f, 0.0f, 0.0f, 0.0f};
         const int reminder = channels_count % 4 == 0 ? 4 : channels_count % 4;
         for (int i = 0; i < reminder; ++i) {
           params.mask[i] = 1.0f;
         }
         const uint8_t* ptr = reinterpret_cast<const uint8_t*>(&params);
         return std::vector<uint8_t>(ptr, ptr + sizeof(uniforms));
       }},
  };

  desc->resize_function = [](const std::map<ValueId, BHWC>& buffers) {
    return std::make_pair(uint3{32u, 1u, 1u}, uint3{1u, 1u, 1u});
  };

  return {desc};
}

}  // namespace metal
}  // namespace gpu