Build failure (Myelin layouts.cpp:2180) or wrong plugin input of TensorRT 10.16 when an IPluginV3 with tactics reads a kHWC8 network output, on Jetson Thor #4882

// Standalone repro: an IPluginV3 (FP16, TensorFormat::kHWC8 in and out) placed after a convolution whose output the
// plugin reads. The plugin is either a pass-through copy, or (--up2) a nearest x2 upsample keeping the first 16
// channels, so its output can be checked exactly against conv_out on the host.
//
//   ./repro [--up2] [--tactics] [--out-hwc8] [--consumer] [--resize] [--slice] [--trunk] [--cat] [--cat-producer]
//     --up2      : plugin output (N, 16, 2H, 2W) = nearest x2 of the first 16 channels of its input
//     --tactics  : the plugin offers 8 tactics (getValidTactics); every tactic computes the same thing
//     --out-hwc8 : conv_out is also a network output restricted to kHWC8 (the layout the plugin asks for)
//     --consumer : conv_out also feeds a second convolution (another output)
//     --resize   : conv_out also feeds an IResizeLayer (linear x2, another output)
//     --slice    : conv_out's first batch item also feeds a 1x1 conv (another output, kHWC8)
//     --trunk    : the plugin output is convolved, concatenated with another conv's output and convolved again
//     --cat      : the plugin output is concatenated directly with another conv's output, then convolved
//     --cat-producer : conv_out's own input comes from a concat -> conv chain
//
// Prints the build status and the number of plugin_out elements that differ from what conv_out implies.
#include <NvInfer.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>

#include <cmath>
#include <cstdio>
#include <cstring>
#include <memory>
#include <random>
#include <string>
#include <vector>

using namespace nvinfer1;

// out[n][2h + a][2w + b][k] = in[n][h][w][k], k < K (both channel-last, channel strides Ci / K)
static __global__ void up2Kernel(const __half* in, __half* out, int N, int Ci, int H, int W, int K) {
    const long long i = blockIdx.x * static_cast<long long>(blockDim.x) + threadIdx.x;
    const long long total = static_cast<long long>(N) * 2 * H * 2 * W * K;
    if (i >= total) return;
    const int k = i % K;
    long long t = i / K;
    const int x = t % (2 * W);
    t /= 2 * W;
    const int y = t % (2 * H);
    const int n = static_cast<int>(t / (2 * H));
    out[i] = in[((static_cast<long long>(n) * H + y / 2) * W + x / 2) * Ci + k];
}

