TorchLean API

NN.Runtime.Autograd.Engine.LibTorch.ConvPool

CUDA Conv/Pool FFI #

Foreign-function declarations for TorchLean's float32 convolution and pooling adapter, implemented in csrc/libtorch/torchlean.cpp.

All buffers are contiguous LibTorch.Buffer values and shape/stride/padding metadata is passed explicitly through the FFI boundary.

@[never_extract, extern torchlean_cuda_convtranspose_fwd]
opaque Runtime.Autograd.LibTorch.torchleanConvTransposeFwdCuda (input kernel bias : Buffer) (inSpatial kernelSpatial stride padding : Array Nat) (inC outC : UInt32) :

Float32 transposed convolution forward (channels-first, no batch).

Shapes/parameters:

  • inSpatial: length d (input spatial dims)
  • kernelSpatial: length d (kernel window)
  • stride: length d
  • padding: length d

All arrays must have the same length 1 ≤ d ≤ 3.

Layout conventions:

  • input: (inC, spatial...)
  • kernel: (inC, outC, kernelSpatial...)
  • bias: (outC)
  • output: (outC, outSpatial...), where outSpatial[i] = (inSpatial[i] - 1) * stride[i] - 2*padding[i] + kernelSpatial[i].
@[never_extract, extern torchlean_cuda_convtranspose_bwd]
opaque Runtime.Autograd.LibTorch.torchleanConvTransposeBwdCuda (input kernel gradOutput : Buffer) (inSpatial kernelSpatial stride padding : Array Nat) (inC outC : UInt32) :

Float32 transposed convolution backward.

Returns (dKernel, dBias, dInput) as device buffers. Array conventions match torchleanConvTransposeFwdCuda.

@[never_extract, extern torchlean_cuda_conv_fwd]
opaque Runtime.Autograd.LibTorch.torchleanConvFwdCuda (input kernel bias : Buffer) (inSpatial kernelSpatial stride padding : Array Nat) (inC outC : UInt32) :

Float32 convolution forward (channels-first, no batch).

Shapes/parameters:

  • inSpatial: length d (spatial dims)
  • kernelSpatial: length d (kernel window)
  • stride: length d
  • padding: length d

All arrays must have the same length 1 ≤ d ≤ 3.

@[never_extract, extern torchlean_cuda_conv_bwd]
opaque Runtime.Autograd.LibTorch.torchleanConvBwdCuda (input kernel gradOutput : Buffer) (inSpatial kernelSpatial stride padding : Array Nat) (inC outC : UInt32) :

Float32 convolution backward.

Returns (dKernel, dBias, dInput) as device buffers. Array conventions match torchleanConvFwdCuda.

@[never_extract, extern torchlean_cuda_maxpool_fwd]
opaque Runtime.Autograd.LibTorch.torchleanMaxPoolFwdCuda (input : Buffer) (inSpatial kernel stride padding : Array Nat) (inC : UInt32) :

Float32 max-pooling forward (channels preserved).

@[never_extract, extern torchlean_cuda_maxpool_bwd]
opaque Runtime.Autograd.LibTorch.torchleanMaxPoolBwdCuda (input gradOutput : Buffer) (inSpatial kernel stride padding : Array Nat) (inC : UInt32) :

Float32 max-pooling backward: returns dInput.

@[never_extract, extern torchlean_cuda_avgpool_fwd]
opaque Runtime.Autograd.LibTorch.torchleanAvgPoolFwdCuda (input : Buffer) (inSpatial kernel stride padding : Array Nat) (inC : UInt32) :

Float32 average-pooling forward (channels preserved).

@[never_extract, extern torchlean_cuda_avgpool_bwd]
opaque Runtime.Autograd.LibTorch.torchleanAvgPoolBwdCuda (gradOutput : Buffer) (inSpatial kernel stride padding : Array Nat) (inC : UInt32) :

Float32 average-pooling backward: returns dInput.

@[never_extract, extern torchlean_cuda_smooth_maxpool_fwd]
opaque Runtime.Autograd.LibTorch.torchleanSmoothMaxPoolFwdCuda (input : Buffer) (beta : Float) (inSpatial kernel stride padding : Array Nat) (inC : UInt32) :

Float32 smooth max-pooling forward with channels preserved.

The native implementation requires finite nonzero beta and uses a maximum input pivot for positive beta or a minimum input pivot for negative beta, matching the two-dimensional path.

@[never_extract, extern torchlean_cuda_smooth_maxpool_bwd]
opaque Runtime.Autograd.LibTorch.torchleanSmoothMaxPoolBwdCuda (input gradOutput : Buffer) (beta : Float) (inSpatial kernel stride padding : Array Nat) (inC : UInt32) :

Float32 smooth max-pooling backward, returning dInput.

It shares the forward operation's finite nonzero-beta contract and sign-aware max/min input pivot.