CUDA PTX inline asm listing for `.func` block fails internal exception for -e stmt_html

[antonysigma]

Description

Related to #9196 : Sometimes, cos/sin functions are lowered to an internal .func NVPTX block having function name *trig_reduction* containing nested braces. The previous bugfix did not cover this corner case because it is .func, not .visible .entry.

Example .func nvptx code triggering the failure

const .align 8 .b8 __cudart_i2opi_d[144] = {...};
.const .align 8 .b8 __cudart_sin_cos_coeffs[128] = {...};
                                        // -- Begin function __internal_trig_reduction_slowpathd
                                        // @__internal_trig_reduction_slowpathd
.func  (.param .b64 func_retval0) __internal_trig_reduction_slowpathd(
    .param .b64 __internal_trig_reduction_slowpathd_param_0,
    .param .b64 __internal_trig_reduction_slowpathd_param_1
)
{

...

// begin inline asm
    {
    .reg .u32 r0, r1, r2, r3, alo, ahi, blo, bhi, clo, chi;
    mov.b64         {alo,ahi}, %rd51;    
    mov.b64         {blo,bhi}, %rd52;    
    mov.b64         {clo,chi}, %rd102;    
    mad.lo.cc.u32   r0, alo, blo, clo;
    madc.hi.cc.u32  r1, alo, blo, chi;
    madc.hi.u32     r2, alo, bhi,   0;
    mad.lo.cc.u32   r1, alo, bhi,  r1;
    madc.hi.cc.u32  r2, ahi, blo,  r2;
    madc.hi.u32     r3, ahi, bhi,   0;
    mad.lo.cc.u32   r1, ahi, blo,  r1;
    madc.lo.cc.u32  r2, ahi, bhi,  r2;
    addc.u32        r3,  r3,   0;     
    mov.b64         %rd49, {r0,r1};      
    mov.b64         %rd102, {r2,r3};      
    }
    // end inline asm
...
}

Reproducing case

const .align 8 .b8 __cudart_i2opi_d[144] = {...};
.const .align 8 .b8 __cudart_sin_cos_coeffs[128] = {...};
                                        // -- Begin function __internal_trig_reduction_slowpathd
                                        // @__internal_trig_reduction_slowpathd
.func  (.param .b64 func_retval0) __internal_trig_reduction_slowpathd(
    .param .b64 __internal_trig_reduction_slowpathd_param_0,
    .param .b64 __internal_trig_reduction_slowpathd_param_1
)
{

...

// begin inline asm
    {
    .reg .u32 r0, r1, r2, r3, alo, ahi, blo, bhi, clo, chi;
    mov.b64         {alo,ahi}, %rd51;    
    mov.b64         {blo,bhi}, %rd52;    
    mov.b64         {clo,chi}, %rd102;    
    mad.lo.cc.u32   r0, alo, blo, clo;
    madc.hi.cc.u32  r1, alo, blo, chi;
    madc.hi.u32     r2, alo, bhi,   0;
    mad.lo.cc.u32   r1, alo, bhi,  r1;
    madc.hi.cc.u32  r2, ahi, blo,  r2;
    madc.hi.u32     r3, ahi, bhi,   0;
    mad.lo.cc.u32   r1, ahi, blo,  r1;
    madc.lo.cc.u32  r2, ahi, bhi,  r2;
    addc.u32        r3,  r3,   0;     
    mov.b64         %rd49, {r0,r1};      
    mov.b64         %rd102, {r2,r3};      
    }
    // end inline asm
...
}

How did you get Halide?

Built from source

Halide version

18.0.0

Halide commit (if known)

No response

Target

x86_64-gnu-linux-cuda

Operating system

No response

Additional context

No response

[alexreinking]

@mcourteaux

[mcourteaux]

@antonysigma before spending any time at thisz you indicated Halide 18. That's ... Really old. Is that correct ? Please update if so.

[antonysigma]

Yes I agree the Halide version that I use is ancient. I can reproduce the error in the latest Halide version if desired.

[mcourteaux]

Please do. Provide some halide code that triggers this, because the "repro" is not really a repro: it's the output of the compiler you pasted here.

[antonysigma]

Minimum reproducible case as follows. The .func section in NVPTX assembly appears when Li2018 auto-scheduler is used.

#include "Halide.h"

namespace {

using namespace Halide;

class MomentOfInertia : public Generator<MomentOfInertia> {
   public:
    Input<Buffer<float, 2>> input{"input"};
    Input<double> theta_in{"theta_in"};

    Output<Buffer<double, 1>> output{"output"};

    void generate();
    void schedule();

   private:
    const Var kx{"kx"};
    const Var ky{"ky"};

    // Placeholder axis to compute the Jacobian vector
    const Var k{"k"};

    // cost value if i == 0, first derivative if i == 1.
    const Var i{"i"};

    Func theta{"theta"};
    Func centered{"centered"};
    Func moment_of_inertia{"moment_of_inertia"};
    Func d_moment_d_theta{"d_moment_d_theta"};

    void setBounds();
};

void
MomentOfInertia::generate() {
    const Expr cx = input.width() / 2;
    const Expr cy = input.height() / 2;

    centered(kx, ky) = input(kx + cx, ky + cy);

    theta(k) = theta_in;

    RDom r{-cx, input.width(), -cy, input.height()};
    const Expr rxx = r.x * cos(theta(k)) + r.y * sin(theta(k));
    moment_of_inertia(k) = Expr{0.0};
    moment_of_inertia(k) += centered(r.x, r.y) * rxx * rxx;

    const auto d_output = propagate_adjoints(moment_of_inertia);
    d_moment_d_theta = d_output(theta);

    output(i) = mux(i, {moment_of_inertia(0), d_moment_d_theta(0)});
}

void
MomentOfInertia::setBounds() {
    const Expr w = input.width() / 32 * 32;
    input.dim(0).set_bounds(0, w).set_stride(1);
    input.dim(1).set_bounds(0, w).set_stride(w);

    output.dim(0).set_bounds(0, 2).set_stride(1);
}

void
MomentOfInertia::schedule() {
    assert(using_autoscheduler() && "Manual GPU implementation not yet implemented.");

    input.set_estimates({{0, 256}, {0, 256}});
    theta_in.set_estimate(0.0);
    output.set_estimates({{0, 2}});
}
}  // namespace

HALIDE_REGISTER_GENERATOR(MomentOfInertia, moment_of_inertia)