# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
#
# This source code is licensed under the BSD license found in the
# LICENSE file in the root directory of this source tree.


import logging
from enum import Enum
from typing import Optional

import torch
import triton
from torch.cuda.amp import custom_bwd, custom_fwd

from xformers.triton.k_softmax import _softmax, _softmax_backward

# CREDITS: This is adapted from the vanilla Triton example. See https://openai.com/blog/triton/
# and https://triton-lang.org/getting-started/tutorials/02-fused-softmax.html


_triton_registered_overflow = False
_triton_registered_warnings = False
_triton_softmax_fp16_enabled = False  # NOTE: PyTorch keeps softmax as fp32


class MaskType(str, Enum):
    ADD = "add"
    MUL = "mul"


# Helper to handle the SPMD launch grid and error cases
class _softmax_triton(torch.autograd.Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float16 if _triton_softmax_fp16_enabled else None)
    def forward(ctx, x, mask, log_outputs, causal):
        """
        Fused softmax implementation, using the Triton programming model.
        This only supports a reduction over the last dimension for now
        """

        # Handle 2D/3D tensors
        x_ = x.unsqueeze(0) if x.ndim == 2 else x

        if not x_.is_contiguous():
            x_ = x_.contiguous()

        y = torch.empty_like(x_)
        assert (
            y.stride(2) == 1 and x_.stride(2) == 1
        ), f"{x.shape} - {x_.shape} - {x_.stride()}"

        # SPMD launch grid
        grid_2d = (
            x_.shape[0],
            x_.shape[1],
        )

        # enqueue GPU kernel
        use_mask = True
        if mask is None:
            #  placeholder, will not be used
            mask = x_
            use_mask = False
        else:
            # Make sure that the mask is binary
            assert mask.dtype == x.dtype, "An additive mask is requested"

        _softmax[grid_2d](
            y,
            x_,
            mask,
            y.stride(0),
            y.stride(1),
            x_.stride(0),
            x_.stride(1),
            mask.stride(0),
            x_.shape[2],
            log=log_outputs,
            use_mask=use_mask,
            causal=causal,
        )

        ctx.save_for_backward(y)
        ctx.log_outputs = log_outputs
        ctx.causal = causal
        return y.reshape_as(x)

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_out):
        (out,) = ctx.saved_tensors

        # Handle 2D/3D tensors
        grad_out_ = grad_out.unsqueeze(0) if grad_out.ndim == 2 else grad_out

        # SPMD launch grid
        grid_2d = (
            grad_out_.shape[0],
            grad_out_.shape[1],
        )

        depth = triton.next_power_of_2(grad_out_.shape[2])
        grad_in = torch.empty_like(
            out
        )  # torch.zeros is measurably slower, we'll zero out in the kernel

        # Make sure that the tensor are contiguous
        grad_in, grad_out, out = map(lambda x: x.contiguous(), [grad_in, grad_out, out])

        # fmt: off
        _softmax_backward[grid_2d](
            grad_in, grad_out_, out,
            grad_in.stride(0), grad_in.stride(1),
            grad_out_.stride(0), grad_out_.stride(1),
            out.stride(0), out.stride(1),
            out.shape[2],
            depth=depth,
            log=ctx.log_outputs,
            causal=ctx.causal
        )
        # fmt: on
        return grad_in.reshape_as(grad_out), None, None, None


def softmax(
    x: torch.Tensor, mask: Optional[torch.Tensor] = None, causal: bool = False
) -> torch.Tensor:
    r"""Applies the Softmax function to an 3-dimensional input Tensor
    rescaling them so that the elements of the n-dimensional output Tensor
    lie in the range [0,1] and sum to 1.

    Softmax is defined as:

    .. math::
        \text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}

    .. warning: softmax is computed on the last dimension of the input tensor.


    Args:
        x: input tensor.
        mask: optional mask, its application will be fused to the softmax computation if triton is used
        causal: optional performance optimization, if triton is used and the attention is causal

    Returns:
        a Tensor of the same dimension and shape as the input with
        values in the range [0, 1] and sum to 1
    """
    return _softmax_dispatch(x, log=False, mask=mask, causal=causal)