namespace {

class Logger : public ILogger {
    void log(Severity s, const char* msg) noexcept override {
        if (s <= Severity::kWARNING) std::printf("[trt] %s\n", msg);
    }
} gLogger;

#define CK(x)                                                                                                      \
    do {                                                                                                           \
        cudaError_t e_ = (x);                                                                                      \
        if (e_ != cudaSuccess) {                                                                                   \
            std::printf("%s: %s\n", #x, cudaGetErrorString(e_));                                                   \
            return 1;                                                                                              \
        }                                                                                                          \
    } while (0)

constexpr char kName[] = "PassHWC8", kVersion[] = "1", kNamespace[] = "repro";
bool gUp2 = false, gTactics = false;
constexpr int kUpK = 16;



// out = in, FP16 kHWC8 (channels padded to 8: N * H * W * roundup(C, 8) elements)
class PassPlugin : public IPluginV3, public IPluginV3OneCore, public IPluginV3OneBuild, public IPluginV3OneRuntime {
public:
    IPluginCapability* getCapabilityInterface(PluginCapabilityType t) noexcept override {
        if (t == PluginCapabilityType::kCORE) return static_cast<IPluginV3OneCore*>(this);
        if (t == PluginCapabilityType::kBUILD) return static_cast<IPluginV3OneBuild*>(this);
        return static_cast<IPluginV3OneRuntime*>(this);
    }
    IPluginV3* clone() noexcept override { return new PassPlugin(*this); }
    AsciiChar const* getPluginName() const noexcept override { return kName; }
    AsciiChar const* getPluginVersion() const noexcept override { return kVersion; }
    AsciiChar const* getPluginNamespace() const noexcept override { return kNamespace; }
    int32_t getNbOutputs() const noexcept override { return 1; }
    int32_t getOutputDataTypes(DataType* o, int32_t, DataType const* in, int32_t) const noexcept override {
        o[0] = in[0];
        return 0;
    }
    int32_t getOutputShapes(DimsExprs const* in, int32_t, DimsExprs const*, int32_t, DimsExprs* out, int32_t,
                            IExprBuilder& eb) noexcept override {
        out[0] = in[0];
        if (gUp2) {
            out[0].d[1] = eb.constant(kUpK);
            out[0].d[2] = eb.operation(DimensionOperation::kPROD, *in[0].d[2], *eb.constant(2));
            out[0].d[3] = eb.operation(DimensionOperation::kPROD, *in[0].d[3], *eb.constant(2));
        }
        return 0;
    }
    int32_t getNbTactics() noexcept override { return gTactics ? 8 : 0; }
    int32_t getValidTactics(int32_t* t, int32_t n) noexcept override {
        for (int32_t i = 0; i < n; ++i) t[i] = i + 1;
        return 0;
    }
    int32_t setTactic(int32_t) noexcept override { return 0; }
    bool supportsFormatCombination(int32_t pos, DynamicPluginTensorDesc const* io, int32_t, int32_t) noexcept override {
        return io[pos].desc.type == DataType::kHALF && io[pos].desc.format == TensorFormat::kHWC8;
    }
    int32_t configurePlugin(DynamicPluginTensorDesc const*, int32_t, DynamicPluginTensorDesc const*,
                            int32_t) noexcept override {
        return 0;
    }
    size_t getWorkspaceSize(DynamicPluginTensorDesc const*, int32_t, DynamicPluginTensorDesc const*,
                            int32_t) const noexcept override {
        return 0;
    }
    int32_t onShapeChange(PluginTensorDesc const*, int32_t, PluginTensorDesc const*, int32_t) noexcept override { return 0; }
    int32_t enqueue(PluginTensorDesc const* inDesc, PluginTensorDesc const*, void const* const* inputs, void* const* outputs,
                    void*, cudaStream_t stream) noexcept override {
        const Dims& d = inDesc[0].dims;
        const int Ci = static_cast<int>((d.d[1] + 7) / 8 * 8);
        if (gUp2) {
            const long long total = d.d[0] * 4 * d.d[2] * d.d[3] * kUpK;
            up2Kernel<<<static_cast<unsigned>((total + 255) / 256), 256, 0, stream>>>(
                static_cast<const __half*>(inputs[0]), static_cast<__half*>(outputs[0]), static_cast<int>(d.d[0]), Ci,
                static_cast<int>(d.d[2]), static_cast<int>(d.d[3]), kUpK);
            return cudaPeekAtLastError() == cudaSuccess ? 0 : 1;
        }
        const size_t n = static_cast<size_t>(d.d[0]) * Ci * d.d[2] * d.d[3];
        return cudaMemcpyAsync(outputs[0], inputs[0], n * sizeof(__half), cudaMemcpyDeviceToDevice, stream) == cudaSuccess
                   ? 0
                   : 1;
    }
    IPluginV3* attachToContext(IPluginResourceContext*) noexcept override { return clone(); }
    PluginFieldCollection const* getFieldsToSerialize() noexcept override {
        fc_.nbFields = 0;
        fc_.fields = nullptr;
        return &fc_;
    }

private:
    PluginFieldCollection fc_{};
};

class PassCreator : public IPluginCreatorV3One {
public:
    IPluginV3* createPlugin(AsciiChar const*, PluginFieldCollection const*, TensorRTPhase) noexcept override {
        return new PassPlugin();
    }
    PluginFieldCollection const* getFieldNames() noexcept override { return &fc_; }
    AsciiChar const* getPluginName() const noexcept override { return kName; }
    AsciiChar const* getPluginVersion() const noexcept override { return kVersion; }
    AsciiChar const* getPluginNamespace() const noexcept override { return kNamespace; }

private:
    PluginFieldCollection fc_{};
};

// FP16 tensor of format f (kLINEAR NCHW, or kHWC8 = NHWC with C padded to 8) -> NCHW floats
std::vector<float> toNCHW(const std::vector<__half>& v, TensorFormat f, int N, int C, int H, int W) {
    std::vector<float> o(static_cast<size_t>(N) * C * H * W);
    const int Cp = (C + 7) / 8 * 8;
    for (int n = 0; n < N; ++n)
        for (int c = 0; c < C; ++c)
            for (int h = 0; h < H; ++h)
                for (int w = 0; w < W; ++w) {
                    const size_t src = f == TensorFormat::kHWC8 ? ((static_cast<size_t>(n) * H + h) * W + w) * Cp + c
                                                                : ((static_cast<size_t>(n) * C + c) * H + h) * W + w;
                    o[((static_cast<size_t>(n) * C + c) * H + h) * W + w] = __half2float(v[src]);
                }
    return o;
}

std::vector<__half> randomHalf(size_t n, float scale, std::mt19937& g) {
    std::uniform_real_distribution<float> u(-scale, scale);
    std::vector<__half> v(n);
    for (auto& x : v) x = __float2half(u(g));
    return v;
}

}  // namespace

int main(int argc, char** argv) {
    bool outHwc8 = false, consumer = false, trunk = false, resize = false, slice = false, cat = false, catProducer = false;
    for (int i = 1; i < argc; ++i) {
        const std::string a = argv[i];
        outHwc8 |= a == "--out-hwc8";
        consumer |= a == "--consumer";
        trunk |= a == "--trunk";
        resize |= a == "--resize";
        slice |= a == "--slice";
        cat |= a == "--cat";
        catProducer |= a == "--cat-producer";
        gUp2 |= a == "--up2";
        gTactics |= a == "--tactics";
    }
    static PassCreator creator;
    getPluginRegistry()->registerCreator(creator, kNamespace);
    std::string var;
    for (int i = 1; i < argc; ++i) var += std::string(" ") + argv[i];
    std::printf("TensorRT %d.%d.%d, variant:%s\n", getInferLibVersion() / 10000, (getInferLibVersion() / 100) % 100,
                getInferLibVersion() % 100, var.c_str());

    // shapes of the model where this was found: batch 2 (two camera views), 64 channels, 128 x 224
    const int N = 2, C = 64, H = 128, W = 224, K = 16;
    std::mt19937 g(1);
    std::vector<std::vector<__half>> keep;
    auto weights = [&](size_t n, float scale) {
        keep.push_back(randomHalf(n, scale, g));
        return Weights{DataType::kHALF, keep.back().data(), static_cast<int64_t>(n)};
    };

    std::unique_ptr<IBuilder> builder(createInferBuilder(gLogger));
    std::unique_ptr<INetworkDefinition> net(
        builder->createNetworkV2(1u << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kSTRONGLY_TYPED)));
    ITensor* x = net->addInput("x", DataType::kHALF, Dims4{N, C, H, W});
    ITensor* xin = x;
    if (catProducer) {   // conv_out = relu(conv(cat(conv(x), x))), as the model's neck output that feeds the plugin
        auto* p0 = net->addConvolutionNd(*x, C, Dims2{3, 3}, weights(size_t(C) * C * 9, 1.f / 24.f), weights(C, 0.f));
        p0->setPaddingNd(Dims2{1, 1});
        ITensor* pc[2] = {net->addActivation(*p0->getOutput(0), ActivationType::kRELU)->getOutput(0), x};
        auto* cc = net->addConcatenation(pc, 2);
        cc->setAxis(1);
        auto* p1 = net->addConvolutionNd(*cc->getOutput(0), C, Dims2{3, 3}, weights(size_t(C) * 2 * C * 9, 1.f / 34.f),
                                          weights(C, 0.f));
        p1->setPaddingNd(Dims2{1, 1});
        xin = net->addActivation(*p1->getOutput(0), ActivationType::kRELU)->getOutput(0);
    }
    auto* conv = net->addConvolutionNd(*xin, C, Dims2{3, 3}, weights(size_t(C) * C * 9, 1.f / 24.f), weights(C, 0.1f));
    conv->setPaddingNd(Dims2{1, 1});
    ITensor* y = net->addActivation(*conv->getOutput(0), ActivationType::kRELU)->getOutput(0);
    y->setName("conv_out");
    net->markOutput(*y);
    if (outHwc8) y->setAllowedFormats(1u << static_cast<uint32_t>(TensorFormat::kHWC8));

