Repository: opencv/opencv
File: modules/dnn/src/layers/affinegrid_layer.cpp (C++)

Code before:

        std::vector<int> outShape = is3d ? std::vector<int>{N, D, H, W, 3}
                                         : std::vector<int>{N, H, W, 2};
        Mat& grid = outputs[0];
        grid.create((int)outShape.size(), outShape.data(), CV_32F);

        auto computeLinspace = [&](int len, float& scale, float& delta){
            if (align_corners) {
                scale = len > 1 ? 2.f / float(len - 1) : 0.f;
                delta = -1.f;
            } else {
                scale = 2.f / float(len);
                delta = scale * 0.5f - 1.f;
            }
        };

        float xs, xd, ys, yd, zs=0.f, zd=0.f;
        computeLinspace(W, xs, xd);
        computeLinspace(H, ys, yd);
        if (is3d) computeLinspace(D, zs, zd);

        Mat thetaF;
        if (theta.depth() == CV_32F) thetaF = theta; else theta.convertTo(thetaF, CV_32F);

        parallel_for_(Range(0, N), [&](const Range& range){
            for (int n = range.start; n < range.end; n++) {
                int i0[3] = {n, 0, 0};
                int i1[3] = {n, 1, 0};
                int i2[3] = {n, 2, 0};
                const float* T0 = thetaF.ptr<float>(i0);
                const float* T1 = thetaF.ptr<float>(i1);
                const float* T2 = is3d ? thetaF.ptr<float>(i2) : nullptr;
                if (!is3d) {
                    for (int y = 0; y < H; y++) {
                        float ny = yd + ys * y;
                        float base0 = T0[1]*ny + T0[2];
                        float base1 = T1[1]*ny + T1[2];
                        float* out = grid.ptr<float>(n, y, 0);
                        for (int x = 0; x < W; x++) {
                            float nx = xd + xs * x;
                            out[2*x + 0] = T0[0]*nx + base0;
                            out[2*x + 1] = T1[0]*nx + base1;
                        }
                    }
                } else {
                    for (int zh = 0; zh < D * H; zh++) {
                        int z = zh / H;
                        int y = zh - z * H;
                        float ny = yd + ys * y;
                        float nz = zd + zs * z;
                        float base0 = T0[1]*ny + T0[2]*nz + T0[3];

Review comment (performance):
it's not good that parallelism is only on the batch. We often have N=1. Parallelism should be over N*H.

```
for (int ny = range.start; ny < range.end; ny++) {
     int n = ny / H;
     int y = ny % H;
     ...
}
```


Code after:
        } else {
            H = sizeVec[2];
            W = sizeVec[3];
            if (H > 0 && W > 0)
                outShape = MatShape{N, H, W, 2};
        }
    }

    bool getMemoryShapes(const std::vector<MatShape> &inputs,
                         const int requiredOutputs,
                         std::vector<MatShape> &outputs,
                         std::vector<MatShape> &internals) const CV_OVERRIDE
    {
        CV_CheckGE((int)inputs.size(), 2, "AffineGrid requires two inputs: theta and size");

        const MatShape thetaShape = inputs[0];
        CV_CheckGE((int)thetaShape.size(), 3, "theta must be at least 3D [N, r, c]");
        int N = thetaShape[0];
        int r = thetaShape[1];
        int c = thetaShape[2];
        CV_CheckTrue((r == 2 && c == 3) || (r == 3 && c == 4), "theta must be [N,2,3] or [N,3,4]");
        const bool is3d = (r == 3);

        MatShape outShape;
        const MatShape sizeShape = inputs[1];
        CV_CheckTrue(sizeShape.size() == 1, "size input must be 1D tensor");
        const int dims = (int)sizeShape[0];
        CV_CheckTrue(dims == (is3d ? 5 : 4), "size input must have 4 (2D) or 5 (3D) elements");
        int D, H, W;
        resolveSizeAndShape(N, is3d, /*explicitSz*/nullptr, outShape, D, H, W);

        outputs.assign(1, outShape);
        return false;
    }

    bool dynamicOutputShapes() const CV_OVERRIDE
    {
        Net::Impl* netimpl_ = getNetImpl(this);
        if (!netimpl_) return true;
        CV_CheckGE((int)this->inputs.size(), 2, "AffineGrid requires two inputs: theta and size");

        if (netimpl_->isConstArg(this->inputs[1])) return false;
        try {
            const Mat& sz = netimpl_->argTensor(this->inputs[1]);
            return sz.empty();
        } catch (const cv::Exception& e) {
            CV_Error(cv::Error::StsError, cv::format("DNN/AffineGrid: cannot query 'size' tensor at runtime: %s", e.what()));
        }
    }

    void getTypes(const std::vector<MatType>& inputs,