def log_softmax(
    x: torch.Tensor, mask: Optional[torch.Tensor] = None, causal: bool = False
) -> torch.Tensor:
    r"""Applies the :math:`\log(\text{Softmax}(x))` function to an 3-dimensional
    input Tensor. The LogSoftmax formulation can be simplified as:

    .. math::
        \text{LogSoftmax}(x_{i}) = \log\left(\frac{\exp(x_i) }{ \sum_j \exp(x_j)} \right)

    Args:
        x: input tensor.

    Returns:
        a Tensor of the same dimension and shape as the input with
        values in the range [-inf, 0)
    """
    return _softmax_dispatch(x, log=True, mask=mask, causal=causal)


def _softmax_dispatch(
    x: torch.Tensor, log: bool, mask: Optional[torch.Tensor], causal: bool = False
) -> torch.Tensor:
    # Triton is used if
    # - CUDA
    # - there's enough data to make it faster than pytorch. This could change over time, Triton is improving
    # - there was no previous failure

    global _triton_registered_overflow
    global _triton_registered_warnings

    try:
        if torch.cuda.is_available() and x.is_cuda and not _triton_registered_overflow:
            return _softmax_triton.apply(x, mask, log, causal)
    except (triton.code_gen.OutOfResources, RuntimeError) as e:
        # Catch cases where the current GPU does not have enough registers to hold a full tensor line
        # fallback to PyTorch's implementation, which streams the tensor in and out
        _triton_registered_overflow = True
        logging.warning(
            "Triton softmax kernel register spillover or invalid image caught."
            "Deactivating this kernel, please file an issue int the xFormers repository"
        )
        logging.warning(e)

    if causal and not _triton_registered_warnings:
        logging.warning(
            "Triton softmax could not be used. \
                The causal flags is being passed but it does not provide any benefit with PyTorch softmax."
        )
        _triton_registered_warnings = True

    if mask is not None:
        x += mask

    if log:
        return torch.log_softmax(x, dim=-1)
    else:
        return torch.softmax(x, dim=-1)

xformers/triton/softmax.py
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
#
# This source code is licensed under the BSD license found in the
# LICENSE file in the root directory of this source tree.


import logging
from enum import Enum
from typing import Optional

import torch
import triton
from torch.cuda.amp import custom_bwd, custom_fwd

from xformers.triton.k_softmax import _softmax, _softmax_backward

# CREDITS: This is adapted from the vanilla Triton example. See https://openai.com/blog/triton/
# and https://triton-lang.org/getting-started/tutorials/02-fused-softmax.html


_triton_registered_overflow = False
_triton_registered_warnings = False
_triton_softmax_fp16_enabled = False  # NOTE: PyTorch keeps softmax as fp32


class MaskType(str, Enum):
    ADD = "add"
    MUL = "mul"


