Skip to content

Refactor to avoid using xnn_define_static_broadcast #10688

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Refactor to avoid using xnn_define_static_broadcast #10688

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions litert/tensor/backends/xnnpack/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ cc_library(
hdrs = ["utils.h"],
deps = [
"//:XNNPACK",
"//litert/tensor/internal:graph",
"//litert/tensor/utils:macros",
"@abseil-cpp//absl/status",
"@abseil-cpp//absl/strings",
Expand Down
120 changes: 109 additions & 11 deletions litert/tensor/backends/xnnpack/arithmetic.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,78 @@ absl::StatusOr<uint32_t> DynamicallyQuantizeInput(
return qd_id;
}

absl::StatusOr<uint32_t> DefineZeroTensor(
XnnpackBuildContext& ctx, const graph::TensorInformation& input_info,
size_t num_dims, const size_t* dims, absl::string_view op_name) {
xnn_datatype datatype = GetXnnpackType(input_info);
if (datatype == xnn_datatype_invalid) {
return absl::UnimplementedError(
absl::StrFormat("%s: unsupported input type %d", op_name,
static_cast<int>(input_info.type)));
}

size_t num_elements = 1;
for (size_t i = 0; i < num_dims; ++i) {
num_elements *= dims[i];
}

uint32_t zero_id = XNN_INVALID_VALUE_ID;
const void* data_ptr = nullptr;

if (input_info.quantization) {
if (auto maybe_pcq =
input_info.quantization->As<PerChannelAffineQuantization>();
maybe_pcq.ok()) {
const auto& pcq = maybe_pcq.value();
if (pcq.scales.size() == 1) {
int32_t zero_point = pcq.zero_points.empty() ? 0 : pcq.zero_points[0];
float scale = pcq.scales[0];

if (datatype == xnn_datatype_qint8) {
std::vector<int8_t> zeros(num_elements,
static_cast<int8_t>(zero_point));
data_ptr = ctx.KeepAlive(std::move(zeros));
} else {
return absl::UnimplementedError(
absl::StrFormat("%s: unsupported quantized type %d", op_name,
static_cast<int>(datatype)));
}

LRT_TENSOR_RETURN_IF_ERROR(xnn_define_quantized_tensor_value(
ctx.subgraph(), datatype, zero_point, scale, num_dims, dims,
data_ptr,
/*external_id=*/XNN_INVALID_VALUE_ID, /*flags=*/0, &zero_id))
<< "Could not define quantized zero tensor.";
} else {
return absl::UnimplementedError(absl::StrFormat(
"%s: per-channel quantized zero tensor not supported", op_name));
}
} else {
return absl::UnimplementedError(
absl::StrFormat("%s: unsupported quantization type", op_name));
}
} else {
if (datatype == xnn_datatype_fp32) {
std::vector<float> zeros(num_elements, 0.0f);
data_ptr = ctx.KeepAlive(std::move(zeros));
} else if (datatype == xnn_datatype_int32) {
std::vector<int32_t> zeros(num_elements, 0);
data_ptr = ctx.KeepAlive(std::move(zeros));
} else {
return absl::UnimplementedError(
absl::StrFormat("%s: unsupported type %d for zero tensor", op_name,
static_cast<int>(datatype)));
}

LRT_TENSOR_RETURN_IF_ERROR(xnn_define_tensor_value(
ctx.subgraph(), datatype, num_dims, dims, data_ptr,
/*external_id=*/XNN_INVALID_VALUE_ID, /*flags=*/0, &zero_id))
<< "Could not define zero tensor.";
}

return zero_id;
}

template <Type... Types>
absl::Status ValidateTensorType(const graph::Tensor& tensor,
absl::string_view op_name) {
Expand Down Expand Up @@ -1316,7 +1388,7 @@ absl::Status OpMixin<TileOperation, XnnpackMixinTag>::ToXnnpack(
input_info.shape.size(), num_dims));
}

std::vector<size_t> new_shape(num_dims);
std::vector<size_t> zero_shape(num_dims);
for (size_t i = 0; i < num_dims; ++i) {
int mult = multiples_data[i];
int dim = input_info.shape[i];
Expand All @@ -1326,13 +1398,21 @@ absl::Status OpMixin<TileOperation, XnnpackMixinTag>::ToXnnpack(
"Dimension %d has size %d but multiples[%d]=%d",
op_name, i, dim, i, mult));
}
new_shape[i] = static_cast<size_t>(dim * mult);
zero_shape[i] = static_cast<size_t>(mult);
}

