Method computeForNormalize

source/geometry/GeometryLRN.cpp:19–125

class GeometryLRN : public GeometryComputer {
public:
    bool computeForNormalize(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
                             Context& context, CommandBuffer& res) const {
        auto normalize      = op->main_as_Normalize();
        auto mAcrossSpatial = normalize->acrossSpatial();
        auto mChannelShared = normalize->channelShared();
        Tensor* eps         = nullptr;
        Tensor* scale       = nullptr;
        auto cache          = context.searchConst(op);
        if (!cache.empty()) {
            eps   = cache[0].get();
            scale = cache[1].get();
        } else {
            auto mEps              = normalize->eps();
            auto epsT              = context.allocConst(op, {}, halide_type_of<float>());
            epsT->host<float>()[0] = mEps;
            eps                    = epsT.get();
            auto mScale = context.allocConst(op, {1, (int)normalize->scale()->size(), 1}, halide_type_of<float>());
            ::memcpy(mScale->host<float>(), normalize->scale()->data(), normalize->scale()->size() * sizeof(float));
            scale = mScale.get();
        }
        auto inputTensor = inputs[0];
        // Across channel
        int inside  = inputTensor->width() * inputTensor->height();
        int axis    = inputTensor->channel();
        int outside = inputTensor->batch();

        {
            // 1, axis, 1 -> outside, axis, inside
            std::shared_ptr<Tensor> broadCastScale(Tensor::createDevice<float>({outside, axis, inside}, Tensor::CAFFE));
            res.extras.emplace_back(broadCastScale);
            auto des = TensorUtils::getDescribe(broadCastScale.get());
            des->memoryType = Tensor::InsideDescribe::MEMORY_VIRTUAL;
            des->regions.resize(1);
            auto& reg = des->regions[0];
            reg.size[0] = outside;
            reg.size[1] = axis;
            reg.size[2] = inside;
            reg.src.offset = 0;
            reg.src.stride[0] = 0;
            reg.src.stride[1] = 1;
            reg.src.stride[2] = 0;
            reg.dst.offset = 0;
            reg.dst.stride[0] = axis * inside;
            reg.dst.stride[1] = inside;
            reg.dst.stride[2] = 1;
            reg.origin = scale;
            scale = broadCastScale.get();
        }

        // Across Spatial
        if (mAcrossSpatial) {
            inside  = 1;
            axis    = inputTensor->width() * inputTensor->height() * inputTensor->channel();
        }
        std::shared_ptr<Tensor> inputRaw(Tensor::createDevice<float>({outside, axis, inside}, Tensor::CAFFE));
        res.extras.emplace_back(inputRaw);
        std::shared_ptr<Tensor> inputRawSquare(Tensor::createDevice<float>({outside, axis, inside}, Tensor::CAFFE));
        res.extras.emplace_back(inputRawSquare);