Neural network total size for allocation is constant but exceeds 2^31 - 1

[dCubelic]

Hi,

I am trying to build a neural network using Halide. I have a network with ~75 layers (mostly convolution, bias add, maxpool and relu layers)

But when I run the network I get the error saying Total size for allocation Conv2D$8 is constant but exceeds 2^31 - 1.

I also tried JIT compiling for target with LargeBuffers feature like this, but I get the same error.

Halide::Target target = Halide::get_host_target();
std::vector< Halide::Target::Feature > features;
features.push_back( Halide::Target::LargeBuffers );
target.set_features( features );
out.compile_jit( target );

Should it be possible to achieve what I am trying, and if so, how? If not, what is the exact problem I am facing?

Let me know if you are missing any info.

Thanks! :)

[dCubelic]

Ok, so I tried bisecting my network a bit, and managed to reproduce the problem with a minimal network of 3 layers: convolution, upsample and convolution.

here are my implementations of those 2 layers:

Conv2D:

Halide::Func Conv2D::forward( Halide::Func input )
{
    Halide::Var n("n"), h("h"), w("w"), c("c");
    Halide::Func output( "Conv2D" );

    auto paddedInput{ Halide::BoundaryConditions::constant_exterior( input, 0, 0, inputSizes[ 0 ], 0, inputSizes[ 1 ], 0, inputSizes[ 2 ], 0, inputSizes[ 3 ] ) };

    auto heightPadding{ ( outputSize( 2 ) - 1 ) * stride_h + kernel.dim( 3 ).max() + 1 - inputSizes[ 2 ] };
    auto widthPadding { ( outputSize( 1 ) - 1 ) * stride_w + kernel.dim( 2 ).max() + 1 - inputSizes[ 1 ] };

    auto padTop { heightPadding / 2 };
    auto padLeft{ widthPadding  / 2 };

    // NCWH
    Halide::RDom r( kernel );

    output( c, w, h, n ) = 0.f;
    output( c, w, h, n ) += paddedInput( r.y, w * stride_w - padLeft + r.z, h * stride_h - padTop + r.w, 0 ) * kernel( c, r.y, r.z, r.w );

    output.compute_root();

    return output;
}

Upsample:

Halide::Func Upsample::forward( Halide::Func input )
{
    Halide::Func output( "Upsample" );
    Halide::Var n, h, w, c;

    output( c, w, h, n ) = input
    (
        c,
        Halide::cast( Halide::UInt( 16 ), Halide::ceil( w / 2.f ) ),
        Halide::cast( Halide::UInt( 16 ), Halide::ceil( h / 2.f ) ),
        n
    );

    return output;
}

And I am propagating through the network the following way:

    auto a1 { conv2d_1.forward( input ) };
    auto a2 { upsample.forward( a1   ) };
    auto a3 { conv2d_2.forward( a2  ) };

What am I doing wrong?

[dCubelic]

I realised that if i remove Halide::ceil calls from my Upsample implementation it actually works (I don't get allocation size errors). What could be the problem with that function?

[abadams]

ceil should work fine in bounds inference, but its presence probably means Halide can't prove the cast to uint16 and back to int32 won't cause overflow. Keep things in int32:

Halide::Func Upsample::forward( Halide::Func input )
{
    Halide::Func output( "Upsample" );
    Halide::Var n, h, w, c;

    output( c, w, h, n ) = input
    (
        c,
        (w + 1)/2, 
        (h + 1)/2,
        n
    );

    return output;
}