    PassPlugin plugin;
    ITensor* in[1] = {y};
    ITensor* z = net->addPluginV3(in, 1, nullptr, 0, plugin)->getOutput(0);
    z->setName("plugin_out");
    net->markOutput(*z);
    z->setAllowedFormats(1u << static_cast<uint32_t>(TensorFormat::kHWC8));

    if (consumer) {   // y also read by a 1x1 conv
        auto* c2 = net->addConvolutionNd(*y, 8, Dims2{1, 1}, weights(size_t(8) * C, 0.1f), weights(8, 0.f));
        c2->getOutput(0)->setName("side");
        net->markOutput(*c2->getOutput(0));
    }
    if (slice) {   // y's first batch item also read by a 1x1 conv (in the model: guide features of the left view)
        auto* sl = net->addSlice(*y, Dims4{0, 0, 0, 0}, Dims4{1, C, H, W}, Dims4{1, 1, 1, 1});
        auto* c3 = net->addConvolutionNd(*sl->getOutput(0), 8, Dims2{1, 1}, weights(size_t(8) * C, 0.1f), weights(8, 0.f));
        c3->getOutput(0)->setName("left");
        net->markOutput(*c3->getOutput(0));
        c3->getOutput(0)->setAllowedFormats(1u << static_cast<uint32_t>(TensorFormat::kHWC8));
    }
    if (resize) {   // y also read by a linear x2 resize
        auto* r = net->addResize(*y);
        const float sc[4] = {1.f, 1.f, 2.f, 2.f};
        r->setScales(sc, 4);
        r->setResizeMode(InterpolationMode::kLINEAR);
        r->setCoordinateTransformation(ResizeCoordinateTransformation::kHALF_PIXEL);
        r->getOutput(0)->setName("resized");
        net->markOutput(*r->getOutput(0));
    }
    if (cat) {   // cat(plugin_out, conv(x or its x2 resize)) -> conv (kHWC8 output), as in the model's neck
        const int Cz = gUp2 ? kUpK : C;
        ITensor* xb = x;
        if (gUp2) {
            auto* r = net->addResize(*x);
            const float sc[4] = {1.f, 1.f, 2.f, 2.f};
            r->setScales(sc, 4);
            xb = r->getOutput(0);
        }
        auto* b = net->addConvolutionNd(*xb, Cz, Dims2{3, 3}, weights(size_t(Cz) * C * 9, 1.f / 24.f), weights(Cz, 0.f));
        b->setPaddingNd(Dims2{1, 1});
        ITensor* cin[2] = {z, net->addActivation(*b->getOutput(0), ActivationType::kRELU)->getOutput(0)};
        auto* cc = net->addConcatenation(cin, 2);
        cc->setAxis(1);
        auto* t = net->addConvolutionNd(*cc->getOutput(0), K, Dims2{3, 3}, weights(size_t(K) * 2 * Cz * 9, 1.f / 17.f),
                                         weights(K, 0.f));
        t->setPaddingNd(Dims2{1, 1});
        t->getOutput(0)->setName("cat_tail");
        net->markOutput(*t->getOutput(0));
        t->getOutput(0)->setAllowedFormats(1u << static_cast<uint32_t>(TensorFormat::kHWC8));
    }
    if (trunk) {   // cat(conv(plugin_out), conv(x or its x2 resize)) -> conv, as in the model's neck
        const int Cz = gUp2 ? kUpK : C;
        auto* a = net->addConvolutionNd(*z, K, Dims2{3, 3}, weights(size_t(K) * Cz * 9, 1.f / 24.f), weights(K, 0.f));
        a->setPaddingNd(Dims2{1, 1});
        ITensor* xb = x;
        if (gUp2) {
            auto* r = net->addResize(*x);
            const float sc[4] = {1.f, 1.f, 2.f, 2.f};
            r->setScales(sc, 4);
            xb = r->getOutput(0);
        }
        auto* b = net->addConvolutionNd(*xb, K, Dims2{3, 3}, weights(size_t(K) * C * 9, 1.f / 24.f), weights(K, 0.f));
        b->setPaddingNd(Dims2{1, 1});
        ITensor* cat[2] = {a->getOutput(0), b->getOutput(0)};
        auto* cc = net->addConcatenation(cat, 2);
        cc->setAxis(1);
        auto* t = net->addConvolutionNd(*cc->getOutput(0), K, Dims2{3, 3}, weights(size_t(K) * 2 * K * 9, 1.f / 17.f),
                                         weights(K, 0.f));
        t->setPaddingNd(Dims2{1, 1});
        t->getOutput(0)->setName("tail");
        net->markOutput(*t->getOutput(0));
    }

