# 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.


# CREDITS: This comes almost as-is from the Triton dropout tutorial
# https://raw.githubusercontent.com/openai/triton/master/python/tutorials/04-low-memory-dropout.py

import triton
import triton.language as tl

_configs = [
    triton.Config({}, num_warps=1),
    triton.Config({}, num_warps=2),
    triton.Config({}, num_warps=4),
    triton.Config({}, num_warps=8),
    triton.Config({}, num_warps=16),
]


@triton.jit
def _get_4_bin_masks(seed, rand_offsets, p):
    rand1, rand2, rand3, rand4 = tl.randint4x(seed.to(tl.int32), rand_offsets)

    # binarize masks, save registers
    # NOTE: We keep the random numbers as is there (integers over int32),
    # and convert the threshold instead, for speed

    # The initial distribution is -2**31 / 2**31 -1
    # and our float threshold  is in between [0, 1]
    # The full computation is: `start_point + full range * p`
    threshold = (-2147483648.0 + 4294967295.0 * p).to(tl.int32)
    rand_mask1 = rand1 > threshold
    rand_mask2 = rand2 > threshold
    rand_mask3 = rand3 > threshold
    rand_mask4 = rand4 > threshold

    return rand_mask1, rand_mask2, rand_mask3, rand_mask4


@triton.jit
def _random_prune_and_scale(x, rand_mask, p, p_scale):
    zero = 0.0

    if p > 0.0:
        # generate all the random numbers for the block at once, then reshape
        keep = tl.reshape(rand_mask, x.shape)

        # prune and normalize in one go
        x = tl.where(keep, (x * p_scale).to(x.dtype), zero.to(x.dtype))

    return x


# fmt: off
@triton.heuristics({"SIZE_RAND_BLOCK": lambda *_, **meta: meta["BLOCK_N"] * meta["BLOCK_M"]})
@triton.autotune(
    configs=_configs,
    key=["M", "N", "is_fp16"],
)
@triton.jit
def k_dropout_fw(
    Y, X, BIAS, SEEDS,
    stride,
    M, N,
    p,
    is_fp16,  # autotune
    **meta,
):
    """
    Apply dropout on an input tensor
    Y : Output  (M, N)
    X : Input   (M, N)
    BIAS        (N,)
    SEEDS       (M,)
    p : dropout probability
    """
    # fmt: on

    BLOCK_M = meta["BLOCK_M"]
    BLOCK_N = meta["BLOCK_N"]
    SIZE_RAND_BLOCK = meta["SIZE_RAND_BLOCK"]

    row_id = tl.program_id(axis=0)
    rows = row_id * BLOCK_M * 4 + tl.arange(0, BLOCK_M)

    col_id = tl.program_id(axis=1)
    cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N)
    seed = SEEDS + col_id

    # pointers starting point
    x_ptrs = X + rows[:, None] * stride + cols[None, :]
    y_ptrs = Y + rows[:, None] * stride + cols[None, :]

    # go over all the tiles, one by one
    rand_offsets = tl.arange(0, SIZE_RAND_BLOCK) + row_id * BLOCK_M * 4
    rand_mask1, rand_mask2, rand_mask3, rand_mask4 = _get_4_bin_masks(seed, rand_offsets, p)

    col_mask = cols[None, :] < N
    p_scale = 1 / (1 - p) if p < 1. else 1.

    if meta["USE_BIAS"]:
        b_ptrs = BIAS + cols[None, :]
        bias = tl.load(b_ptrs, mask=cols[None, :] < N, other=0.)

    for i in range(4):
        # cycle through the binary masks (workaround / no indexing)
        if i == 0:
            rand_mask = rand_mask1
        elif i == 1:
            rand_mask = rand_mask2
        elif i == 2:
            rand_mask = rand_mask3
        else:
            rand_mask = rand_mask4

        block_mask = (rows[:, None] < M) & col_mask
        x = tl.load(x_ptrs, mask=block_mask, other=0.)

        # optionally apply a fused bias
        if meta["USE_BIAS"]:
            x += bias

        # optional: fused activation (while the data is in shared memory)
        if meta["ACTIVATION"]:
            x = meta["ACTIVATION"](x)

        # randomly prune (and scale) the resulting buffer, possibly a no-op
        output = _random_prune_and_scale(x, rand_mask, p, p_scale)

        tl.store(y_ptrs, output, mask=block_mask)

        # Update the pointers
        rows += BLOCK_M  # needs to be updated for the mask to be correct
        x_ptrs += BLOCK_M * stride
        y_ptrs += BLOCK_M * stride