# Helper to handle the SPMD launch grid and error cases
class _softmax_triton(torch.autograd.Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float16 if _triton_softmax_fp16_enabled else None)
    def forward(ctx, x, mask, log_outputs, causal):
        """
        Fused softmax implementation, using the Triton programming model.
        This only supports a reduction over the last dimension for now
        """

        # Handle 2D/3D tensors
        x_ = x.unsqueeze(0) if x.ndim == 2 else x

        if not x_.is_contiguous():
            x_ = x_.contiguous()

        y = torch.empty_like(x_)
        assert (
            y.stride(2) == 1 and x_.stride(2) == 1
        ), f"{x.shape} - {x_.shape} - {x_.stride()}"

        # SPMD launch grid
        grid_2d = (
            x_.shape[0],
            x_.shape[1],
        )

        # enqueue GPU kernel
        use_mask = True
        if mask is None:
            #  placeholder, will not be used
            mask = x_
            use_mask = False
        else:
            # Make sure that the mask is binary
            assert mask.dtype == x.dtype, "An additive mask is requested"

        _softmax[grid_2d](
            y,
            x_,
            mask,
            y.stride(0),
            y.stride(1),
            x_.stride(0),
            x_.stride(1),
            mask.stride(0),
            x_.shape[2],
            log=log_outputs,
            use_mask=use_mask,
            causal=causal,
        )

        ctx.save_for_backward(y)
        ctx.log_outputs = log_outputs
        ctx.causal = causal
        return y.reshape_as(x)

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_out):
        (out,) = ctx.saved_tensors

        # Handle 2D/3D tensors
        grad_out_ = grad_out.unsqueeze(0) if grad_out.ndim == 2 else grad_out

        # SPMD launch grid
        grid_2d = (
            grad_out_.shape[0],
            grad_out_.shape[1],
        )

        depth = triton.next_power_of_2(grad_out_.shape[2])
        grad_in = torch.empty_like(
            out
        )  # torch.zeros is measurably slower, we'll zero out in the kernel

        # Make sure that the tensor are contiguous
        grad_in, grad_out, out = map(lambda x: x.contiguous(), [grad_in, grad_out, out])

        # fmt: off
        _softmax_backward[grid_2d](
            grad_in, grad_out_, out,
            grad_in.stride(0), grad_in.stride(1),
            grad_out_.stride(0), grad_out_.stride(1),
            out.stride(0), out.stride(1),
            out.shape[2],
            depth=depth,
            log=ctx.log_outputs,
            causal=ctx.causal
        )
        # fmt: on
        return grad_in.reshape_as(grad_out), None, None, None


def softmax(
    x: torch.Tensor, mask: Optional[torch.Tensor] = None, causal: bool = False
) -> torch.Tensor:
    r"""Applies the Softmax function to an 3-dimensional input Tensor
    rescaling them so that the elements of the n-dimensional output Tensor
    lie in the range [0,1] and sum to 1.

    Softmax is defined as:

    .. math::
        \text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}

    .. warning: softmax is computed on the last dimension of the input tensor.


    Args:
        x: input tensor.
        mask: optional mask, its application will be fused to the softmax computation if triton is used
        causal: optional performance optimization, if triton is used and the attention is causal

    Returns:
        a Tensor of the same dimension and shape as the input with
        values in the range [0, 1] and sum to 1
    """
    return _softmax_dispatch(x, log=False, mask=mask, causal=causal)


def log_softmax(
    x: torch.Tensor, mask: Optional[torch.Tensor] = None, causal: bool = False
) -> torch.Tensor:
    r"""Applies the :math:`\log(\text{Softmax}(x))` function to an 3-dimensional
    input Tensor. The LogSoftmax formulation can be simplified as:

    .. math::
        \text{LogSoftmax}(x_{i}) = \log\left(\frac{\exp(x_i) }{ \sum_j \exp(x_j)} \right)

    Args:
        x: input tensor.

    Returns:
        a Tensor of the same dimension and shape as the input with
        values in the range [-inf, 0)
    """
    return _softmax_dispatch(x, log=True, mask=mask, causal=causal)


def _softmax_dispatch(
    x: torch.Tensor, log: bool, mask: Optional[torch.Tensor], causal: bool = False
) -> torch.Tensor:
    # Triton is used if
    # - CUDA
    # - there's enough data to make it faster than pytorch. This could change over time, Triton is improving
    # - there was no previous failure

    global _triton_registered_overflow
    global _triton_registered_warnings

    try:
        if torch.cuda.is_available() and x.is_cuda and not _triton_registered_overflow:
            return _softmax_triton.apply(x, mask, log, causal)
    except (triton.code_gen.OutOfResources, RuntimeError) as e:
        # Catch cases where the current GPU does not have enough registers to hold a full tensor line
        # fallback to PyTorch's implementation, which streams the tensor in and out
        _triton_registered_overflow = True
        logging.warning(
            "Triton softmax kernel register spillover or invalid image caught."
            "Deactivating this kernel, please file an issue int the xFormers repository"
        )
        logging.warning(e)

    if causal and not _triton_registered_warnings:
        logging.warning(
            "Triton softmax could not be used. \
                The causal flags is being passed but it does not provide any benefit with PyTorch softmax."
        )
        _triton_registered_warnings = True

    if mask is not None:
        x = x + mask

    if log:
        return torch.log_softmax(x, dim=-1)
    else:
        return torch.softmax(x, dim=-1)

