|
| 1 | +/* |
| 2 | + * Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | + * All rights reserved. |
| 4 | + * |
| 5 | + * This source code is licensed under the BSD-style license found in the |
| 6 | + * LICENSE file in the root directory of this source tree. |
| 7 | + */ |
| 8 | +#pragma once |
| 9 | + |
| 10 | +#include "cortex_m_ops_common.h" |
| 11 | +extern "C" { |
| 12 | +#include "arm_nnfunctions.h" |
| 13 | +} |
| 14 | + |
| 15 | +namespace cortex_m { |
| 16 | +namespace native { |
| 17 | + |
| 18 | +// During AOT phase, quantized_linear_fusion_pass allocates total buffer |
| 19 | +// and passes in as 'Tensor'. (Total buffer = 8-byte header + x bytes) |
| 20 | +// ┌─────────────────┬─────────────────────────────────────┐ |
| 21 | +// │ KernelSum Header│ CMSIS Workspace │ |
| 22 | +// │ (8 bytes) │ (x bytes) │ |
| 23 | +// └─────────────────┴─────────────────────────────────────┘ |
| 24 | +// │ │ |
| 25 | +// │ └─> Passed to CMSIS API |
| 26 | +// │ |
| 27 | +// └─> State for kernel sum |
| 28 | + |
| 29 | +// C++ Runtime: |
| 30 | +// ┌─────────────────┬─────────────────────────────────────┐ |
| 31 | +// │ KernelSum Header│ CMSIS Workspace │ |
| 32 | +// │ (8 bytes) │ (x bytes) │ |
| 33 | +// └─────────────────┴─────────────────────────────────────┘ |
| 34 | +// ^ ^ |
| 35 | +// │ │ |
| 36 | +// scratch_ptr cmsis_workspace_ptr |
| 37 | +// │ │ |
| 38 | +// ▼ ▼ |
| 39 | +// arm_vector_sum_s8() writes kernel sums (with bias if avail): |
| 40 | +// [sum₀+bias₀][sum₁+bias₁][sum₂+bias₂]...[sum_{n-1}+bias_{n-1}] |
| 41 | +// (n * 4-byte int32_t values = x bytes) |
| 42 | +// |
| 43 | +// - n = out_features (number of output features) |
| 44 | +// - x = n * 4 bytes (total CMSIS buffer size) |
| 45 | +// - Total buffer = 8 + x bytes |
| 46 | + |
| 47 | +class CMSISScratchBufferContext final { |
| 48 | + public: |
| 49 | + CMSISScratchBufferContext( |
| 50 | + Tensor& scratch_buffer, |
| 51 | + const Tensor& weights, |
| 52 | + const Tensor& weight_zero_point, |
| 53 | + const torch::executor::optional<Tensor>& bias) |
| 54 | + : scratch_ptr_(scratch_buffer.mutable_data_ptr<int8_t>()), |
| 55 | + total_size_(scratch_buffer.size(0)), |
| 56 | + base_ptr_(reinterpret_cast<uint8_t*>(scratch_ptr_)), |
| 57 | + in_features_(weights.size(1)), |
| 58 | + out_features_(weights.size(0)), |
| 59 | + is_per_channel_(weight_zero_point.numel() > 1), |
| 60 | + weight_data_offset_(calculate_offset(weights.const_data_ptr<int8_t>())), |
| 61 | + weight_zp_data_offset_( |
| 62 | + calculate_offset(weight_zero_point.const_data_ptr<int32_t>())), |
| 63 | + bias_data_offset_( |
| 64 | + bias.has_value() |
| 65 | + ? calculate_offset(bias.value().const_data_ptr<int32_t>()) |
| 66 | + : 0), |
| 67 | + header_(reinterpret_cast<KernelSumHeader*>(scratch_ptr_)), |
| 68 | + cmsis_workspace_ptr_(scratch_ptr_ + KERNEL_SUM_HEADER_SIZE) { |
| 69 | + cmsis_nn_dims filter_dims = {in_features_, 1, 1, out_features_}; |
| 70 | + validate_size(filter_dims); |
| 71 | + } |
| 72 | + |
| 73 | + cmsis_nn_context get_cmsis_ctx() const { |
| 74 | + cmsis_nn_context ctx; |
| 75 | + ET_CHECK_MSG( |
| 76 | + reinterpret_cast<uintptr_t>(cmsis_workspace_ptr_) % 4 == 0, |
| 77 | + "CMSIS workspace not 4-byte aligned"); |
| 78 | + ctx.buf = cmsis_workspace_ptr_; |
| 79 | + ctx.size = get_cmsis_workspace_size(); |
| 80 | + return ctx; |
| 81 | + } |
| 82 | + |
| 83 | + bool is_kernel_sum_updated() const { |
| 84 | + return header_->updated; |
| 85 | + } |
| 86 | + |
| 87 | + void compute_kernel_sums_if_needed() { |
| 88 | + if (!header_->updated) { |
| 89 | + arm_vector_sum_s8( |
| 90 | + reinterpret_cast<int32_t*>(cmsis_workspace_ptr_), |
| 91 | + in_features_, |
| 92 | + out_features_, |
| 93 | + get_weight_data(), |
| 94 | + get_weight_zp_data()[0], |
| 95 | + 0, |
| 96 | + get_bias_data()); |
| 97 | + header_->updated = true; |
| 98 | + ET_LOG( |
| 99 | + Info, |
| 100 | + "Computed kernel sums. [required_bytes : %d]", |
| 101 | + header_->required_size); |
| 102 | + } |
| 103 | + } |
| 104 | + |
| 105 | + const int8_t* get_weight_data() const { |
| 106 | + return reinterpret_cast<const int8_t*>(base_ptr_ + weight_data_offset_); |
| 107 | + } |
| 108 | + |
| 109 | + const int32_t* get_weight_zp_data() const { |
| 110 | + return reinterpret_cast<const int32_t*>(base_ptr_ + weight_zp_data_offset_); |
| 111 | + } |
| 112 | + |
| 113 | + const int32_t* get_bias_data() const { |
| 114 | + return bias_data_offset_ == 0 |
| 115 | + ? nullptr |
| 116 | + : reinterpret_cast<const int32_t*>(base_ptr_ + bias_data_offset_); |
| 117 | + } |
| 118 | + |
| 119 | + bool is_per_channel_quant() const { |
| 120 | + return is_per_channel_; |
| 121 | + } |
| 122 | + int32_t get_in_features() const { |
| 123 | + return in_features_; |
| 124 | + } |
| 125 | + int32_t get_out_features() const { |
| 126 | + return out_features_; |
| 127 | + } |
| 128 | + |
| 129 | + private: |
| 130 | + static constexpr size_t KERNEL_SUM_HEADER_SIZE = 8; |
| 131 | + |
| 132 | + // Header for kernel sum computation state only |
| 133 | + struct KernelSumHeader { |
| 134 | + bool updated = false; |
| 135 | + int32_t required_size = 0; |
| 136 | + }; |
| 137 | + static_assert( |
| 138 | + sizeof(KernelSumHeader) == KERNEL_SUM_HEADER_SIZE, |
| 139 | + "KernelSumHeader must be exactly 8 bytes"); |
| 140 | + |
| 141 | + int8_t* scratch_ptr_; |
| 142 | + size_t total_size_; |
| 143 | + uint8_t* base_ptr_; |
| 144 | + |
| 145 | + // Context members |
| 146 | + const int32_t in_features_; |
| 147 | + const int32_t out_features_; |
| 148 | + const bool is_per_channel_; |
| 149 | + const uint32_t weight_data_offset_; |
| 150 | + const uint32_t weight_zp_data_offset_; |
| 151 | + const uint32_t bias_data_offset_; |
| 152 | + |
| 153 | + KernelSumHeader* header_; |
| 154 | + int8_t* cmsis_workspace_ptr_; |
| 155 | + |
| 156 | + uint32_t calculate_offset(const void* ptr) const { |
| 157 | + if (ptr == nullptr) |
| 158 | + return 0; |
| 159 | + |
| 160 | + const uint8_t* ptr_bytes = reinterpret_cast<const uint8_t*>(ptr); |
| 161 | + ET_CHECK_MSG(ptr_bytes >= base_ptr_, "Pointer is before base address"); |
| 162 | + |
| 163 | + const std::ptrdiff_t offset = ptr_bytes - base_ptr_; |
| 164 | + ET_CHECK_MSG( |
| 165 | + offset >= 0 && offset <= UINT32_MAX, "Offset out of valid range"); |
| 166 | + return static_cast<uint32_t>(offset); |
| 167 | + } |
| 168 | + |
| 169 | + size_t get_cmsis_workspace_size() const { |
| 170 | + return total_size_ - KERNEL_SUM_HEADER_SIZE; |
| 171 | + } |
| 172 | + |
| 173 | + void validate_size(const cmsis_nn_dims& filter_dims) const { |
| 174 | + header_->required_size = |
| 175 | + arm_fully_connected_s8_get_buffer_size(&filter_dims); |
| 176 | + |
| 177 | + ET_CHECK_MSG( |
| 178 | + get_cmsis_workspace_size() >= |
| 179 | + static_cast<size_t>(header_->required_size), |
| 180 | + "Scratch buffer size %zu insufficient for required size %d", |
| 181 | + get_cmsis_workspace_size(), |
| 182 | + header_->required_size); |
| 183 | + } |
| 184 | +}; |
| 185 | + |
| 186 | +} // namespace native |
| 187 | +} // namespace cortex_m |
0 commit comments