# fmt: off
@triton.heuristics({"SIZE_RAND_BLOCK": lambda *_, **meta: meta["BLOCK_N"] * meta["BLOCK_M"]})
@triton.autotune(
    configs=_configs,
    key=["M", "N", "is_fp16"],
)
@triton.jit
def k_dropout_bw(
    GRAD_IN, GRAD_BIAS, GRAD_OUT,
    INPUTS, BIAS, SEEDS,
    stride_grad, stride_inputs,
    M, N,
    p,
    is_fp16,  # autotune
    **meta,
):
    """
    Apply dropout on an input tensor
    GRAD_OUT    (M, N)
    GRAD_BIAS   (N,)
    GRAD_IN     (M, N)
    BIAS        (N,)
    SEEDS       (N,)
    p : dropout probability
    """
    # fmt: on

    BLOCK_M = meta["BLOCK_M"]
    BLOCK_N = meta["BLOCK_N"]
    SIZE_RAND_BLOCK = meta["SIZE_RAND_BLOCK"]
    TRAINABLE_BIAS = meta["TRAINABLE_BIAS"]

    rows = tl.arange(0, BLOCK_M)
    row_id = tl.program_id(axis=0)
    rows = row_id * BLOCK_M * 4 + tl.arange(0, BLOCK_M)

    col_id = tl.program_id(axis=1)
    cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N)
    seed = SEEDS + col_id  # FIXME index the seed properly

    # pointers starting point
    grad_out_ptrs = GRAD_OUT + rows[:, None] * stride_grad + cols[None, :]
    grad_in_ptrs = GRAD_IN + rows[:, None] * stride_grad + cols[None, :]
    input_ptrs = INPUTS + rows[:, None] * stride_inputs + cols[None, :]

    # random binary masks, save registers
    rand_offsets = tl.arange(0, SIZE_RAND_BLOCK) + row_id * BLOCK_M * 4
    rand_mask1, rand_mask2, rand_mask3, rand_mask4 = _get_4_bin_masks(seed, rand_offsets, p)

    # now go over the tiles
    grad_bias = tl.zeros((BLOCK_N,), dtype=tl.float32)
    col_mask = cols[None, :] < N
    p_scale = 1 / (1 - p) if p < 1. else 1.

    if meta["USE_BIAS"]:
        b_ptrs = BIAS + cols[None, :]
        bias = tl.load(b_ptrs, mask=col_mask, other=0.)

    for i in range(4):
        # cycle through the binary masks (workaround / no indexing)
        if i == 0:
            rand_mask = rand_mask1
        elif i == 1:
            rand_mask = rand_mask2
        elif i == 2:
            rand_mask = rand_mask3
        else:
            rand_mask = rand_mask4

        block_mask = (rows[:, None] < M) & col_mask
        grad_out = tl.load(grad_out_ptrs, mask=block_mask, other=0.)

        # optional: fused activation (while the data is in shared memory)
        if meta["ACTIVATION_GRAD"]:
            inputs = tl.load(input_ptrs, mask=block_mask, other=0.)

            # optionally apply a fused bias
            if meta["USE_BIAS"]:
                inputs += bias

            act_grad = meta["ACTIVATION_GRAD"](inputs).to(grad_out.dtype)
            grad_out *= act_grad

        # randomly prune (and scale) the resulting buffer, possibly a no-op
        # note that even if we did not save the mask from the FW pass, it is generated
        # from the same seeds, so the same drop mask is applied here
        output = _random_prune_and_scale(grad_out, rand_mask, p, p_scale)

        # write-back
        tl.store(grad_in_ptrs, output, mask=block_mask)

        # optionally accumulate the bias gradient
        if TRAINABLE_BIAS:
            grad_bias += tl.sum(output, axis=0)

        # Update the pointers
        rows += BLOCK_M  # needs to be updated for the mask to be correct
        grad_out_ptrs += BLOCK_M * stride_grad
        input_ptrs += BLOCK_M * stride_inputs
        grad_in_ptrs += BLOCK_M * stride_grad

    if TRAINABLE_BIAS:
        grad_bias_ptr = GRAD_BIAS + row_id * N + cols
        tl.store(grad_bias_ptr, grad_bias, mask=cols < N)

