-
Notifications
You must be signed in to change notification settings - Fork 74
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
#6938: Implement softplus as a single kernel
- Loading branch information
Showing
13 changed files
with
208 additions
and
46 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
50 changes: 50 additions & 0 deletions
50
tt_metal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/ckernel_sfpu_softplus.h
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,50 @@ | ||
| // SPDX-FileCopyrightText: © 2023 Tenstorrent Inc. | ||
| // | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| #pragma once | ||
|
|
||
| #include "ckernel.h" | ||
| #include "ckernel_defs.h" | ||
| #include "ckernel_sfpu_converter.h" | ||
| #include "ckernel_sfpu_exp.h" | ||
| #include "ckernel_sfpu_log.h" | ||
|
|
||
| using namespace sfpi; | ||
|
|
||
| namespace ckernel { | ||
| namespace sfpu { | ||
|
|
||
| template <bool APPROXIMATION_MODE> | ||
| inline void calculate_softplus_body(vFloat beta, vFloat beta_reciprocal, vFloat threshold) { | ||
| vFloat a = dst_reg[0]; | ||
| vFloat a_beta = a * beta; | ||
| v_if(a_beta < threshold) { | ||
| exp_init<APPROXIMATION_MODE, false>(); | ||
| a = calculate_exponential_body<APPROXIMATION_MODE>(a_beta) + 1.0f; | ||
|
|
||
| log_init<APPROXIMATION_MODE>(); | ||
| dst_reg[0] = a; | ||
| calculate_log_body<false>(0); | ||
| a = beta_reciprocal * dst_reg[0]; | ||
| } | ||
| v_endif; | ||
| dst_reg[0] = a; | ||
| } | ||
|
|
||
| template <bool APPROXIMATION_MODE, int ITERATIONS = 8> | ||
| inline void calculate_softplus(uint param0, uint param1, uint param2) { | ||
| vFloat beta = Converter::to_float(param0); | ||
| vFloat beta_reciprocal = Converter::to_float(param1); | ||
| vFloat threshold = Converter::to_float(param2); | ||
| for (int d = 0; d < ITERATIONS; d++) { | ||
| calculate_softplus_body<APPROXIMATION_MODE>(beta, beta_reciprocal, threshold); | ||
| dst_reg++; | ||
| } | ||
| } | ||
|
|
||
| template <bool APPROXIMATION_MODE> | ||
| void softplus_init() {} | ||
|
|
||
| } // namespace sfpu | ||
| } // namespace ckernel |
29 changes: 29 additions & 0 deletions
29
...tal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/llk_math_eltwise_unary_sfpu_softplus.h
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,29 @@ | ||
| // SPDX-FileCopyrightText: © 2023 Tenstorrent Inc. | ||
| // | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| #pragma once | ||
|
|
||
| #include "ckernel_sfpu_softplus.h" | ||
| #include "llk_math_eltwise_unary_sfpu_3_param.h" | ||
| #include "llk_math_eltwise_unary_sfpu_init.h" | ||
|
|
||
| namespace ckernel { | ||
|
|
||
| template <bool APPROXIMATE> | ||
| inline void llk_math_eltwise_unary_sfpu_softplus_init() { | ||
| llk_math_eltwise_unary_sfpu_init<SfpuType::softplus, APPROXIMATE>(); | ||
| } | ||
|
|
||
| template <bool APPROXIMATE> | ||
| inline void llk_math_eltwise_unary_sfpu_softplus( | ||
| uint dst_index, uint param0, uint param1, uint param2, int vector_mode = (int)VectorMode::RC) { | ||
| llk_math_eltwise_unary_sfpu_3_param<APPROXIMATE>( | ||
| ckernel::sfpu::calculate_softplus<APPROXIMATE>, | ||
| ckernel::sfpu::calculate_softplus<APPROXIMATE>, | ||
| dst_index, | ||
| vector_mode, | ||
| param0, param1, param2); | ||
| } | ||
|
|
||
| } // namespace ckernel |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -73,6 +73,7 @@ enum SfpuType { | |
| unary_ne, | ||
| unary_gt, | ||
| unary_lt, | ||
| softplus, | ||
| tiled_prod, | ||
| unused, | ||
| }; | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
47 changes: 47 additions & 0 deletions
47
tt_metal/include/compute_kernel_api/eltwise_unary/softplus.h
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,47 @@ | ||
| // SPDX-FileCopyrightText: © 2023 Tenstorrent Inc. | ||
| // | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| #pragma once | ||
|
|
||
|
|
||
| #include "compute_kernel_api/common_globals.h" | ||
| #ifdef TRISC_MATH | ||
| #include "llk_math_eltwise_unary_sfpu_softplus.h" | ||
| #define MAIN math_main() | ||
| #define MATH(x) x | ||
| #else | ||
| #define MATH(x) | ||
| #endif | ||
|
|
||
|
|
||
|
|
||
| namespace ckernel { | ||
|
|
||
| /** | ||
| * Performs element-wise computation of softplus (`1/beta * log(1 + exp(beta * x))`) on each element | ||
| * of a tile in DST register at index tile_index. Any input value greater than the provided threshold | ||
| * with return itself. The DST register buffer must be in acquired state via *acquire_dst* call. This | ||
| * call is blocking and is only available on the compute engine. | ||
| * | ||
| * Return value: None | ||
| * | ||
| * | Argument | Description | Type | Valid Range | Required | | ||
| * |-----------------|----------------------------------------------------------------------------|----------|-------------------------------------------------------|----------| | ||
| * | tile_index | The index of the tile in DST register buffer to perform the computation on | uint32_t | Must be less than the size of the DST register buffer | True | | ||
| * | beta | Beta used in softplus calculation | uint32_t | Greater than 0 | True | | ||
| * | beta_reciprocal | Reciprocal of beta (1/beta) used in softplus calculation | uint32_t | Greater than 0 | True | | ||
| * | threshold | Threshold used in softplus calculation | uint32_t | Greater than 0 | True | | ||
| */ | ||
| ALWI void softplus_tile(uint32_t idst, uint32_t beta, uint32_t beta_reciprocal, uint32_t threshold) { | ||
| MATH(( llk_math_eltwise_unary_sfpu_softplus<APPROX>(idst, beta, beta_reciprocal, threshold) )); | ||
| } | ||
|
|
||
| /** | ||
| * Please refer to documentation for any_init. | ||
| */ | ||
| ALWI void softplus_tile_init() { | ||
| MATH(( llk_math_eltwise_unary_sfpu_softplus_init<APPROX>() )); | ||
| } | ||
|
|
||
| } // namespace ckernel |
Oops, something went wrong.