Method Run

executor/operator/arm64/softmax_float.cpp:222–723

    }

    bool Run(Node* node)
    {
        Tensor* input_tensor = node->GetInputTensor(0);
        Tensor* output_tensor = node->GetOutputTensor(0);
        const std::vector<int>& dims = input_tensor->GetShape().GetDim();
        Softmax* softmax_op = dynamic_cast<Softmax*>(node->GetOp());
        SoftmaxParam* param_ = softmax_op->GetParam();
        float* input = ( float* )get_tensor_mem(input_tensor);
        float* output = ( float* )get_tensor_mem(output_tensor);


        int dim_size = dims.size();
        int axis = param_->axis;

        if (dim_size == 1) // axis == 0
        {
            int w = dims[0];
            float* ptr = input;
            float* output_ptr = output;

            // get max
            float max = -__FLT_MAX__;
            float32x4_t _max4 = vdupq_n_f32(max);
            for(int i = 0; i < (w & -4); i += 4)
            {
                float32x4_t _in = vld1q_f32(ptr + i);
                _max4 = vmaxq_f32(_max4, _in);
            }
            max = vmaxvq_f32(_max4);
            float32x2_t _max2 = vdup_n_f32(max);
            for (int j = (w & ~3);j<(w & -2);j+=2)
            {
                float32x2_t _in = vld1_f32(ptr + j);
                _max2 = vmax_f32(_max2, _in);
            }
            max = vpmaxs_f32(_max2);
            for (int j=(w & ~1); j<w; j++)
            {
                max = std::max(max, ptr[j]);
            }

            // get sum
            float sum = 0.f;
            float32x4_t _sum4 = vdupq_n_f32(0.0f);
            _max4 = vdupq_n_f32(max);
            for(int i = 0; i < (w & -4); i += 4)
            {
                float32x4_t _in = vld1q_f32(ptr + i);
                float32x4_t _out = vexpq10_f32(vsubq_f32(_in, _max4));
                vst1q_f32(output_ptr+i, _out);
                _sum4 = vaddq_f32(_sum4, _out);
            }
            float32x2_t _sum2 = vdup_n_f32(0.0f);
            _max2 = vdup_n_f32(max);
            for(int i = (w & ~3); i < (w & -2); i += 2)
            {
                float32x2_t _in = vld1_f32(ptr + i);
                float32x2_t _out = vexp10_f32(vsub_f32(_in, _max2));