HI, I am trying to call the native functions. However, I met ‘CUDA error: an illegal memory access was encountered’ when I ran the CUDA version and it gave ‘Segmentation fault’ when I switched to the CPU version.
I tested the code on pytorch 1.5/1.6/1.7 with cuda 9.2, all pytorch version gave the same error.
What I was trying to do:
First wrap the functions:
#include <torch/extension.h>
#include <ATen/NativeFunctions.h>
#include <ATen/Config.h>

std::tuple<at::Tensor, at::Tensor, at::Tensor> layer_norm_forward_cpu(
    const at::Tensor & input,
    const at::Tensor & weight,
    const at::Tensor & bias,
    int64_t M, int64_t N, double eps) {
    return at::native::layer_norm_cpu(input, weight, bias, M, N, eps);
}

std::tuple<at::Tensor, at::Tensor, at::Tensor> backward_layer_norm_cpu(
    const at::Tensor & grad_out,
    const at::Tensor & input,
    const at::Tensor & mean,
    const at::Tensor & rstd,
    const at::Tensor & weight,
    int64_t M, int64_t N, std::array<bool,3> output_mask) {
    return at::native::layer_norm_backward_cpu(grad_out, input, mean, rstd, weight, M, N, output_mask);
}

std::tuple<at::Tensor, at::Tensor, at::Tensor> layer_norm_forward_cuda(
    const at::Tensor & input,
    const at::Tensor & weight,
    const at::Tensor & bias,
    int64_t M, int64_t N, double eps) {
    return at::native::layer_norm_cuda(input, weight, bias, M, N, eps);
}

std::tuple<at::Tensor, at::Tensor, at::Tensor> backward_layer_norm_cuda(
    const at::Tensor & grad_out,
    const at::Tensor & input,
    const at::Tensor & mean,
    const at::Tensor & rstd,
    const at::Tensor & weight,
    int64_t M, int64_t N, std::array<bool,3> output_mask) {
    return at::native::layer_norm_backward_cuda(grad_out, input, mean, rstd, weight, M, N, output_mask);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("layer_norm_forward_cpu",  &layer_norm_forward_cpu, "layer norm forward (cpu version)");
    m.def("layer_norm_backward_cpu", &backward_layer_norm_cpu, "layer norm backward (cpu version)");
    m.def("layer_norm_forward_cuda", &layer_norm_forward_cuda, "layer norm forward (cuda version)");
    m.def("layer_norm_backward_cuda",&backward_layer_norm_cuda, "layer norm backward (cuda version)");
}

and then call the functions in
import torch
import torch.nn as nn
import torch.nn.functional as F
import pdb

import native

class layer_norm(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, normalized_shape, weight, bias, eps, training):
        N = 1
        if isinstance(normalized_shape, int):
            N = normalized_shape
        elif isinstance(normalized_shape, (list, tuple)):
            for i in normalized_shape:
                N *= i
        else:
            raise RuntimeError("unexpected type of normalized_shape".format(type(normalized_shape)))
        M = x.nelement() // N

        if x.is_cuda:

            y, mean, rstd = native.layer_norm_forward_cuda(x, weight, bias, M, N, eps)
        else:
            y, mean, rstd = native.layer_norm_forward_cpu(x, weight, bias, M, N, eps)

        if training:
            ctx.layer_norm_input = x
            ctx.layer_norm_parameters = (mean, rstd, weight, M, N)
        return y

    @staticmethod
    def backward(ctx, grad_output):
        x = ctx.layer_norm_input
        mean, rstd, weight, M, N = ctx.layer_norm_parameters

        output_mask = [True, True, True]
        if grad_output.is_cuda:
            grad_input, grad_weight, grad_bias = native.layer_norm_backward_cuda(grad_output, x, mean, rstd, weight, M, N, output_mask)
        else:
            grad_input, grad_weight, grad_bias = native.layer_norm_backward_cpu(grad_output, x, mean, rstd, weight, M, N, output_mask)
        ctx.layer_norm_input = None
        ctx.layer_norm_parameters = None
        return grad_input, None, grad_weight, grad_bias, None, None, None, None

class LayerNorm(nn.LayerNorm):
    def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True):
        nn.LayerNorm.__init__(self, normalized_shape, eps=eps, elementwise_affine=elementwise_affine)

    def forward(self, x):
        y = layer_norm.apply(x, self.normalized_shape, self.weight, self.bias, self.eps, self.training)
        return y


if __name__ == "__main__":
    seed = 2809
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.deterministic=True #https://github.com/pytorch/pytorch/issues/8019

    model = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            LayerNorm([64,56,56])
            )
    print(model)

    #model = model.cuda()
    model.train()
    iteration = 10
    for i in range(iteration):
        print("index: ", i)
        x = torch.rand(512,64,56,56)
        x = x - 0.5
        #x = x.cuda()

        y = model(x)
        z = y.sum()
        z.backward()


I also uploaded all the code to GitHub - irving-qin/nativefunctions
Just run bash install.sh for testing.
The forward process of the wrapped layer norm seemed to be normal. However, it throwed errors in the backward function.
Thanks so much if anyone could give me some tips.