[Question] Correct implementation of Fused Conv3d + Bias on SM80 Tensor Cores

[Crinton]

I am trying to implement a fused Conv3d (NDHWC) + Bias operator using CUTLASS on an Ampere (SM80) GPU with Tensor Cores.

The goal is to achieve: Output = alpha * Conv3d(Input, Weight) + Bias.

When I don't use bias, the results are exactly the same as with torch.conv3d. However, when I use bias, a calculation error occurs, and this is not due to accumulated error.

The Problem

1. Implicit Broadcast (Stride 0): Setting Tensor C's stride to 0 to mimic a bias vector. This works for Conv2d in official examples but fails for Conv3d, likely due to 5D coordinate mapping issues.

Request
Could the maintainers provide a clear example or confirm the correct Arguments structure for Conv3d + Bias on SM80 Tensor Cores? Specifically:

Is LinearCombination fully supported for OpClassTensorOp in 3D convolutions?

What is the correct way to pass the Bias vector?

``

[Arthur031221]

A zero-stride C tensor does work as a bias for Conv3d fprop on SM80 tensor cores. The program below uses the pieces you mention (LinearCombination, OpClassTensorOp, SM80, NDHWC) with DefaultConv3dFprop, half inputs and float accumulate and output, and checks every output element against cutlass::reference::host::Conv3d with the bias added on the host:

optimized, NoBetaScaling,     alpha=1        : y[0,0,0,0,0:3] = 77 98 61, max |y - (conv + bias)| = 0
analytic,  NoBetaScaling,     alpha=1        : y[0,0,0,0,0:3] = 77 98 61, max |y - (conv + bias)| = 0
optimized, Default,           alpha=1, beta=1: y[0,0,0,0,0:3] = 77 98 61, max |y - (conv + bias)| = 0
optimized, Default,           alpha=1, beta=0: y[0,0,0,0,0:3] = -23 -3 -41, max |y - (conv + bias)| = 163
odd shape, optimized, NoBeta, alpha=1        : y[0,0,0,0,0:3] = 105 120 116, max |y - (conv + bias)| = 0
odd shape, analytic,  NoBeta, alpha=1        : y[0,0,0,0,0:3] = 105 120 116, max |y - (conv + bias)| = 0

The inputs are small integers from fixed seeds, so the comparison is exact and 0 means every element matches. I got the same six lines with the headers from main (0b55a2f) and from v4.3.5, the release current when you asked. I built it from a CUTLASS checkout with nvcc -std=c++17 -O2 -gencode arch=compute_80,code=compute_80 -Iinclude -Itools/util/include conv3d_bias.cu -o conv3d_bias (nvcc 13.2) and ran it on an RTX 5090, where the compute_80 PTX is JIT-compiled. I don't have an A100 to run it on.

On your questions:

Is LinearCombination supported for OpClassTensorOp in 3D convolutions? Yes. Both the optimized and analytic iterators give exact results, including a shape where neither N*Z*P*Q = 315 nor K = 40 is a multiple of the 128x128 threadblock tile.

How should the bias be passed? Pass it as tensor C with all strides zero, as examples/17_fprop_per_channel_bias does for Conv2d: {bias_ptr, cutlass::layout::TensorNDHWC::Stride(0)}. Stride(0) sets all four NDHWC strides to 0. With the default template arguments the fprop epilogue treats C as a row-major (N*Z*P*Q) x K matrix and takes its row stride from stride(0) of C (ConvOutputIteratorParameter in include/cutlass/epilogue/threadblock/output_iterator_parameter.h), so the 5D coordinates never come into it.

One setting drops the bias without any error. With the default ScaleType::Default, LinearCombination::is_source_needed() returns beta != 0, so with beta = 0 C is never read. That is the beta=0 line above: the error of 163 is the largest bias value (100 + 63). Either pass beta = 1, or use ScaleType::NoBetaScaling as example 17 does, which computes alpha * acc + C and ignores beta.

The other easy mistake is passing bias.device_ref() for a HostTensor of extent {1, 1, 1, 1, K}. Its strides are packed, so the row stride is K rather than 0. When I swapped that into the program, the first output row came out right and the later rows were wrong, because they read past the end of the K-element bias buffer.

If your output still differs, the kernel type and the Arguments construction from your code would show where it departs from this one.

conv3d_bias.cu

// Conv3d fprop (NDHWC) + per-channel bias on SM80 tensor cores.
// Bias is passed as tensor C with all-zero strides, as in examples/17_fprop_per_channel_bias.
#include <iostream>
#include <cmath>

#include "cutlass/cutlass.h"
#include "cutlass/conv/kernel/default_conv3d_fprop.h"
#include "cutlass/conv/device/implicit_gemm_convolution.h"
#include "cutlass/epilogue/thread/linear_combination.h"

#include "cutlass/util/host_tensor.h"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/convolution.h"

using ElementA = cutlass::half_t;
using ElementB = cutlass::half_t;
using ElementC = float;
using ElementAcc = float;
using Layout = cutlass::layout::TensorNDHWC;

template <cutlass::conv::IteratorAlgorithm Algo, cutlass::epilogue::thread::ScaleType::Kind Scale>
using Conv3dFprop = cutlass::conv::device::ImplicitGemmConvolution<
  typename cutlass::conv::kernel::DefaultConv3dFprop<
    ElementA, Layout, ElementB, Layout, ElementC, Layout, ElementAcc,
    cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80,
    cutlass::gemm::GemmShape<128, 128, 32>,
    cutlass::gemm::GemmShape<64, 64, 32>,
    cutlass::gemm::GemmShape<16, 8, 16>,
    cutlass::epilogue::thread::LinearCombination<ElementC, 4, ElementAcc, ElementAcc, Scale>,
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
    3, cutlass::arch::OpMultiplyAdd, Algo
  >::Kernel>;

