Skip to content

Commit

Permalink
softmax: Correct to include exponential
Browse files Browse the repository at this point in the history
  • Loading branch information
colluca committed Oct 24, 2023
1 parent 68a927c commit 051aaa1
Showing 1 changed file with 15 additions and 9 deletions.
24 changes: 15 additions & 9 deletions sw/dnn/src/softmax.h
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
// Copyright 2020 ETH Zurich and University of Bologna.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0
//
// Viviane Potocnik, ETH Zurich, <[email protected]>

#pragma once

#include "math.h"
#include "snrt.h"
// #include "printf.h"
#include "utils.h"

/**
Expand Down Expand Up @@ -41,6 +42,10 @@ typedef struct softmax_layer_struct {
precision_t dtype;
} softmax_layer_t;

/* Dump function for single-cycle printf. */
dump_float(input, 7);
dump_float(output, 8);

/**
* Implementation of the SoftMax layer.
*/
Expand All @@ -50,15 +55,19 @@ static inline void softmax_fp32(float *input, float *output, int32_t ldI,
float max_core = 0.0; // max value of the current core
float sum = 0.0; // sum of the exp values of the current core

// uint32_t compute_id = snrt_global_core_idx();
// uint32_t num_cores = snrt_cluster_compute_core_num();
uint32_t compute_id = snrt_global_core_idx();
uint32_t num_cores = snrt_cluster_compute_core_num();

// printf("Hello from core [%d/%d]\n", compute_id, num_cores);
// dump_softmax(compute_id);

for (int32_t b = 0; b < batch_size; b++) {
for (int32_t s = 0; s < seq_len; s++) {
max_core = -INFINITY;
sum = 0.0;

for (int32_t i = 0; i < input_samples; i++) {
// dump_input(input[b * batch_offset + s * ldI + i]);
if (input[b * batch_offset + s * ldI + i] > max_core) {
max_core = input[b * batch_offset + s * ldI + i];
}
Expand All @@ -69,18 +78,15 @@ static inline void softmax_fp32(float *input, float *output, int32_t ldI,
output[b * batch_offset + s * ldI + i] =
// FIXME: Below code is erroring due to the standard math
// lib conflict
// TODO: Try out with musl lib
// expf(input[b * batch_offset + s * ldI + i] - max_core);
// FIXME: actually there should be an exponentiation
input[b * batch_offset + s * ldI + i] - max_core;
expf(input[b * batch_offset + s * ldI + i] - max_core);
// input[b * batch_offset + s * ldI + i] - max_core;
sum += output[b * batch_offset + s * ldI + i];
}

// compute the softmax value of the current row
for (int32_t i = 0; i < input_samples; i++) {
// INFO: DIVSQRT unit MUST be activated in the cluster
// configuration
output[b * batch_offset + s * ldI + i] /= sum;
// dump_output(output[b * batch_offset + s * ldI + i]);
// printf("output[%d] = %f\n", compute_id * input_samples + b *
// batch_offset + s * ldI + i,
// output[b * batch_offset + s * ldI + i]);
Expand Down

0 comments on commit 051aaa1

Please sign in to comment.