    std::unique_ptr<IBuilderConfig> cfg(builder->createBuilderConfig());
    std::unique_ptr<IHostMemory> plan(builder->buildSerializedNetwork(*net, *cfg));
    if (!plan) {
        std::printf("RESULT: engine build FAILED\n");
        return 2;
    }
    std::unique_ptr<IRuntime> runtime(createInferRuntime(gLogger));
    std::unique_ptr<ICudaEngine> engine(runtime->deserializeCudaEngine(plan->data(), plan->size()));
    std::unique_ptr<IExecutionContext> ctx(engine->createExecutionContext());
    std::vector<void*> bufs;
    size_t nXY = 0;
    for (int i = 0; i < engine->getNbIOTensors(); ++i) {
        const char* name = engine->getIOTensorName(i);
        const Dims d = engine->getTensorShape(name);
        const int vec = engine->getTensorVectorizedDim(name) >= 0 ? engine->getTensorComponentsPerElement(name) : 1;
        size_t n = 1;
        for (int k = 0; k < d.nbDims; ++k) n *= k == 1 ? (d.d[k] + vec - 1) / vec * vec : d.d[k];
        void* p = nullptr;
        CK(cudaMalloc(&p, n * sizeof(__half)));
        CK(cudaMemset(p, 0, n * sizeof(__half)));
        if (std::strcmp(name, "x") == 0) {
            auto hx = randomHalf(n, 1.f, g);
            CK(cudaMemcpy(p, hx.data(), n * sizeof(__half), cudaMemcpyHostToDevice));
        }
        if (std::strcmp(name, "conv_out") == 0) nXY = n;
        std::printf("io %-10s format %s\n", name, engine->getTensorFormatDesc(name));
        ctx->setTensorAddress(name, p);
        bufs.push_back(p);
    }
    cudaStream_t s;
    CK(cudaStreamCreate(&s));
    if (!ctx->enqueueV3(s)) {
        std::printf("RESULT: enqueue failed\n");
        return 3;
    }
    CK(cudaStreamSynchronize(s));
    std::vector<__half> hy(nXY), hz(nXY);
    CK(cudaMemcpy(hy.data(), ctx->getTensorAddress("conv_out"), nXY * sizeof(__half), cudaMemcpyDeviceToHost));
    CK(cudaMemcpy(hz.data(), ctx->getTensorAddress("plugin_out"), nXY * sizeof(__half), cudaMemcpyDeviceToHost));
    const auto fy = toNCHW(hy, engine->getTensorFormat("conv_out"), N, C, H, W);
    std::vector<float> expect = fy, fz;
    if (gUp2) {   // nearest x2 of the first kUpK channels
        expect.assign(static_cast<size_t>(N) * kUpK * 4 * H * W, 0.f);
        for (int n = 0; n < N; ++n)
            for (int k = 0; k < kUpK; ++k)
                for (int yy = 0; yy < 2 * H; ++yy)
                    for (int xx = 0; xx < 2 * W; ++xx)
                        expect[((static_cast<size_t>(n) * kUpK + k) * 2 * H + yy) * 2 * W + xx] =
                            fy[((static_cast<size_t>(n) * C + k) * H + yy / 2) * W + xx / 2];
        std::vector<__half> hz2(static_cast<size_t>(N) * kUpK * 4 * H * W);
        CK(cudaMemcpy(hz2.data(), ctx->getTensorAddress("plugin_out"), hz2.size() * sizeof(__half), cudaMemcpyDeviceToHost));
        fz = toNCHW(hz2, engine->getTensorFormat("plugin_out"), N, kUpK, 2 * H, 2 * W);
    } else {
        fz = toNCHW(hz, engine->getTensorFormat("plugin_out"), N, C, H, W);
    }
    float maxd = 0.f;
    size_t bad = 0;
    for (size_t i = 0; i < expect.size(); ++i) {
        const float d = std::fabs(expect[i] - fz[i]);
        maxd = std::fmax(maxd, d);
        bad += d > 0.f;
    }
    nXY = expect.size();
    std::printf("RESULT: max |plugin_out - expected from conv_out| = %g, %zu of %zu elements differ -> %s\n", maxd, bad, nXY,
                bad ? "WRONG (the plugin was given a different input than conv_out)" : "OK");
    for (void* p : bufs) cudaFree(p);
    return bad ? 4 : 0;
}