wehub-resource-sync
232 lines
8.8 KiB
C++
232 lines
8.8 KiB
C++
// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#pragma once
|
|
|
|
#include "paddle/phi/kernels/fake_quantize_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/fake_quantize_functor.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void FakeQuantizeRangeAbsMaxKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
const DenseTensor &in_scale,
|
|
const optional<DenseTensor> &iter,
|
|
int window_size,
|
|
int bit_length,
|
|
bool is_test,
|
|
int round_type,
|
|
DenseTensor *out,
|
|
DenseTensor *out_scale,
|
|
DenseTensor *out_scales) {
|
|
dev_ctx.template Alloc<T>(out);
|
|
int bin_cnt = std::pow(2, bit_length - 1) - 1;
|
|
|
|
// testing
|
|
if (is_test) {
|
|
funcs::ClipAndFakeQuantFunctor<Context, T>()(
|
|
dev_ctx, x, in_scale, bin_cnt, round_type, out);
|
|
return;
|
|
}
|
|
|
|
// training
|
|
dev_ctx.template Alloc<T>(out_scale);
|
|
|
|
DenseTensor cur_scale;
|
|
cur_scale.Resize({1});
|
|
T *cur_scale_data = dev_ctx.template Alloc<T>(&cur_scale);
|
|
funcs::FindAbsMaxFunctor<Context, T>()(
|
|
dev_ctx, x.data<T>(), x.numel(), cur_scale_data);
|
|
funcs::FindRangeAbsMaxFunctor<Context, T>()(dev_ctx,
|
|
cur_scale,
|
|
in_scale,
|
|
iter.get(),
|
|
window_size,
|
|
out_scales,
|
|
out_scale);
|
|
funcs::ClipAndFakeQuantFunctor<Context, T>()(
|
|
dev_ctx, x, *out_scale, bin_cnt, round_type, out);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FakeQuantizeAbsMaxKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
int bit_length,
|
|
int round_type,
|
|
DenseTensor *out,
|
|
DenseTensor *out_scale) {
|
|
T *out_s = dev_ctx.template Alloc<T>(out_scale);
|
|
int bin_cnt = std::pow(2, bit_length - 1) - 1;
|
|
const T *in_data = x.data<T>();
|
|
funcs::FindAbsMaxFunctor<Context, T> find_abs_max_functor;
|
|
find_abs_max_functor(dev_ctx, in_data, x.numel(), out_s);
|
|
|
|
funcs::ClipAndFakeQuantFunctor<Context, T> clip_and_fake_quant_functor;
|
|
clip_and_fake_quant_functor(dev_ctx, x, *out_scale, bin_cnt, round_type, out);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FakeQuantOrWithDequantMovingAverageAbsMaxKernel(
|
|
const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
const DenseTensor &in_scale,
|
|
const optional<DenseTensor> &in_accum,
|
|
const optional<DenseTensor> &in_state,
|
|
float moving_rate,
|
|
int bit_length,
|
|
bool is_test,
|
|
int round_type,
|
|
DenseTensor *out,
|
|
DenseTensor *out_scale,
|
|
DenseTensor *out_state,
|
|
DenseTensor *out_accum) {
|
|
int bin_cnt = std::pow(2, bit_length - 1) - 1;
|
|
|
|
// testing
|
|
if (is_test) {
|
|
funcs::ClipAndFakeQuantFunctor<Context, T>()(
|
|
dev_ctx, x, in_scale, bin_cnt, round_type, out);
|
|
return;
|
|
}
|
|
|
|
// training
|
|
DenseTensor tmp_scale;
|
|
tmp_scale.Resize(common::make_dim(1));
|
|
T *cur_scale_data = dev_ctx.template Alloc<T>(&tmp_scale);
|
|
|
|
funcs::FindAbsMaxFunctor<Context, T>()(
|
|
dev_ctx, x.data<T>(), x.numel(), cur_scale_data);
|
|
|
|
funcs::FindMovingAverageAbsMaxFunctor<Context, T>()(dev_ctx,
|
|
in_accum.get(),
|
|
in_state.get(),
|
|
cur_scale_data,
|
|
moving_rate,
|
|
out_state,
|
|
out_accum,
|
|
out_scale);
|
|
|
|
funcs::ClipAndFakeQuantFunctor<Context, T>()(
|
|
dev_ctx, x, *out_scale, bin_cnt, round_type, out);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FakeChannelWiseQuantizeAbsMaxKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
int bit_length,
|
|
int round_type,
|
|
int quant_axis,
|
|
bool is_test,
|
|
DenseTensor *out,
|
|
DenseTensor *out_scale) {
|
|
dev_ctx.template Alloc<T>(out);
|
|
int bin_cnt = std::pow(2, bit_length - 1) - 1;
|
|
|
|
if (!is_test) {
|
|
T *out_scale_data = dev_ctx.template Alloc<T>(out_scale);
|
|
funcs::FindChannelAbsMaxFunctor<Context, T>()(
|
|
dev_ctx, x, quant_axis, out_scale_data);
|
|
}
|
|
funcs::ChannelClipAndFakeQuantFunctor<Context, T>()(
|
|
dev_ctx, x, *out_scale, bin_cnt, round_type, quant_axis, out);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FakeChannelWiseQuantizeDequantizeAbsMaxKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
int bit_length,
|
|
int round_type,
|
|
int quant_axis,
|
|
DenseTensor *out,
|
|
DenseTensor *out_scale) {
|
|
T *out_scale_data = dev_ctx.template Alloc<T>(out_scale);
|
|
dev_ctx.template Alloc<T>(out);
|
|
int bin_cnt = std::pow(2, bit_length - 1) - 1;
|
|
|
|
funcs::FindChannelAbsMaxFunctor<Context, T>()(
|
|
dev_ctx, x, quant_axis, out_scale_data);
|
|
|
|
funcs::ChannelClipFakeQuantDequantFunctor<Context, T>()(
|
|
dev_ctx, x, *out_scale, bin_cnt, round_type, quant_axis, out);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FakeQuantizeDequantizeMovingAverageAbsMaxKernel(
|
|
const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
const DenseTensor &in_scale,
|
|
const optional<DenseTensor> &in_accum,
|
|
const optional<DenseTensor> &in_state,
|
|
float moving_rate,
|
|
int bit_length,
|
|
bool is_test,
|
|
int round_type,
|
|
DenseTensor *out,
|
|
DenseTensor *out_scale,
|
|
DenseTensor *out_state,
|
|
DenseTensor *out_accum) {
|
|
dev_ctx.template Alloc<T>(out);
|
|
int bin_cnt = std::pow(2, bit_length - 1) - 1;
|
|
|
|
// testing
|
|
if (is_test) {
|
|
funcs::ClipAndFakeQuantDequantFunctor<Context, T>()(
|
|
dev_ctx, x, in_scale, bin_cnt, round_type, out);
|
|
return;
|
|
}
|
|
|
|
// training
|
|
|
|
DenseTensor tmp_scale;
|
|
tmp_scale.Resize(common::make_dim(1));
|
|
T *cur_scale_data = dev_ctx.template Alloc<T>(&tmp_scale);
|
|
funcs::FindAbsMaxFunctor<Context, T>()(
|
|
dev_ctx, x.data<T>(), x.numel(), cur_scale_data);
|
|
|
|
dev_ctx.template Alloc<T>(out_state);
|
|
dev_ctx.template Alloc<T>(out_accum);
|
|
dev_ctx.template Alloc<T>(out_scale);
|
|
|
|
funcs::FindMovingAverageAbsMaxFunctor<Context, T>()(dev_ctx,
|
|
in_accum.get(),
|
|
in_state.get(),
|
|
cur_scale_data,
|
|
moving_rate,
|
|
out_state,
|
|
out_accum,
|
|
out_scale);
|
|
|
|
funcs::ClipAndFakeQuantDequantFunctor<Context, T>()(
|
|
dev_ctx, x, *out_scale, bin_cnt, round_type, out);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FakeQuantizeDequantizeAbsMaxKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
int bit_length,
|
|
int round_type,
|
|
DenseTensor *out,
|
|
DenseTensor *out_scale) {
|
|
T *out_s = dev_ctx.template Alloc<T>(out_scale);
|
|
int bin_cnt = std::pow(2, bit_length - 1) - 1;
|
|
const T *in_data = x.data<T>();
|
|
funcs::FindAbsMaxFunctor<Context, T>()(dev_ctx, in_data, x.numel(), out_s);
|
|
|
|
funcs::ClipAndFakeQuantDequantFunctor<Context, T>()(
|
|
dev_ctx, x, *out_scale, bin_cnt, round_type, out);
|
|
}
|
|
|
|
} // namespace phi
|