xformers/triton/k_dropout.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.


# CREDITS: This comes almost as-is from the Triton dropout tutorial
# https://raw.githubusercontent.com/openai/triton/master/python/tutorials/04-low-memory-dropout.py

import triton
import triton.language as tl

_configs = [
    triton.Config({}, num_warps=1),
    triton.Config({}, num_warps=2),
    triton.Config({}, num_warps=4),
    triton.Config({}, num_warps=8),
    triton.Config({}, num_warps=16),
]


@triton.jit
def _get_4_bin_masks(seed, rand_offsets, p):
    seed = tl.load(seed)
    rand1, rand2, rand3, rand4 = tl.randint4x(seed, rand_offsets)

    # binarize masks, save registers
    # NOTE: We keep the random numbers as is there (integers over int32),
    # and convert the threshold instead, for speed

    # The initial distribution is -2**31 / 2**31 -1
    # and our float threshold  is in between [0, 1]
    # The full computation is: `start_point + full range * p`
    threshold = (-2147483648.0 + 4294967295.0 * p).to(tl.int32)
    rand_mask1 = rand1 > threshold
    rand_mask2 = rand2 > threshold
    rand_mask3 = rand3 > threshold
    rand_mask4 = rand4 > threshold

    return rand_mask1, rand_mask2, rand_mask3, rand_mask4


@triton.jit
def _random_prune_and_scale(x, rand_mask, p, p_scale):
    zero = 0.0

    if p > 0.0:
        # generate all the random numbers for the block at once, then reshape
        keep = tl.reshape(rand_mask, x.shape)

        # prune and normalize in one go
        x = tl.where(keep, (x * p_scale).to(x.dtype), zero.to(x.dtype))

    return x


# fmt: off
@triton.heuristics({"SIZE_RAND_BLOCK": lambda *_, **meta: meta["BLOCK_N"] * meta["BLOCK_M"]})
@triton.autotune(
    configs=_configs,
    key=["M", "N", "is_fp16"],
)
@triton.jit
def k_dropout_fw(
    Y, X, BIAS, SEEDS,
    stride,
    M, N,
    p,
    is_fp16,  # autotune
    **meta,
):
    """
    Apply dropout on an input tensor
    Y : Output  (M, N)
    X : Input   (M, N)
    BIAS        (N,)
    SEEDS       (M,)
    p : dropout probability
    """
    # fmt: on

    BLOCK_M = meta["BLOCK_M"]
    BLOCK_N = meta["BLOCK_N"]
    SIZE_RAND_BLOCK = meta["SIZE_RAND_BLOCK"]

    row_id = tl.program_id(axis=0)
    rows = row_id * BLOCK_M * 4 + tl.arange(0, BLOCK_M)

    col_id = tl.program_id(axis=1)
    cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N)
    seed = SEEDS + col_id

    # pointers starting point
    x_ptrs = X + rows[:, None] * stride + cols[None, :]
    y_ptrs = Y + rows[:, None] * stride + cols[None, :]

    # go over all the tiles, one by one
    rand_offsets = tl.arange(0, SIZE_RAND_BLOCK) + row_id * BLOCK_M * 4
    rand_mask1, rand_mask2, rand_mask3, rand_mask4 = _get_4_bin_masks(seed, rand_offsets, p)

    col_mask = cols[None, :] < N
    p_scale = 1 / (1 - p) if p < 1. else 1.

    if meta["USE_BIAS"]:
        b_ptrs = BIAS + cols[None, :]
        bias = tl.load(b_ptrs, mask=cols[None, :] < N, other=0.)

    for i in range(4):
        # cycle through the binary masks (workaround / no indexing)
        if i == 0:
            rand_mask = rand_mask1
        elif i == 1:
            rand_mask = rand_mask2
        elif i == 2:
            rand_mask = rand_mask3
        else:
            rand_mask = rand_mask4

        block_mask = (rows[:, None] < M) & col_mask
        x = tl.load(x_ptrs, mask=block_mask, other=0.)

        # optionally apply a fused bias
        if meta["USE_BIAS"]:
            x += bias

        # optional: fused activation (while the data is in shared memory)
        if meta["ACTIVATION"]:
            x = meta["ACTIVATION"](x)

        # randomly prune (and scale) the resulting buffer, possibly a no-op
        output = _random_prune_and_scale(x, rand_mask, p, p_scale)

        tl.store(y_ptrs, output, mask=block_mask)

        # Update the pointers
        rows += BLOCK_M  # needs to be updated for the mask to be correct
        x_ptrs += BLOCK_M * stride
        y_ptrs += BLOCK_M * stride