template <typename Conv>
void run(char const *name, float alpha, float beta,
         cutlass::Tensor5DCoord in = cutlass::Tensor5DCoord(2, 4, 8, 8, 32), int K = 64) {
  // N, D, H, W, C / K, T, R, S, C / pad / stride / dilation / Z, P, Q
  cutlass::conv::Conv3dProblemSize ps(
    in, cutlass::Tensor5DCoord(K, 3, 3, 3, in[4]),
    cutlass::make_Coord(1, 1, 1), cutlass::make_Coord(1, 1, 1), cutlass::make_Coord(1, 1, 1), cutlass::Tensor5DCoord(in[0], in[1], in[2], in[3], K),
    cutlass::conv::Mode::kCrossCorrelation, 1);

  cutlass::HostTensor<ElementA, Layout> x(ps.activation_extent());
  cutlass::HostTensor<ElementB, Layout> w(ps.filter_extent());
  cutlass::HostTensor<ElementC, Layout> bias({1, 1, 1, 1, ps.K});
  cutlass::HostTensor<ElementC, Layout> y(ps.output_extent());
  cutlass::HostTensor<ElementC, Layout> y_ref(ps.output_extent());
  cutlass::HostTensor<ElementC, Layout> zeros(ps.output_extent());

  // Small integers, fixed seeds: every product and sum is exact in fp32.
  cutlass::reference::host::TensorFillRandomUniform(x.host_view(), 1, ElementA(2), ElementA(-2), 0);
  cutlass::reference::host::TensorFillRandomUniform(w.host_view(), 2, ElementB(2), ElementB(-2), 0);
  for (int k = 0; k < ps.K; ++k) bias.at({0, 0, 0, 0, k}) = float(100 + k);
  cutlass::reference::host::TensorFill(y.host_view());
  cutlass::reference::host::TensorFill(zeros.host_view());
  x.sync_device(); w.sync_device(); bias.sync_device(); y.sync_device();

  typename Conv::Arguments args{
    ps, x.device_ref(), w.device_ref(),
    {bias.device_data(), Layout::Stride(0)},   // C = bias, all four strides zero
    y.device_ref(),
    {alpha, beta}};

  Conv op;
  cutlass::Status st = op.can_implement(args);
  if (st == cutlass::Status::kSuccess) st = op.initialize(args);
  if (st == cutlass::Status::kSuccess) st = op();
  cudaError_t ce = cudaDeviceSynchronize();
  if (st != cutlass::Status::kSuccess || ce != cudaSuccess) {
    std::cout << name << ": status " << cutlassGetStatusString(st) << ", " << cudaGetErrorString(ce) << "\n";
    return;
  }
  y.sync_host();

  // Reference: alpha * conv(x, w), then + bias[k] on the host.
  cutlass::reference::host::Conv3d<ElementA, Layout, ElementB, Layout, ElementC, Layout, ElementAcc, ElementAcc>(
    cutlass::conv::Operator::kFprop, ps, x.host_ref(), w.host_ref(), zeros.host_ref(), y_ref.host_ref(), alpha, 0.f);

  double max_err = 0;
  auto ext = ps.output_extent();
  for (int n = 0; n < ext[0]; ++n)
  for (int z = 0; z < ext[1]; ++z)
  for (int p = 0; p < ext[2]; ++p)
  for (int q = 0; q < ext[3]; ++q)
  for (int k = 0; k < ext[4]; ++k) {
    double want = y_ref.at({n, z, p, q, k}) + bias.at({0, 0, 0, 0, k});
    max_err = std::fmax(max_err, std::fabs(double(y.at({n, z, p, q, k})) - want));
  }
  std::cout << name << ": y[0,0,0,0,0:3] = " << y.at({0,0,0,0,0}) << " " << y.at({0,0,0,0,1}) << " " << y.at({0,0,0,0,2})
            << ", max |y - (conv + bias)| = " << max_err << "\n";
}

int main() {
  using SK = cutlass::epilogue::thread::ScaleType;
  using CA = cutlass::conv::IteratorAlgorithm;
  run<Conv3dFprop<CA::kOptimized, SK::NoBetaScaling>>("optimized, NoBetaScaling,     alpha=1        ", 1.f, 0.f);
  run<Conv3dFprop<CA::kAnalytic,  SK::NoBetaScaling>>("analytic,  NoBetaScaling,     alpha=1        ", 1.f, 0.f);
  run<Conv3dFprop<CA::kOptimized, SK::Default>>      ("optimized, Default,           alpha=1, beta=1", 1.f, 1.f);
  run<Conv3dFprop<CA::kOptimized, SK::Default>>      ("optimized, Default,           alpha=1, beta=0", 1.f, 0.f);
  // Odd shape: M = 1*5*7*9 = 315 and K = 40, neither a multiple of the 128x128 tile.
  run<Conv3dFprop<CA::kOptimized, SK::NoBetaScaling>>("odd shape, optimized, NoBeta, alpha=1        ", 1.f, 0.f,
                                                      cutlass::Tensor5DCoord(1, 5, 7, 9, 16), 40);
  run<Conv3dFprop<CA::kAnalytic,  SK::NoBetaScaling>>("odd shape, analytic,  NoBeta, alpha=1        ", 1.f, 0.f,
                                                      cutlass::Tensor5DCoord(1, 5, 7, 9, 16), 40);
  return 0;
}