xformers/triton/softmax.py
diff --git a/tests/test_triton_softmax.py b/tests/test_triton_softmax.py
index 34abcc8488..83cea8943d 100644
--- a/tests/test_triton_softmax.py
+++ b/tests/test_triton_softmax.py
@@ -103,3 +103,61 @@ def test_softmax_fp16(dtype):
 
     a = torch.rand(b, s, d, device="cuda", dtype=dtype)
     triton_softmax(a)
+
+
+@pytest.mark.skipif(not _triton_available, reason="Triton is not available")
+@pytest.mark.parametrize("log", [False, True])
+@pytest.mark.parametrize("masking", [True, False])
+@pytest.mark.parametrize("causal", [True, False])
+@pytest.mark.parametrize("contiguous", [True, False])
+def test_softmax_parity_fallback(log, masking, causal, contiguous):
+    """Check that the fallback paths are correct"""
+    torch.random.manual_seed(0)
+
+    shape = (16, 16)
+
+    # Check the result of a FW pass
+    X = torch.normal(0, 1, size=shape, device="cpu", requires_grad=False)
+
+    if not contiguous:
+        # Make sure that the buffer is not contiguous
+        X = X.transpose(-2, -1).contiguous().transpose(-2, -1)
+
+    X_ = X.clone()
+    X.requires_grad = True
+    X_.requires_grad = True
+
+    seq = shape[1]
+    mask = torch.zeros((seq, seq))
+    if masking:
+        mask[torch.rand((seq, seq)) > 0.8] = -float("inf")
+
+    if causal:
+        mask[~torch.tril(torch.ones_like(mask)).bool()] = -float("inf")
+
+    y_torch = (
+        torch.log_softmax(X + mask, dim=-1) if log else torch.softmax(X + mask, dim=-1)
+    )
+    y_triton = (
+        triton_log_softmax(X_, mask, causal)
+        if log
+        else triton_softmax(X_, mask, causal)
+    )
+
+    assert torch.allclose(y_torch, y_triton, equal_nan=True)
+
+    print(y_torch)
+
+    # Check that BW also gives the same result
+    loss_torch = torch.norm(y_torch.transpose(-2, -1) @ y_torch)
+    loss_torch.backward()
+
+    loss_triton = torch.norm(y_triton.transpose(-2, -1) @ y_triton)
+    loss_triton.backward()
+
+    print(X.grad)
+    print(X_.grad)
+
+    assert torch.allclose(
+        torch.norm(X.grad), torch.norm(X_.grad), equal_nan=True, atol=1e-5
+    ), f"{torch.norm(X.grad)}, {torch.norm(X_.grad)}"
diff --git a/xformers/triton/softmax.py b/xformers/triton/softmax.py
index 008d0fb4b5..e1ccd1dcfa 100644
--- a/xformers/triton/softmax.py
+++ b/xformers/triton/softmax.py
@@ -200,7 +200,7 @@ def _softmax_dispatch(
         _triton_registered_warnings = True
 
     if mask is not None:
-        x += mask
+        x = x + mask
 
     if log:
         return torch.log_softmax(x, dim=-1)

Follow up on #141 with a proper unit test
# 🐛 Bug
Make sure that #141 cannot happen anymore, add a small unit test to guard parity

Follow up on #141 with a proper unit test
# 🐛 Bug
Make sure that #141 cannot happen anymore, add a small unit test to guard parity

