wehub-resource-sync
71 lines
2.5 KiB
C++
Executable File
71 lines
2.5 KiB
C++
Executable File
//
|
|
// MetalSharedGather.hpp
|
|
// MNN
|
|
|
|
#ifndef MetalSharedGather_hpp
|
|
#define MetalSharedGather_hpp
|
|
|
|
#import "MetalExecution.hpp"
|
|
#import "MNN_generated.h"
|
|
|
|
#if MNN_METAL_ENABLED
|
|
namespace MNN {
|
|
|
|
// SharedGather implementation on Metal backend.
|
|
// It reuses quantized 1x1 convolution weights and dequantization parameters
|
|
// to gather selected output-channel rows into a dense floating-point matrix.
|
|
class MetalSharedGather : public MetalExecution {
|
|
public:
|
|
MetalSharedGather(Backend *backend,
|
|
int oc,
|
|
std::shared_ptr<Tensor> weight,
|
|
std::shared_ptr<Tensor> dequantScaleBias,
|
|
int dequantBits,
|
|
float scaleCoef);
|
|
virtual ~MetalSharedGather() = default;
|
|
|
|
virtual ErrorCode onResize(const std::vector<Tensor *> &inputs,
|
|
const std::vector<Tensor *> &outputs) override;
|
|
|
|
virtual void onEncode(const std::vector<Tensor *> &inputs,
|
|
const std::vector<Tensor *> &outputs,
|
|
id<MTLComputeCommandEncoder> encoder) override;
|
|
|
|
virtual bool onClone(Backend *bn, const Op *op, Execution **dst) override;
|
|
|
|
private:
|
|
// conv1x1_constants host-side mirror, used by conv1x1_w_dequant and shared_gather kernels
|
|
struct Conv1x1Constants {
|
|
int input_size; // repurposed as ic for SharedGather
|
|
int input_slice; // ic_4
|
|
int output_width; // selectSize
|
|
int output_height; // unused for SharedGather
|
|
int output_size; // total elements (selectSize * ic)
|
|
int output_slice; // oc_4
|
|
int output_channel; // oc
|
|
int batch; // unused for SharedGather
|
|
int block_size; // quant block size along K axis
|
|
int activation; // not used (no activation)
|
|
float scale_coef; // scale normalization factor
|
|
};
|
|
|
|
private:
|
|
int mOc = 0; // number of output channels (rows in weight)
|
|
std::shared_ptr<Tensor> mWeight; // quantized int8/4 weight
|
|
std::shared_ptr<Tensor> mDequantScaleBias; // packed scale + bias
|
|
int mDequantBits = 0; // 4 or 8
|
|
float mScaleCoef = 1.0f;
|
|
|
|
// direct int4/int8 quant gather pipeline (preferred path)
|
|
id<MTLComputePipelineState> mQuantPipeline = nil;
|
|
std::pair<MTLSize, MTLSize> mQuantThreads;
|
|
|
|
// constant buffer shared by dequant and gather kernels
|
|
id<MTLBuffer> mConstBuffer = nil;
|
|
};
|
|
|
|
} // namespace MNN
|
|
#endif /* MNN_METAL_ENABLED */
|
|
|
|
#endif /* MetalSharedGather_hpp */
|