LRT_TENSOR_RETURN_IF_ERROR(xnn_define_static_broadcast(
ctx.subgraph(), num_dims, new_shape.data(), input_id, output_id,
LRT_TENSOR_ASSIGN_OR_RETURN(
uint32_t zero_id, DefineZeroTensor(ctx, input_info, zero_shape.size(),
zero_shape.data(), op_name));

LRT_TENSOR_ASSIGN_OR_RETURN(auto params,
BuildBinaryParams(kActNone, op_name));

LRT_TENSOR_RETURN_IF_ERROR(xnn_define_binary(
ctx.subgraph(), xnn_binary_add, &params, input_id, zero_id, output_id,
/*flags=*/0))
<< op_name;

return absl::OkStatus();
}

Expand Down Expand Up @@ -1511,17 +1591,24 @@ OpMixin<ResizeNearestNeighborOperation, XnnpackMixinTag>::ToXnnpack(
static_cast<size_t>(scale_h), static_cast<size_t>(input_w),
static_cast<size_t>(scale_w), static_cast<size_t>(channels)};

xnn_datatype datatype = GetXnnpackType(input_info);
if (datatype == xnn_datatype_invalid) {
return absl::UnimplementedError(
absl::StrFormat("%s: unsupported input type %d", op_name,
static_cast<int>(input_info.type)));
}

uint32_t reshape_id = XNN_INVALID_VALUE_ID;
LRT_TENSOR_RETURN_IF_ERROR(xnn_define_tensor_value(
ctx.subgraph(), xnn_datatype_fp32, reshape_dims.size(),
reshape_dims.data(), /*data=*/nullptr, XNN_INVALID_VALUE_ID,
ctx.subgraph(), datatype, reshape_dims.size(), reshape_dims.data(),
/*data=*/nullptr, XNN_INVALID_VALUE_ID,
/*flags=*/0, &reshape_id))
<< op_name;

uint32_t broadcast_id = XNN_INVALID_VALUE_ID;
LRT_TENSOR_RETURN_IF_ERROR(xnn_define_tensor_value(
ctx.subgraph(), xnn_datatype_fp32, broadcast_dims.size(),
broadcast_dims.data(), /*data=*/nullptr, XNN_INVALID_VALUE_ID,
ctx.subgraph(), datatype, broadcast_dims.size(), broadcast_dims.data(),
/*data=*/nullptr, XNN_INVALID_VALUE_ID,
/*flags=*/0, &broadcast_id))
<< op_name;

Expand All @@ -1530,9 +1617,20 @@ OpMixin<ResizeNearestNeighborOperation, XnnpackMixinTag>::ToXnnpack(
reshape_id, /*flags=*/0))
<< op_name;

LRT_TENSOR_RETURN_IF_ERROR(xnn_define_static_broadcast(
ctx.subgraph(), broadcast_dims.size(), broadcast_dims.data(), reshape_id,
broadcast_id, /*flags=*/0))
const std::array<size_t, 6> zero_dims = {
1, 1, static_cast<size_t>(scale_h), 1, static_cast<size_t>(scale_w), 1};

LRT_TENSOR_ASSIGN_OR_RETURN(
uint32_t zero_id, DefineZeroTensor(ctx, input_info, zero_dims.size(),
zero_dims.data(), op_name));

LRT_TENSOR_ASSIGN_OR_RETURN(auto params,
BuildBinaryParams(kActNone, op_name));

LRT_TENSOR_RETURN_IF_ERROR(xnn_define_binary(ctx.subgraph(), xnn_binary_add,
&params, reshape_id, zero_id,
broadcast_id,
/*flags=*/0))
<< op_name;