# fmt: off
@triton.heuristics({"SIZE_RAND_BLOCK": lambda *_, **meta: meta["BLOCK_N"] * meta["BLOCK_M"]})
@triton.autotune(
    configs=_configs,
    key=["M", "N", "is_fp16"],
)
@triton.jit
def k_dropout_bw(
    GRAD_IN, GRAD_BIAS, GRAD_OUT,
    INPUTS, BIAS, SEEDS,
    stride_grad, stride_inputs,
    M, N,
    p,
    is_fp16,  # autotune
    **meta,
):
    """
    Apply dropout on an input tensor
    GRAD_OUT    (M, N)
    GRAD_BIAS   (N,)
    GRAD_IN     (M, N)
    BIAS        (N,)
    SEEDS       (N,)
    p : dropout probability
    """
    # fmt: on

    BLOCK_M = meta["BLOCK_M"]
    BLOCK_N = meta["BLOCK_N"]
    SIZE_RAND_BLOCK = meta["SIZE_RAND_BLOCK"]
    TRAINABLE_BIAS = meta["TRAINABLE_BIAS"]

    rows = tl.arange(0, BLOCK_M)
    row_id = tl.program_id(axis=0)
    rows = row_id * BLOCK_M * 4 + tl.arange(0, BLOCK_M)

    col_id = tl.program_id(axis=1)
    cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N)
    seed = SEEDS + col_id  # FIXME index the seed properly

    # pointers starting point
    grad_out_ptrs = GRAD_OUT + rows[:, None] * stride_grad + cols[None, :]
    grad_in_ptrs = GRAD_IN + rows[:, None] * stride_grad + cols[None, :]
    input_ptrs = INPUTS + rows[:, None] * stride_inputs + cols[None, :]

    # random binary masks, save registers
    rand_offsets = tl.arange(0, SIZE_RAND_BLOCK) + row_id * BLOCK_M * 4
    rand_mask1, rand_mask2, rand_mask3, rand_mask4 = _get_4_bin_masks(seed, rand_offsets, p)

    # now go over the tiles
    grad_bias = tl.zeros((BLOCK_N,), dtype=tl.float32)
    col_mask = cols[None, :] < N
    p_scale = 1 / (1 - p) if p < 1. else 1.

    if meta["USE_BIAS"]:
        b_ptrs = BIAS + cols[None, :]
        bias = tl.load(b_ptrs, mask=col_mask, other=0.)

    for i in range(4):
        # cycle through the binary masks (workaround / no indexing)
        if i == 0:
            rand_mask = rand_mask1
        elif i == 1:
            rand_mask = rand_mask2
        elif i == 2:
            rand_mask = rand_mask3
        else:
            rand_mask = rand_mask4

        block_mask = (rows[:, None] < M) & col_mask
        grad_out = tl.load(grad_out_ptrs, mask=block_mask, other=0.)

        # optional: fused activation (while the data is in shared memory)
        if meta["ACTIVATION_GRAD"]:
            inputs = tl.load(input_ptrs, mask=block_mask, other=0.)

            # optionally apply a fused bias
            if meta["USE_BIAS"]:
                inputs += bias

            act_grad = meta["ACTIVATION_GRAD"](inputs).to(grad_out.dtype)
            grad_out *= act_grad

        # randomly prune (and scale) the resulting buffer, possibly a no-op
        # note that even if we did not save the mask from the FW pass, it is generated
        # from the same seeds, so the same drop mask is applied here
        output = _random_prune_and_scale(grad_out, rand_mask, p, p_scale)

        # write-back
        tl.store(grad_in_ptrs, output, mask=block_mask)

        # optionally accumulate the bias gradient
        if TRAINABLE_BIAS:
            grad_bias += tl.sum(output, axis=0)

        # Update the pointers
        rows += BLOCK_M  # needs to be updated for the mask to be correct
        grad_out_ptrs += BLOCK_M * stride_grad
        input_ptrs += BLOCK_M * stride_inputs
        grad_in_ptrs += BLOCK_M * stride_grad

    if TRAINABLE_BIAS:
        grad_bias_ptr = GRAD_BIAS + row_id * N + cols
        tl.store(grad_bias_ptr, grad_bias, mask=cols < N)

xformers/triton/k_dropout.py
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b8902f02f3..02bef01fe3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## TBD
 ### Fixed
 - Much faster fused dropout [#164]
+- Fused dropout repeatability [#173]
 
 ### Added
 - Embedding weight tying option [#172]
diff --git a/tests/test_triton_dropout.py b/tests/test_triton_dropout.py
index b65f7ff459..7a1fe45f38 100644
--- a/tests/test_triton_dropout.py
+++ b/tests/test_triton_dropout.py
@@ -113,6 +113,17 @@ def test_dropout(shape, amp, bias, p):
         drop_p = (y.numel() - y.count_nonzero()) / y.numel()
         assert abs(drop_p - p) < 0.01
 
+        # Check that the same seeds lead to the same dropout
+        torch.manual_seed(0)
+        torch.cuda.manual_seed(0)
+        y_1 = triton_dropout(x, p=0.5)
+
+        torch.manual_seed(0)
+        torch.cuda.manual_seed(0)
+        y_2 = triton_dropout(x, p=0.5)
+
+        assert torch.allclose(y_1, y_2)
+
 
 @pytest.mark.skipif(not _triton_available, reason="Triton is not available")
 @pytest.mark.skipif(
diff --git a/xformers/triton/k_dropout.py b/xformers/triton/k_dropout.py
index 2af660ffd9..d39d66022b 100644
--- a/xformers/triton/k_dropout.py
+++ b/xformers/triton/k_dropout.py
@@ -21,7 +21,8 @@
 
 @triton.jit
 def _get_4_bin_masks(seed, rand_offsets, p):
-    rand1, rand2, rand3, rand4 = tl.randint4x(seed.to(tl.int32), rand_offsets)
+    seed = tl.load(seed)
+    rand1, rand2, rand3, rand4 = tl.randint4x(seed, rand_offsets)
 
     # binarize masks, save registers
     # NOTE: We keep the random numbers as is there (integers over int32),

Dropout uses the *memory address* of seeds instead of reading seeds from memory
# 🐛 Bug

From reading the code for `k_dropout_fw` and `k_dropout_bw`, it seems to me that the seeds are never read from memory and the code simply uses the *memory address* of the seed.
For [example](https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_dropout.py):
```def _get_4_bin_masks(seed, rand_offsets, p):
    rand1, rand2, rand3, rand4 = tl.randint4x(seed.to(tl.int32), rand_offsets)
```
Here seed is still a memory address and not an integer.

As a result, when `k_dropout_fw` is passed in two identical seed tensors with different memory addresses, it produces different results.  

## To Reproduce

Setting the Pytorch seed should produce the same `seed` used in dropout, and should produce the same dropout mask.
However, that's not the case

```
import torch
from xformers.triton.dropout import dropout

x = torch.randn(3, 5, device='cuda')
print(x)

torch.manual_seed(0)
torch.cuda.manual_seed(0)
print(dropout(x, 0.5))

torch.manual_seed(0)
torch.cuda.manual_seed(0)
print(dropout(x, 0.5))
```

```
 tensor([[ 0.4821, -1.6949,  0.8196,  1.9093, -1.0018],
        [ 0.4030, -1.5175, -0.3187, -0.0959,  2.7204],
        [ 1.0645, -0.1254,  0.3978, -2.9882,  0.2232]], device='cuda:0')

tensor([[ 0.9642,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.8059, -3.0350,  0.0000, -0.1918,  5.4409],
        [ 0.0000, -0.2507,  0.7955,  0.0000,  0.0000]], device='cuda:0')

tensor([[ 0.9642, -3.3897,  0.0000,  3.8186, -2.0037],
        [ 0.0000,  0.0000, -0.6374,  0.0000,  5.4409],
        [ 2.1290, -0.2507,  0.7955,  0.0000,  0.4464]], device='cuda:0')
```


- PyTorch Version (e.g., 1.0): 1.10.1
- OS (e.g., Linux): Ubuntu 20.04
- How you installed PyTorch (`conda`, `pip`, source): conda
- Build command you used (if compiling from source): pip install -e . (from master)
- Python version: 3.8
- CUDA/cuDNN version: 11.3
- GPU models and configuration: V100

Dropout uses the *memory address* of seeds instead of reading seeds from memory
# 🐛 Bug

From reading the code for `k_dropout_fw` and `k_dropout_bw`, it seems to me that the seeds are never read from memory and the code simply uses the *memory address* of the seed.
For [example](https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_dropout.py):
```def _get_4_bin_masks(seed, rand_offsets, p):
    rand1, rand2, rand3, rand4 = tl.randint4x(seed.to(tl.int32), rand_offsets)
```
Here seed is still a memory address and not an integer.

As a result, when `k_dropout_fw` is passed in two identical seed tensors with different memory addresses, it produces different results.  

## To Reproduce

Setting the Pytorch seed should produce the same `seed` used in dropout, and should produce the same dropout mask.
However, that's not the case

```
import torch
from xformers.triton.dropout import dropout

x = torch.randn(3, 5, device='cuda')
print(x)

torch.manual_seed(0)
torch.cuda.manual_seed(0)
print(dropout(x, 0.5))

torch.manual_seed(0)
torch.cuda.manual_seed(0)
print(dropout(x, 0.5))
```

```
 tensor([[ 0.4821, -1.6949,  0.8196,  1.9093, -1.0018],
        [ 0.4030, -1.5175, -0.3187, -0.0959,  2.7204],
        [ 1.0645, -0.1254,  0.3978, -2.9882,  0.2232]], device='cuda:0')

tensor([[ 0.9642,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.8059, -3.0350,  0.0000, -0.1918,  5.4409],
        [ 0.0000, -0.2507,  0.7955,  0.0000,  0.0000]], device='cuda:0')

tensor([[ 0.9642, -3.3897,  0.0000,  3.8186, -2.0037],
        [ 0.0000,  0.0000, -0.6374,  0.0000,  5.4409],
        [ 2.1290, -0.2507,  0.7955,  0.0000,  0.4464]], device='cuda:0')
```


- PyTorch Version (e.g., 1.0): 1.10.1
- OS (e.g., Linux): Ubuntu 20.04
- How you installed PyTorch (`conda`, `pip`, source): conda
- Build command you used (if compiling from source): pip install -e . (from master)
- Python version: 3.8
- CUDA/cuDNN version: 11.3
- GPU models and configuration: V100