const std::array<size_t, 4> output_dims = {
Expand Down
25 changes: 25 additions & 0 deletions litert/tensor/backends/xnnpack/arithmetic.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ limitations under the License.
#include <cstddef>
#include <cstdint>
#include <memory>
#include <utility>
#include <vector>

#include "include/xnnpack.h"
Expand All @@ -37,6 +38,21 @@ struct xnn_subgraph;

namespace litert::tensor {

class BufferHolder {
public:
virtual ~BufferHolder() = default;
};

template <typename T>
class TypedBufferHolder : public BufferHolder {
public:
explicit TypedBufferHolder(std::vector<T>&& vec) : vec_(std::move(vec)) {}
const T* data() const { return vec_.data(); }

private:
std::vector<T> vec_;
};

// Tag to identify the XNNPACK mixin.
struct XnnpackMixinTag {};

Expand Down Expand Up @@ -65,6 +81,14 @@ class XnnpackBuildContext {
// Returns the XNNPACK subgraph.
::xnn_subgraph* subgraph();

template <typename T>
const T* KeepAlive(std::vector<T>&& buffer) {
auto holder = std::make_unique<TypedBufferHolder<T>>(std::move(buffer));
const T* ptr = holder->data();
custom_buffers_.push_back(std::move(holder));
return ptr;
}

private:
xnn_subgraph* subgraph_ = nullptr;
std::vector<graph::Tensor> outputs_;
Expand All @@ -74,6 +98,7 @@ class XnnpackBuildContext {
absl::flat_hash_map<graph::Tensor, uint32_t> external_ids_;
std::vector<std::vector<float>> dequantized_buffers_;
std::vector<std::vector<fp16_t>> fp16_buffers_;
std::vector<std::unique_ptr<BufferHolder>> custom_buffers_;

friend absl::StatusOr<std::unique_ptr<XnnpackGraph>> BuildXnnpackGraph(
std::vector<TensorHandle> outputs);
Expand Down
68 changes: 13 additions & 55 deletions litert/tensor/backends/xnnpack/conversion.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,50 +56,6 @@ absl::Status EnsureXnnInitialized() {
return *g_xnn_init_status;
}

xnn_datatype GetXnnpackType(const XnnpackValue& value) {
switch (value.info.type) {
case Type::kUnknown:
case Type::kBOOL:
case Type::kI2:
case Type::kI4:
if (value.info.quantization) {
if (value.info.quantization->As<PerChannelAffineQuantization>().ok()) {
return xnn_datatype_qcint4;
} else if (value.info.quantization->As<BlockwiseQuantization>().ok()) {
return xnn_datatype_qbint4;
}
}
break;
case Type::kI8:
if (value.info.quantization) {
if (auto it =
value.info.quantization->As<PerChannelAffineQuantization>();
it.ok()) {
return it->scales.size() > 1 ? xnn_datatype_qcint8
: xnn_datatype_qint8;
}
}
break;
case Type::kI16:
case Type::kI64:
case Type::kU4:
case Type::kU8:
case Type::kU16:
case Type::kU32:
case Type::kU64:
case Type::kFP16:
return xnn_datatype_fp16;
case Type::kI32:
return xnn_datatype_int32;
case Type::kFP32:
return xnn_datatype_fp32;
case Type::kFP64:
break;
case Type::kBF16:
return xnn_datatype_bf16;
}
return xnn_datatype_invalid;
}

// TODO: b/493560478 - Decide whether to delete this from here.
[[maybe_unused]]
Expand Down Expand Up @@ -225,13 +181,15 @@ XnnpackGraph::XnnpackGraph(
absl::flat_hash_map<graph::Tensor, size_t> tensor_index,
absl::flat_hash_set<graph::Tensor> external_outputs,
std::vector<std::vector<float>> dequantized_buffers,
std::vector<std::vector<fp16_t>> fp16_buffers)
std::vector<std::vector<fp16_t>> fp16_buffers,
std::vector<std::unique_ptr<BufferHolder>> custom_buffers)
: subgraph_(subgraph),
values_(std::move(values)),
tensor_index_(std::move(tensor_index)),
external_outputs_(std::move(external_outputs)),
dequantized_buffers_(std::move(dequantized_buffers)),
fp16_buffers_(std::move(fp16_buffers)) {}
fp16_buffers_(std::move(fp16_buffers)),
custom_buffers_(std::move(custom_buffers)) {}

XnnpackGraph::~XnnpackGraph() {
if (subgraph_ != nullptr) {
Expand Down Expand Up @@ -285,7 +243,7 @@ absl::StatusOr<std::unique_ptr<XnnpackGraph>> XnnpackBuildContext::Finalize() {
return std::make_unique<XnnpackGraph>(
subgraph, std::move(values_), std::move(tensor_index_),
std::move(external_outputs_), std::move(dequantized_buffers_),
std::move(fp16_buffers_));
std::move(fp16_buffers_), std::move(custom_buffers_));
}

absl::StatusOr<std::unique_ptr<XnnpackGraph>> BuildXnnpackGraph(
Expand Down Expand Up @@ -370,18 +328,18 @@ absl::StatusOr<uint32_t> XnnpackBuildContext::DefineValue(
}

if (!info.quantization) {
LRT_TENSOR_RETURN_IF_ERROR(
xnn_define_tensor_value(subgraph_, GetXnnpackType(value), dims.size(),
dims.empty() ? nullptr : dims.data(), data_ptr,
external_id, value.flags, &value.id))
LRT_TENSOR_RETURN_IF_ERROR(xnn_define_tensor_value(
subgraph_, GetXnnpackType(value.info), dims.size(),
dims.empty() ? nullptr : dims.data(), data_ptr, external_id,
value.flags, &value.id))
<< "Could not define a new tensor value.";
} else if (auto maybe_pcq =
info.quantization->As<PerChannelAffineQuantization>();
maybe_pcq.ok()) {
const auto& pcq = maybe_pcq.value();
if (pcq.scales.size() == 1) {
LRT_TENSOR_RETURN_IF_ERROR(xnn_define_quantized_tensor_value(
subgraph_, GetXnnpackType(value),
subgraph_, GetXnnpackType(value.info),
pcq.zero_points.empty() ? 0 : pcq.zero_points[0], pcq.scales[0],
dims.size(), dims.empty() ? nullptr : dims.data(), data_ptr,
external_id, value.flags, &value.id))
Expand All @@ -408,7 +366,7 @@ absl::StatusOr<uint32_t> XnnpackBuildContext::DefineValue(
} else {
LRT_TENSOR_RETURN_IF_ERROR(
xnn_define_channelwise_quantized_tensor_value_v3(
subgraph_, GetXnnpackType(value), /*zero_point=*/0,
subgraph_, GetXnnpackType(value.info), /*zero_point=*/0,
pcq.scales.data(), dims.size(), pcq.quantized_dimension,
dims.empty() ? nullptr : dims.data(), data_ptr, external_id,
value.flags, &value.id, /*channelwise_zero_point=*/nullptr))
Expand All @@ -422,8 +380,8 @@ absl::StatusOr<uint32_t> XnnpackBuildContext::DefineValue(
const void* scale_ptr = fp16_buffers_.back().data();
int32_t zero_point = bwq.zero_points.empty() ? 0 : bwq.zero_points[0];
LRT_TENSOR_RETURN_IF_ERROR(xnn_define_blockwise_quantized_tensor_value_v2(
subgraph_, GetXnnpackType(value), zero_point, scale_ptr, dims.size(),
bwq.quantized_dimension, bwq.block_size,
subgraph_, GetXnnpackType(value.info), zero_point, scale_ptr,
dims.size(), bwq.quantized_dimension, bwq.block_size,
dims.empty() ? nullptr : dims.data(), data_ptr, external_id,
value.flags, xnn_datatype_fp16, &value.id))
<< "Could not define a new blockwise quantized tensor value.";
Expand Down
4 changes: 3 additions & 1 deletion litert/tensor/backends/xnnpack/conversion.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,8 @@ class XnnpackGraph {
absl::flat_hash_map<graph::Tensor, size_t> tensor_index,
absl::flat_hash_set<graph::Tensor> external_outputs,
std::vector<std::vector<float>> dequantized_buffers = {},
std::vector<std::vector<fp16_t>> fp16_buffers = {});
std::vector<std::vector<fp16_t>> fp16_buffers = {},
std::vector<std::unique_ptr<BufferHolder>> custom_buffers = {});
~XnnpackGraph();

// Returns the XNNPACK subgraph.
Expand All @@ -66,6 +67,7 @@ class XnnpackGraph {
absl::flat_hash_set<graph::Tensor> external_outputs_;
std::vector<std::vector<float>> dequantized_buffers_;
std::vector<std::vector<fp16_t>> fp16_buffers_;
std::vector<std::unique_ptr<BufferHolder>> custom_buffers_;
};

// Builds an XNNPACK graph from the given outputs.
Expand Down
Loading
Loading