[ { "path": ".clang-format", "content": "# Start with a built-in style and modify it\nBasedOnStyle: Google\n\n# Overrides\nColumnLimit: 120\n" }, { "path": ".git-blame-ignore-revs", "content": "# Commits to ignore in git-blame\n# These commits are bulk formatting or refactoring changes that should be skipped when viewing blame history\n\n# Add pre-commit and GitHub Actions workflow for it (#1949)\n1f20398756f0eeba37d6887a2d3f65e0687ec94f\n# Remove github actions config of pre-commit in favor of pre-commit ci (#1958)\n27e0e8951352d9d58c88b2895cd8f2c752bda963\n# Enable Ruff pre-commit hooks (#1957)\n16fadfe71c0d57312351c2d8b056251a0c8ce1ef\n" }, { "path": ".github/ISSUE_TEMPLATE/bug_report.md", "content": "---\nname: Bug report\nabout: Create a report to help us improve apex\ntitle: ''\nlabels: bug\nassignees: ''\n\n---\n\n**Describe the Bug**\n\n**Minimal Steps/Code to Reproduce the Bug**\n \n\n**Expected Behavior**\n\n\n**Environment**\n\n" }, { "path": ".gitignore", "content": "apex.egg-info\ndist\nbuild\ndocs/build\n*~\n__pycache__\n.vscode\n\n# Copied from https://raw.githubusercontent.com/github/gitignore/master/Python.gitignore\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\ncover/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\n.pybuilder/\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n# For a library or package, you might want to ignore these files since the code is\n# intended to run in multiple environments; otherwise, check them in:\n# .python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# pytype static type analyzer\n.pytype/\n\n# Cython debug symbols\ncython_debug/\n" }, { "path": ".gitmodules", "content": "[submodule \"apex/contrib/csrc/multihead_attn/cutlass\"]\n\tpath = apex/contrib/csrc/multihead_attn/cutlass\n\turl = https://github.com/NVIDIA/cutlass.git\n\tbranch = v1.2.0\n[submodule \"apex/contrib/csrc/cudnn-frontend\"]\n\tpath = apex/contrib/csrc/cudnn-frontend\n\turl = https://github.com/NVIDIA/cudnn-frontend.git\n" }, { "path": ".nojekyll", "content": "" }, { "path": ".pre-commit-config.yaml", "content": "repos:\n- repo: https://github.com/pre-commit/mirrors-clang-format\n rev: v22.1.1 # Or pin to your preferred clang-format version\n hooks:\n - id: clang-format\n files: \\.(c|h|cpp|hpp|proto|cu|cuh)$\n exclude: ^(apex/contrib/csrc/multihead_attn/cutlass|apex/contrib/csrc/cudnn-frontend)/\n\n- repo: https://github.com/astral-sh/ruff-pre-commit\n rev: v0.15.6\n hooks:\n - id: ruff-check\n args: [\"--fix\"]\n - id: ruff-format\n types_or: [python]\n exclude: \"examples\"\n" }, { "path": "LICENSE", "content": "All rights reserved.\n\nRedistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:\n\n1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.\n\n2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.\n\n3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.\n\nTHIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE." }, { "path": "README.md", "content": "# Introduction\n\nThis repository holds NVIDIA-maintained utilities to streamline mixed precision and distributed training in Pytorch.\nSome of the code here will be included in upstream Pytorch eventually.\nThe intent of Apex is to make up-to-date utilities available to users as quickly as possible.\n\n# Installation\nEach [`apex.contrib`](./apex/contrib) module requires one or more install options other than `--cpp_ext` and `--cuda_ext`.\nNote that contrib modules do not necessarily support stable PyTorch releases, some of them might only be compatible with nightlies.\n\n## Containers\nNVIDIA PyTorch Containers are available on NGC: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch.\nThe containers come with all the custom extensions available at the moment. \n\nSee [the NGC documentation](https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/index.html) for details such as:\n- how to pull a container\n- how to run a pulled container\n- release notes\n\n## From Source\n\nTo install Apex from source, we recommend using the nightly Pytorch obtainable from https://github.com/pytorch/pytorch.\n\nThe latest stable release obtainable from https://pytorch.org should also work.\n\nWe recommend installing [`Ninja`](https://ninja-build.org/) to make compilation faster.\n\n### Linux\n\nFor performance and full functionality, we recommend installing Apex with CUDA and C++ extensions using environment variables:\n\n#### Using Environment Variables (Recommended)\n\n```bash\ngit clone https://github.com/NVIDIA/apex\ncd apex\n# Build with core extensions (cpp and cuda)\nAPEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .\n\n# To build with additional extensions, specify them with environment variables\nAPEX_CPP_EXT=1 APEX_CUDA_EXT=1 APEX_FAST_MULTIHEAD_ATTN=1 APEX_FUSED_CONV_BIAS_RELU=1 pip install -v --no-build-isolation .\n\n# To build all contrib extensions at once\nAPEX_CPP_EXT=1 APEX_CUDA_EXT=1 APEX_ALL_CONTRIB_EXT=1 pip install -v --no-build-isolation .\n```\n\nTo reduce the build time, parallel building can be enabled:\n\n```bash\nNVCC_APPEND_FLAGS=\"--threads 4\" APEX_PARALLEL_BUILD=8 APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .\n```\n\nWhen CPU cores or memory are limited, the `--parallel` option is generally preferred over `--threads`. See [pull#1882](https://github.com/NVIDIA/apex/pull/1882) for more details.\n\n#### Using Command-Line Flags (Legacy Method)\n\nThe traditional command-line flags are still supported:\n\n```bash\n# Using pip config-settings (pip >= 23.1)\npip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings \"--build-option=--cpp_ext\" --config-settings \"--build-option=--cuda_ext\" ./\n\n# For older pip versions\npip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./\n\n# To build with additional extensions\npip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" --global-option=\"--fast_multihead_attn\" ./\n```\n\n#### Python-Only Build\n\nAPEX also supports a Python-only build via:\n```bash\npip install -v --disable-pip-version-check --no-build-isolation --no-cache-dir ./\n```\nA Python-only build omits:\n- Fused kernels required to use `apex.optimizers.FusedAdam`.\n- Fused kernels required to use `apex.normalization.FusedLayerNorm` and `apex.normalization.FusedRMSNorm`.\n- Fused kernels that improve the performance and numerical stability of `apex.parallel.SyncBatchNorm`.\n- Fused kernels that improve the performance of `apex.parallel.DistributedDataParallel` and `apex.amp`.\n`DistributedDataParallel`, `amp`, and `SyncBatchNorm` will still be usable, but they may be slower.\n\n\n### [Experimental] Windows\n`pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings \"--build-option=--cpp_ext\" --config-settings \"--build-option=--cuda_ext\" .` may work if you were able to build Pytorch from source\non your system. A Python-only build via `pip install -v --no-cache-dir .` is more likely to work. \nIf you installed Pytorch in a Conda environment, make sure to install Apex in that same environment.\n\n\n## Custom C++/CUDA Extensions and Install Options\n\nIf a requirement of a module is not met, then it will not be built.\n\n| Module Name | Environment Variable | Install Option | Misc |\n|---------------|------------------------|------------------|--------|\n| `apex_C` | `APEX_CPP_EXT=1` | `--cpp_ext` | |\n| `amp_C` | `APEX_CUDA_EXT=1` | `--cuda_ext` | |\n| `syncbn` | `APEX_CUDA_EXT=1` | `--cuda_ext` | |\n| `fused_layer_norm_cuda` | `APEX_CUDA_EXT=1` | `--cuda_ext` | [`apex.normalization`](./apex/normalization) |\n| `mlp_cuda` | `APEX_CUDA_EXT=1` | `--cuda_ext` | |\n| `scaled_upper_triang_masked_softmax_cuda` | `APEX_CUDA_EXT=1` | `--cuda_ext` | |\n| `generic_scaled_masked_softmax_cuda` | `APEX_CUDA_EXT=1` | `--cuda_ext` | |\n| `scaled_masked_softmax_cuda` | `APEX_CUDA_EXT=1` | `--cuda_ext` | |\n| `fused_weight_gradient_mlp_cuda` | `APEX_CUDA_EXT=1` | `--cuda_ext` | Requires CUDA>=11 |\n| `permutation_search_cuda` | `APEX_PERMUTATION_SEARCH=1` | `--permutation_search` | [`apex.contrib.sparsity`](./apex/contrib/sparsity) |\n| `bnp` | `APEX_BNP=1` | `--bnp` | [`apex.contrib.groupbn`](./apex/contrib/groupbn) |\n| `xentropy` | `APEX_XENTROPY=1` | `--xentropy` | [`apex.contrib.xentropy`](./apex/contrib/xentropy) |\n| `focal_loss_cuda` | `APEX_FOCAL_LOSS=1` | `--focal_loss` | [`apex.contrib.focal_loss`](./apex/contrib/focal_loss) |\n| `fused_index_mul_2d` | `APEX_INDEX_MUL_2D=1` | `--index_mul_2d` | [`apex.contrib.index_mul_2d`](./apex/contrib/index_mul_2d) |\n| `fused_adam_cuda` | `APEX_DEPRECATED_FUSED_ADAM=1` | `--deprecated_fused_adam` | [`apex.contrib.optimizers`](./apex/contrib/optimizers) |\n| `fused_lamb_cuda` | `APEX_DEPRECATED_FUSED_LAMB=1` | `--deprecated_fused_lamb` | [`apex.contrib.optimizers`](./apex/contrib/optimizers) |\n| `fast_layer_norm` | `APEX_FAST_LAYER_NORM=1` | `--fast_layer_norm` | [`apex.contrib.layer_norm`](./apex/contrib/layer_norm). different from `fused_layer_norm` |\n| `fmhalib` | `APEX_FMHA=1` | `--fmha` | [`apex.contrib.fmha`](./apex/contrib/fmha) |\n| `fast_multihead_attn` | `APEX_FAST_MULTIHEAD_ATTN=1` | `--fast_multihead_attn` | [`apex.contrib.multihead_attn`](./apex/contrib/multihead_attn) |\n| `transducer_joint_cuda` | `APEX_TRANSDUCER=1` | `--transducer` | [`apex.contrib.transducer`](./apex/contrib/transducer) |\n| `transducer_loss_cuda` | `APEX_TRANSDUCER=1` | `--transducer` | [`apex.contrib.transducer`](./apex/contrib/transducer) |\n| `cudnn_gbn_lib` | `APEX_CUDNN_GBN=1` | `--cudnn_gbn` | Requires cuDNN>=8.5, [`apex.contrib.cudnn_gbn`](./apex/contrib/cudnn_gbn) |\n| `peer_memory_cuda` | `APEX_PEER_MEMORY=1` | `--peer_memory` | [`apex.contrib.peer_memory`](./apex/contrib/peer_memory) |\n| `nccl_p2p_cuda` | `APEX_NCCL_P2P=1` | `--nccl_p2p` | Requires NCCL >= 2.10, [`apex.contrib.nccl_p2p`](./apex/contrib/nccl_p2p) |\n| `fast_bottleneck` | `APEX_FAST_BOTTLENECK=1` | `--fast_bottleneck` | Requires `peer_memory_cuda` and `nccl_p2p_cuda`, [`apex.contrib.bottleneck`](./apex/contrib/bottleneck) |\n| `fused_conv_bias_relu` | `APEX_FUSED_CONV_BIAS_RELU=1` | `--fused_conv_bias_relu` | Requires cuDNN>=8.4, [`apex.contrib.conv_bias_relu`](./apex/contrib/conv_bias_relu) |\n| `distributed_adam_cuda` | `APEX_DISTRIBUTED_ADAM=1` | `--distributed_adam` | [`apex.contrib.optimizers`](./apex/contrib/optimizers) |\n| `distributed_lamb_cuda` | `APEX_DISTRIBUTED_LAMB=1` | `--distributed_lamb` | [`apex.contrib.optimizers`](./apex/contrib/optimizers) |\n| `_apex_nccl_allocator` | `APEX_NCCL_ALLOCATOR=1` | `--nccl_allocator` | Requires NCCL >= 2.19, [`apex.contrib.nccl_allocator`](./apex/contrib/nccl_allocator) |\n| `_apex_gpu_direct_storage` | `APEX_GPU_DIRECT_STORAGE=1` | `--gpu_direct_storage` | [`apex.contrib.gpu_direct_storage`](./apex/contrib/gpu_direct_storage) |\n\nYou can also build all contrib extensions at once by setting `APEX_ALL_CONTRIB_EXT=1`.\n" }, { "path": "apex/__init__.py", "content": "import logging\nimport warnings\n\n# May help avoid undefined symbol errors https://pytorch.org/cppdocs/notes/faq.html#undefined-symbol-errors-from-pytorch-aten\nimport torch\n\n# For optimizers and normalization there is no Python fallback.\n# Absence of cuda backend is a hard error.\n# I would like the errors from importing fused_adam_cuda or fused_layer_norm_cuda\n# to be triggered lazily, because if someone has installed with --cpp_ext and --cuda_ext\n# so they expect those backends to be available, but for some reason they actually aren't\n# available (for example because they built improperly in a way that isn't revealed until\n# load time) the error message is timely and visible.\nfrom . import optimizers\nfrom . import normalization\n\n\n__all__ = [\"optimizers\", \"normalization\"]\n\n\ndef check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int) -> bool:\n cudnn_available = torch.backends.cudnn.is_available()\n cudnn_version = torch.backends.cudnn.version() if cudnn_available else None\n if not (cudnn_available and (cudnn_version >= required_cudnn_version)):\n warnings.warn(\n f\"`{global_option}` depends on cuDNN {required_cudnn_version} or later, \"\n f\"but {'cuDNN is not available' if not cudnn_available else cudnn_version}\"\n )\n return False\n return True\n\n\nclass DeprecatedFeatureWarning(FutureWarning):\n pass\n\n\ndef deprecated_warning(msg: str) -> None:\n if (\n not torch.distributed.is_available\n or not torch.distributed.is_initialized()\n or (torch.distributed.is_initialized() and torch.distributed.get_rank() == 0)\n ):\n warnings.warn(msg, DeprecatedFeatureWarning)\n" }, { "path": "apex/_autocast_utils.py", "content": "from typing import Optional, Sequence\n\nimport torch\n\n\n__all__ = [\"_cast_if_autocast_enabled\"]\n\n\ndef _get_autocast_dtypes() -> Sequence[torch.dtype]:\n if torch.cuda.is_bf16_supported():\n return [torch.half, torch.bfloat16]\n return [torch.half]\n\n\ndef _get_current_dtype(dtype: Optional[torch.dtype] = None) -> torch.dtype:\n if not torch.is_autocast_enabled():\n return torch.float or dtype\n else:\n return torch.get_autocast_gpu_dtype()\n\n\ndef _cast_if_autocast_enabled(*args):\n if not torch.is_autocast_enabled():\n return args\n else:\n return torch.cuda.amp.autocast_mode._cast(args, torch.get_autocast_gpu_dtype())\n" }, { "path": "apex/contrib/__init__.py", "content": "" }, { "path": "apex/contrib/bottleneck/__init__.py", "content": "from .bottleneck import Bottleneck, SpatialBottleneck\nfrom .halo_exchangers import (\n HaloExchangerNoComm,\n HaloExchangerAllGather,\n HaloExchangerSendRecv,\n HaloExchangerPeer,\n)\n" }, { "path": "apex/contrib/bottleneck/bottleneck.py", "content": "import functools as func\n\nimport torch\nfrom torch import nn\n\nfrom apex import check_cudnn_version_and_warn\nimport fast_bottleneck\nimport nccl_p2p_cuda as inc\n\n\nassert check_cudnn_version_and_warn(__name__, 8400)\n\n\ndef kaiming_uniform_(tensor, a=0, mode=\"fan_in\", nonlinearity=\"leaky_relu\"):\n weight_tensor_nchw = tensor\n nn.init.kaiming_uniform_(weight_tensor_nchw, a=a, mode=mode, nonlinearity=nonlinearity)\n\n\ndef compute_scale_bias_one(nhwc, weight, bias, running_mean, running_var, w_scale, w_bias):\n scale = weight * running_var.rsqrt()\n bias = bias - running_mean * scale\n w_scale.copy_(scale)\n w_bias.copy_(bias)\n\n\ndef compute_scale_bias_method(nhwc, args):\n for arg in args:\n # arg is tuple of (weight, bias, running_mean, running_var, w_scale, w_bias)\n compute_scale_bias_one(nhwc, *arg)\n\n\nclass FrozenBatchNorm2d(torch.jit.ScriptModule):\n \"\"\"\n BatchNorm2d where the batch statistics and the affine parameters are fixed\n \"\"\"\n\n def __init__(self, n):\n super(FrozenBatchNorm2d, self).__init__()\n self.register_buffer(\"weight\", torch.ones(n))\n self.register_buffer(\"bias\", torch.zeros(n))\n self.register_buffer(\"running_mean\", torch.zeros(n))\n self.register_buffer(\"running_var\", torch.ones(n))\n\n @torch.jit.script_method\n def get_scale_bias(self, nhwc):\n # type: (bool) -> List[torch.Tensor]\n scale = self.weight * self.running_var.rsqrt()\n bias = self.bias - self.running_mean * scale\n if nhwc:\n scale = scale.reshape(1, 1, 1, -1)\n bias = bias.reshape(1, 1, 1, -1)\n else:\n scale = scale.reshape(1, -1, 1, 1)\n bias = bias.reshape(1, -1, 1, 1)\n return scale, bias\n\n @torch.jit.script_method\n def forward(self, x):\n scale, bias = self.get_scale_bias(False)\n return x * scale + bias\n\n\n@torch.jit.script\ndef drelu_dscale1(grad_o, output, scale1):\n relu_mask = output > 0\n dx_relu = relu_mask * grad_o\n g1 = dx_relu * scale1\n return g1, dx_relu\n\n\n@torch.jit.script\ndef drelu_dscale2(grad_o, output, scale1, scale2):\n relu_mask = output > 0\n dx_relu = relu_mask * grad_o\n g1 = dx_relu * scale1\n g2 = dx_relu * scale2\n return g1, g2\n\n\nclass BottleneckFunction(torch.autograd.Function):\n @staticmethod\n def forward(ctx, nhwc, stride_1x1, scale, bias, x, *conv):\n # TODO: clean up order of tensors\n args = [x, *conv[0:3], *scale[0:3], *bias[0:3]]\n ctx.downsample = len(conv) > 3\n if ctx.downsample:\n args.append(conv[3])\n args.append(scale[3])\n args.append(bias[3])\n\n # weight buffers are always in nhwc while shape can be nhwc or channels_last\n # here we pass in flag and let c++ handle it\n # alternatively, we can put all sizes into a fixed format and pass it in\n outputs = fast_bottleneck.forward(nhwc, stride_1x1, args)\n ctx.save_for_backward(*(args + outputs))\n # save relu outputs for drelu\n ctx.nhwc = nhwc\n ctx.stride_1x1 = stride_1x1\n return outputs[2]\n\n # backward relu is not exposed, MUL with mask used now\n # only support dgrad\n @staticmethod\n def backward(ctx, grad_o):\n outputs = ctx.saved_tensors[-3:]\n\n if ctx.downsample:\n grad_conv3, grad_conv4 = drelu_dscale2(\n grad_o, outputs[2], ctx.saved_tensors[6], ctx.saved_tensors[11]\n )\n else:\n grad_conv3, grad_conv4 = drelu_dscale1(grad_o, outputs[2], ctx.saved_tensors[6])\n\n # create input vector for backward\n t_list = [*ctx.saved_tensors[0:10]]\n t_list.append(grad_conv3)\n t_list.append(grad_conv4)\n\n # outputs used for wgrad and generating drelu mask\n t_list.append(outputs[0])\n t_list.append(outputs[1])\n\n # in case there is downsample\n if ctx.downsample:\n t_list.append(ctx.saved_tensors[10])\n\n grads = fast_bottleneck.backward(ctx.nhwc, ctx.stride_1x1, t_list)\n\n return (None, None, None, None, *grads)\n\n\nbottleneck_function = BottleneckFunction.apply\n\n\ndef conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):\n \"\"\"3x3 convolution with padding\"\"\"\n return nn.Conv2d(\n in_planes,\n out_planes,\n kernel_size=3,\n stride=stride,\n padding=dilation,\n groups=groups,\n bias=False,\n dilation=dilation,\n )\n\n\ndef conv1x1(in_planes, out_planes, stride=1):\n \"\"\"1x1 convolution\"\"\"\n return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n\n\nclass Bottleneck(torch.nn.Module):\n # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)\n # while original implementation places the stride at the first 1x1 convolution(self.conv1)\n # according to \"Deep residual learning for image recognition\"https://arxiv.org/abs/1512.03385.\n # This variant is also known as ResNet V1.5 and improves accuracy according to\n # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.\n # here we put it at 1x1\n\n def __init__(\n self,\n in_channels,\n bottleneck_channels,\n out_channels,\n stride=1,\n groups=1,\n dilation=1,\n norm_func=None,\n use_cudnn=False,\n explicit_nhwc=False,\n ):\n super(Bottleneck, self).__init__()\n if groups != 1:\n raise RuntimeError(\"Only support groups == 1\")\n if dilation != 1:\n raise RuntimeError(\"Only support dilation == 1\")\n if norm_func == None:\n norm_func = FrozenBatchNorm2d\n else:\n raise RuntimeError(\"Only support frozen BN now.\")\n\n if stride != 1 or in_channels != out_channels:\n self.downsample = nn.Sequential(\n conv1x1(in_channels, out_channels, stride),\n norm_func(out_channels),\n )\n else:\n self.downsample = None\n\n # Both self.conv2 and self.downsample layers downsample the input when stride != 1\n self.conv1 = conv1x1(in_channels, bottleneck_channels, stride)\n self.conv2 = conv3x3(bottleneck_channels, bottleneck_channels)\n self.conv3 = conv1x1(bottleneck_channels, out_channels)\n self.relu = nn.ReLU(inplace=True)\n self.stride = stride\n\n self.bn1 = norm_func(bottleneck_channels)\n self.bn2 = norm_func(bottleneck_channels)\n self.bn3 = norm_func(out_channels)\n self.w_scale = None\n\n self.use_cudnn = use_cudnn\n\n # setup conv weights\n self.w_conv = [self.conv1.weight, self.conv2.weight, self.conv3.weight]\n if self.downsample is not None:\n self.w_conv.append(self.downsample[0].weight)\n\n # init weight in nchw format before possible transpose\n for w in self.w_conv:\n kaiming_uniform_(w, a=1)\n\n # TODO: prevent unsupported case usage\n # support cases\n # native cudnn\n # normal yes no\n # channel_last yes yes\n # explicit_nhwc no yes\n self.explicit_nhwc = explicit_nhwc\n if self.explicit_nhwc:\n for p in self.parameters():\n with torch.no_grad():\n p.data = p.data.permute(0, 2, 3, 1).contiguous()\n\n return\n\n # Returns single callable that recomputes scale and bias for all frozen batch-norms.\n # This method must be called before cuda graphing.\n # The callable it returns can be called anytime.\n # Calling this method will prevent these from being computed every forward call.\n def get_scale_bias_callable(self):\n self.w_scale, self.w_bias, args = [], [], []\n batch_norms = [self.bn1, self.bn2, self.bn3]\n if self.downsample is not None:\n batch_norms.append(self.downsample[1])\n for bn in batch_norms:\n s = torch.empty_like(bn.weight)\n b = torch.empty_like(s)\n args.append((bn.weight, bn.bias, bn.running_mean, bn.running_var, s, b))\n if self.explicit_nhwc:\n self.w_scale.append(s.reshape(1, 1, 1, -1))\n self.w_bias.append(b.reshape(1, 1, 1, -1))\n else:\n self.w_scale.append(s.reshape(1, -1, 1, 1))\n self.w_bias.append(b.reshape(1, -1, 1, 1))\n return func.partial(compute_scale_bias_method, self.explicit_nhwc, args)\n\n def forward(self, x):\n if self.use_cudnn:\n if self.w_scale is None:\n # calculate scale/bias from registered buffers\n # TODO: make this better\n s1, b1 = self.bn1.get_scale_bias(self.explicit_nhwc)\n s2, b2 = self.bn2.get_scale_bias(self.explicit_nhwc)\n s3, b3 = self.bn3.get_scale_bias(self.explicit_nhwc)\n w_scale = [s1, s2, s3]\n w_bias = [b1, b2, b3]\n if self.downsample is not None:\n s4, b4 = self.downsample[1].get_scale_bias(self.explicit_nhwc)\n w_scale.append(s4)\n w_bias.append(b4)\n out = bottleneck_function(\n self.explicit_nhwc, self.stride, w_scale, w_bias, x, *self.w_conv\n )\n else:\n out = bottleneck_function(\n self.explicit_nhwc,\n self.stride,\n self.w_scale,\n self.w_bias,\n x,\n *self.w_conv,\n )\n return out\n\n if self.explicit_nhwc:\n raise RuntimeError(\"explicit nhwc with native ops is not supported.\")\n\n # fallback to native ops\n identity = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out += identity\n out = self.relu(out)\n\n return out\n\n\nclass SpatialBottleneckFunction(torch.autograd.Function):\n @staticmethod\n def forward(\n ctx,\n spatial_group_size,\n spatial_group_rank,\n spatial_communicator,\n spatial_halo_exchanger,\n spatial_method,\n use_delay_kernel,\n explicit_nhwc,\n stride_1x1,\n scale,\n bias,\n thresholdTop,\n thresholdBottom,\n x,\n *conv,\n ):\n if spatial_group_size > 1:\n stream1 = spatial_halo_exchanger.stream1\n stream2 = spatial_halo_exchanger.stream2\n stream3 = spatial_halo_exchanger.stream3\n\n # TODO: clean up order of tensors\n args = [x, *conv[0:3], *scale[0:3], *bias[0:3]]\n ctx.downsample = len(conv) > 3\n if ctx.downsample:\n args.append(conv[3])\n args.append(scale[3])\n args.append(bias[3])\n\n # weight buffers are always in explicit_nhwc while shape can be explicit_nhwc or channels_last\n # here we pass in flag and let c++ handle it\n # alternatively, we can put all sizes into a fixed format and pass it in\n outputs = fast_bottleneck.forward_init(explicit_nhwc, stride_1x1, args)\n fast_bottleneck.forward_out1(explicit_nhwc, stride_1x1, args, outputs)\n\n if spatial_group_size > 1:\n out1 = outputs[0]\n if explicit_nhwc:\n N, Hs, W, C = list(out1.shape)\n memory_format = torch.contiguous_format\n out1_pad = torch.empty([N, Hs + 2, W, C], dtype=out1.dtype, device=\"cuda\")\n else:\n N, C, Hs, W = list(out1.shape)\n memory_format = (\n torch.channels_last\n if out1.is_contiguous(memory_format=torch.channels_last)\n else torch.contiguous_format\n )\n out1_pad = torch.empty(\n [N, C, Hs + 2, W],\n dtype=out1.dtype,\n device=\"cuda\",\n memory_format=memory_format,\n )\n stream1.wait_stream(torch.cuda.current_stream())\n if spatial_method != 2:\n stream3.wait_stream(torch.cuda.current_stream())\n with torch.cuda.stream(stream1):\n if explicit_nhwc:\n top_out1_halo = out1_pad[:, :1, :, :]\n btm_out1_halo = out1_pad[:, Hs + 1 : Hs + 2, :, :]\n spatial_halo_exchanger.left_right_halo_exchange(\n out1[:, :1, :, :],\n out1[:, Hs - 1 :, :, :],\n top_out1_halo,\n btm_out1_halo,\n )\n else:\n top_out1_halo = out1_pad[:, :, :1, :]\n btm_out1_halo = out1_pad[:, :, Hs + 1 : Hs + 2, :]\n spatial_halo_exchanger.left_right_halo_exchange(\n out1[:, :, :1, :],\n out1[:, :, Hs - 1 :, :],\n top_out1_halo,\n btm_out1_halo,\n )\n if spatial_method == 1:\n # overlap mid convolution with halo transfer\n if spatial_group_rank < spatial_group_size - 1:\n stream2.wait_stream(stream1)\n with torch.cuda.stream(stream2):\n if explicit_nhwc:\n btm_fat_halo = torch.empty(\n (N, 3, W, C), dtype=out1.dtype, device=out1.device\n )\n btm_fat_halo[:, 0:2, :, :].copy_(out1[:, Hs - 2 :, :, :])\n btm_fat_halo[:, 2:, :, :].copy_(btm_out1_halo)\n else:\n btm_fat_halo = torch.empty(\n (N, C, 3, W), dtype=out1.dtype, device=out1.device\n )\n btm_fat_halo[:, :, 0:2, :].copy_(out1[:, :, Hs - 2 :, :])\n btm_fat_halo[:, :, 2:, :].copy_(btm_out1_halo)\n btm_out2 = fast_bottleneck.forward_out2_halo(\n explicit_nhwc, btm_fat_halo, args\n )\n if spatial_group_rank > 0:\n with torch.cuda.stream(stream1):\n if explicit_nhwc:\n top_fat_halo = torch.empty(\n (N, 3, W, C), dtype=out1.dtype, device=out1.device\n )\n top_fat_halo[:, :1, :, :].copy_(top_out1_halo)\n top_fat_halo[:, 1:3, :, :].copy_(out1[:, :2, :, :])\n else:\n top_fat_halo = torch.empty(\n (N, C, 3, W), dtype=out1.dtype, device=out1.device\n )\n top_fat_halo[:, :, :1, :].copy_(top_out1_halo)\n top_fat_halo[:, :, 1:3, :].copy_(out1[:, :, :2, :])\n top_out2 = fast_bottleneck.forward_out2_halo(\n explicit_nhwc, top_fat_halo, args\n )\n if use_delay_kernel:\n inc.add_delay(10)\n elif spatial_method != 2 and spatial_method != 3:\n assert False, \"spatial_method must be 1, 2 or 3\"\n\n if spatial_group_size <= 1:\n fast_bottleneck.forward_out2(explicit_nhwc, stride_1x1, args, outputs)\n elif spatial_method == 1:\n fast_bottleneck.forward_out2(explicit_nhwc, stride_1x1, args, outputs)\n with torch.cuda.stream(stream3):\n if explicit_nhwc:\n out1_pad[:, 1 : Hs + 1, :, :].copy_(out1)\n else:\n out1_pad[:, :, 1 : Hs + 1, :].copy_(out1)\n elif spatial_method == 2:\n # wait for halo transfer to finish before doing a full convolution of padded x\n if explicit_nhwc:\n out1_pad[:, 1 : Hs + 1, :, :].copy_(out1)\n else:\n out1_pad[:, :, 1 : Hs + 1, :].copy_(out1)\n torch.cuda.current_stream().wait_stream(stream1)\n fast_bottleneck.forward_out2_pad(explicit_nhwc, stride_1x1, args, outputs, out1_pad)\n elif spatial_method == 3:\n fast_bottleneck.forward_out2_mask(\n explicit_nhwc, stride_1x1, args, outputs, thresholdTop, thresholdBottom\n )\n with torch.cuda.stream(stream3):\n if explicit_nhwc:\n out1_pad[:, 1 : Hs + 1, :, :].copy_(out1)\n else:\n out1_pad[:, :, 1 : Hs + 1, :].copy_(out1)\n\n # compute halo cells for outputs[1] (out2)\n if spatial_group_size > 1:\n out2 = outputs[1]\n if explicit_nhwc:\n top_out2_halo = out2[:, :1, :, :]\n btm_out2_halo = out2[:, Hs - 1 :, :, :]\n else:\n top_out2_halo = out2[:, :, :1, :]\n btm_out2_halo = out2[:, :, Hs - 1 :, :]\n if spatial_method == 1:\n if spatial_group_rank > 0:\n torch.cuda.current_stream().wait_stream(stream1)\n top_out2_halo.copy_(top_out2)\n if spatial_group_rank < spatial_group_size - 1:\n torch.cuda.current_stream().wait_stream(stream2)\n btm_out2_halo.copy_(btm_out2)\n elif spatial_method == 3:\n # Note\n # out2 halo correction cannot overlap with anything since it has\n # to wait for out2_mask to finish, but itself has to finish before\n # the first kernel of _forward_rest can launch.\n # At least we can overlap the two halo correction kernels.\n if spatial_group_rank < spatial_group_size - 1:\n stream2.wait_stream(stream1) # wait for halo transfers to finish\n stream2.wait_stream(\n torch.cuda.current_stream()\n ) # wait for *_out2_mask to finish\n with torch.cuda.stream(stream2):\n w1by3 = args[2][:, 2:3, :, :].clone()\n btm_out1_halo = btm_out1_halo.clone()\n btm_out2 = fast_bottleneck.forward_out2_halo_corr(\n explicit_nhwc,\n btm_out1_halo,\n args,\n w1by3,\n btm_out2_halo.clone(),\n )\n btm_out2_halo.copy_(btm_out2)\n if spatial_group_rank > 0:\n stream1.wait_stream(\n torch.cuda.current_stream()\n ) # wait for *_out2_mask to finish\n with torch.cuda.stream(stream1):\n w1by3 = args[2][:, :1, :, :].clone()\n top_out1_halo = top_out1_halo.clone()\n top_out2 = fast_bottleneck.forward_out2_halo_corr(\n explicit_nhwc,\n top_out1_halo,\n args,\n w1by3,\n top_out2_halo.clone(),\n )\n top_out2_halo.copy_(top_out2)\n if spatial_group_rank < spatial_group_size - 1:\n torch.cuda.current_stream().wait_stream(stream2)\n if spatial_group_rank > 0:\n torch.cuda.current_stream().wait_stream(stream1)\n\n fast_bottleneck.forward_rest(explicit_nhwc, stride_1x1, args, outputs)\n # save halos for backward pass\n if spatial_group_size > 1:\n if spatial_method != 2:\n # make sure copy of mid-section of out1 into out1_pad is done before exiting\n torch.cuda.current_stream().wait_stream(stream3)\n ctx.save_for_backward(\n *(\n args\n + outputs\n + [\n out1_pad,\n ]\n )\n )\n else:\n ctx.save_for_backward(*(args + outputs))\n # save relu outputs for drelu\n ctx.explicit_nhwc = explicit_nhwc\n ctx.stride_1x1 = stride_1x1\n ctx.spatial_group_size = spatial_group_size\n if spatial_group_size > 1:\n ctx.spatial_group_rank = spatial_group_rank\n ctx.spatial_halo_exchanger = spatial_halo_exchanger\n ctx.spatial_method = spatial_method\n ctx.use_delay_kernel = use_delay_kernel\n ctx.thresholdTop = thresholdTop\n ctx.thresholdBottom = thresholdBottom\n ctx.stream1 = stream1\n ctx.stream2 = stream2\n ctx.stream3 = stream3\n return outputs[2]\n\n # backward relu is not exposed, MUL with mask used now\n # only support dgrad\n @staticmethod\n def backward(ctx, grad_o):\n if ctx.spatial_group_size > 1:\n out1_pad = ctx.saved_tensors[-1]\n outputs = ctx.saved_tensors[-4:-1]\n else:\n outputs = ctx.saved_tensors[-3:]\n\n if ctx.downsample:\n grad_conv3, grad_conv4 = drelu_dscale2(\n grad_o, outputs[2], ctx.saved_tensors[6], ctx.saved_tensors[11]\n )\n else:\n grad_conv3, grad_conv4 = drelu_dscale1(grad_o, outputs[2], ctx.saved_tensors[6])\n\n # create input vector for backward\n t_list = [*ctx.saved_tensors[0:10]]\n t_list.append(grad_conv3)\n t_list.append(grad_conv4)\n\n # outputs used for wgrad and generating drelu mask\n t_list.append(outputs[0])\n t_list.append(outputs[1])\n\n # in case there is downsample\n if ctx.downsample:\n t_list.append(ctx.saved_tensors[10])\n\n grads = fast_bottleneck.backward_init(ctx.explicit_nhwc, ctx.stride_1x1, t_list)\n wgrad3_stream = torch.cuda.Stream()\n wgrad3_stream.wait_stream(torch.cuda.current_stream())\n grad_out2 = fast_bottleneck.backward_grad_out2(\n ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads\n )\n wgrad2_stream = torch.cuda.Stream()\n wgrad2_stream.wait_stream(torch.cuda.current_stream())\n # do halo exchange of grad_out2 here\n # compute halo cells for grad_out1\n if ctx.spatial_group_size > 1:\n if ctx.explicit_nhwc:\n N, Hs, W, C = list(grad_out2.shape)\n else:\n N, C, Hs, W = list(grad_out2.shape)\n relu1 = t_list[12]\n ctx.stream1.wait_stream(torch.cuda.current_stream())\n with torch.cuda.stream(ctx.stream1):\n top_halo, btm_halo = ctx.spatial_halo_exchanger.left_right_halo_exchange(\n grad_out2[:, :1, :, :], grad_out2[:, Hs - 1 :, :, :]\n )\n # copy halos to send buffer\n if ctx.spatial_method == 1 or ctx.spatial_method == 2:\n # 1 -> halo recompute approach\n # 2 -> wait for concatenated halos, then do single conv on full input (not implemented yet for bprop)\n if ctx.spatial_group_rank < ctx.spatial_group_size - 1:\n ctx.stream2.wait_stream(ctx.stream1)\n with torch.cuda.stream(ctx.stream2):\n if ctx.explicit_nhwc:\n btm_fat_halo = torch.empty(\n (N, 3, W, C),\n dtype=grad_out2.dtype,\n device=grad_out2.device,\n )\n btm_fat_halo[:, :2, :, :].copy_(grad_out2[:, Hs - 2 :, :, :])\n btm_fat_halo[:, 2:, :, :].copy_(btm_halo)\n btm_fat_relu_halo = torch.empty(\n (N, 3, W, C),\n dtype=grad_out2.dtype,\n device=grad_out2.device,\n )\n btm_fat_relu_halo[:, :2, :, :].copy_(relu1[:, Hs - 2 :, :, :])\n btm_fat_relu_halo[:, 2:, :, :].zero_()\n else:\n btm_fat_halo = torch.empty(\n (N, C, 3, W),\n dtype=grad_out2.dtype,\n device=grad_out2.device,\n )\n btm_fat_halo[:, :, :2, :].copy_(grad_out2[:, :, Hs - 2 :, :])\n btm_fat_halo[:, :, 2:, :].copy_(btm_halo)\n btm_fat_relu_halo = torch.empty(\n (N, C, 3, W),\n dtype=grad_out2.dtype,\n device=grad_out2.device,\n )\n btm_fat_relu_halo[:, :, :2, :].copy_(relu1[:, :, Hs - 2 :, :])\n btm_fat_relu_halo[:, :, 2:, :].zero_()\n btm_grad_out1_halo = fast_bottleneck.backward_grad_out1_halo(\n ctx.explicit_nhwc,\n ctx.stride_1x1,\n t_list,\n grads,\n btm_fat_halo,\n btm_fat_relu_halo,\n )\n if ctx.explicit_nhwc:\n btm_grad_out1_halo = btm_grad_out1_halo[:, 1:2, :, :]\n else:\n btm_grad_out1_halo = btm_grad_out1_halo[:, :, 1:2, :]\n if ctx.spatial_group_rank > 0:\n with torch.cuda.stream(ctx.stream1):\n if ctx.explicit_nhwc:\n top_fat_halo = torch.empty(\n (N, 3, W, C),\n dtype=grad_out2.dtype,\n device=grad_out2.device,\n )\n top_fat_halo[:, :1, :, :].copy_(top_halo)\n top_fat_halo[:, 1:, :, :].copy_(grad_out2[:, :2, :, :])\n top_fat_relu_halo = torch.empty(\n (N, 3, W, C),\n dtype=grad_out2.dtype,\n device=grad_out2.device,\n )\n top_fat_relu_halo[:, :1, :, :].zero_()\n top_fat_relu_halo[:, 1:, :, :].copy_(relu1[:, :2, :, :])\n else:\n top_fat_halo = torch.empty(\n (N, C, 3, W),\n dtype=grad_out2.dtype,\n device=grad_out2.device,\n )\n top_fat_halo[:, :, :1, :].copy_(top_halo)\n top_fat_halo[:, :, 1:, :].copy_(grad_out2[:, :, :2, :])\n top_fat_relu_halo = torch.empty(\n (N, C, 3, W),\n dtype=grad_out2.dtype,\n device=grad_out2.device,\n )\n top_fat_relu_halo[:, :, :1, :].zero_()\n top_fat_relu_halo[:, :, 1:, :].copy_(relu1[:, :, :2, :])\n top_grad_out1_halo = fast_bottleneck.backward_grad_out1_halo(\n ctx.explicit_nhwc,\n ctx.stride_1x1,\n t_list,\n grads,\n top_fat_halo,\n top_fat_relu_halo,\n )\n if ctx.explicit_nhwc:\n top_grad_out1_halo = top_grad_out1_halo[:, 1:2, :, :]\n else:\n top_grad_out1_halo = top_grad_out1_halo[:, :, 1:2, :]\n if ctx.use_delay_kernel:\n inc.add_delay(10)\n elif ctx.spatial_method != 3:\n assert False, \"spatial_method must be 1, 2 or 3\"\n\n # compute grad_out1 for internal cells\n if ctx.spatial_group_size <= 1 or ctx.spatial_method == 1 or ctx.spatial_method == 2:\n grad_out1 = fast_bottleneck.backward_grad_out1(\n ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads, grad_out2\n )\n elif ctx.spatial_group_size > 1 and ctx.spatial_method == 3:\n grad_out1 = fast_bottleneck.backward_grad_out1_mask(\n ctx.explicit_nhwc,\n ctx.stride_1x1,\n t_list,\n grads,\n grad_out2,\n ctx.thresholdTop,\n ctx.thresholdBottom,\n )\n\n # apply halo cells to grad_out1\n if ctx.spatial_group_size > 1:\n w = t_list[2]\n z = t_list[4]\n relu1 = t_list[12]\n # print(\"w.shape = %s, z.shape = %s, relu1.shape = %s\" % (str(list(w.shape)), str(list(z.shape)), str(list(relu1.shape))))\n if ctx.spatial_method == 1 or ctx.spatial_method == 2:\n if ctx.spatial_group_rank < ctx.spatial_group_size - 1:\n torch.cuda.current_stream().wait_stream(ctx.stream2)\n if ctx.explicit_nhwc:\n grad_out1[:, Hs - 1 :, :, :].copy_(btm_grad_out1_halo)\n else:\n grad_out1[:, :, Hs - 1 :, :].copy_(btm_grad_out1_halo)\n # print(\"ctx.spatial_group_rank = %d, apply grad_out1 btm halo (grad_out1.shape = %s)\" % (ctx.spatial_group_rank, str(list(grad_out1.shape))))\n if ctx.spatial_group_rank > 0:\n torch.cuda.current_stream().wait_stream(ctx.stream1)\n if ctx.explicit_nhwc:\n grad_out1[:, :1, :, :].copy_(top_grad_out1_halo)\n else:\n grad_out1[:, :, :1, :].copy_(top_grad_out1_halo)\n # print(\"ctx.spatial_group_rank = %d, apply grad_out1 top halo (grad_out1.shape = %s)\" % (ctx.spatial_group_rank, str(list(grad_out1.shape))))\n elif ctx.spatial_method == 3:\n if ctx.spatial_group_rank < ctx.spatial_group_size - 1:\n if ctx.explicit_nhwc:\n btm_relu_halo = relu1[:, Hs - 1 :, :, :].clone()\n btm_grad_out1 = grad_out1[:, Hs - 1 :, :, :]\n else:\n btm_relu_halo = relu1[:, :, Hs - 1 :, :].clone()\n btm_grad_out1 = grad_out1[:, :, Hs - 1 :, :]\n w1by3 = w[:, :1, :, :].clone()\n ctx.stream2.wait_stream(ctx.stream1) # wait for halo transfers to finish\n ctx.stream2.wait_stream(\n torch.cuda.current_stream()\n ) # wait for backward_grad_out1_mask to finish before launching halo correction kernel\n with torch.cuda.stream(ctx.stream2):\n btm_grad_out1_halo = fast_bottleneck.backward_grad_out1_halo_corr(\n ctx.explicit_nhwc,\n ctx.stride_1x1,\n t_list,\n w1by3,\n grads,\n btm_halo,\n btm_relu_halo,\n btm_grad_out1.clone(),\n )\n btm_grad_out1.copy_(btm_grad_out1_halo)\n if ctx.spatial_group_rank > 0:\n if ctx.explicit_nhwc:\n top_relu_halo = relu1[:, :1, :, :].clone()\n top_grad_out1 = grad_out1[:, :1, :, :]\n else:\n top_relu_halo = relu1[:, :, :1, :].clone()\n top_grad_out1 = grad_out1[:, :, :1, :]\n w1by3 = w[:, 2:, :, :].clone()\n ctx.stream1.wait_stream(\n torch.cuda.current_stream()\n ) # wait for backward_grad_out1_mask to finish before launching halo correction kernel\n with torch.cuda.stream(ctx.stream1):\n top_grad_out1_halo = fast_bottleneck.backward_grad_out1_halo_corr(\n ctx.explicit_nhwc,\n ctx.stride_1x1,\n t_list,\n w1by3,\n grads,\n top_halo,\n top_relu_halo,\n top_grad_out1.clone(),\n )\n top_grad_out1.copy_(top_grad_out1_halo)\n if ctx.spatial_group_rank < ctx.spatial_group_size - 1:\n torch.cuda.current_stream().wait_stream(\n ctx.stream2\n ) # wait for halo correction to finish\n if ctx.spatial_group_rank > 0:\n torch.cuda.current_stream().wait_stream(ctx.stream1)\n\n wgrad1_stream = torch.cuda.Stream()\n wgrad1_stream.wait_stream(torch.cuda.current_stream())\n fast_bottleneck.backward_rest(\n ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads, grad_out2, grad_out1\n )\n with torch.cuda.stream(wgrad3_stream):\n fast_bottleneck.backward_wgrad3(ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads)\n with torch.cuda.stream(wgrad2_stream):\n if ctx.spatial_group_size > 1:\n fast_bottleneck.backward_wgrad2_pad(\n ctx.explicit_nhwc,\n ctx.stride_1x1,\n t_list,\n grads,\n out1_pad,\n grad_out2,\n )\n else:\n fast_bottleneck.backward_wgrad2(\n ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads, grad_out2\n )\n with torch.cuda.stream(wgrad1_stream):\n fast_bottleneck.backward_wgrad1(\n ctx.explicit_nhwc, ctx.stride_1x1, t_list, grads, grad_out1\n )\n torch.cuda.current_stream().wait_stream(wgrad3_stream)\n torch.cuda.current_stream().wait_stream(wgrad2_stream)\n torch.cuda.current_stream().wait_stream(wgrad1_stream)\n\n return (\n None,\n None,\n None,\n None,\n None,\n None,\n None,\n None,\n None,\n None,\n None,\n None,\n *grads,\n )\n\n\nspatial_bottleneck_function = SpatialBottleneckFunction.apply\n\n\nclass SpatialBottleneck(torch.nn.Module):\n # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)\n # while original implementation places the stride at the first 1x1 convolution(self.conv1)\n # according to \"Deep residual learning for image recognition\"https://arxiv.org/abs/1512.03385.\n # This variant is also known as ResNet V1.5 and improves accuracy according to\n # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.\n # here we put it at 1x1\n\n def __init__(\n self,\n in_channels,\n bottleneck_channels,\n out_channels,\n stride=1,\n groups=1,\n dilation=1,\n norm_func=None,\n use_cudnn=False,\n explicit_nhwc=False,\n spatial_parallel_args=None,\n ):\n super(SpatialBottleneck, self).__init__()\n if groups != 1:\n raise RuntimeError(\"Only support groups == 1\")\n if dilation != 1:\n raise RuntimeError(\"Only support dilation == 1\")\n if norm_func == None:\n norm_func = FrozenBatchNorm2d\n else:\n raise RuntimeError(\"Only support frozen BN now.\")\n\n if stride != 1 or in_channels != out_channels:\n self.downsample = nn.Sequential(\n conv1x1(in_channels, out_channels, stride),\n norm_func(out_channels),\n )\n else:\n self.downsample = None\n\n # Both self.conv2 and self.downsample layers downsample the input when stride != 1\n self.conv1 = conv1x1(in_channels, bottleneck_channels, stride)\n self.conv2 = conv3x3(bottleneck_channels, bottleneck_channels)\n self.conv3 = conv1x1(bottleneck_channels, out_channels)\n self.relu = nn.ReLU(inplace=True)\n self.stride = stride\n\n self.bn1 = norm_func(bottleneck_channels)\n self.bn2 = norm_func(bottleneck_channels)\n self.bn3 = norm_func(out_channels)\n self.w_scale = None\n\n self.use_cudnn = use_cudnn\n\n # setup conv weights\n self.w_conv = [self.conv1.weight, self.conv2.weight, self.conv3.weight]\n if self.downsample is not None:\n self.w_conv.append(self.downsample[0].weight)\n\n # init weight in nchw format before possible transpose\n for w in self.w_conv:\n kaiming_uniform_(w, a=1)\n\n self.thresholdTop, self.thresholdBottom = None, None\n\n # TODO: prevent unsupported case usage\n # support cases\n # native cudnn\n # normal yes no\n # channel_last yes yes\n # explicit_nhwc no yes\n self.explicit_nhwc = explicit_nhwc\n if self.explicit_nhwc:\n for p in self.parameters():\n with torch.no_grad():\n p.data = p.data.permute(0, 2, 3, 1).contiguous()\n\n # spatial communicator\n if spatial_parallel_args is None:\n self.spatial_parallel_args = (1, 0, None, None, 0, False)\n else:\n self.spatial_parallel_args = spatial_parallel_args\n return\n\n # Returns single callable that recomputes scale and bias for all frozen batch-norms.\n # This method must be called before cuda graphing.\n # The callable it returns can be called anytime.\n # Calling this method will prevent these from being computed every forward call.\n def get_scale_bias_callable(self):\n self.w_scale, self.w_bias, args = [], [], []\n batch_norms = [self.bn1, self.bn2, self.bn3]\n if self.downsample is not None:\n batch_norms.append(self.downsample[1])\n for bn in batch_norms:\n s = torch.empty_like(bn.weight)\n b = torch.empty_like(s)\n args.append((bn.weight, bn.bias, bn.running_mean, bn.running_var, s, b))\n if self.explicit_nhwc:\n self.w_scale.append(s.reshape(1, 1, 1, -1))\n self.w_bias.append(b.reshape(1, 1, 1, -1))\n else:\n self.w_scale.append(s.reshape(1, -1, 1, 1))\n self.w_bias.append(b.reshape(1, -1, 1, 1))\n return func.partial(compute_scale_bias_method, self.explicit_nhwc, args)\n\n def forward(self, x):\n if self.use_cudnn:\n if self.thresholdTop is None:\n spatial_group_size, spatial_group_rank, _, _, _, _ = self.spatial_parallel_args\n if self.explicit_nhwc:\n N, H, W, C = list(x.shape)\n else:\n N, C, H, W = list(x.shape)\n self.thresholdTop = torch.tensor(\n [1 if spatial_group_rank > 0 else 0],\n dtype=torch.int32,\n device=\"cuda\",\n )\n self.thresholdBottom = torch.tensor(\n [H - 2 if spatial_group_rank < spatial_group_size - 1 else H - 1],\n dtype=torch.int32,\n device=\"cuda\",\n )\n\n if self.w_scale is None:\n # calculate scale/bias from registered buffers\n # TODO: make this better\n s1, b1 = self.bn1.get_scale_bias(self.explicit_nhwc)\n s2, b2 = self.bn2.get_scale_bias(self.explicit_nhwc)\n s3, b3 = self.bn3.get_scale_bias(self.explicit_nhwc)\n w_scale = [s1, s2, s3]\n w_bias = [b1, b2, b3]\n if self.downsample is not None:\n s4, b4 = self.downsample[1].get_scale_bias(self.explicit_nhwc)\n w_scale.append(s4)\n w_bias.append(b4)\n out = spatial_bottleneck_function(\n *self.spatial_parallel_args,\n self.explicit_nhwc,\n self.stride,\n w_scale,\n w_bias,\n self.thresholdTop,\n self.thresholdBottom,\n x,\n *self.w_conv,\n )\n else:\n out = spatial_bottleneck_function(\n *self.spatial_parallel_args,\n self.explicit_nhwc,\n self.stride,\n self.w_scale,\n self.w_bias,\n self.thresholdTop,\n self.thresholdBottom,\n x,\n *self.w_conv,\n )\n return out\n\n if self.explicit_nhwc:\n raise RuntimeError(\"explicit nhwc with native ops is not supported.\")\n\n # fallback to native ops\n identity = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out += identity\n out = self.relu(out)\n\n return out\n" }, { "path": "apex/contrib/bottleneck/halo_exchangers.py", "content": "import torch\nimport nccl_p2p_cuda as inc\nimport peer_memory_cuda as pm\n\n\n# Communication free halo exchanger.\n# NB! This halo exchanger does not exchange halos with neighbors as it should, it merely swaps the inputs\n# NB! This is only useful for performance testing.\n# NB! Do not use for actual production runs\nclass HaloExchanger(object):\n def __init__(self, ranks, rank_in_group):\n self.stream1 = torch.cuda.Stream()\n self.stream2 = torch.cuda.Stream()\n self.stream3 = torch.cuda.Stream()\n self.group_size = len(ranks)\n self.ranks = ranks\n self.rank_in_group = rank_in_group\n self.wrap_around_left_rank_in_group = (\n rank_in_group + self.group_size - 1\n ) % self.group_size\n self.wrap_around_right_rank_in_group = (rank_in_group + 1) % self.group_size\n self.left_rank = ranks[rank_in_group - 1] if rank_in_group > 0 else -1\n self.left_zero = True if rank_in_group == 0 else False\n self.right_rank = ranks[rank_in_group + 1] if rank_in_group < self.group_size - 1 else -1\n self.right_zero = True if rank_in_group == self.group_size - 1 else False\n\n\nclass HaloExchangerNoComm(HaloExchanger):\n def __init__(self, ranks, rank_in_group):\n super(HaloExchangerNoComm, self).__init__(ranks, rank_in_group)\n\n def left_right_halo_exchange(\n self,\n left_output_halo,\n right_output_halo,\n left_input_halo=None,\n right_input_halo=None,\n ):\n if left_input_halo is None:\n return right_output_halo, left_output_halo\n else:\n left_input_halo.copy_(right_output_halo)\n right_input_halo.copy_(left_output_halo)\n\n\nclass HaloExchangerAllGather(HaloExchanger):\n def __init__(self, ranks, rank_in_group, comm):\n super(HaloExchangerAllGather, self).__init__(ranks, rank_in_group)\n # self.comm must be NCCL process_group created with torch.distributed.new_group(ranks=ranks)\n self.comm = comm\n\n def left_right_halo_exchange(\n self,\n left_output_halo,\n right_output_halo,\n left_input_halo=None,\n right_input_halo=None,\n ):\n N, Hh, W, C = list(left_output_halo.shape)\n send_halos = torch.empty(\n (N, 2 * Hh, W, C),\n dtype=left_output_halo.dtype,\n device=left_output_halo.device,\n )\n send_halos[:, :Hh, :, :].copy_(left_output_halo)\n send_halos[:, Hh:, :, :].copy_(right_output_halo)\n all_halos = torch.empty(\n (N, 2 * Hh * self.group_size, W, C),\n dtype=left_output_halo.dtype,\n device=left_output_halo.device,\n )\n all_halos = [\n all_halos[:, i * 2 * Hh : (i + 1) * 2 * Hh, :, :] for i in range(self.group_size)\n ]\n torch.distributed.all_gather(all_halos, send_halos, group=self.comm, no_copy=True)\n ag_left_input_halo = all_halos[self.wrap_around_left_rank_in_group][:, Hh:, :, :]\n ag_right_input_halo = all_halos[self.wrap_around_right_rank_in_group][:, :Hh, :, :]\n if left_input_halo is None:\n if self.left_zero:\n ag_left_input_halo.zero_()\n if self.right_zero:\n ag_right_input_halo.zero_()\n return ag_left_input_halo, ag_right_input_halo\n else:\n if self.left_zero:\n left_input_halo.zero_()\n else:\n left_input_halo.copy_(ag_left_input_halo)\n if self.right_zero:\n right_input_halo.zero_()\n else:\n right_input_halo.copy_(ag_right_input_halo)\n\n\nclass HaloExchangerSendRecv(HaloExchanger):\n def __init__(self, ranks, rank_in_group):\n super(HaloExchangerSendRecv, self).__init__(ranks, rank_in_group)\n nccl_id = inc.get_unique_nccl_id(1).cuda()\n torch.distributed.broadcast(nccl_id, 0)\n nccl_id = nccl_id.cpu()\n print(\"%d :: nccl_id = %s\" % (torch.distributed.get_rank(), str(nccl_id)))\n # Create another global nccl communicator in addition to the one created by torch.distributed.init_process_group(\"nccl\")\n # This is unavoidable because the underlying NCCL communicator torch.distributed creates is a protected variable, hence\n # it cannot be accessed from another class.\n # TODO: Figure out a way to avoid creating a second global communicator\n assert torch.distributed.get_rank() == self.ranks[self.rank_in_group], (\n \"ranks[%d](%d) != torch.distributed.get_rank()(%d)\"\n % (\n self.rank_in_group,\n self.ranks[self.rank_in_group],\n torch.distributed.get_rank(),\n )\n )\n self.handle = inc.init_nccl_comm(\n nccl_id, torch.distributed.get_rank(), torch.distributed.get_world_size()\n )\n\n def left_right_halo_exchange(\n self,\n left_output_halo,\n right_output_halo,\n left_input_halo=None,\n right_input_halo=None,\n ):\n if left_input_halo is None:\n left_input_halo, right_input_halo = inc.left_right_halo_exchange(\n self.handle,\n self.left_rank,\n self.right_rank,\n left_output_halo,\n right_output_halo,\n )\n return left_input_halo, right_input_halo\n else:\n inc.left_right_halo_exchange_inplace(\n self.handle,\n self.left_rank,\n self.right_rank,\n left_output_halo,\n right_output_halo,\n left_input_halo,\n right_input_halo,\n )\n\n\nclass HaloExchangerPeer(HaloExchanger):\n def __init__(self, ranks, rank_in_group, peer_pool, explicit_nhwc, numSM=0):\n super(HaloExchangerPeer, self).__init__(ranks, rank_in_group)\n self.diagnostics = False\n self.explicit_nhwc = explicit_nhwc\n self.numSM = numSM\n self.peer_pool = peer_pool\n\n def _allocate_peer_tensor(self, halo):\n # Compute size in bytes\n # Note: Pad buffer so each CUDA block gets required buffer size\n size = 4 * halo.numel() * halo.element_size()\n size_per_block = 128 * 2 * 16 # 128 threads each require two 128b buffers\n size = (size + size_per_block - 1) // size_per_block * size_per_block\n\n # Construct dtype peer buffer with desired size\n shape = [1, 1, 1, size // halo.element_size()]\n return self.peer_pool.allocate_peer_tensors(shape, halo.dtype, False, True)\n\n def left_right_halo_exchange(\n self,\n left_output_halo,\n right_output_halo,\n left_input_halo=None,\n right_input_halo=None,\n ):\n inplace = False if left_input_halo is None and right_input_halo is None else True\n if not inplace:\n left_input_halo = torch.empty_like(right_output_halo)\n right_input_halo = torch.empty_like(left_output_halo)\n channels_last = (\n left_output_halo.is_contiguous(memory_format=torch.channels_last)\n and not self.explicit_nhwc\n )\n left_tx = self._allocate_peer_tensor(left_input_halo)\n right_tx = self._allocate_peer_tensor(right_input_halo)\n pm.push_pull_halos_1d(\n self.diagnostics,\n self.explicit_nhwc,\n self.numSM,\n self.rank_in_group,\n self.left_zero,\n left_output_halo,\n left_tx[self.rank_in_group],\n right_tx[self.wrap_around_left_rank_in_group],\n left_input_halo,\n self.right_zero,\n right_output_halo,\n right_tx[self.rank_in_group],\n left_tx[self.wrap_around_right_rank_in_group],\n right_input_halo,\n )\n if not inplace:\n return left_input_halo, right_input_halo\n\n\n# Class that combines input volume with halos from neighbors (1d).\nclass HaloPadder:\n def __init__(self, halo_ex):\n self.halo_ex = halo_ex\n self.stream1 = torch.cuda.Stream()\n self.stream2 = torch.cuda.Stream()\n\n def __call__(self, y, half_halo, explicit_nhwc, H_split):\n channels_last = not explicit_nhwc and y.is_contiguous(memory_format=torch.channels_last)\n if explicit_nhwc:\n N, H, W, C = list(y.shape)\n if H_split:\n padded_shape = [N, H + 2 * half_halo, W, C]\n ypad = torch.empty(\n shape=padded_shape,\n dtype=y.dtype,\n device=y.device,\n memory_format=torch.contiguous_format,\n )\n yleft = ypad[:, :half_halo, :, :]\n ymid = ypad[:, half_halo : H + half_halo, :, :]\n yright = ypad[:, H + half_halo : H + 2 * half_halo, :, :]\n oleft = y[:, :half_halo, :, :]\n oright = y[:, H - half_halo :, :, :]\n else:\n padded_shape = [N, H, W + 2 * half_halo, C]\n ypad = torch.empty(\n shape=padded_shape,\n dtype=y.dtype,\n device=y.device,\n memory_format=torch.contiguous_format,\n )\n yleft = ypad[:, :, :half_halo, :]\n ymid = ypad[:, :, half_halo : W + half_halo, :]\n yright = ypad[:, :, W + half_halo : W + 2 * half_halo, :]\n oleft = y[:, :, :half_halo, :]\n oright = y[:, :, W - half_halo :, :]\n else:\n N, C, H, W = list(y.shape)\n if H_split:\n padded_shape = [N, C, H + 2 * half_halo, W]\n ypad = torch.empty(\n shape=padded_shape,\n dtype=y.dtype,\n device=y.device,\n memory_format=torch.channels_last,\n )\n yleft = ypad[:, :, :half_halo, :]\n ymid = ypad[:, :, half_halo : H + half_halo, :]\n yright = ypad[:, :, H + half_halo : H + 2 * half_halo, :]\n oleft = y[:, :, :half_halo, :]\n oright = y[:, :, H - half_halo :, :]\n else:\n padded_shape = [N, C, H, W + 2 * half_halo]\n ypad = torch.empty(\n shape=padded_shape,\n dtype=y.dtype,\n device=y.device,\n memory_format=torch.channels_last,\n )\n yleft = ypad[:, :, :, :half_halo]\n ymid = ypad[:, :, :, half_halo : W + half_halo]\n yright = ypad[:, :, :, W + half_halo : W + 2 * half_halo]\n oleft = y[:, :, :, :half_halo]\n oright = y[:, :, :, W - half_halo :]\n with torch.cuda.stream(self.stream1):\n self.halo_ex(oleft, oright, yleft, yright)\n with torch.cuda.stream(self.stream2):\n ymid.copy_(y)\n return ypad\n\n def wait(self):\n current_stream = torch.cuda.current_stream()\n current_stream.wait_stream(self.stream1)\n current_stream.wait_stream(self.stream2)\n" }, { "path": "apex/contrib/bottleneck/test.py", "content": "import torch\nfrom bottleneck import Bottleneck\n\ntorch.manual_seed(23337)\n\n# use True to print layerwise sum for all outputs in reference code path\nDEBUG = False # True\n\nfor stride, o_channel in [(1, 32), (1, 128), (2, 32)]:\n print(\"testing stride ==\", stride, \", in_channel == 32 , out_channel ==\", o_channel)\n a_ = torch.randn(17, 32, 28, 28)\n\n a = a_.cuda().half().to(memory_format=torch.channels_last).requires_grad_()\n model = (\n Bottleneck(32, 8, o_channel, stride=stride)\n .cuda()\n .half()\n .to(memory_format=torch.channels_last)\n )\n\n # test model\n b = model(a)\n b.mean().backward()\n d_grad = a.grad.float()\n a.grad = None\n torch.cuda.synchronize()\n\n if DEBUG:\n print(\"[DEBUG] ref dx :\", d_grad.sum().item())\n # print wgrad. we don't need to reset since later cpp print before accumulation\n for i, w in enumerate(model.w_conv):\n print(\"[DEBUG] ref wgrad{} :\".format(i + 1), w.grad.sum().item())\n\n wgrads = []\n for w in model.w_conv:\n wgrads.append(w.grad.float())\n\n model.use_cudnn = True\n model.zero_grad()\n c = model(a)\n c.mean().backward()\n\n torch.cuda.synchronize()\n print(\"comparing native and channels_last:\")\n print(\n \"max error fprop:\",\n (b - c).abs().max().item(),\n \"max elem:\",\n b.abs().max().item(),\n )\n print(\n \"max error dgrad:\",\n (d_grad - a.grad.float()).abs().max().item(),\n \"max elem:\",\n d_grad.abs().max().item(),\n )\n for i, (w, wgrad) in enumerate(zip(model.w_conv, wgrads)):\n print(\n \"max error wgrad{}:\".format(i + 1),\n (wgrad - w.grad.float()).abs().max().item(),\n \"max elem:\",\n wgrad.abs().max().item(),\n )\n\n nhwc_a = a_.permute(0, 2, 3, 1).contiguous().cuda().half().requires_grad_()\n nhwc_model = (\n Bottleneck(32, 8, o_channel, stride=stride, explicit_nhwc=True, use_cudnn=True)\n .cuda()\n .half()\n )\n for p, q in zip(model.parameters(), nhwc_model.parameters()):\n # model's storage is already in nhwc, we clone and assign to explicit nhwc model\n q.data.copy_(p.data.permute(0, 2, 3, 1).contiguous())\n for p, q in zip(model.buffers(), nhwc_model.buffers()):\n q.data.copy_(p.data)\n\n d = nhwc_model(nhwc_a)\n d.mean().backward()\n torch.cuda.synchronize()\n\n # reset reference to cudnn channels_last permute\n # c_s = c.storage().tolist()\n # d_s = d.storage().tolist()\n # print(max([x-y for x,y in zip(c_s,d_s)]))\n c = c.contiguous(memory_format=torch.contiguous_format).permute(0, 2, 3, 1).contiguous()\n d_grad = a.grad.float().permute(0, 2, 3, 1).contiguous()\n wgrads = []\n for w in model.w_conv:\n wgrads.append(w.grad.float().permute(0, 2, 3, 1).contiguous())\n\n torch.cuda.synchronize()\n print(\"comparing nhwc and channels_last:\")\n print(\n \"max error fprop:\",\n (d - c).abs().max().item(),\n \"max elem:\",\n c.abs().max().item(),\n )\n print(\n \"max error dgrad:\",\n (d_grad - nhwc_a.grad.float()).abs().max().item(),\n \"max elem:\",\n d_grad.abs().max().item(),\n )\n for i, (w, wgrad) in enumerate(zip(nhwc_model.w_conv, wgrads)):\n print(\n \"max error wgrad{}:\".format(i + 1),\n (wgrad - w.grad.float()).abs().max().item(),\n \"max elem:\",\n wgrad.abs().max().item(),\n )\n" }, { "path": "apex/contrib/clip_grad/__init__.py", "content": "from .clip_grad import clip_grad_norm_\n" }, { "path": "apex/contrib/clip_grad/clip_grad.py", "content": "from typing import Union, Iterable\n\nimport torch\n\n_kernel_import_succeeded = False\ntry:\n import amp_C\n from apex.multi_tensor_apply import multi_tensor_applier\n\n _kernel_import_succeeded = True\nexcept ImportError:\n _kernel_import_succeeded = False\n\n_tensor_or_tensors = Union[torch.Tensor, Iterable[torch.Tensor]]\n\n\ndef clip_grad_norm_(\n parameters: _tensor_or_tensors,\n max_norm: float,\n norm_type: float = 2.0,\n error_if_nonfinite: bool = False,\n) -> torch.Tensor:\n r\"\"\"Clips gradient norm of an iterable of parameters.\n\n The norm is computed over all gradients together, as if they were\n concatenated into a single vector. Gradients are modified in-place.\n\n This is identical to torch.nn.utils.clip_grad_norm_, except it\n uses a fused CUDA kernel when computing the 2-norm of GPU tensors\n in float32 and float16.\n\n Args:\n parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a\n single Tensor that will have gradients normalized\n max_norm (float or int): max norm of the gradients\n norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for\n infinity norm.\n error_if_nonfinite (bool): if True, an error is thrown if the total\n norm of the gradients from :attr:`parameters` is ``nan``,\n ``inf``, or ``-inf``. Default: False (will switch to True in the future)\n\n Returns:\n Total norm of the parameters (viewed as a single vector).\n\n \"\"\"\n if isinstance(parameters, torch.Tensor):\n parameters = [parameters]\n parameters = [p for p in parameters if p.grad is not None]\n max_norm = float(max_norm)\n norm_type = float(norm_type)\n\n # Trivial case\n if len(parameters) == 0:\n return torch.tensor(0.0)\n\n # Fallback implementation\n if not (_kernel_import_succeeded and norm_type == 2.0 and any(p.is_cuda for p in parameters)):\n return torch.nn.utils.clip_grad_norm_(\n parameters,\n max_norm,\n norm_type=norm_type,\n error_if_nonfinite=error_if_nonfinite,\n )\n\n # Find fp32 and fp16 gradients on GPU\n device = next(p.device for p in parameters if p.is_cuda)\n grads_fp32, grads_fp16, grads_misc = [], [], []\n for p in parameters:\n grad = p.grad.detach()\n if p.dtype == torch.float32 and p.device == device:\n grads_fp32.append(grad)\n elif p.dtype == torch.float16 and p.device == device:\n grads_fp16.append(grad)\n else:\n grads_misc.append(grad)\n\n # Compute gradient L2 norms\n norms = []\n dummy_overflow_buf = torch.zeros([1], dtype=torch.int32, device=device)\n if grads_fp32:\n norms.append(\n multi_tensor_applier(\n amp_C.multi_tensor_l2norm,\n dummy_overflow_buf,\n [grads_fp32],\n False,\n )[0]\n )\n if grads_fp16:\n norms.append(\n multi_tensor_applier(\n amp_C.multi_tensor_l2norm,\n dummy_overflow_buf,\n [grads_fp16],\n False,\n )[0],\n )\n for g in grads_misc:\n norms.append(torch.linalg.norm(g).unsqueeze(0).to(device))\n total_norm = torch.linalg.norm(torch.cat(norms))\n\n # Check for non-finite values\n if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()):\n raise RuntimeError(\n f\"The total norm of order {norm_type} for gradients from \"\n \"`parameters` is non-finite, so it cannot be clipped. To disable \"\n \"this error and scale the gradients by the non-finite norm anyway, \"\n \"set `error_if_nonfinite=False`\"\n )\n\n # Scale gradients\n clip_coef = max_norm / (total_norm + 1e-6)\n clip_coef_clamped = torch.clamp(clip_coef, max=1.0)\n if grads_fp32:\n multi_tensor_applier(\n amp_C.multi_tensor_scale,\n dummy_overflow_buf,\n [grads_fp32, grads_fp32],\n clip_coef_clamped,\n )\n if grads_fp16:\n multi_tensor_applier(\n amp_C.multi_tensor_scale,\n dummy_overflow_buf,\n [grads_fp16, grads_fp16],\n clip_coef_clamped,\n )\n for g in grads_misc:\n g.mul_(clip_coef_clamped.to(g.device))\n\n return total_norm\n" }, { "path": "apex/contrib/conv_bias_relu/__init__.py", "content": "from .conv_bias_relu import (\n ConvBiasReLU,\n ConvBias,\n ConvBiasMaskReLU,\n ConvFrozenScaleBiasReLU,\n)\n" }, { "path": "apex/contrib/conv_bias_relu/conv_bias_relu.py", "content": "import torch\n\nfrom apex import check_cudnn_version_and_warn\nimport fused_conv_bias_relu\n\ncheck_cudnn_version_and_warn(__name__, 8400)\n\n\nclass ConvBiasReLU_(torch.autograd.Function):\n @staticmethod\n @torch.amp.custom_fwd(cast_inputs=torch.half, device_type=\"cuda\")\n def forward(ctx, x, weight, bias, padding, stride):\n outputs = fused_conv_bias_relu.forward([x, weight, bias], padding, stride)\n ctx.save_for_backward(x, weight, outputs[0])\n ctx.padding = padding\n ctx.stride = stride\n\n return outputs[0]\n\n @staticmethod\n @torch.amp.custom_bwd(device_type=\"cuda\")\n def backward(ctx, grad_output):\n bwd_args = [*ctx.saved_tensors, grad_output]\n padding = ctx.padding\n stride = ctx.stride\n grads = fused_conv_bias_relu.backward(bwd_args, padding, stride)\n\n return grads[0], grads[1], grads[2], None, None\n\n\nclass ConvBiasMaskReLU_(torch.autograd.Function):\n @staticmethod\n @torch.amp.custom_fwd(cast_inputs=torch.half, device_type=\"cuda\")\n def forward(ctx, x, weight, bias, mask, padding, stride):\n outputs = fused_conv_bias_relu.forward_mask([x, weight, bias, mask], padding, stride)\n ctx.save_for_backward(x, weight, outputs[0])\n ctx.padding = padding\n ctx.stride = stride\n\n return outputs[0]\n\n @staticmethod\n @torch.amp.custom_bwd(device_type=\"cuda\")\n def backward(ctx, grad_output):\n bwd_args = [*ctx.saved_tensors, grad_output]\n padding = ctx.padding\n stride = ctx.stride\n grads = fused_conv_bias_relu.backward(bwd_args, padding, stride)\n\n return grads[0], grads[1], grads[2], None, None, None\n\n\nclass ConvBias_(torch.autograd.Function):\n @staticmethod\n @torch.amp.custom_fwd(cast_inputs=torch.half, device_type=\"cuda\")\n def forward(ctx, x, weight, bias, padding, stride):\n outputs = fused_conv_bias_relu.forward_no_relu([x, weight, bias], padding, stride)\n ctx.save_for_backward(x, weight)\n ctx.padding = padding\n ctx.stride = stride\n\n return outputs[0]\n\n @staticmethod\n @torch.amp.custom_bwd(device_type=\"cuda\")\n def backward(ctx, grad_output):\n bwd_args = [*ctx.saved_tensors, grad_output]\n padding = ctx.padding\n stride = ctx.stride\n grads = fused_conv_bias_relu.backward_no_relu(bwd_args, padding, stride)\n\n return grads[0], grads[1], grads[2], None, None\n\n\nclass ConvFrozenScaleBiasReLU_(torch.autograd.Function):\n @staticmethod\n @torch.amp.custom_fwd(cast_inputs=torch.half, device_type=\"cuda\")\n def forward(ctx, x, weight, scale, bias, padding, stride):\n output = fused_conv_bias_relu.forward_cscale_cbias_relu(\n [x, weight, scale, bias], padding, stride\n )\n ctx.save_for_backward(x, weight, scale, output)\n ctx.padding = padding\n ctx.stride = stride\n\n return output\n\n @staticmethod\n @torch.amp.custom_bwd(device_type=\"cuda\")\n def backward(ctx, grad_output):\n bwd_args = [*ctx.saved_tensors, grad_output]\n padding = ctx.padding\n stride = ctx.stride\n grads = fused_conv_bias_relu.backward_cscale_cbias_relu(bwd_args, padding, stride)\n\n return grads[0], grads[1], None, None, None, None\n\n\nConvBiasReLU = ConvBiasReLU_.apply\nConvBiasMaskReLU = ConvBiasMaskReLU_.apply\nConvBias = ConvBias_.apply\nConvFrozenScaleBiasReLU = ConvFrozenScaleBiasReLU_.apply\n" }, { "path": "apex/contrib/csrc/bottleneck/bottleneck.cpp", "content": "#include \n#include // for getcudnnhandle\n#include \n#include \n#include \n\n#include \n#include \n\n#ifdef DEBUG\n#define DEBUG_MSG(str) \\\n do { \\\n std::cout << str << std::endl; \\\n } while (false)\n#else\n#define DEBUG_MSG(str) \\\n do { \\\n } while (false)\n#endif\n\n#ifdef DEBUG_CUDNN\n#define DEBUG_CUDNN_MSG(buf, str) \\\n do { \\\n buf << str << std::endl; \\\n } while (false)\n#else\n#define DEBUG_CUDNN_MSG(buf, str) \\\n do { \\\n } while (false)\n#endif\n\n#define checkCudnnErr(...) \\\n do { \\\n int err = checkCudnnError(__VA_ARGS__, #__VA_ARGS__, __FILE__, __LINE__); \\\n if (err) { \\\n return; \\\n } \\\n } while (0)\n\nint checkCudnnError(cudnnStatus_t code, const char* expr, const char* file, int line) {\n if (code) {\n printf(\"CUDNN error at %s:%d, code=%d (%s) in '%s'\\n\", file, line, (int)code, cudnnGetErrorString(code), expr);\n return 1;\n }\n return 0;\n}\n\nvoid checkError(cudaError_t code, char const* func, const char* file, const int line, bool abort = true);\n#define checkCUDAError(val) \\\n { \\\n checkError((val), #val, __FILE__, __LINE__); \\\n } // in-line regular function\n\nvoid checkError(cudaError_t code, char const* func, const char* file, const int line, bool abort) {\n if (code != cudaSuccess) {\n const char* errorMessage = cudaGetErrorString(code);\n fprintf(stderr, \"CUDA error returned from \\\"%s\\\" at %s:%d, Error code: %d (%s)\\n\", func, file, line, code,\n errorMessage);\n if (abort) {\n cudaDeviceReset();\n exit(code);\n }\n }\n}\n\nvoid generateStrides(const int64_t* dimA, int64_t* strideA, int nbDims, cudnnTensorFormat_t filterFormat) {\n // For INT8x4 and INT8x32 we still compute standard strides here to input\n // into the cuDNN functions. We will manually scale by resizeFactor in the cpu ref.\n if (filterFormat == CUDNN_TENSOR_NCHW) {\n strideA[nbDims - 1] = 1;\n for (int64_t d = nbDims - 2; d >= 0; d--) {\n strideA[d] = strideA[d + 1] * dimA[d + 1];\n }\n } else {\n // Here we assume that the format is CUDNN_TENSOR_NHWC\n strideA[1] = 1;\n strideA[nbDims - 1] = strideA[1] * dimA[1];\n for (int64_t d = nbDims - 2; d >= 2; d--) {\n strideA[d] = strideA[d + 1] * dimA[d + 1];\n }\n strideA[0] = strideA[2] * dimA[2];\n }\n}\n\nint getFwdConvDilatedFilterDim(int filterDim, int dilation) { return ((filterDim - 1) * dilation) + 1; }\n\nint getFwdConvPaddedImageDim(int tensorDim, int pad) { return tensorDim + (2 * pad); }\n\nint getFwdConvOutputDim(int tensorDim, int pad, int filterDim, int stride, int dilation) {\n int p = (getFwdConvPaddedImageDim(tensorDim, pad) - getFwdConvDilatedFilterDim(filterDim, dilation)) / stride + 1;\n return (p);\n}\n\nenum {\n X_TENSOR,\n Y_TENSOR,\n W_TENSOR,\n Z_TENSOR,\n B_TENSOR,\n AFTERADD_TENSOR,\n AFTERBIAS_TENSOR,\n AFTERCONV_TENSOR,\n OPTIONAL,\n AFTEROPT_TENSOR,\n};\n\nusing common_conv_descriptors =\n std::tuple;\n\ncommon_conv_descriptors create_common_descriptors(int64_t* x_dim_padded, int64_t* padA, int64_t* convstrideA,\n int64_t* dilationA, int64_t* w_dim_padded, int64_t* y_dim_padded,\n cudnnDataType_t dataType, cudnnConvolutionMode_t mode) {\n const int convDim = 2;\n\n int64_t strideA_padded[4];\n int64_t outstrideA_padded[4];\n int64_t filterstrideA_padded[4];\n\n generateStrides(w_dim_padded, filterstrideA_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(x_dim_padded, strideA_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(y_dim_padded, outstrideA_padded, 4, CUDNN_TENSOR_NHWC);\n\n return common_conv_descriptors(cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, strideA_padded)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, outstrideA_padded)\n .setId('y')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, w_dim_padded)\n .setStrides(4, filterstrideA_padded)\n .setId('w')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(mode)\n .setNDims(convDim)\n .setStrides(convDim, convstrideA)\n .setPrePadding(convDim, padA)\n .setPostPadding(convDim, padA)\n .setDilation(convDim, dilationA)\n .build());\n}\n\nusing common_convbias_descriptors =\n std::tuple;\n\ncommon_convbias_descriptors create_conv_bias_add_act_descriptors(int64_t* x_dim_padded, int64_t* padA,\n int64_t* convstrideA, int64_t* dilationA,\n int64_t* w_dim_padded, int64_t* y_dim_padded,\n cudnnDataType_t dataType) {\n const int convDim = 2;\n\n int64_t b_dim_padded[4];\n b_dim_padded[0] = 1;\n b_dim_padded[1] = y_dim_padded[1];\n b_dim_padded[2] = 1;\n b_dim_padded[3] = 1;\n\n int64_t x_stride_padded[4];\n int64_t y_stride_padded[4];\n int64_t w_stride_padded[4];\n int64_t b_stride_padded[4];\n\n generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);\n\n return common_convbias_descriptors(cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('y')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, w_dim_padded)\n .setStrides(4, w_stride_padded)\n .setId('w')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim_padded)\n .setStrides(4, b_stride_padded)\n .setId('z')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim_padded)\n .setStrides(4, b_stride_padded)\n .setId('b')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setVirtual()\n .setId('A') // after add\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setVirtual()\n .setId('B') // after bias\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('C') // after conv\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('i')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('D') // after optional add\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_FLOAT)\n .build());\n}\n\n// tensor descriptors used for dgrad\nenum {\n X_OR_DX_TENSOR,\n DY_TENSOR,\n W_OR_DW_TENSOR,\n SCALE_TENSOR,\n RELU_TENSOR,\n AFTER_DCONV_TENSOR,\n AFTER_DRELU_TENSOR,\n};\n\nusing dconv_descriptors =\n std::tuple;\n\ndconv_descriptors create_dconv_descriptors(int64_t* x_dim_padded, int64_t* padA, int64_t* convstrideA,\n int64_t* dilationA, int64_t* w_dim_padded, int64_t* y_dim_padded,\n cudnnDataType_t dataType) {\n const int convDim = 2;\n\n int64_t b_dim_padded[4];\n b_dim_padded[0] = 1;\n b_dim_padded[1] = x_dim_padded[1];\n b_dim_padded[2] = 1;\n b_dim_padded[3] = 1;\n\n int64_t x_stride_padded[4];\n int64_t y_stride_padded[4];\n int64_t w_stride_padded[4];\n int64_t b_stride_padded[4];\n\n generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);\n\n return dconv_descriptors(cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('y')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, w_dim_padded)\n .setStrides(4, w_stride_padded)\n .setId('w')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim_padded)\n .setStrides(4, b_stride_padded)\n .setId('s')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setId('r')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setVirtual()\n .setId('A') // after dconv\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setVirtual()\n .setId('B') // after drelu\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build());\n}\n\n// create a cache for plan\nstd::unordered_map plan_cache;\n\n// TODO: better name\nstd::string getConvFusionString(int64_t* x_dim_padded, int64_t* padA, int64_t* convstrideA, int64_t* dilationA,\n int64_t* w_dim_padded, cudnnDataType_t dataType, std::string fusion_string) {\n for (int i = 0; i < 4; i++) {\n fusion_string += 'X';\n fusion_string += std::to_string(x_dim_padded[i]);\n }\n for (int i = 0; i < 4; i++) {\n fusion_string += 'W';\n fusion_string += std::to_string(w_dim_padded[i]);\n }\n for (int i = 0; i < 2; i++) {\n fusion_string += 'P';\n fusion_string += std::to_string(padA[i]);\n }\n for (int i = 0; i < 2; i++) {\n fusion_string += 'S';\n fusion_string += std::to_string(convstrideA[i]);\n }\n for (int i = 0; i < 2; i++) {\n fusion_string += 'D';\n fusion_string += std::to_string(dilationA[i]);\n }\n fusion_string += 'T';\n fusion_string += std::to_string(dataType);\n return fusion_string;\n}\n\ncudnn_frontend::ExecutionPlan& getOrCreatePlan(cudnnHandle_t handle_, std::stringstream& log_buf,\n cudnn_frontend::OperationGraph& opGraph, std::string cache_string,\n bool use_heuristic = true) {\n auto it = plan_cache.find(cache_string);\n if (it != plan_cache.end()) {\n DEBUG_CUDNN_MSG(log_buf, \"Found plan in cache\");\n return it->second;\n } else {\n if (use_heuristic) {\n // TODO: confirm which mode to use\n auto heuristics = cudnn_frontend::EngineHeuristicsBuilder()\n .setOperationGraph(opGraph)\n .setHeurMode(CUDNN_HEUR_MODE_INSTANT)\n .build();\n // try 3 times for now as WAR for no heuristic training\n int max_tries = 3, count = 0;\n auto& engine_configs = heuristics.getEngineConfig(max_tries);\n while (true) {\n try {\n plan_cache.emplace(cache_string, std::move(cudnn_frontend::ExecutionPlanBuilder()\n .setHandle(handle_)\n .setEngineConfig(engine_configs[count], opGraph.getTag())\n .build()));\n break;\n } catch (cudnn_frontend::cudnnException e) {\n if (++count == max_tries) throw e;\n }\n }\n } else {\n DEBUG_CUDNN_MSG(log_buf, \"No plan in cache\");\n // How many engines support this operation graph ?\n auto total_engines = opGraph.getEngineCount();\n DEBUG_CUDNN_MSG(log_buf, opGraph.describe() << \" has \" << total_engines << \" engines.\");\n // We have to randomly pick one engine from [0, total_engines)\n // Selecting \"0\" by default\n auto engine = cudnn_frontend::EngineBuilder().setGlobalEngineIdx(0).setOperationGraph(opGraph).build();\n DEBUG_CUDNN_MSG(log_buf, engine.describe());\n auto& knobs = engine.getSupportedKnobs();\n for (auto it = std::begin(knobs); it != std::end(knobs); ++it) {\n DEBUG_CUDNN_MSG(log_buf, it->describe());\n }\n if (knobs.begin() != knobs.end()) {\n DEBUG_CUDNN_MSG(log_buf, \"Updated knob choice\");\n knobs.begin()->setChoice(knobs.begin()->getMinValue() + 1);\n DEBUG_CUDNN_MSG(log_buf, knobs.begin()->describe());\n }\n\n // Createmplacee the requisite engine config\n auto engine_config = cudnn_frontend::EngineConfigBuilder().setEngine(engine).build();\n DEBUG_CUDNN_MSG(log_buf, engine_config.describe());\n plan_cache.emplace(\n cache_string,\n std::move(cudnn_frontend::ExecutionPlanBuilder().setHandle(handle_).setEngineConfig(engine_config).build()));\n }\n\n return plan_cache.find(cache_string)->second;\n }\n}\n\nvoid run_conv_scale_bias_add_activation(int64_t* x_dim_padded, int64_t* pad, int64_t* convstride, int64_t* dilation,\n int64_t* w_dim_padded, int64_t* y_dim_padded, cudnnDataType_t dataType,\n at::Half* devPtrX, at::Half* devPtrW, at::Half* devPtrY, at::Half* devPtrZ,\n at::Half* devPtrB, at::Half* devPtrI) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n try {\n int convDim = 2;\n\n // Creates the necessary tensor descriptors\n common_convbias_descriptors tensors = create_conv_bias_add_act_descriptors(x_dim_padded, pad, convstride, dilation,\n w_dim_padded, y_dim_padded, dataType);\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n\n // Define the add operation\n auto scaleDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_MUL).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());\n\n // Define the bias operation\n auto biasDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());\n\n // optional add\n auto addDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, addDesc.describe());\n\n // Define the activation operation\n auto actDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_RELU_FWD)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, actDesc.describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(convDim)\n .setStrides(convDim, convstride)\n .setPrePadding(convDim, pad)\n .setPostPadding(convDim, pad)\n .setDilation(convDim, dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n float alpha = 1.0f;\n float beta = 0.0f;\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setwDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // Create a Add Node with scaling parameters.\n auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(conv_op.getOutputTensor())\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(scaleDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, scale_op.describe());\n\n // Create a Bias Node.\n auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(scale_op.getOutputTensor())\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(biasDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bias_op.describe());\n\n // Create a optional add Node.\n auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(bias_op.getOutputTensor())\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(addDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, add_op.describe());\n\n // Create an Activation Node.\n auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(devPtrI ? add_op.getOutputTensor() : bias_op.getOutputTensor())\n .setyDesc(std::get(tensors))\n .setpwDesc(actDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, act_op.describe());\n\n // Create an Operation Graph. In this case it is convolution add bias activation\n std::array ops = {&conv_op, &scale_op, &bias_op, devPtrI ? &add_op : &act_op,\n &act_op};\n\n auto opGraph = cudnn_frontend::OperationGraphBuilder()\n .setHandle(handle_)\n .setOperationGraph(devPtrI ? ops.size() : 4, ops.data())\n .build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrB, devPtrI};\n int64_t uids[] = {'x', 'y', 'w', 'z', 'b', 'i'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(devPtrI ? 6 : 5, data_ptrs)\n .setUids(devPtrI ? 6 : 5, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_conv_scale_bias(int64_t* x_dim_padded, int64_t* pad, int64_t* convstride, int64_t* dilation,\n int64_t* w_dim_padded, int64_t* y_dim_padded, cudnnDataType_t dataType, at::Half* devPtrX,\n at::Half* devPtrW, at::Half* devPtrY, at::Half* devPtrZ, at::Half* devPtrB) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n try {\n int convDim = 2;\n\n // Creates the necessary tensor descriptors\n common_convbias_descriptors tensors = create_conv_bias_add_act_descriptors(x_dim_padded, pad, convstride, dilation,\n w_dim_padded, y_dim_padded, dataType);\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n\n // Define the add operation\n auto scaleDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_MUL).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());\n\n // Define the bias operation\n auto addDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, addDesc.describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(convDim)\n .setStrides(convDim, convstride)\n .setPrePadding(convDim, pad)\n .setPostPadding(convDim, pad)\n .setDilation(convDim, dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n float alpha = 1.0f;\n float beta = 0.0f;\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setwDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // Create a Add Node with scaling parameters.\n auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(conv_op.getOutputTensor())\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors)) // TODO: change enum to aftermul\n .setpwDesc(scaleDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, scale_op.describe());\n\n // Create a Bias Node.\n auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(scale_op.getOutputTensor())\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(addDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, add_op.describe());\n\n // Create an Operation Graph. In this case it is convolution add bias activation\n std::array ops = {&conv_op, &scale_op, &add_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrB};\n int64_t uids[] = {'x', 'y', 'w', 'z', 'b'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(5, data_ptrs)\n .setUids(5, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_dconv_drelu_dscale(int64_t* x_dim_padded, int64_t* pad, int64_t* convstride, int64_t* dilation,\n int64_t* w_dim_padded, int64_t* y_dim_padded, cudnnDataType_t dataType, at::Half* devPtrX,\n at::Half* devPtrW, at::Half* devPtrY, at::Half* devPtrZ, at::Half* devPtrR) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n try {\n int convDim = 2;\n\n // Creates the necessary tensor descriptors\n dconv_descriptors tensors =\n create_dconv_descriptors(x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(convDim)\n .setStrides(convDim, convstride)\n .setPrePadding(convDim, pad)\n .setPostPadding(convDim, pad)\n .setDilation(convDim, dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n // Define the activation backward operation\n auto actDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_RELU_BWD)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, actDesc.describe());\n\n // Define the scale backward operation\n auto scaleDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_MUL).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());\n\n float alpha = 1.0f;\n float beta = 0.0f;\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR)\n .setdxDesc(std::get(tensors))\n .setwDesc(std::get(tensors))\n .setdyDesc(std::get(tensors))\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // TODO: do we need getOutputTensor(), and what it returns in backward case?\n // Create an relu backward Node.\n auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setdyDesc(std::get(tensors))\n .setxDesc(std::get(tensors))\n .setdxDesc(std::get(tensors))\n .setpwDesc(actDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, act_op.describe());\n\n // Create a Scale Node.\n auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(scaleDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, scale_op.describe());\n\n // Create an Operation Graph. In this case it is convolution add bias activation\n std::array ops = {&conv_op, &act_op, &scale_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrR};\n int64_t uids[] = {'x', 'y', 'w', 's', 'r'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(5, data_ptrs)\n .setUids(5, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_dconv(int64_t* x_dim_padded, int64_t* pad, int64_t* convstride, int64_t* dilation, int64_t* w_dim_padded,\n int64_t* y_dim_padded, cudnnDataType_t dataType, at::Half* devPtrX, at::Half* devPtrW, at::Half* devPtrY,\n cudnnBackendDescriptorType_t mode) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n try {\n int convDim = 2;\n\n // Creates the necessary tensor descriptors\n dconv_descriptors tensors =\n create_dconv_descriptors(x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(convDim)\n .setStrides(convDim, convstride)\n .setPrePadding(convDim, pad)\n .setPostPadding(convDim, pad)\n .setDilation(convDim, dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n float alpha = 1.0f;\n float beta = 0.0f;\n\n // Create a convolution Node\n // mode should be one of following\n // CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR\n // CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR\n auto conv_op_builder = cudnn_frontend::OperationBuilder(mode);\n if (mode == CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR) {\n conv_op_builder.setdxDesc(std::get(tensors))\n .setwDesc(std::get(tensors))\n .setdyDesc(std::get(tensors))\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta);\n } else {\n conv_op_builder.setxDesc(std::get(tensors))\n .setdwDesc(std::get(tensors))\n .setdyDesc(std::get(tensors))\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta);\n }\n auto conv_op = conv_op_builder.build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // Create an Operation Graph. In this case it is convolution add bias activation\n std::array ops = {&conv_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrY, devPtrW};\n int64_t uids[] = {'x', 'y', 'w'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(3, data_ptrs)\n .setUids(3, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_dconv_add(int64_t* x_dim_padded, int64_t* pad, int64_t* convstride, int64_t* dilation, int64_t* w_dim_padded,\n int64_t* y_dim_padded, cudnnDataType_t dataType, at::Half* devPtrX, at::Half* devPtrW,\n at::Half* devPtrY, at::Half* devPtrR) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n try {\n int convDim = 2;\n\n // Creates the necessary tensor descriptors\n dconv_descriptors tensors =\n create_dconv_descriptors(x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(convDim)\n .setStrides(convDim, convstride)\n .setPrePadding(convDim, pad)\n .setPostPadding(convDim, pad)\n .setDilation(convDim, dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n // Define the add backward operation\n auto addDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, addDesc.describe());\n\n float alpha = 1.0f;\n float beta = 0.0f;\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR)\n .setdxDesc(std::get(tensors))\n .setwDesc(std::get(tensors))\n .setdyDesc(std::get(tensors))\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // TODO: do we need getOutputTensor(), and what it returns in backward case?\n // Create add Node.\n auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(addDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, add_op.describe());\n\n // Create an Operation Graph. In this case it is convolution add bias activation\n std::array ops = {&conv_op, &add_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrR};\n int64_t uids[] = {'x', 'y', 'w', 'r'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(4, data_ptrs)\n .setUids(4, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\n// inputs contains x,w,z,b,(i)\nstd::vector bottleneck_forward(bool explicit_nhwc, int stride_1X1, std::vector inputs) {\n std::cout << std::fixed;\n // create output vector\n std::vector outputs;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // setup dimensions\n int64_t dimA[] = {0, 0, 0, 0};\n int64_t filterdimA1[] = {0, 0, 0, 0};\n int64_t filterdimA2[] = {0, 0, 0, 0};\n int64_t filterdimA3[] = {0, 0, 0, 0};\n int64_t filterdimA4[] = {0, 0, 0, 0};\n\n // All dim calculation after this order of n,c,h,w\n int axis[]{0, 1, 2, 3};\n if (explicit_nhwc) {\n axis[0] = 0;\n axis[1] = 3;\n axis[2] = 1;\n axis[3] = 2;\n }\n for (int dim = 0; dim < 4; dim++) {\n dimA[dim] = inputs[0].size(axis[dim]);\n filterdimA1[dim] = inputs[1].size(axis[dim]);\n filterdimA2[dim] = inputs[2].size(axis[dim]);\n filterdimA3[dim] = inputs[3].size(axis[dim]);\n }\n if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {\n for (int dim = 0; dim < 4; dim++) {\n filterdimA4[dim] = inputs[10].size(axis[dim]);\n }\n }\n\n // output dim in n,c,h,w used by backend\n int64_t outdimA1[] = {0, 0, 0, 0}; // Computed Below\n int64_t outdimA2[] = {0, 0, 0, 0}; // Computed Below\n int64_t outdimA3[] = {0, 0, 0, 0}; // Computed Below\n\n // use these fixed value for test run\n int64_t padA[] = {0, 0};\n int64_t padA1[] = {1, 1};\n int64_t dilationA[] = {1, 1};\n int64_t convstrideA[] = {1, 1};\n int64_t convstride1X1[] = {stride_1X1, stride_1X1};\n\n // compute output from pad/stride/dilation\n outdimA1[0] = dimA[0];\n outdimA1[1] = filterdimA1[0];\n for (int dim = 0; dim < 2; dim++) {\n outdimA1[dim + 2] =\n getFwdConvOutputDim(dimA[dim + 2], padA[dim], filterdimA1[dim + 2], convstride1X1[dim], dilationA[dim]);\n }\n\n outdimA2[0] = outdimA1[0];\n outdimA2[1] = filterdimA2[0];\n for (int dim = 0; dim < 2; dim++) {\n outdimA2[dim + 2] =\n getFwdConvOutputDim(outdimA1[dim + 2], padA1[dim], filterdimA2[dim + 2], convstrideA[dim], dilationA[dim]);\n }\n\n outdimA3[0] = outdimA2[0];\n outdimA3[1] = filterdimA3[0];\n for (int dim = 0; dim < 2; dim++) {\n outdimA3[dim + 2] =\n getFwdConvOutputDim(outdimA2[dim + 2], padA[dim], filterdimA3[dim + 2], convstrideA[dim], dilationA[dim]);\n }\n\n // Create output tensor in the correct shape in pytorch's view\n int64_t outdim1[] = {0, 0, 0, 0};\n int64_t outdim2[] = {0, 0, 0, 0};\n int64_t outdim3[] = {0, 0, 0, 0};\n if (explicit_nhwc) {\n axis[0] = 0;\n axis[1] = 2;\n axis[2] = 3;\n axis[3] = 1;\n }\n for (int dim = 0; dim < 4; dim++) {\n outdim1[dim] = outdimA1[axis[dim]];\n outdim2[dim] = outdimA2[axis[dim]];\n outdim3[dim] = outdimA3[axis[dim]];\n }\n\n // run\n at::Half* x = inputs[0].data_ptr();\n at::Half* w = inputs[1].data_ptr();\n at::Half* z = inputs[4].data_ptr();\n at::Half* b = inputs[7].data_ptr();\n auto out1 = at::empty(outdim1, inputs[0].type(), output_format);\n at::Half* y1 = out1.data_ptr();\n\n run_conv_scale_bias_add_activation(dimA, padA, convstride1X1, dilationA, filterdimA1, outdimA1, CUDNN_DATA_HALF, x, w,\n y1, z, b, nullptr);\n\n DEBUG_MSG(\"[DEBUG] new relu1 : \" << out1.to(at::kFloat).sum().item());\n\n w = inputs[2].data_ptr();\n z = inputs[5].data_ptr();\n b = inputs[8].data_ptr();\n auto out2 = at::empty(outdim2, inputs[0].type(), output_format);\n at::Half* y2 = out2.data_ptr();\n\n run_conv_scale_bias_add_activation(outdimA1, padA1, convstrideA, dilationA, filterdimA2, outdimA2, CUDNN_DATA_HALF,\n y1, w, y2, z, b, nullptr);\n DEBUG_MSG(\"[DEBUG] new relu2 : \" << out2.to(at::kFloat).sum().item());\n\n // create output of conv3\n auto out3 = at::empty(outdim3, inputs[0].type(), output_format);\n at::Half* y3 = out3.data_ptr();\n\n // create output of conv4 that may exist\n auto identity = at::empty_like(out3);\n at::Half* yi = identity.data_ptr();\n\n if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {\n w = inputs[10].data_ptr();\n z = inputs[11].data_ptr();\n b = inputs[12].data_ptr();\n run_conv_scale_bias(dimA, padA, convstride1X1, dilationA, filterdimA4, outdimA3, CUDNN_DATA_HALF, x, w, yi, z, b);\n DEBUG_MSG(\"[DEBUG] new downsample : \" << identity.to(at::kFloat).sum().item());\n } else {\n yi = x;\n }\n\n w = inputs[3].data_ptr();\n z = inputs[6].data_ptr();\n b = inputs[9].data_ptr();\n\n run_conv_scale_bias_add_activation(outdimA2, padA, convstrideA, dilationA, filterdimA3, outdimA3, CUDNN_DATA_HALF, y2,\n w, y3, z, b, yi);\n DEBUG_MSG(\"[DEBUG] new relu3 : \" << out3.to(at::kFloat).sum().item());\n\n outputs.push_back(out1);\n outputs.push_back(out2);\n outputs.push_back(out3);\n\n return outputs;\n}\n\nstd::vector bottleneck_backward(bool explicit_nhwc, int stride_1X1, std::vector inputs) {\n bool requires_grad = inputs[0].requires_grad();\n\n std::cout << std::fixed;\n // create output vector\n std::vector outputs;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // setup dimensions\n int64_t dimA[] = {0, 0, 0, 0};\n int64_t filterdimA1[] = {0, 0, 0, 0};\n int64_t filterdimA2[] = {0, 0, 0, 0};\n int64_t filterdimA3[] = {0, 0, 0, 0};\n int64_t filterdimA4[] = {0, 0, 0, 0};\n\n // All dim calculation after this order of n,c,h,w\n int axis[]{0, 1, 2, 3};\n if (explicit_nhwc) {\n axis[0] = 0;\n axis[1] = 3;\n axis[2] = 1;\n axis[3] = 2;\n }\n for (int dim = 0; dim < 4; dim++) {\n dimA[dim] = inputs[0].size(axis[dim]);\n filterdimA1[dim] = inputs[1].size(axis[dim]);\n filterdimA2[dim] = inputs[2].size(axis[dim]);\n filterdimA3[dim] = inputs[3].size(axis[dim]);\n }\n if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {\n for (int dim = 0; dim < 4; dim++) {\n filterdimA4[dim] = inputs[14].size(axis[dim]);\n }\n }\n\n // output dim in n,c,h,w used by backend\n int64_t outdimA1[] = {0, 0, 0, 0}; // Computed Below\n int64_t outdimA2[] = {0, 0, 0, 0}; // Computed Below\n int64_t outdimA3[] = {0, 0, 0, 0}; // Computed Below\n\n // use these fixed value for test run\n int64_t padA[] = {0, 0};\n int64_t padA1[] = {1, 1};\n int64_t dilationA[] = {1, 1};\n int64_t convstrideA[] = {1, 1};\n int64_t convstride1X1[] = {stride_1X1, stride_1X1};\n\n // compute output from pad/stride/dilation\n outdimA1[0] = dimA[0];\n outdimA1[1] = filterdimA1[0];\n for (int dim = 0; dim < 2; dim++) {\n outdimA1[dim + 2] =\n getFwdConvOutputDim(dimA[dim + 2], padA[dim], filterdimA1[dim + 2], convstride1X1[dim], dilationA[dim]);\n }\n\n outdimA2[0] = outdimA1[0];\n outdimA2[1] = filterdimA2[0];\n for (int dim = 0; dim < 2; dim++) {\n outdimA2[dim + 2] =\n getFwdConvOutputDim(outdimA1[dim + 2], padA1[dim], filterdimA2[dim + 2], convstrideA[dim], dilationA[dim]);\n }\n\n outdimA3[0] = outdimA2[0];\n outdimA3[1] = filterdimA3[0];\n for (int dim = 0; dim < 2; dim++) {\n outdimA3[dim + 2] =\n getFwdConvOutputDim(outdimA2[dim + 2], padA[dim], filterdimA3[dim + 2], convstrideA[dim], dilationA[dim]);\n }\n\n // Create output tensor in the correct shape in pytorch's view\n int64_t outdim1[] = {0, 0, 0, 0};\n int64_t outdim2[] = {0, 0, 0, 0};\n int64_t outdim3[] = {0, 0, 0, 0};\n if (explicit_nhwc) {\n axis[0] = 0;\n axis[1] = 2;\n axis[2] = 3;\n axis[3] = 1;\n }\n for (int dim = 0; dim < 4; dim++) {\n outdim1[dim] = outdimA1[axis[dim]];\n outdim2[dim] = outdimA2[axis[dim]];\n outdim3[dim] = outdimA3[axis[dim]];\n }\n\n // dconv3+drelu2+dscale2\n at::Half* conv_in = inputs[13].data_ptr();\n at::Half* dy3 = inputs[10].data_ptr();\n\n DEBUG_MSG(\"[DEBUG] new dconv3 : \" << inputs[10].to(at::kFloat).sum().item());\n\n // wgrad\n auto wgrad3 = at::empty_like(inputs[3]);\n at::Half* dw3 = wgrad3.data_ptr();\n run_dconv(outdimA2, padA, convstrideA, dilationA, filterdimA3, outdimA3, CUDNN_DATA_HALF, conv_in, dw3, dy3,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n\n // dgrad\n auto grad_out2 = at::empty(outdim2, inputs[0].type(), output_format);\n at::Half* dy2 = grad_out2.data_ptr();\n at::Half* w = inputs[3].data_ptr();\n at::Half* z = inputs[5].data_ptr();\n\n at::Half* relu2 = inputs[13].data_ptr();\n\n run_dconv_drelu_dscale(outdimA2, padA, convstrideA, dilationA, filterdimA3, outdimA3, CUDNN_DATA_HALF, dy2, w, dy3, z,\n relu2);\n\n DEBUG_MSG(\"[DEBUG] new dconv2 : \" << grad_out2.to(at::kFloat).sum().item());\n\n // dconv2+drelu1+dscale1\n conv_in = inputs[12].data_ptr();\n\n // wgrad\n auto wgrad2 = at::empty_like(inputs[2]);\n at::Half* dw2 = wgrad2.data_ptr();\n run_dconv(outdimA1, padA1, convstrideA, dilationA, filterdimA2, outdimA2, CUDNN_DATA_HALF, conv_in, dw2, dy2,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n\n // dgrad\n auto grad_out1 = at::empty(outdim1, inputs[0].type(), output_format);\n at::Half* dy1 = grad_out1.data_ptr();\n w = inputs[2].data_ptr();\n z = inputs[4].data_ptr();\n\n at::Half* relu1 = inputs[12].data_ptr();\n // fused dgrad\n run_dconv_drelu_dscale(outdimA1, padA1, convstrideA, dilationA, filterdimA2, outdimA2, CUDNN_DATA_HALF, dy1, w, dy2,\n z, relu1);\n\n /*\n // backward strided conv cannot be fused\n // if stride == 1 but channel changes, we can fuse here\n if (stride_1X1 != 1){\n // dgrad\n run_dconv(outdimA1,\n padA1,\n convstride1X1,\n dilationA,\n filterdimA2,\n outdimA2,\n CUDNN_DATA_HALF,\n dy1,\n w,\n dy2,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);\n\n // mul fused mask\n grad_out1.mul_(inputs[15]);\n }\n else {\n at::Half* relu1 = inputs[12].data_ptr();\n // fused dgrad\n run_dconv_drelu_dscale(outdimA1,\n padA1,\n convstride1X1,\n dilationA,\n filterdimA2,\n outdimA2,\n CUDNN_DATA_HALF,\n dy1,\n w,\n dy2,\n z,\n relu1);\n }\n */\n DEBUG_MSG(\"[DEBUG] new dconv1 : \" << grad_out1.to(at::kFloat).sum().item());\n\n // create grads of conv4 that may exist\n auto grad_x_conv4 = at::empty_like(inputs[0]);\n at::Half* dx_conv4 = grad_x_conv4.data_ptr();\n at::Tensor wgrad4;\n\n // x used for dconv1 and dconv4 wgrad\n at::Half* x = inputs[0].data_ptr();\n\n if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {\n w = inputs[14].data_ptr();\n at::Half* dy_conv4 = inputs[11].data_ptr();\n if (requires_grad) {\n run_dconv(dimA, padA, convstride1X1, dilationA, filterdimA4, outdimA3, CUDNN_DATA_HALF, dx_conv4, w, dy_conv4,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);\n // we don't print here since we can't hook out this grad in pytorch alone to compare, due to addition with dx\n // DEBUG_MSG(\"[DEBUG] new dx_identity : \" << grad_x_conv4.to(at::kFloat).sum().item());\n }\n // wgrad\n wgrad4 = at::empty_like(inputs[14]);\n at::Half* dw4 = wgrad4.data_ptr();\n run_dconv(dimA, padA, convstride1X1, dilationA, filterdimA4, outdimA3, CUDNN_DATA_HALF, x, dw4, dy_conv4,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n } else {\n // if there is no downsample, dx_conv4 is fork of drelu3\n dx_conv4 = inputs[11].data_ptr();\n }\n\n // dconv1+add\n // wgrad\n auto wgrad1 = at::empty_like(inputs[1]);\n at::Half* dw1 = wgrad1.data_ptr();\n run_dconv(dimA, padA, convstride1X1, dilationA, filterdimA1, outdimA1, CUDNN_DATA_HALF, x, dw1, dy1,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n\n // dgrad\n w = inputs[1].data_ptr();\n auto grad_x = at::empty_like(inputs[0]);\n at::Half* dx = grad_x.data_ptr();\n\n // backward strided conv cannot be fused\n // if stride == 1 but channel changes, we can fuse here\n if (requires_grad) {\n if (stride_1X1 != 1) {\n run_dconv(dimA, padA, convstride1X1, dilationA, filterdimA1, outdimA1, CUDNN_DATA_HALF, dx, w, dy1,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);\n // add 2 together\n grad_x.add_(grad_x_conv4);\n } else {\n run_dconv_add(dimA, padA, convstride1X1, dilationA, filterdimA1, outdimA1, CUDNN_DATA_HALF, dx, w, dy1, dx_conv4);\n }\n }\n\n DEBUG_MSG(\"[DEBUG] new dx : \" << grad_x.to(at::kFloat).sum().item());\n DEBUG_MSG(\"[DEBUG] new wgrad1 : \" << wgrad1.to(at::kFloat).sum().item());\n DEBUG_MSG(\"[DEBUG] new wgrad2 : \" << wgrad2.to(at::kFloat).sum().item());\n DEBUG_MSG(\"[DEBUG] new wgrad3 : \" << wgrad3.to(at::kFloat).sum().item());\n outputs.push_back(grad_x);\n outputs.push_back(wgrad1);\n outputs.push_back(wgrad2);\n outputs.push_back(wgrad3);\n\n if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {\n DEBUG_MSG(\"[DEBUG] new wgrad4 : \" << wgrad4.to(at::kFloat).sum().item());\n outputs.push_back(wgrad4);\n }\n\n return outputs;\n}\n\nnamespace {\n\nenum {\n X_TENSOR,\n Y_TENSOR,\n W_TENSOR,\n Z_TENSOR,\n B_TENSOR,\n AFTERADD_TENSOR,\n AFTERBIAS_TENSOR,\n AFTERCONV_TENSOR,\n OPTIONAL,\n AFTEROPT_TENSOR,\n AFTERACT_TENSOR,\n GEN_INDEX_TENSOR,\n MASK_TOP_TENSOR,\n MASK_BOTTOM_TENSOR,\n MASK_TENSOR,\n THRESHOLD_TOP_TENSOR,\n THRESHOLD_BOTTOM_TENSOR,\n};\n\nusing masked_convbias_descriptors =\n std::tuple;\n\nmasked_convbias_descriptors create_conv_bias_add_act_mask_descriptors(int64_t* x_dim_padded, int64_t* padA,\n int64_t* convstrideA, int64_t* dilationA,\n int64_t* w_dim_padded, int64_t* y_dim_padded,\n int64_t* threshold_dim,\n cudnnDataType_t dataType) {\n const int convDim = 2;\n\n int64_t b_dim_padded[4];\n b_dim_padded[0] = 1;\n b_dim_padded[1] = y_dim_padded[1];\n b_dim_padded[2] = 1;\n b_dim_padded[3] = 1;\n\n int64_t x_stride_padded[4];\n int64_t y_stride_padded[4];\n int64_t w_stride_padded[4];\n int64_t b_stride_padded[4];\n int64_t threshold_stride[4];\n\n generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(threshold_dim, threshold_stride, 4, CUDNN_TENSOR_NHWC);\n\n return masked_convbias_descriptors(cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('y')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, w_dim_padded)\n .setStrides(4, w_stride_padded)\n .setId('w')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim_padded)\n .setStrides(4, b_stride_padded)\n .setId('z')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim_padded)\n .setStrides(4, b_stride_padded)\n .setId('b')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setVirtual()\n .setId('A') // after add\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setVirtual()\n .setId('B') // after bias\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('C') // after conv\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('i')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('D') // after optional add\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('E') // after act for masked\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('I') // output of the gen index operation\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_INT32)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('m') // top half of the mask created after the less than\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_BOOLEAN)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('n') // bottom half of the mask\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_BOOLEAN)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('M') // OR of the top and bottom masks\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_BOOLEAN)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, threshold_dim)\n .setStrides(4, threshold_stride)\n .setId('t') // threshold for creating the top mask\n .setAlignment(16)\n .setDataType(CUDNN_DATA_INT32)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, threshold_dim)\n .setStrides(4, threshold_stride)\n .setId('u') // threshold for creating the bottom mask\n .setAlignment(16)\n .setDataType(CUDNN_DATA_INT32)\n .build());\n}\n\n// tensor descriptors used for dgrad\nenum {\n X_OR_DX_TENSOR,\n DY_TENSOR,\n W_OR_DW_TENSOR,\n SCALE_TENSOR,\n RELU_TENSOR,\n AFTER_DCONV_TENSOR,\n AFTER_DRELU_TENSOR,\n DGRAD_INPUT_TENSOR,\n DGRAD_OPTIONAL_TENSOR,\n DGRAD_GEN_INDEX_TENSOR,\n DGRAD_MASK_TOP_TENSOR,\n DGRAD_MASK_BOTTOM_TENSOR,\n DGRAD_MASK_TENSOR,\n DGRAD_THRESHOLD_TOP_TENSOR,\n DGRAD_THRESHOLD_BOTTOM_TENSOR,\n};\n\nusing dconv_add_descriptors = std::tuple;\n\ndconv_add_descriptors create_dconv_add_descriptors(int64_t* x_dim_padded, int64_t* padA, int64_t* convstrideA,\n int64_t* dilationA, int64_t* w_dim_padded, int64_t* y_dim_padded,\n cudnnDataType_t dataType) {\n const int convDim = 2;\n\n int64_t b_dim_padded[4];\n b_dim_padded[0] = 1;\n b_dim_padded[1] = x_dim_padded[1];\n b_dim_padded[2] = 1;\n b_dim_padded[3] = 1;\n\n int64_t x_stride_padded[4];\n int64_t y_stride_padded[4];\n int64_t w_stride_padded[4];\n int64_t b_stride_padded[4];\n\n generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);\n\n return dconv_add_descriptors(cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('y')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, w_dim_padded)\n .setStrides(4, w_stride_padded)\n .setId('w')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim_padded)\n .setStrides(4, b_stride_padded)\n .setId('s')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setId('r')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setVirtual()\n .setId('A') // after dconv\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setVirtual()\n .setId('B') // after drelu\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('i')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('D') // after optional add\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_FLOAT)\n .build());\n}\n\nusing dconv_mask_descriptors =\n std::tuple;\n\ndconv_mask_descriptors create_dconv_mask_descriptors(int64_t* x_dim_padded, int64_t* padA, int64_t* convstrideA,\n int64_t* dilationA, int64_t* w_dim_padded, int64_t* y_dim_padded,\n int64_t* threshold_dim, cudnnDataType_t dataType) {\n const int convDim = 2;\n\n int64_t b_dim_padded[4];\n b_dim_padded[0] = 1;\n b_dim_padded[1] = x_dim_padded[1];\n b_dim_padded[2] = 1;\n b_dim_padded[3] = 1;\n\n int64_t x_stride_padded[4];\n int64_t y_stride_padded[4];\n int64_t w_stride_padded[4];\n int64_t b_stride_padded[4];\n int64_t threshold_stride[4];\n\n generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);\n generateStrides(threshold_dim, threshold_stride, 4, CUDNN_TENSOR_NHWC);\n\n return dconv_mask_descriptors(cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('y')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, w_dim_padded)\n .setStrides(4, w_stride_padded)\n .setId('w')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim_padded)\n .setStrides(4, b_stride_padded)\n .setId('s')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setId('r')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setVirtual()\n .setId('A') // after dconv\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim_padded)\n .setStrides(4, x_stride_padded)\n .setVirtual()\n .setId('B') // after drelu\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('i')\n .setAlignment(16)\n .setDataType(dataType)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('D') // after optional add\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_FLOAT)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('I') // output of the gen index operation\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_INT32)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('m') // top half of the mask created after the less than\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_BOOLEAN)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('n') // bottom half of the mask\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_BOOLEAN)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim_padded)\n .setStrides(4, y_stride_padded)\n .setId('M') // OR of the top and bottom masks\n .setAlignment(16)\n .setVirtual()\n .setDataType(CUDNN_DATA_BOOLEAN)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, threshold_dim)\n .setStrides(4, threshold_stride)\n .setId('t') // threshold for creating the top mask\n .setAlignment(16)\n .setDataType(CUDNN_DATA_INT32)\n .build(),\n cudnn_frontend::TensorBuilder()\n .setDim(4, threshold_dim)\n .setStrides(4, threshold_stride)\n .setId('u') // threshold for creating the bottom mask\n .setAlignment(16)\n .setDataType(CUDNN_DATA_INT32)\n .build());\n}\n\nvoid run_conv_add_scale_bias_activation(int64_t* x_dim_padded, int64_t* pad, int64_t* convstride, int64_t* dilation,\n int64_t* w_dim_padded, int64_t* y_dim_padded, cudnnDataType_t dataType,\n at::Half* devPtrX, at::Half* devPtrW, at::Half* devPtrY, at::Half* devPtrZ,\n at::Half* devPtrB, at::Half* devPtrI) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n try {\n int convDim = 2;\n\n // Creates the necessary tensor descriptors\n common_convbias_descriptors tensors = create_conv_bias_add_act_descriptors(x_dim_padded, pad, convstride, dilation,\n w_dim_padded, y_dim_padded, dataType);\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n\n // Define the add operation\n auto scaleDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_MUL).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());\n\n // Define the bias operation\n auto biasDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());\n\n // optional add\n auto addDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, addDesc.describe());\n\n // Define the activation operation\n auto actDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_RELU_FWD)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, actDesc.describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(convDim)\n .setStrides(convDim, convstride)\n .setPrePadding(convDim, pad)\n .setPostPadding(convDim, pad)\n .setDilation(convDim, dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n float alpha = 1.0f;\n float beta = 0.0f;\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setwDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // create an add node.\n auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(conv_op.getOutputTensor())\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(addDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, add_op.describe());\n\n // Create a Add Node with scaling parameters.\n auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(add_op.getOutputTensor())\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(scaleDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, scale_op.describe());\n\n // Create a Bias Node.\n auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(scale_op.getOutputTensor())\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(biasDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bias_op.describe());\n\n // Create an Activation Node.\n auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(bias_op.getOutputTensor())\n .setyDesc(std::get(tensors))\n .setpwDesc(actDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, act_op.describe());\n\n // Create an Operation Graph. In this case it is convolution add bias activation\n std::array ops = {&conv_op, &add_op, &scale_op, &bias_op, &act_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrB, devPtrI};\n int64_t uids[] = {'x', 'y', 'w', 'z', 'b', 'i'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(6, data_ptrs)\n .setUids(6, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_conv_scale_bias_add_activation_mask(int64_t* x_dim_padded, int64_t* pad, int64_t* convstride,\n int64_t* dilation, int64_t* w_dim_padded, int64_t* y_dim_padded,\n int64_t* threshold_dim, cudnnDataType_t dataType, at::Half* devPtrX,\n at::Half* devPtrW, at::Half* devPtrY, at::Half* devPtrZ, at::Half* devPtrB,\n at::Half* devPtrI, int* devPtrT, int* devPtrU, int axis) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n try {\n int convDim = 2;\n\n // Creates the necessary tensor descriptors\n masked_convbias_descriptors tensors = create_conv_bias_add_act_mask_descriptors(\n x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, threshold_dim, dataType);\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n\n // Define the add operation\n auto scaleDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_MUL).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());\n\n // Define the bias operation\n auto biasDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());\n\n // optional add\n auto addDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, addDesc.describe());\n\n // Define the activation operation\n auto actDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_RELU_FWD)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, actDesc.describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(convDim)\n .setStrides(convDim, convstride)\n .setPrePadding(convDim, pad)\n .setPostPadding(convDim, pad)\n .setDilation(convDim, dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n // Define the genIndex descriptor\n auto genIndexDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_GEN_INDEX)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .setAxis(axis)\n .build();\n DEBUG_CUDNN_MSG(log_buf, genIndexDesc.describe());\n\n // Define the lessThan descriptor\n auto lessThanDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_CMP_LT)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, lessThanDesc.describe());\n\n // Define the greaterThan descriptor\n auto greaterThanDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_CMP_GT)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, greaterThanDesc.describe());\n\n // Define the logical_or descriptor\n auto logicalOrDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_LOGICAL_OR)\n .setMathPrecision(CUDNN_DATA_BOOLEAN)\n .build();\n DEBUG_CUDNN_MSG(log_buf, logicalOrDesc.describe());\n\n // Define the binary_selection descriptor\n auto selectionDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_BINARY_SELECT)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, selectionDesc.describe());\n\n float alpha = 1.0f;\n float beta = 0.0f;\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setwDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // Create a Add Node with scaling parameters.\n auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(conv_op.getOutputTensor())\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(scaleDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, scale_op.describe());\n\n // Create a Bias Node.\n auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(scale_op.getOutputTensor())\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(biasDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bias_op.describe());\n\n // Create a optional add Node.\n auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(bias_op.getOutputTensor())\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(addDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, add_op.describe());\n\n // Create an Activation Node.\n auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(devPtrI ? add_op.getOutputTensor() : bias_op.getOutputTensor())\n .setyDesc(std::get(tensors))\n .setpwDesc(actDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, act_op.describe());\n\n // Create a Gen_Index Node.\n auto genIndex_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(genIndexDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, genIndex_op.describe());\n\n // Create a LessThan Node.\n auto lessThan_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(lessThanDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, lessThan_op.describe());\n\n // Create a GreaterThan Node.\n auto greaterThan_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(greaterThanDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, greaterThan_op.describe());\n\n // Create a LogicalOr Node.\n auto logicalOr_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(logicalOrDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, logicalOr_op.describe());\n\n // Create a Binary_Selection Node.\n auto selection_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .settDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(selectionDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, selection_op.describe());\n\n // Create an Operation Graph. In this case it is convolution add bias activation\n if (devPtrI) {\n std::array ops = {\n &conv_op, &scale_op, &bias_op, &add_op, &act_op,\n &genIndex_op, &lessThan_op, &greaterThan_op, &logicalOr_op, &selection_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrB, devPtrI, devPtrT, devPtrU};\n int64_t uids[] = {'x', 'y', 'w', 'z', 'b', 'i', 't', 'u'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(8, data_ptrs)\n .setUids(8, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } else {\n std::array ops = {&conv_op, &scale_op, &bias_op,\n &act_op, &genIndex_op, &lessThan_op,\n &greaterThan_op, &logicalOr_op, &selection_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrB, devPtrT, devPtrU};\n int64_t uids[] = {'x', 'y', 'w', 'z', 'b', 't', 'u'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(7, data_ptrs)\n .setUids(7, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n }\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_dconv_add_drelu_dscale(int64_t* x_dim_padded, int64_t* pad, int64_t* convstride, int64_t* dilation,\n int64_t* w_dim_padded, int64_t* y_dim_padded, cudnnDataType_t dataType,\n at::Half* devPtrX, at::Half* devPtrW, at::Half* devPtrY, at::Half* devPtrZ,\n at::Half* devPtrR, at::Half* devPtrI) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n try {\n int convDim = 2;\n\n // Creates the necessary tensor descriptors\n dconv_add_descriptors tensors =\n create_dconv_add_descriptors(x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(convDim)\n .setStrides(convDim, convstride)\n .setPrePadding(convDim, pad)\n .setPostPadding(convDim, pad)\n .setDilation(convDim, dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n // optional add\n auto addDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, addDesc.describe());\n\n // Define the activation backward operation\n auto actDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_RELU_BWD)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, actDesc.describe());\n\n // Define the scale backward operation\n auto scaleDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_MUL).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());\n\n float alpha = 1.0f;\n float beta = 0.0f;\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR)\n .setdxDesc(std::get(tensors))\n .setwDesc(std::get(tensors))\n .setdyDesc(std::get(tensors))\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // Create add Node.\n auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(addDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, add_op.describe());\n\n // TODO: do we need getOutputTensor(), and what it returns in backward case?\n // Create an relu backward Node.\n auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setdyDesc(std::get(tensors))\n .setxDesc(std::get(tensors))\n .setdxDesc(std::get(tensors))\n .setpwDesc(actDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, act_op.describe());\n\n // Create a Scale Node.\n auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(scaleDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, scale_op.describe());\n\n // Create an Operation Graph. In this case it is convolution add bias activation\n std::array ops = {&conv_op, &add_op, &act_op, &scale_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrR, devPtrI};\n int64_t uids[] = {'x', 'y', 'w', 's', 'r', 'i'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(6, data_ptrs)\n .setUids(6, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_dconv_drelu_dscale_mask(int64_t* x_dim_padded, int64_t* pad, int64_t* convstride, int64_t* dilation,\n int64_t* w_dim_padded, int64_t* y_dim_padded, int64_t* threshold_dim,\n cudnnDataType_t dataType, at::Half* devPtrX, at::Half* devPtrW, at::Half* devPtrY,\n at::Half* devPtrZ, at::Half* devPtrR, int* devPtrT, int* devPtrU, int axis) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n try {\n int convDim = 2;\n\n // Creates the necessary tensor descriptors\n dconv_mask_descriptors tensors = create_dconv_mask_descriptors(x_dim_padded, pad, convstride, dilation,\n w_dim_padded, y_dim_padded, threshold_dim, dataType);\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n DEBUG_CUDNN_MSG(log_buf, std::get(tensors).describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(convDim)\n .setStrides(convDim, convstride)\n .setPrePadding(convDim, pad)\n .setPostPadding(convDim, pad)\n .setDilation(convDim, dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n // Define the activation backward operation\n auto actDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_RELU_BWD)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, actDesc.describe());\n\n // Define the scale backward operation\n auto scaleDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_MUL).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());\n\n // Define the genIndex descriptor\n auto genIndexDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_GEN_INDEX)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .setAxis(axis)\n .build();\n DEBUG_CUDNN_MSG(log_buf, genIndexDesc.describe());\n\n // Define the lessThan descriptor\n auto lessThanDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_CMP_LT)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, lessThanDesc.describe());\n\n // Define the greaterThan descriptor\n auto greaterThanDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_CMP_GT)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, greaterThanDesc.describe());\n\n // Define the logical_or descriptor\n auto logicalOrDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_LOGICAL_OR)\n .setMathPrecision(CUDNN_DATA_BOOLEAN)\n .build();\n DEBUG_CUDNN_MSG(log_buf, logicalOrDesc.describe());\n\n // Define the binary_selection descriptor\n auto selectionDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_BINARY_SELECT)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, selectionDesc.describe());\n\n float alpha = 1.0f;\n float beta = 0.0f;\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR)\n .setdxDesc(std::get(tensors))\n .setwDesc(std::get(tensors))\n .setdyDesc(std::get(tensors))\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // TODO: do we need getOutputTensor(), and what it returns in backward case?\n // Create an relu backward Node.\n auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setdyDesc(std::get(tensors))\n .setxDesc(std::get(tensors))\n .setdxDesc(std::get(tensors))\n .setpwDesc(actDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, act_op.describe());\n\n // Create a Scale Node.\n auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(scaleDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, scale_op.describe());\n\n // Create a Gen_Index Node.\n auto genIndex_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(genIndexDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, genIndex_op.describe());\n\n // Create a LessThan Node.\n auto lessThan_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(lessThanDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, lessThan_op.describe());\n\n // Create a GreaterThan Node.\n auto greaterThan_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(greaterThanDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, greaterThan_op.describe());\n\n // Create a LogicalOr Node.\n auto logicalOr_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(logicalOrDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, logicalOr_op.describe());\n\n // Create a Binary_Selection Node.\n auto selection_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(std::get(tensors))\n .setbDesc(std::get(tensors))\n .settDesc(std::get(tensors))\n .setyDesc(std::get(tensors))\n .setpwDesc(selectionDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, selection_op.describe());\n\n // Create an Operation Graph. In this case it is convolution add bias activation\n std::array ops = {&conv_op, &act_op, &scale_op, &genIndex_op,\n &lessThan_op, &greaterThan_op, &logicalOr_op, &selection_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrR, devPtrT, devPtrU};\n int64_t uids[] = {'x', 'y', 'w', 's', 'r', 't', 'u'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(7, data_ptrs)\n .setUids(7, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nstruct bottleneck_forward_status {\n int64_t dimA[4];\n int64_t filterdimA1[4];\n int64_t filterdimA2[4];\n int64_t filterdimA2hh[4];\n int64_t filterdimA3[4];\n int64_t filterdimA4[4];\n\n int64_t threshdim[4];\n\n int axis[4];\n\n int64_t outdimA0[4];\n int64_t outdimA1[4];\n int64_t outdimA1b[4]; // out1_pad\n int64_t outdimA2[4];\n int64_t outdimA3[4];\n int64_t outdimA4[4];\n\n int64_t padA[2];\n int64_t padA1[2];\n int64_t padA2[2]; // halo padding\n int64_t dilationA[2];\n int64_t convstrideA[2];\n int64_t convstride1X1[2];\n\n int64_t outdim0[4]; // halo input shape\n int64_t outdim1[4];\n int64_t outdim1b[4];\n int64_t outdim2[4];\n int64_t outdim3[4];\n int64_t outdim4[4]; // halo output shape\n\n void init(bool explicit_nhwc, int stride_1X1, std::vector inputs) {\n dimA[0] = dimA[1] = dimA[2] = dimA[3] = 0;\n filterdimA1[0] = filterdimA1[1] = filterdimA1[2] = filterdimA1[3] = 0;\n filterdimA2[0] = filterdimA2[1] = filterdimA2[2] = filterdimA2[3] = 0;\n filterdimA2hh[0] = filterdimA2hh[1] = filterdimA2hh[2] = filterdimA2hh[3] = 0;\n filterdimA3[0] = filterdimA3[1] = filterdimA3[2] = filterdimA3[3] = 0;\n filterdimA4[0] = filterdimA4[1] = filterdimA4[2] = filterdimA4[3] = 0;\n threshdim[0] = threshdim[1] = threshdim[2] = threshdim[3] = 1;\n\n // All dim calculation after this order of n,c,h,w\n if (explicit_nhwc) {\n axis[0] = 0;\n axis[1] = 3;\n axis[2] = 1;\n axis[3] = 2;\n } else {\n axis[0] = 0;\n axis[1] = 1;\n axis[2] = 2;\n axis[3] = 3;\n }\n\n for (int dim = 0; dim < 4; dim++) {\n dimA[dim] = inputs[0].size(axis[dim]);\n filterdimA1[dim] = inputs[1].size(axis[dim]);\n filterdimA2[dim] = inputs[2].size(axis[dim]);\n filterdimA3[dim] = inputs[3].size(axis[dim]);\n }\n if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {\n for (int dim = 0; dim < 4; dim++) {\n filterdimA4[dim] = inputs[10].size(axis[dim]);\n }\n }\n for (int dim = 0; dim < 4; dim++) {\n if (dim == 2) {\n filterdimA2hh[dim] = 1;\n } else {\n filterdimA2hh[dim] = filterdimA2[dim];\n }\n }\n\n // output dim in n,c,h,w used by backend\n outdimA0[0] = outdimA0[1] = outdimA0[2] = outdimA0[3] = 0;\n outdimA1[0] = outdimA1[1] = outdimA1[2] = outdimA1[3] = 0;\n outdimA1b[0] = outdimA1b[1] = outdimA1b[2] = outdimA1b[3] = 0;\n outdimA2[0] = outdimA2[1] = outdimA2[2] = outdimA2[3] = 0;\n outdimA3[0] = outdimA3[1] = outdimA3[2] = outdimA3[3] = 0;\n outdimA4[0] = outdimA4[1] = outdimA4[2] = outdimA4[3] = 0;\n\n // use these fixed value for test run\n padA[0] = 0;\n padA[1] = 0;\n padA1[0] = 1;\n padA1[1] = 1;\n padA2[0] = 0;\n padA2[1] = 1;\n dilationA[0] = 1;\n dilationA[1] = 1;\n convstrideA[0] = 1;\n convstrideA[1] = 1;\n convstride1X1[0] = stride_1X1;\n convstride1X1[1] = stride_1X1;\n\n // compute output from pad/stride/dilation\n outdimA1[0] = dimA[0];\n outdimA1[1] = filterdimA1[0];\n for (int dim = 0; dim < 2; dim++) {\n outdimA1[dim + 2] =\n getFwdConvOutputDim(dimA[dim + 2], padA[dim], filterdimA1[dim + 2], convstride1X1[dim], dilationA[dim]);\n }\n for (int dim = 0; dim < 4; dim++) {\n if (dim == 2) {\n outdimA1b[dim] = outdimA1[dim] + 2;\n } else {\n outdimA1b[dim] = outdimA1[dim];\n }\n }\n\n outdimA2[0] = outdimA1[0];\n outdimA2[1] = filterdimA2[0];\n for (int dim = 0; dim < 2; dim++) {\n outdimA2[dim + 2] =\n getFwdConvOutputDim(outdimA1[dim + 2], padA1[dim], filterdimA2[dim + 2], convstrideA[dim], dilationA[dim]);\n }\n\n for (int dim = 0; dim < 4; dim++) {\n if (dim == 2) {\n outdimA0[dim] = 3;\n outdimA4[dim] = 1;\n } else {\n outdimA0[dim] = outdimA1[dim];\n outdimA4[dim] = outdimA2[dim];\n }\n }\n\n outdimA3[0] = outdimA2[0];\n outdimA3[1] = filterdimA3[0];\n for (int dim = 0; dim < 2; dim++) {\n outdimA3[dim + 2] =\n getFwdConvOutputDim(outdimA2[dim + 2], padA[dim], filterdimA3[dim + 2], convstrideA[dim], dilationA[dim]);\n }\n\n // Create output tensor in the correct shape in pytorch's view\n outdim1[0] = outdim1[1] = outdim1[2] = outdim1[3] = 0;\n outdim1b[0] = outdim1b[1] = outdim1b[2] = outdim1b[3] = 0;\n outdim2[0] = outdim2[1] = outdim2[2] = outdim2[3] = 0;\n outdim3[0] = outdim3[1] = outdim3[2] = outdim3[3] = 0;\n if (explicit_nhwc) {\n axis[0] = 0;\n axis[1] = 2;\n axis[2] = 3;\n axis[3] = 1;\n }\n for (int dim = 0; dim < 4; dim++) {\n outdim0[dim] = outdimA0[axis[dim]];\n outdim1[dim] = outdimA1[axis[dim]];\n outdim1b[dim] = outdimA1b[axis[dim]];\n outdim2[dim] = outdimA2[axis[dim]];\n outdim3[dim] = outdimA3[axis[dim]];\n outdim4[dim] = outdimA4[axis[dim]];\n }\n }\n};\n\nbottleneck_forward_status forward_state;\n\n} // end of anonymous namespace\n\nstd::vector bottleneck_forward_init(bool explicit_nhwc, int stride_1X1, std::vector inputs) {\n // NB! Bottleneck_forward and bottleneck_backward are NOT thread safe method.\n // NB! We use a global object to store state.\n forward_state.init(explicit_nhwc, stride_1X1, inputs);\n\n // create output vector\n std::vector outputs;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // printf(\"outdim1 =\n // (%d,%d,%d,%d)\\n\",forward_state.outdim1[0],forward_state.outdim1[1],forward_state.outdim1[2],forward_state.outdim1[3]);\n auto out1 = at::empty(forward_state.outdim1, inputs[0].type(), output_format);\n auto out2 = at::empty(forward_state.outdim2, inputs[0].type(), output_format);\n auto out3 = at::empty(forward_state.outdim3, inputs[0].type(), output_format);\n\n outputs.push_back(out1);\n outputs.push_back(out2);\n outputs.push_back(out3);\n\n return outputs;\n}\n\n// inputs contains x,w,z,b,(i)\nvoid bottleneck_forward_out1(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs) {\n std::cout << std::fixed;\n\n // run\n at::Half* x = inputs[0].data_ptr();\n at::Half* w = inputs[1].data_ptr();\n at::Half* z = inputs[4].data_ptr();\n at::Half* b = inputs[7].data_ptr();\n auto out1 = outputs[0];\n at::Half* y1 = out1.data_ptr();\n\n run_conv_scale_bias_add_activation(forward_state.dimA, forward_state.padA, forward_state.convstride1X1,\n forward_state.dilationA, forward_state.filterdimA1, forward_state.outdimA1,\n CUDNN_DATA_HALF, x, w, y1, z, b, nullptr);\n\n DEBUG_MSG(\"[DEBUG] new relu1 : \" << out1.to(at::kFloat).sum().item());\n}\n\n// computes halo (top or bottom) from fat halo input.\n// fat halo input is 3 pixels wide in H.\nat::Tensor bottleneck_forward_out2_halo(bool explicit_nhwc, at::Tensor fat_halo_y1, std::vector inputs) {\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // run\n at::Half* w = inputs[2].data_ptr();\n at::Half* z = inputs[5].data_ptr();\n at::Half* b = inputs[8].data_ptr();\n\n at::Half* y1 = fat_halo_y1.data_ptr();\n\n auto halo_y2 = at::empty(forward_state.outdim4, inputs[0].type(), output_format);\n at::Half* y2 = halo_y2.data_ptr();\n\n run_conv_scale_bias_add_activation(forward_state.outdimA0, forward_state.padA2, forward_state.convstrideA,\n forward_state.dilationA, forward_state.filterdimA2, forward_state.outdimA4,\n CUDNN_DATA_HALF, y1, w, y2, z, b, nullptr);\n\n return halo_y2;\n}\n\n// compute halo correction term (top or bottom) from slim halo input (N,C,1,W).\n// slim halo input is 1 pixel wide in H.\nat::Tensor bottleneck_forward_out2_halo_corr(bool explicit_nhwc, at::Tensor slim_halo_y1,\n std::vector inputs, at::Tensor w1by3,\n at::Tensor out2_part_halo) {\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // run\n at::Half* w = w1by3.data_ptr(); // C,C,1,3\n at::Half* z = inputs[5].data_ptr();\n at::Half* b = inputs[8].data_ptr();\n\n at::Half* y1 = slim_halo_y1.data_ptr();\n\n at::Half* prev_out2 = out2_part_halo.data_ptr();\n\n auto halo_y2 = at::empty(forward_state.outdim4, inputs[0].type(), output_format);\n at::Half* y2 = halo_y2.data_ptr();\n\n run_conv_add_scale_bias_activation(forward_state.outdimA4, forward_state.padA2, forward_state.convstrideA,\n forward_state.dilationA, forward_state.filterdimA2hh, forward_state.outdimA4,\n CUDNN_DATA_HALF, y1, w, y2, z, b, prev_out2);\n\n return halo_y2;\n}\n\nvoid bottleneck_forward_out2(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs) {\n std::cout << std::fixed;\n\n // from _out1 method\n at::Half* x = inputs[0].data_ptr();\n auto out1 = outputs[0];\n at::Half* y1 = out1.data_ptr();\n\n // run\n at::Half* w = inputs[2].data_ptr();\n at::Half* z = inputs[5].data_ptr();\n at::Half* b = inputs[8].data_ptr();\n auto out2 = outputs[1];\n at::Half* y2 = out2.data_ptr();\n\n // printf(\"forward_state.outdimA1 =\n // {%d,%d,%d,%d}\\n\",forward_state.outdimA1[0],forward_state.outdimA1[1],forward_state.outdimA1[2],forward_state.outdimA1[3]);\n // printf(\"forward_state.padA1 = {%d,%d}\\n\",forward_state.padA1[0],forward_state.padA1[1]);\n // printf(\"forward_state.convstrideA = {%d,%d}\\n\",forward_state.convstrideA[0],forward_state.convstrideA[1]);\n // printf(\"forward_state.dilationA = {%d,%d}\\n\",forward_state.dilationA[0],forward_state.dilationA[1]);\n // printf(\"forward_state.filterdimA2 =\n // {%d,%d,%d,%d}\\n\",forward_state.filterdimA2[0],forward_state.filterdimA2[1],forward_state.filterdimA2[2],forward_state.filterdimA2[3]);\n // printf(\"forward_state.outdimA2 =\n // {%d,%d,%d,%d}\\n\",forward_state.outdimA2[0],forward_state.outdimA2[1],forward_state.outdimA2[2],forward_state.outdimA2[3]);\n run_conv_scale_bias_add_activation(forward_state.outdimA1, forward_state.padA1, forward_state.convstrideA,\n forward_state.dilationA, forward_state.filterdimA2, forward_state.outdimA2,\n CUDNN_DATA_HALF, y1, w, y2, z, b, nullptr);\n DEBUG_MSG(\"[DEBUG] new relu2 : \" << out2.to(at::kFloat).sum().item());\n}\n\nvoid bottleneck_forward_out2_mask(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs, at::Tensor thresholdTop,\n at::Tensor thresholdBottom) {\n std::cout << std::fixed;\n\n // from _out1 method\n at::Half* x = inputs[0].data_ptr();\n auto out1 = outputs[0];\n at::Half* y1 = out1.data_ptr();\n\n // run\n at::Half* w = inputs[2].data_ptr();\n at::Half* z = inputs[5].data_ptr();\n at::Half* b = inputs[8].data_ptr();\n auto out2 = outputs[1];\n at::Half* y2 = out2.data_ptr();\n\n // printf(\"forward_state.outdimA1 =\n // {%d,%d,%d,%d}\\n\",forward_state.outdimA1[0],forward_state.outdimA1[1],forward_state.outdimA1[2],forward_state.outdimA1[3]);\n // printf(\"forward_state.padA1 = {%d,%d}\\n\",forward_state.padA1[0],forward_state.padA1[1]);\n // printf(\"forward_state.convstrideA = {%d,%d}\\n\",forward_state.convstrideA[0],forward_state.convstrideA[1]);\n // printf(\"forward_state.dilationA = {%d,%d}\\n\",forward_state.dilationA[0],forward_state.dilationA[1]);\n // printf(\"forward_state.filterdimA2 =\n // {%d,%d,%d,%d}\\n\",forward_state.filterdimA2[0],forward_state.filterdimA2[1],forward_state.filterdimA2[2],forward_state.filterdimA2[3]);\n // printf(\"forward_state.outdimA2 =\n // {%d,%d,%d,%d}\\n\",forward_state.outdimA2[0],forward_state.outdimA2[1],forward_state.outdimA2[2],forward_state.outdimA2[3]);\n run_conv_scale_bias_add_activation_mask(forward_state.outdimA1, forward_state.padA1, forward_state.convstrideA,\n forward_state.dilationA, forward_state.filterdimA2, forward_state.outdimA2,\n forward_state.threshdim, CUDNN_DATA_HALF, y1, w, y2, z, b, nullptr,\n thresholdTop.data_ptr(), thresholdBottom.data_ptr(),\n 2); // axis == 1 -> Does this assume explicit NHWC?\n DEBUG_MSG(\"[DEBUG] new relu2 : \" << out2.to(at::kFloat).sum().item());\n}\n\nvoid bottleneck_forward_out2_pad(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs, at::Tensor out1_pad) {\n std::cout << std::fixed;\n\n // from _out1 method\n at::Half* x = inputs[0].data_ptr();\n auto out1 = outputs[0];\n at::Half* y1 = out1_pad.data_ptr();\n\n // run\n at::Half* w = inputs[2].data_ptr();\n at::Half* z = inputs[5].data_ptr();\n at::Half* b = inputs[8].data_ptr();\n auto out2 = outputs[1];\n at::Half* y2 = out2.data_ptr();\n\n // printf(\"forward_state.outdimA1 =\n // {%d,%d,%d,%d}\\n\",forward_state.outdimA1[0],forward_state.outdimA1[1],forward_state.outdimA1[2],forward_state.outdimA1[3]);\n // printf(\"forward_state.padA1 = {%d,%d}\\n\",forward_state.padA1[0],forward_state.padA1[1]);\n // printf(\"forward_state.convstrideA = {%d,%d}\\n\",forward_state.convstrideA[0],forward_state.convstrideA[1]);\n // printf(\"forward_state.dilationA = {%d,%d}\\n\",forward_state.dilationA[0],forward_state.dilationA[1]);\n // printf(\"forward_state.filterdimA2 =\n // {%d,%d,%d,%d}\\n\",forward_state.filterdimA2[0],forward_state.filterdimA2[1],forward_state.filterdimA2[2],forward_state.filterdimA2[3]);\n // printf(\"forward_state.outdimA2 =\n // {%d,%d,%d,%d}\\n\",forward_state.outdimA2[0],forward_state.outdimA2[1],forward_state.outdimA2[2],forward_state.outdimA2[3]);\n run_conv_scale_bias_add_activation(forward_state.outdimA1b, forward_state.padA2, forward_state.convstrideA,\n forward_state.dilationA, forward_state.filterdimA2, forward_state.outdimA2,\n CUDNN_DATA_HALF, y1, w, y2, z, b, nullptr);\n DEBUG_MSG(\"[DEBUG] new relu2 : \" << out2.to(at::kFloat).sum().item());\n}\n\nvoid bottleneck_forward_rest(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs) {\n std::cout << std::fixed;\n\n // from _out1 method\n at::Half* x = inputs[0].data_ptr();\n\n // create output of conv3\n auto out3 = outputs[2];\n at::Half* y3 = out3.data_ptr();\n\n // create output of conv4 that may exist\n auto identity = at::empty_like(out3);\n at::Half* yi = identity.data_ptr();\n\n at::Half *w, *z, *b;\n\n if (stride_1X1 != 1 || forward_state.filterdimA3[0] != forward_state.dimA[1]) {\n w = inputs[10].data_ptr();\n z = inputs[11].data_ptr();\n b = inputs[12].data_ptr();\n run_conv_scale_bias(forward_state.dimA, forward_state.padA, forward_state.convstride1X1, forward_state.dilationA,\n forward_state.filterdimA4, forward_state.outdimA3, CUDNN_DATA_HALF, x, w, yi, z, b);\n DEBUG_MSG(\"[DEBUG] new downsample : \" << identity.to(at::kFloat).sum().item());\n } else {\n yi = x;\n }\n\n auto out2 = outputs[1];\n at::Half* y2 = out2.data_ptr();\n\n w = inputs[3].data_ptr();\n z = inputs[6].data_ptr();\n b = inputs[9].data_ptr();\n\n run_conv_scale_bias_add_activation(forward_state.outdimA2, forward_state.padA, forward_state.convstrideA,\n forward_state.dilationA, forward_state.filterdimA3, forward_state.outdimA3,\n CUDNN_DATA_HALF, y2, w, y3, z, b, yi);\n DEBUG_MSG(\"[DEBUG] new relu3 : \" << out3.to(at::kFloat).sum().item());\n}\n\nnamespace {\n\nstruct bottleneck_backward_state {\n int64_t dimA[4];\n int64_t filterdimA1[4];\n int64_t filterdimA2[4];\n int64_t filterdimA3[4];\n int64_t filterdimA4[4];\n int64_t filterdimA2hh[4]; // Cin,Cout,1,3\n int64_t threshdim[4];\n\n int axis[4];\n\n int64_t outdimA1[4]; // grad_out1\n int64_t outdimA1b[4]; // out1_pad\n int64_t outdimA2[4]; // grad_out2\n int64_t outdimA3[4];\n int64_t outdimA1h[4]; // output: grad_out1 halo (H=3)\n int64_t outdimA2h[4]; // input : grad_out2 halo cells (H=3)\n int64_t outdimA1hh[4]; // input: grad_out2 halo (H=1)\n int64_t outdimA2hh[4]; // input: out1 halo (H=1)\n\n int64_t padA[2];\n int64_t padA1[2];\n int64_t padA2[2];\n int64_t dilationA[2];\n int64_t convstrideA[2];\n int64_t convstride1X1[2];\n\n int64_t filterdim2hh[4]; // Cin,1,3,Cout\n\n int64_t outdim1[4];\n int64_t outdim1b[4];\n int64_t outdim2[4];\n int64_t outdim3[4];\n int64_t outdim1h[4];\n int64_t outdim1hh[4];\n\n void init(bool explicit_nhwc, int stride_1X1, std::vector inputs) {\n // setup dimensions\n dimA[0] = dimA[1] = dimA[2] = dimA[3] = 0;\n filterdimA1[0] = filterdimA1[1] = filterdimA1[2] = filterdimA1[3] = 0;\n filterdimA2[0] = filterdimA2[1] = filterdimA2[2] = filterdimA2[3] = 0;\n filterdimA3[0] = filterdimA3[1] = filterdimA3[2] = filterdimA3[3] = 0;\n filterdimA4[0] = filterdimA4[1] = filterdimA4[2] = filterdimA4[3] = 0;\n filterdimA2hh[0] = filterdimA2hh[1] = filterdimA2hh[2] = filterdimA2hh[3] = 0;\n threshdim[0] = threshdim[1] = threshdim[2] = threshdim[3] = 1;\n\n // All dim calculation after this order of n,c,h,w\n if (explicit_nhwc) {\n axis[0] = 0;\n axis[1] = 3;\n axis[2] = 1;\n axis[3] = 2;\n } else {\n axis[0] = 0;\n axis[1] = 1;\n axis[2] = 2;\n axis[3] = 3;\n }\n\n for (int dim = 0; dim < 4; dim++) {\n dimA[dim] = inputs[0].size(axis[dim]);\n filterdimA1[dim] = inputs[1].size(axis[dim]);\n filterdimA2[dim] = inputs[2].size(axis[dim]);\n filterdimA3[dim] = inputs[3].size(axis[dim]);\n }\n if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {\n for (int dim = 0; dim < 4; dim++) {\n filterdimA4[dim] = inputs[14].size(axis[dim]);\n }\n }\n\n for (int dim = 0; dim < 4; dim++) {\n if (dim == 2) {\n filterdimA2hh[dim] = 1;\n } else {\n filterdimA2hh[dim] = filterdimA2[dim];\n }\n }\n\n // output dim in n,c,h,w used by backend\n outdimA1[0] = outdimA1[1] = outdimA1[2] = outdimA1[3] = 0;\n outdimA1b[0] = outdimA1b[1] = outdimA1b[2] = outdimA1b[3] = 0;\n outdimA2[0] = outdimA2[1] = outdimA2[2] = outdimA2[3] = 0;\n outdimA3[0] = outdimA3[1] = outdimA3[2] = outdimA3[3] = 0;\n outdimA1h[0] = outdimA1h[1] = outdimA1h[2] = outdimA1h[3] = 0;\n outdimA2h[0] = outdimA2h[1] = outdimA2h[2] = outdimA2h[3] = 0;\n outdimA1hh[0] = outdimA1hh[1] = outdimA1hh[2] = outdimA1hh[3] = 0;\n outdimA2hh[0] = outdimA2hh[1] = outdimA2hh[2] = outdimA2hh[3] = 0;\n\n // use these fixed value for test run\n padA[0] = 0;\n padA[1] = 0;\n padA1[0] = 1;\n padA1[1] = 1;\n padA2[0] = 0;\n padA2[1] = 1;\n dilationA[0] = 1;\n dilationA[1] = 1;\n convstrideA[0] = 1;\n convstrideA[1] = 1;\n convstride1X1[0] = stride_1X1;\n convstride1X1[1] = stride_1X1;\n\n // compute output from pad/stride/dilation\n outdimA1[0] = dimA[0];\n outdimA1[1] = filterdimA1[0];\n for (int dim = 0; dim < 2; dim++) {\n outdimA1[dim + 2] =\n getFwdConvOutputDim(dimA[dim + 2], padA[dim], filterdimA1[dim + 2], convstride1X1[dim], dilationA[dim]);\n }\n for (int dim = 0; dim < 4; dim++) {\n if (dim == 2) {\n outdimA1b[dim] = outdimA1[dim] + 2;\n } else {\n outdimA1b[dim] = outdimA1[dim];\n }\n }\n\n outdimA2[0] = outdimA1[0];\n outdimA2[1] = filterdimA2[0];\n for (int dim = 0; dim < 2; dim++) {\n outdimA2[dim + 2] =\n getFwdConvOutputDim(outdimA1[dim + 2], padA1[dim], filterdimA2[dim + 2], convstrideA[dim], dilationA[dim]);\n }\n\n outdimA3[0] = outdimA2[0];\n outdimA3[1] = filterdimA3[0];\n for (int dim = 0; dim < 2; dim++) {\n outdimA3[dim + 2] =\n getFwdConvOutputDim(outdimA2[dim + 2], padA[dim], filterdimA3[dim + 2], convstrideA[dim], dilationA[dim]);\n }\n\n for (int dim = 0; dim < 4; dim++) {\n if (dim == 2) {\n outdimA1h[dim] = 3;\n outdimA2h[dim] = 3;\n outdimA1hh[dim] = 1;\n outdimA2hh[dim] = 1;\n } else {\n outdimA1h[dim] = outdimA1[dim];\n outdimA2h[dim] = outdimA2[dim];\n outdimA1hh[dim] = outdimA1[dim];\n outdimA2hh[dim] = outdimA2[dim];\n }\n }\n\n // Create output tensor in the correct shape in pytorch's view\n outdim1[0] = outdim1[1] = outdim1[2] = outdim1[3] = 0;\n outdim1b[0] = outdim1b[1] = outdim1b[2] = outdim1b[3] = 0;\n outdim2[0] = outdim2[1] = outdim2[2] = outdim2[3] = 0;\n outdim3[0] = outdim3[1] = outdim3[2] = outdim3[3] = 0;\n outdim1h[0] = outdim1h[1] = outdim1h[2] = outdim1h[3] = 0;\n outdim1hh[0] = outdim1hh[1] = outdim1hh[2] = outdim1hh[3] = 0;\n filterdim2hh[0] = filterdim2hh[1] = filterdim2hh[2] = filterdim2hh[3] = 0;\n if (explicit_nhwc) {\n axis[0] = 0;\n axis[1] = 2;\n axis[2] = 3;\n axis[3] = 1;\n }\n for (int dim = 0; dim < 4; dim++) {\n outdim1[dim] = outdimA1[axis[dim]];\n outdim1b[dim] = outdimA1b[axis[dim]];\n outdim2[dim] = outdimA2[axis[dim]];\n outdim3[dim] = outdimA3[axis[dim]];\n outdim1h[dim] = outdimA1h[axis[dim]];\n outdim1hh[dim] = outdimA1hh[axis[dim]];\n filterdim2hh[dim] = filterdimA2hh[axis[dim]];\n }\n }\n};\n\nbottleneck_backward_state backward_state;\n\n} // namespace\n\nstd::vector bottleneck_backward_init(bool explicit_nhwc, int stride_1X1, std::vector inputs) {\n std::cout << std::fixed;\n\n backward_state.init(explicit_nhwc, stride_1X1, inputs);\n\n // create output vector\n std::vector outputs;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n auto grad_x = at::empty_like(inputs[0]);\n auto wgrad1 = at::empty_like(inputs[1]);\n auto wgrad2 = at::empty_like(inputs[2]);\n auto wgrad3 = at::empty_like(inputs[3]);\n\n outputs.push_back(grad_x);\n outputs.push_back(wgrad1);\n outputs.push_back(wgrad2);\n outputs.push_back(wgrad3);\n if (stride_1X1 != 1 || backward_state.filterdimA3[0] != backward_state.dimA[1]) {\n auto wgrad4 = at::empty_like(inputs[14]);\n outputs.push_back(wgrad4);\n }\n\n return outputs;\n}\n\nvoid bottleneck_backward_wgrad3(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs) {\n // dconv3+drelu2+dscale2\n at::Half* conv_in = inputs[13].data_ptr();\n at::Half* dy3 = inputs[10].data_ptr();\n\n // wgrad\n auto wgrad3 = outputs[3];\n at::Half* dw3 = wgrad3.data_ptr();\n run_dconv(backward_state.outdimA2, backward_state.padA, backward_state.convstrideA, backward_state.dilationA,\n backward_state.filterdimA3, backward_state.outdimA3, CUDNN_DATA_HALF, conv_in, dw3, dy3,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n DEBUG_MSG(\"[DEBUG] new wgrad3 : \" << wgrad3.to(at::kFloat).sum().item());\n}\n\nat::Tensor bottleneck_backward_grad_out2(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs) {\n bool requires_grad = inputs[0].requires_grad();\n\n std::cout << std::fixed;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // dconv3+drelu2+dscale2\n at::Half* conv_in = inputs[13].data_ptr();\n at::Half* dy3 = inputs[10].data_ptr();\n\n DEBUG_MSG(\"[DEBUG] new dconv3 : \" << inputs[10].to(at::kFloat).sum().item());\n\n // dgrad\n auto grad_out2 = at::empty(backward_state.outdim2, inputs[0].type(), output_format);\n at::Half* dy2 = grad_out2.data_ptr();\n at::Half* w = inputs[3].data_ptr();\n at::Half* z = inputs[5].data_ptr();\n\n at::Half* relu2 = inputs[13].data_ptr();\n\n run_dconv_drelu_dscale(backward_state.outdimA2, backward_state.padA, backward_state.convstrideA,\n backward_state.dilationA, backward_state.filterdimA3, backward_state.outdimA3, CUDNN_DATA_HALF,\n dy2, w, dy3, z, relu2);\n\n // do halo exchange of dy2 here\n\n DEBUG_MSG(\"[DEBUG] new dconv2 : \" << grad_out2.to(at::kFloat).sum().item());\n\n return grad_out2;\n}\n\nat::Tensor bottleneck_backward_grad_out1(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs, at::Tensor grad_out2) {\n bool requires_grad = inputs[0].requires_grad();\n\n std::cout << std::fixed;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // dgrad\n at::Half* dy2 = grad_out2.data_ptr();\n\n // dgrad\n auto grad_out1 = at::empty(backward_state.outdim1, inputs[0].type(), output_format);\n at::Half* dy1 = grad_out1.data_ptr();\n at::Half* w = inputs[2].data_ptr();\n at::Half* z = inputs[4].data_ptr();\n\n at::Half* relu1 = inputs[12].data_ptr();\n // printf(\"relu.shape = [%d,%d,%d,%d]\\n\",inputs[12].size(0),inputs[12].size(1),inputs[12].size(2),inputs[12].size(3));\n\n // fused dgrad\n // printf(\"backward_state.outdim1 =\n // {%d,%d,%d,%d}\\n\",backward_state.outdim1[0],backward_state.outdim1[1],backward_state.outdim1[2],backward_state.outdim1[3]);\n run_dconv_drelu_dscale(backward_state.outdimA1, backward_state.padA1, backward_state.convstrideA,\n backward_state.dilationA, backward_state.filterdimA2, backward_state.outdimA2, CUDNN_DATA_HALF,\n dy1, w, dy2, z, relu1);\n\n return grad_out1;\n}\n\nat::Tensor bottleneck_backward_grad_out1_mask(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs, at::Tensor grad_out2,\n at::Tensor thresholdTop, at::Tensor thresholdBottom) {\n bool requires_grad = inputs[0].requires_grad();\n\n std::cout << std::fixed;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // dgrad\n at::Half* dy2 = grad_out2.data_ptr();\n\n // dgrad\n auto grad_out1 = at::empty(backward_state.outdim1, inputs[0].type(), output_format);\n at::Half* dy1 = grad_out1.data_ptr();\n at::Half* w = inputs[2].data_ptr();\n at::Half* z = inputs[4].data_ptr();\n\n at::Half* relu1 = inputs[12].data_ptr();\n // printf(\"relu.shape = [%d,%d,%d,%d]\\n\",inputs[12].size(0),inputs[12].size(1),inputs[12].size(2),inputs[12].size(3));\n\n // fused dgrad\n run_dconv_drelu_dscale_mask(backward_state.outdimA1, backward_state.padA1, backward_state.convstrideA,\n backward_state.dilationA, backward_state.filterdimA2, backward_state.outdimA2,\n backward_state.threshdim, CUDNN_DATA_HALF, dy1, w, dy2, z, relu1,\n thresholdTop.data_ptr(), thresholdBottom.data_ptr(), 2);\n\n return grad_out1;\n}\n\n// perform backward data 1x3 convolution (grad_out * w_rot180) on grad_out2 input of shape [N,1,W,C] with padding=(0,1)\n// to produce output of shape [N,1,W,C]\nat::Tensor bottleneck_backward_grad_out1_halo_corr(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n at::Tensor w1by3, std::vector outputs,\n at::Tensor grad_out2_halo, at::Tensor relu1_halo,\n at::Tensor part_grad_out1) {\n bool requires_grad = inputs[0].requires_grad();\n\n std::cout << std::fixed;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // dgrad\n at::Half* dy2h = grad_out2_halo.data_ptr();\n\n // dgrad\n auto grad_out1_halo = at::empty(backward_state.outdim1hh, inputs[0].type(), output_format);\n at::Half* dy1h = grad_out1_halo.data_ptr();\n // at::Half* w = inputs[2].data_ptr(); // use w1by3 instead, which is a sliced version of inputs[2]\n at::Half* w = w1by3.data_ptr();\n at::Half* z = inputs[4].data_ptr();\n at::Half* relu1h = relu1_halo.data_ptr();\n at::Half* pdy1h = part_grad_out1.data_ptr();\n\n // printf(\"relu.shape = [%d,%d,%d,%d]\\n\",relu1_halo.size(0),relu1_halo.size(1),relu1_halo.size(2),relu1_halo.size(3));\n // fused dgrad\n // printf(\"backward_state.outdimA1h =\n // {%d,%d,%d,%d}\\n\",backward_state.outdimA1h[0],backward_state.outdimA1h[1],backward_state.outdimA1h[2],backward_state.outdimA1h[3]);\n // printf(\"backward_state.outdimA2h =\n // {%d,%d,%d,%d}\\n\",backward_state.outdimA2h[0],backward_state.outdimA2h[1],backward_state.outdimA2h[2],backward_state.outdimA2h[3]);\n // printf(\"backward_state.filterdimA2 =\n // {%d,%d,%d,%d}\\n\",backward_state.filterdimA2[0],backward_state.filterdimA2[1],backward_state.filterdimA2[2],backward_state.filterdimA2[3]);\n run_dconv_add_drelu_dscale(backward_state.outdimA1hh,\n backward_state.padA2, // 0,1\n backward_state.convstrideA, backward_state.dilationA,\n backward_state.filterdimA2hh, // C,1,3,C\n backward_state.outdimA2hh, CUDNN_DATA_HALF, dy1h, w, dy2h, z, relu1h, pdy1h);\n\n return grad_out1_halo;\n}\n\n// perform backward data 3x3 convolution (grad_out * w_rot180) on grad_out2 input of shape [N,3,W,C] with padding=(1,1)\n// to produce output of shape [N,3,W,C]\nat::Tensor bottleneck_backward_grad_out1_halo(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs, at::Tensor grad_out2_halo,\n at::Tensor relu1_halo) {\n bool requires_grad = inputs[0].requires_grad();\n\n std::cout << std::fixed;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // dgrad\n at::Half* dy2h = grad_out2_halo.data_ptr();\n\n // dgrad\n auto grad_out1_halo = at::empty(backward_state.outdim1h, inputs[0].type(), output_format);\n at::Half* dy1h = grad_out1_halo.data_ptr();\n at::Half* w = inputs[2].data_ptr();\n at::Half* z = inputs[4].data_ptr();\n\n at::Half* relu1h = relu1_halo.data_ptr();\n // printf(\"relu.shape = [%d,%d,%d,%d]\\n\",relu1_halo.size(0),relu1_halo.size(1),relu1_halo.size(2),relu1_halo.size(3));\n // fused dgrad\n // printf(\"backward_state.outdimA1h =\n // {%d,%d,%d,%d}\\n\",backward_state.outdimA1h[0],backward_state.outdimA1h[1],backward_state.outdimA1h[2],backward_state.outdimA1h[3]);\n // printf(\"backward_state.outdimA2h =\n // {%d,%d,%d,%d}\\n\",backward_state.outdimA2h[0],backward_state.outdimA2h[1],backward_state.outdimA2h[2],backward_state.outdimA2h[3]);\n // printf(\"backward_state.filterdimA2 =\n // {%d,%d,%d,%d}\\n\",backward_state.filterdimA2[0],backward_state.filterdimA2[1],backward_state.filterdimA2[2],backward_state.filterdimA2[3]);\n run_dconv_drelu_dscale(backward_state.outdimA1h, backward_state.padA1, backward_state.convstrideA,\n backward_state.dilationA, backward_state.filterdimA2, backward_state.outdimA2h,\n CUDNN_DATA_HALF, dy1h, w, dy2h, z, relu1h);\n\n return grad_out1_halo;\n}\n\nvoid bottleneck_backward_wgrad2_pad(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs, at::Tensor input, at::Tensor grad_out2) {\n std::cout << std::fixed;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // dgrad\n at::Half* dy2 = grad_out2.data_ptr();\n\n // dconv2+drelu1+dscale1\n at::Half* conv_in = input.data_ptr();\n\n // wgrad\n auto wgrad2 = outputs[2];\n at::Half* dw2 = wgrad2.data_ptr();\n\n // printf(\"outdimA1b =\n // (%d,%d,%d,%d)\\n\",backward_state.outdimA1b[0],backward_state.outdimA1b[1],backward_state.outdimA1b[2],backward_state.outdimA1b[3]);\n // printf(\"backward_state.padA2 = {%d,%d}\\n\",backward_state.padA2[0],backward_state.padA2[1]);\n run_dconv(backward_state.outdimA1b, // conv_in.shape (including H halos)\n backward_state.padA2, // 0, 1\n backward_state.convstrideA, backward_state.dilationA,\n backward_state.filterdimA2, // dw2.shape\n backward_state.outdimA2, // dy2.shape\n CUDNN_DATA_HALF, conv_in, dw2, dy2, CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n DEBUG_MSG(\"[DEBUG] new wgrad2 : \" << wgrad2.to(at::kFloat).sum().item());\n}\n\nvoid bottleneck_backward_wgrad2(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs, at::Tensor grad_out2) {\n bool requires_grad = inputs[0].requires_grad();\n\n std::cout << std::fixed;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // dgrad\n at::Half* dy2 = grad_out2.data_ptr();\n\n // dconv2+drelu1+dscale1\n at::Half* conv_in = inputs[12].data_ptr();\n\n // wgrad\n auto wgrad2 = outputs[2];\n at::Half* dw2 = wgrad2.data_ptr();\n\n // printf(\"outdimA1 =\n // (%d,%d,%d,%d)\\n\",backward_state.outdimA1[0],backward_state.outdimA1[1],backward_state.outdimA1[2],backward_state.outdimA1[3]);\n run_dconv(backward_state.outdimA1, backward_state.padA1, backward_state.convstrideA, backward_state.dilationA,\n backward_state.filterdimA2, backward_state.outdimA2, CUDNN_DATA_HALF, conv_in, dw2, dy2,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n DEBUG_MSG(\"[DEBUG] new wgrad2 : \" << wgrad2.to(at::kFloat).sum().item());\n}\n\n// compute halo cells for input volume of dimension [N,1,W,C] with padding=(0,1) to produce output volume of dimension\n// [N,1,W,C] input and grad_out2_halo tensors are all of same shape output tensor is of shape [Cin,1,3,Cout] (regular\n// filter dims are [Cin,3,3,Cout]\nat::Tensor bottleneck_backward_wgrad2_halo(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs, at::Tensor input,\n at::Tensor grad_out2_halo) {\n bool requires_grad = inputs[0].requires_grad();\n\n std::cout << std::fixed;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // dgrad\n at::Half* dy2 = grad_out2_halo.data_ptr();\n\n // dconv2+drelu1+dscale1\n at::Half* conv_in = input.data_ptr();\n\n // wgrad\n auto wgrad2_halo = at::empty(backward_state.filterdim2hh, input.type(), output_format);\n at::Half* dw2 = wgrad2_halo.data_ptr();\n\n // printf(\"backward_state.outdimA1hh =\n // {%d,%d,%d,%d}\\n\",backward_state.outdimA1hh[0],backward_state.outdimA1hh[1],backward_state.outdimA1hh[2],backward_state.outdimA1hh[3]);\n // printf(\"backward_state.outdimA2hh =\n // {%d,%d,%d,%d}\\n\",backward_state.outdimA2hh[0],backward_state.outdimA2hh[1],backward_state.outdimA2hh[2],backward_state.outdimA2hh[3]);\n // printf(\"backward_state.filterdim2hh =\n // {%d,%d,%d,%d}\\n\",backward_state.filterdim2hh[0],backward_state.filterdim2hh[1],backward_state.filterdim2hh[2],backward_state.filterdim2hh[3]);\n // printf(\"backward_state.filterdimA2hh =\n // {%d,%d,%d,%d}\\n\",backward_state.filterdimA2hh[0],backward_state.filterdimA2hh[1],backward_state.filterdimA2hh[2],backward_state.filterdimA2hh[3]);\n // printf(\"backward_state.padA2 = {%d,%d}\\n\",backward_state.padA2[0],backward_state.padA2[1]);\n run_dconv(backward_state.outdimA1hh, // N,C,1,W\n backward_state.padA2, // 0, 1\n backward_state.convstrideA, backward_state.dilationA,\n backward_state.filterdimA2hh, // Cin,Cout,1,3\n backward_state.outdimA2hh, // N,C,1,W\n CUDNN_DATA_HALF, conv_in, dw2, dy2, CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n\n return wgrad2_halo;\n}\n\nvoid bottleneck_backward_wgrad1(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs, at::Tensor grad_out1) {\n at::Half* x = inputs[0].data_ptr();\n at::Half* dy1 = grad_out1.data_ptr();\n\n // dconv1+add\n // wgrad\n auto wgrad1 = outputs[1];\n at::Half* dw1 = wgrad1.data_ptr();\n run_dconv(backward_state.dimA, backward_state.padA, backward_state.convstride1X1, backward_state.dilationA,\n backward_state.filterdimA1, backward_state.outdimA1, CUDNN_DATA_HALF, x, dw1, dy1,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n}\n\nvoid bottleneck_backward_rest(bool explicit_nhwc, int stride_1X1, std::vector inputs,\n std::vector outputs, at::Tensor grad_out2, at::Tensor grad_out1) {\n bool requires_grad = inputs[0].requires_grad();\n\n std::cout << std::fixed;\n auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;\n\n // dgrad\n at::Half* dy2 = grad_out2.data_ptr();\n at::Half* dy1 = grad_out1.data_ptr();\n\n /*\n // backward strided conv cannot be fused\n // if stride == 1 but channel changes, we can fuse here\n if (stride_1X1 != 1){\n // dgrad\n run_dconv(outdimA1,\n padA1,\n convstride1X1,\n dilationA,\n filterdimA2,\n outdimA2,\n CUDNN_DATA_HALF,\n dy1,\n w,\n dy2,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);\n\n // mul fused mask\n grad_out1.mul_(inputs[15]);\n }\n else {\n at::Half* relu1 = inputs[12].data_ptr();\n // fused dgrad\n run_dconv_drelu_dscale(outdimA1,\n padA1,\n convstride1X1,\n dilationA,\n filterdimA2,\n outdimA2,\n CUDNN_DATA_HALF,\n dy1,\n w,\n dy2,\n z,\n relu1);\n }\n */\n DEBUG_MSG(\"[DEBUG] new dconv1 : \" << grad_out1.to(at::kFloat).sum().item());\n\n // create grads of conv4 that may exist\n auto grad_x_conv4 = at::empty_like(inputs[0]);\n at::Half* dx_conv4 = grad_x_conv4.data_ptr();\n at::Tensor wgrad4;\n\n // x used for dconv1 and dconv4 wgrad\n at::Half* x = inputs[0].data_ptr();\n\n at::Half* w = NULL;\n\n if (stride_1X1 != 1 || backward_state.filterdimA3[0] != backward_state.dimA[1]) {\n w = inputs[14].data_ptr();\n at::Half* dy_conv4 = inputs[11].data_ptr();\n if (requires_grad) {\n run_dconv(backward_state.dimA, backward_state.padA, backward_state.convstride1X1, backward_state.dilationA,\n backward_state.filterdimA4, backward_state.outdimA3, CUDNN_DATA_HALF, dx_conv4, w, dy_conv4,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);\n // we don't print here since we can't hook out this grad in pytorch alone to compare, due to addition with dx\n // DEBUG_MSG(\"[DEBUG] new dx_identity : \" << grad_x_conv4.to(at::kFloat).sum().item());\n }\n // wgrad\n wgrad4 = outputs[4];\n at::Half* dw4 = wgrad4.data_ptr();\n run_dconv(backward_state.dimA, backward_state.padA, backward_state.convstride1X1, backward_state.dilationA,\n backward_state.filterdimA4, backward_state.outdimA3, CUDNN_DATA_HALF, x, dw4, dy_conv4,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n } else {\n // if there is no downsample, dx_conv4 is fork of drelu3\n dx_conv4 = inputs[11].data_ptr();\n }\n\n // dgrad\n w = inputs[1].data_ptr();\n auto grad_x = outputs[0];\n at::Half* dx = grad_x.data_ptr();\n\n // backward strided conv cannot be fused\n // if stride == 1 but channel changes, we can fuse here\n if (requires_grad) {\n if (stride_1X1 != 1) {\n run_dconv(backward_state.dimA, backward_state.padA, backward_state.convstride1X1, backward_state.dilationA,\n backward_state.filterdimA1, backward_state.outdimA1, CUDNN_DATA_HALF, dx, w, dy1,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);\n // add 2 together\n grad_x.add_(grad_x_conv4);\n } else {\n run_dconv_add(backward_state.dimA, backward_state.padA, backward_state.convstride1X1, backward_state.dilationA,\n backward_state.filterdimA1, backward_state.outdimA1, CUDNN_DATA_HALF, dx, w, dy1, dx_conv4);\n }\n }\n\n DEBUG_MSG(\"[DEBUG] new dx : \" << grad_x.to(at::kFloat).sum().item());\n DEBUG_MSG(\"[DEBUG] new wgrad1 : \" << wgrad1.to(at::kFloat).sum().item());\n\n if (stride_1X1 != 1 || backward_state.filterdimA3[0] != backward_state.dimA[1]) {\n DEBUG_MSG(\"[DEBUG] new wgrad4 : \" << wgrad4.to(at::kFloat).sum().item());\n }\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"forward\", &bottleneck_forward, \"Bottleneck block forward\", py::call_guard());\n m.def(\"backward\", &bottleneck_backward, \"Bottleneck block backward\", py::call_guard());\n m.def(\"forward_init\", &bottleneck_forward_init, \"Bottleneck block init\", py::call_guard());\n m.def(\"forward_out1\", &bottleneck_forward_out1, \"Bottleneck block forward\", py::call_guard());\n m.def(\"forward_out2\", &bottleneck_forward_out2, \"Bottleneck block forward\", py::call_guard());\n m.def(\"forward_out2_mask\", &bottleneck_forward_out2_mask, \"Bottleneck block forward\",\n py::call_guard());\n m.def(\"forward_out2_halo\", &bottleneck_forward_out2_halo, \"Bottleneck block forward\",\n py::call_guard());\n m.def(\"forward_out2_halo_corr\", &bottleneck_forward_out2_halo_corr, \"Bottleneck block forward\",\n py::call_guard());\n m.def(\"forward_out2_pad\", &bottleneck_forward_out2_pad, \"Bottleneck block forward\",\n py::call_guard());\n m.def(\"forward_rest\", &bottleneck_forward_rest, \"Bottleneck block forward\", py::call_guard());\n m.def(\"backward_init\", &bottleneck_backward_init, \"Bottleneck block backward init\",\n py::call_guard());\n m.def(\"backward_grad_out2\", &bottleneck_backward_grad_out2, \"Bottleneck block backward\",\n py::call_guard());\n m.def(\"backward_grad_out1\", &bottleneck_backward_grad_out1, \"Bottleneck block backward\",\n py::call_guard());\n m.def(\"backward_grad_out1_mask\", &bottleneck_backward_grad_out1_mask, \"Bottleneck block backward\",\n py::call_guard());\n m.def(\"backward_grad_out1_halo\", &bottleneck_backward_grad_out1_halo, \"Bottleneck block backward\",\n py::call_guard());\n m.def(\"backward_grad_out1_halo_corr\", &bottleneck_backward_grad_out1_halo_corr, \"Bottleneck block backward\",\n py::call_guard());\n m.def(\"backward_wgrad2_pad\", &bottleneck_backward_wgrad2_pad, \"Bottleneck block backward\",\n py::call_guard());\n m.def(\"backward_wgrad2\", &bottleneck_backward_wgrad2, \"Bottleneck block backward\",\n py::call_guard());\n m.def(\"backward_wgrad2_halo\", &bottleneck_backward_wgrad2_halo, \"Bottleneck block backward\",\n py::call_guard());\n m.def(\"backward_wgrad3\", &bottleneck_backward_wgrad3, \"Bottleneck block backward\",\n py::call_guard());\n m.def(\"backward_wgrad1\", &bottleneck_backward_wgrad1, \"Bottleneck block backward\",\n py::call_guard());\n m.def(\"backward_rest\", &bottleneck_backward_rest, \"Bottleneck block backward\",\n py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/conv_bias_relu/conv_bias_relu.cpp", "content": "#include \n#include // for getcudnnhandle\n#include \n#include \n#include \n\n#include \n#include \n\n#ifdef DEBUG\n#define DEBUG_MSG(str) \\\n do { \\\n std::cout << str << std::endl; \\\n } while (false)\n#else\n#define DEBUG_MSG(str) \\\n do { \\\n } while (false)\n#endif\n\n#ifdef DEBUG_CUDNN\n#define DEBUG_CUDNN_MSG(buf, str) \\\n do { \\\n buf << str << std::endl; \\\n } while (false)\n#else\n#define DEBUG_CUDNN_MSG(buf, str) \\\n do { \\\n } while (false)\n#endif\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(at::MemoryFormat::ChannelsLast), #x \" must be contiguous\")\n#define CHECK_INPUT(x) \\\n CHECK_CUDA(x); \\\n CHECK_CONTIGUOUS(x)\n\n#define checkCudnnErr(...) \\\n do { \\\n int err = checkCudnnError(__VA_ARGS__, #__VA_ARGS__, __FILE__, __LINE__); \\\n if (err) { \\\n return; \\\n } \\\n } while (0)\n\nint checkCudnnError(cudnnStatus_t code, const char* expr, const char* file, int line) {\n if (code) {\n printf(\"CUDNN error at %s:%d, code=%d (%s) in '%s'\\n\", file, line, (int)code, cudnnGetErrorString(code), expr);\n return 1;\n }\n return 0;\n}\n\nvoid checkError(cudaError_t code, char const* func, const char* file, const int line, bool abort = true);\n#define checkCUDAError(val) \\\n { \\\n checkError((val), #val, __FILE__, __LINE__); \\\n } // in-line regular function\n\nvoid checkError(cudaError_t code, char const* func, const char* file, const int line, bool abort) {\n if (code != cudaSuccess) {\n const char* errorMessage = cudaGetErrorString(code);\n fprintf(stderr, \"CUDA error returned from \\\"%s\\\" at %s:%d, Error code: %d (%s)\\n\", func, file, line, code,\n errorMessage);\n if (abort) {\n cudaDeviceReset();\n exit(code);\n }\n }\n}\n\nvoid generateStrides(const int64_t* dimA, int64_t* strideA, int nbDims, cudnnTensorFormat_t filterFormat) {\n // For INT8x4 and INT8x32 we still compute standard strides here to input\n // into the cuDNN functions. We will manually scale by resizeFactor in the cpu ref.\n if (filterFormat == CUDNN_TENSOR_NCHW) {\n strideA[nbDims - 1] = 1;\n for (int64_t d = nbDims - 2; d >= 0; d--) {\n strideA[d] = strideA[d + 1] * dimA[d + 1];\n }\n } else {\n // Here we assume that the format is CUDNN_TENSOR_NHWC\n strideA[1] = 1;\n strideA[nbDims - 1] = strideA[1] * dimA[1];\n for (int64_t d = nbDims - 2; d >= 2; d--) {\n strideA[d] = strideA[d + 1] * dimA[d + 1];\n }\n strideA[0] = strideA[2] * dimA[2];\n }\n}\n\nint getFwdConvDilatedFilterDim(int filterDim, int dilation) { return ((filterDim - 1) * dilation) + 1; }\n\nint getFwdConvPaddedImageDim(int tensorDim, int pad) { return tensorDim + (2 * pad); }\n\nint getFwdConvOutputDim(int tensorDim, int pad, int filterDim, int stride, int dilation) {\n int p = (getFwdConvPaddedImageDim(tensorDim, pad) - getFwdConvDilatedFilterDim(filterDim, dilation)) / stride + 1;\n return (p);\n}\n\n// create a cache for plan\nstd::unordered_map plan_cache;\n\nstd::string getConvFusionString(int64_t* x_dim_padded, int64_t* padA, int64_t* convstrideA, int64_t* dilationA,\n int64_t* w_dim_padded, cudnnDataType_t dataType, std::string fusion_string) {\n for (int i = 0; i < 4; i++) {\n fusion_string += 'X';\n fusion_string += std::to_string(x_dim_padded[i]);\n }\n for (int i = 0; i < 4; i++) {\n fusion_string += 'W';\n fusion_string += std::to_string(w_dim_padded[i]);\n }\n for (int i = 0; i < 2; i++) {\n fusion_string += 'P';\n fusion_string += std::to_string(padA[i]);\n }\n for (int i = 0; i < 2; i++) {\n fusion_string += 'S';\n fusion_string += std::to_string(convstrideA[i]);\n }\n for (int i = 0; i < 2; i++) {\n fusion_string += 'D';\n fusion_string += std::to_string(dilationA[i]);\n }\n fusion_string += 'T';\n fusion_string += std::to_string(dataType);\n return fusion_string;\n}\n\ncudnn_frontend::ExecutionPlan& getOrCreatePlan(cudnnHandle_t handle_, std::stringstream& log_buf,\n cudnn_frontend::OperationGraph& opGraph, std::string cache_string,\n bool use_heuristic = true) {\n auto it = plan_cache.find(cache_string);\n if (it != plan_cache.end()) {\n DEBUG_CUDNN_MSG(log_buf, \"Found plan in cache\");\n return it->second;\n } else {\n DEBUG_CUDNN_MSG(log_buf, \"No plan in cache\");\n if (use_heuristic) {\n // TODO: confirm which mode to use\n auto heuristics = cudnn_frontend::EngineHeuristicsBuilder()\n .setOperationGraph(opGraph)\n .setHeurMode(CUDNN_HEUR_MODE_INSTANT)\n .build();\n auto engine_config_count = heuristics.getEngineConfigCount();\n auto& engine_configs = heuristics.getEngineConfig(engine_config_count);\n for (int64_t count = 0; count < engine_config_count; count++) {\n try {\n plan_cache.emplace(cache_string, std::move(cudnn_frontend::ExecutionPlanBuilder()\n .setHandle(handle_)\n .setEngineConfig(engine_configs[count], opGraph.getTag())\n .build()));\n break;\n } catch (cudnn_frontend::cudnnException e) {\n // Throw exception if all engines failed\n if (count == (engine_config_count - 1)) {\n throw e;\n } else {\n continue;\n }\n }\n }\n } else {\n // How many engines support this operation graph ?\n auto total_engines = opGraph.getEngineCount();\n DEBUG_CUDNN_MSG(log_buf, opGraph.describe() << \" has \" << total_engines << \" engines.\");\n // We have to randomly pick one engine from [0, total_engines)\n // Selecting \"0\" by default\n auto engine = cudnn_frontend::EngineBuilder().setGlobalEngineIdx(0).setOperationGraph(opGraph).build();\n DEBUG_CUDNN_MSG(log_buf, engine.describe());\n auto& knobs = engine.getSupportedKnobs();\n for (auto it = std::begin(knobs); it != std::end(knobs); ++it) {\n DEBUG_CUDNN_MSG(log_buf, it->describe());\n }\n if (knobs.begin() != knobs.end()) {\n DEBUG_CUDNN_MSG(log_buf, \"Updated knob choice\");\n knobs.begin()->setChoice(knobs.begin()->getMinValue() + 1);\n DEBUG_CUDNN_MSG(log_buf, knobs.begin()->describe());\n }\n\n // Createmplacee the requisite engine config\n auto engine_config = cudnn_frontend::EngineConfigBuilder().setEngine(engine).build();\n DEBUG_CUDNN_MSG(log_buf, engine_config.describe());\n plan_cache.emplace(\n cache_string,\n std::move(cudnn_frontend::ExecutionPlanBuilder().setHandle(handle_).setEngineConfig(engine_config).build()));\n }\n\n return plan_cache.find(cache_string)->second;\n }\n}\n\nvoid run_conv_bias(int64_t* x_dim, int64_t* w_dim, int64_t* y_dim, int64_t* conv_pad, int64_t* convstride,\n int64_t* dilation, cudnnDataType_t dataType, at::Half* devPtrX, at::Half* devPtrW, at::Half* devPtrB,\n at::Half* devPtrY) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n\n try {\n int convDim = 2;\n float alpha = 1.0f;\n float beta = 0.0f;\n int64_t b_dim[] = {1, y_dim[1], 1, 1};\n\n // Creates the necessary tensor descriptors\n int64_t stride[4];\n generateStrides(x_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto xTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim)\n .setStrides(4, stride)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, xTensor.describe());\n\n generateStrides(w_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto wTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, w_dim)\n .setStrides(4, stride)\n .setId('w')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, wTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterConvTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('c')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .setVirtual()\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterConvTensor.describe());\n\n generateStrides(b_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto bTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim)\n .setStrides(4, stride)\n .setId('b')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterBiasTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('y')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterBiasTensor.describe());\n\n // Define the bias operation\n auto biasDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(convDim)\n .setStrides(convDim, convstride)\n .setPrePadding(convDim, conv_pad)\n .setPostPadding(convDim, conv_pad)\n .setDilation(convDim, dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)\n .setxDesc(xTensor)\n .setwDesc(wTensor)\n .setyDesc(afterConvTensor)\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // Create a Bias Node.\n auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(conv_op.getOutputTensor())\n .setbDesc(bTensor)\n .setyDesc(afterBiasTensor)\n .setpwDesc(biasDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bias_op.describe());\n\n // Create an Operation Graph. In this case it is convolution bias activation\n std::array ops = {&conv_op, &bias_op};\n\n auto opGraph = cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(2, ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string = getConvFusionString(x_dim, conv_pad, convstride, dilation, w_dim, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrW, devPtrB, devPtrY};\n int64_t uids[] = {'x', 'w', 'b', 'y'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(4, data_ptrs)\n .setUids(4, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_conv_bias_mask_relu(int64_t* x_dim, int64_t* w_dim, int64_t* y_dim, int64_t* conv_pad, int64_t* conv_stride,\n int64_t* conv_dilation, cudnnDataType_t dataType, at::Half* devPtrX, at::Half* devPtrW,\n at::Half* devPtrB, int8_t* devPtrM, at::Half* devPtrY) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n\n try {\n int conv_dim = 2;\n float alpha = 1.0f;\n float beta = 0.0f;\n int64_t b_dim[] = {1, y_dim[1], 1, 1};\n\n // Creates the necessary tensor descriptors\n int64_t stride[4];\n generateStrides(x_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto xTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim)\n .setStrides(4, stride)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, xTensor.describe());\n\n generateStrides(w_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto wTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, w_dim)\n .setStrides(4, stride)\n .setId('w')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, wTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto mTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('m')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_INT8)\n .build();\n DEBUG_CUDNN_MSG(log_buf, wTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterConvTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('c')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .setVirtual()\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterConvTensor.describe());\n\n generateStrides(b_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto bTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim)\n .setStrides(4, stride)\n .setId('b')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterBiasTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('B')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .setVirtual()\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterBiasTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterMaskTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('M')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .setVirtual()\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterBiasTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterReLUTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('y')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterReLUTensor.describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(conv_dim)\n .setStrides(conv_dim, conv_stride)\n .setPrePadding(conv_dim, conv_pad)\n .setPostPadding(conv_dim, conv_pad)\n .setDilation(conv_dim, conv_dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n // Define the bias operation\n auto biasDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());\n\n // Define the mask operation\n auto maskDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_MUL).setMathPrecision(CUDNN_DATA_FLOAT).build();\n\n // Define the activation operation\n auto actDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_RELU_FWD)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, actDesc.describe());\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)\n .setxDesc(xTensor)\n .setwDesc(wTensor)\n .setyDesc(afterConvTensor)\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // Create a Bias Node\n auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(conv_op.getOutputTensor())\n .setbDesc(bTensor)\n .setyDesc(afterBiasTensor)\n .setpwDesc(biasDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bias_op.describe());\n\n // create a Mask Node\n auto mask_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(bias_op.getOutputTensor())\n .setbDesc(mTensor)\n .setyDesc(afterMaskTensor)\n .setpwDesc(maskDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, mask_op.describe());\n\n // Create an Activation Node\n auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(mask_op.getOutputTensor())\n .setyDesc(afterReLUTensor)\n .setpwDesc(actDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, act_op.describe());\n\n // Create an Operation Graph. In this case it is convolution bias activation\n std::array ops = {&conv_op, &bias_op, &mask_op, &act_op};\n\n auto opGraph = cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(4, ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim, conv_pad, conv_stride, conv_dilation, w_dim, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrW, devPtrB, devPtrM, devPtrY};\n int64_t uids[] = {'x', 'w', 'b', 'm', 'y'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(5, data_ptrs)\n .setUids(5, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_conv_cscale_cbias_relu(int64_t* x_dim, int64_t* w_dim, int64_t* y_dim, int64_t* conv_pad, int64_t* conv_stride,\n int64_t* conv_dilation, cudnnDataType_t dataType, at::Half* devPtrX, at::Half* devPtrW,\n at::Half* devPtrS, at::Half* devPtrB, at::Half* devPtrY) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n\n try {\n int conv_dim = 2;\n float alpha = 1.0f;\n float beta = 0.0f;\n int64_t s_dim[] = {1, y_dim[1], 1, 1};\n int64_t b_dim[] = {1, y_dim[1], 1, 1};\n\n // Creates the necessary tensor descriptors\n int64_t stride[4];\n generateStrides(x_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto xTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim)\n .setStrides(4, stride)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, xTensor.describe());\n\n generateStrides(w_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto wTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, w_dim)\n .setStrides(4, stride)\n .setId('w')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, wTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterConvTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('c')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .setVirtual()\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterConvTensor.describe());\n\n generateStrides(s_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto sTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, s_dim)\n .setStrides(4, stride)\n .setId('s')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, sTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterScaleTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('S')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .setVirtual()\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterScaleTensor.describe());\n\n generateStrides(b_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto bTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim)\n .setStrides(4, stride)\n .setId('b')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterBiasTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('B')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .setVirtual()\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterBiasTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterReLUTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('y')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterReLUTensor.describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(conv_dim)\n .setStrides(conv_dim, conv_stride)\n .setPrePadding(conv_dim, conv_pad)\n .setPostPadding(conv_dim, conv_pad)\n .setDilation(conv_dim, conv_dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n // Define the scale operation\n auto scaleDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_MUL).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());\n\n // Define the bias operation\n auto biasDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());\n\n // Define the activation operation\n auto actDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_RELU_FWD)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, actDesc.describe());\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)\n .setxDesc(xTensor)\n .setwDesc(wTensor)\n .setyDesc(afterConvTensor)\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // Create a scale Node.\n auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(conv_op.getOutputTensor())\n .setbDesc(sTensor)\n .setyDesc(afterScaleTensor)\n .setpwDesc(scaleDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, scale_op.describe());\n\n // Create a Bias Node.\n auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(scale_op.getOutputTensor())\n .setbDesc(bTensor)\n .setyDesc(afterBiasTensor)\n .setpwDesc(biasDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bias_op.describe());\n\n // Create an Activation Node.\n auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(bias_op.getOutputTensor())\n .setyDesc(afterReLUTensor)\n .setpwDesc(actDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, act_op.describe());\n\n // Create an Operation Graph. In this case it is convolution bias activation\n std::array ops = {&conv_op, &scale_op, &bias_op, &act_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim, conv_pad, conv_stride, conv_dilation, w_dim, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrW, devPtrS, devPtrB, devPtrY};\n int64_t uids[] = {'x', 'w', 's', 'b', 'y'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(5, data_ptrs)\n .setUids(5, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_conv_bias_relu(int64_t* x_dim, int64_t* w_dim, int64_t* y_dim, int64_t* conv_pad, int64_t* conv_stride,\n int64_t* conv_dilation, cudnnDataType_t dataType, at::Half* devPtrX, at::Half* devPtrW,\n at::Half* devPtrB, at::Half* devPtrY) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n\n try {\n int conv_dim = 2;\n float alpha = 1.0f;\n float beta = 0.0f;\n int64_t b_dim[] = {1, y_dim[1], 1, 1};\n\n // Creates the necessary tensor descriptors\n int64_t stride[4];\n generateStrides(x_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto xTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim)\n .setStrides(4, stride)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, xTensor.describe());\n\n generateStrides(w_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto wTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, w_dim)\n .setStrides(4, stride)\n .setId('w')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, wTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterConvTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('c')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .setVirtual()\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterConvTensor.describe());\n\n generateStrides(b_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto bTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim)\n .setStrides(4, stride)\n .setId('b')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterBiasTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('B')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .setVirtual()\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterBiasTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto afterReLUTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('y')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, afterReLUTensor.describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(conv_dim)\n .setStrides(conv_dim, conv_stride)\n .setPrePadding(conv_dim, conv_pad)\n .setPostPadding(conv_dim, conv_pad)\n .setDilation(conv_dim, conv_dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n // Define the bias operation\n auto biasDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_ADD).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());\n\n // Define the activation operation\n auto actDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_RELU_FWD)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, actDesc.describe());\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)\n .setxDesc(xTensor)\n .setwDesc(wTensor)\n .setyDesc(afterConvTensor)\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // Create a Bias Node.\n auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(conv_op.getOutputTensor())\n .setbDesc(bTensor)\n .setyDesc(afterBiasTensor)\n .setpwDesc(biasDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bias_op.describe());\n\n // Create an Activation Node.\n auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(bias_op.getOutputTensor())\n .setyDesc(afterReLUTensor)\n .setpwDesc(actDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, act_op.describe());\n\n // Create an Operation Graph. In this case it is convolution bias activation\n std::array ops = {&conv_op, &bias_op, &act_op};\n\n auto opGraph = cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(3, ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim, conv_pad, conv_stride, conv_dilation, w_dim, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrW, devPtrB, devPtrY};\n int64_t uids[] = {'x', 'w', 'b', 'y'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(4, data_ptrs)\n .setUids(4, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_drelu_dscale(int64_t* dy_dim, cudnnDataType_t dataType, at::Half* devPtrDY, at::Half* devPtrR,\n at::Half* devPtrS, at::Half* devPtrDX) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n\n try {\n int convDim = 2;\n float alpha = 1.0f;\n float beta = 0.0f;\n int64_t s_dim[] = {1, dy_dim[1], 1, 1};\n\n // Creates the necessary tensor descriptors\n int64_t stride[4];\n generateStrides(dy_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto dyTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, dy_dim)\n .setStrides(4, stride)\n .setId('y')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, dyTensor.describe());\n\n generateStrides(dy_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto rTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, dy_dim)\n .setStrides(4, stride)\n .setId('r')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, rTensor.describe());\n\n generateStrides(dy_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto inActGradTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, dy_dim)\n .setStrides(4, stride)\n .setId('R')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .setVirtual()\n .build();\n DEBUG_CUDNN_MSG(log_buf, inActGradTensor.describe());\n\n generateStrides(s_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto scaleTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, s_dim)\n .setStrides(4, stride)\n .setId('s')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, scaleTensor.describe());\n\n generateStrides(dy_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto dxTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, dy_dim)\n .setStrides(4, stride)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, dxTensor.describe());\n\n // Define the activation backward operation\n auto actDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_RELU_BWD)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, actDesc.describe());\n\n // Define the bias backward operation\n auto scaleDesc =\n cudnn_frontend::PointWiseDescBuilder().setMode(CUDNN_POINTWISE_MUL).setMathPrecision(CUDNN_DATA_FLOAT).build();\n DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());\n\n // Create an relu backward Node\n auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setdyDesc(dyTensor)\n .setxDesc(rTensor)\n .setdxDesc(inActGradTensor)\n .setpwDesc(actDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, act_op.describe());\n\n // Create bias node\n auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setxDesc(inActGradTensor)\n .setbDesc(scaleTensor)\n .setyDesc(dxTensor)\n .setpwDesc(scaleDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, scale_op.describe());\n\n // Create an Operation Graph. In this case it is bias only\n std::array ops = {&act_op, &scale_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n // creating unique dummy values\n int64_t pad_dummy[] = {40, 40};\n int64_t stride_dummy[] = {40, 40};\n int64_t dilation_dummy[] = {40, 40};\n auto cache_string =\n getConvFusionString(dy_dim, pad_dummy, stride_dummy, dilation_dummy, s_dim, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrDY, devPtrR, devPtrS, devPtrDX};\n int64_t uids[] = {'y', 'r', 's', 'x'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(4, data_ptrs)\n .setUids(4, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_drelu_dbias(int64_t* dy_dim, cudnnDataType_t dataType, at::Half* devPtrDY, at::Half* devPtrR,\n at::Half* devPtrDR, float* devPtrDB) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n\n try {\n int convDim = 2;\n float alpha = 1.0f;\n float beta = 0.0f;\n int64_t b_dim[] = {1, dy_dim[1], 1, 1};\n\n // Creates the necessary tensor descriptors\n int64_t stride[4];\n generateStrides(dy_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto dyTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, dy_dim)\n .setStrides(4, stride)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, dyTensor.describe());\n\n generateStrides(dy_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto rTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, dy_dim)\n .setStrides(4, stride)\n .setId('r')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, rTensor.describe());\n\n generateStrides(dy_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto inActGradTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, dy_dim)\n .setStrides(4, stride)\n .setId('R')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, inActGradTensor.describe());\n\n generateStrides(b_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto biasGradTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim)\n .setStrides(4, stride)\n .setId('y')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, biasGradTensor.describe());\n\n // Define the activation backward operation\n auto actDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_RELU_BWD)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, actDesc.describe());\n\n // Define the bias backward operation\n auto biasDesc = cudnn_frontend::ReductionDescBuilder()\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .setReductionOp(CUDNN_REDUCE_TENSOR_ADD)\n .build();\n DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());\n\n // Create an relu backward Node\n auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setdyDesc(dyTensor)\n .setxDesc(rTensor)\n .setdxDesc(inActGradTensor)\n .setpwDesc(actDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, act_op.describe());\n\n // Create bias node\n auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_REDUCTION_DESCRIPTOR)\n .setxDesc(inActGradTensor)\n .setyDesc(biasGradTensor)\n .setreductionDesc(biasDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bias_op.describe());\n\n // Create an Operation Graph. In this case it is bias only\n std::array ops = {&act_op, &bias_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n // creating unique dummy values\n int64_t pad_dummy[] = {20, 20};\n int64_t stride_dummy[] = {20, 20};\n int64_t dilation_dummy[] = {20, 20};\n auto cache_string =\n getConvFusionString(dy_dim, pad_dummy, stride_dummy, dilation_dummy, b_dim, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrDY, devPtrR, devPtrDR, devPtrDB};\n int64_t uids[] = {'x', 'r', 'R', 'y'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(4, data_ptrs)\n .setUids(4, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_dconv_drelu_dbias(int64_t* x_dim, int64_t* w_dim, int64_t* y_dim, int64_t* pad, int64_t* convstride,\n int64_t* dilation, cudnnDataType_t dataType, at::Half* devPtrX, at::Half* devPtrW,\n at::Half* devPtrR, at::Half* devPtrRg, float* devPtrY) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n try {\n int convDim = 2;\n float alpha = 1.0f;\n float beta = 0.0f;\n int64_t b_dim[] = {1, x_dim[1], 1, 1};\n\n int64_t stride[4];\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto outConvGradTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, outConvGradTensor.describe());\n\n generateStrides(w_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto wTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, w_dim)\n .setStrides(4, stride)\n .setId('w')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, wTensor.describe());\n\n generateStrides(x_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto inConvGradTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim)\n .setStrides(4, stride)\n .setId('A')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .setVirtual()\n .build();\n DEBUG_CUDNN_MSG(log_buf, inConvGradTensor.describe());\n\n generateStrides(x_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto rTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim)\n .setStrides(4, stride)\n .setId('r')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, rTensor.describe());\n\n generateStrides(x_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto inReLUGradTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim)\n .setStrides(4, stride)\n .setId('R')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, inReLUGradTensor.describe());\n\n generateStrides(b_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto inBiasGradTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim)\n .setStrides(4, stride)\n .setId('y')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, inBiasGradTensor.describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(convDim)\n .setStrides(convDim, convstride)\n .setPrePadding(convDim, pad)\n .setPostPadding(convDim, pad)\n .setDilation(convDim, dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n // Define the activation backward operation\n auto actDesc = cudnn_frontend::PointWiseDescBuilder()\n .setMode(CUDNN_POINTWISE_RELU_BWD)\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, actDesc.describe());\n\n // Define the bias backward operation\n auto biasDesc = cudnn_frontend::ReductionDescBuilder()\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .setReductionOp(CUDNN_REDUCE_TENSOR_ADD)\n .build();\n DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());\n\n // Create a convolution Node\n auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR)\n .setdyDesc(outConvGradTensor)\n .setwDesc(wTensor)\n .setdxDesc(inConvGradTensor)\n .setcDesc(convDesc)\n .setAlpha(alpha)\n .setBeta(beta)\n .build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // Create an relu backward Node\n auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)\n .setdyDesc(inConvGradTensor)\n .setxDesc(rTensor)\n .setdxDesc(inReLUGradTensor)\n .setpwDesc(actDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, act_op.describe());\n\n // Create bias node\n auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_REDUCTION_DESCRIPTOR)\n .setxDesc(inReLUGradTensor)\n .setyDesc(inBiasGradTensor)\n .setreductionDesc(biasDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bias_op.describe());\n\n // Create an Operation Graph. In this case it is bias only\n std::array ops = {&conv_op, &act_op, &bias_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string = getConvFusionString(x_dim, pad, convstride, dilation, w_dim, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrW, devPtrR, devPtrRg, devPtrY};\n int64_t uids[] = {'x', 'w', 'r', 'R', 'y'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(5, data_ptrs)\n .setUids(5, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_dconv(int64_t* x_dim, int64_t* w_dim, int64_t* y_dim, int64_t* conv_pad, int64_t* conv_stride,\n int64_t* conv_dilation, cudnnDataType_t dataType, at::Half* devPtrX, at::Half* devPtrW,\n at::Half* devPtrY, cudnnBackendDescriptorType_t mode) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n\n try {\n int conv_dim = 2;\n float alpha = 1.0f;\n float beta = 0.0f;\n\n // Define the convolution problem\n int64_t stride[4];\n generateStrides(x_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto xTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim)\n .setStrides(4, stride)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, xTensor.describe());\n\n generateStrides(w_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto wTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, w_dim)\n .setStrides(4, stride)\n .setId('w')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, wTensor.describe());\n\n generateStrides(y_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto yTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, y_dim)\n .setStrides(4, stride)\n .setId('y')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, yTensor.describe());\n\n // Define the convolution problem\n auto convDesc = cudnn_frontend::ConvDescBuilder()\n .setDataType(CUDNN_DATA_FLOAT)\n .setMathMode(CUDNN_CROSS_CORRELATION)\n .setNDims(conv_dim)\n .setStrides(conv_dim, conv_stride)\n .setPrePadding(conv_dim, conv_pad)\n .setPostPadding(conv_dim, conv_pad)\n .setDilation(conv_dim, conv_dilation)\n .build();\n DEBUG_CUDNN_MSG(log_buf, convDesc.describe());\n\n // Create a convolution node\n // mode should be one of following\n // CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR\n // CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR\n auto conv_op_builder = cudnn_frontend::OperationBuilder(mode);\n if (mode == CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR) {\n conv_op_builder.setdxDesc(xTensor).setwDesc(wTensor).setdyDesc(yTensor).setcDesc(convDesc);\n } else {\n conv_op_builder.setxDesc(xTensor).setdwDesc(wTensor).setdyDesc(yTensor).setcDesc(convDesc);\n }\n auto conv_op = conv_op_builder.setAlpha(alpha).setBeta(beta).build();\n DEBUG_CUDNN_MSG(log_buf, conv_op.describe());\n\n // Create an Operation Graph. In this case it is convolution add bias activation\n std::array ops = {&conv_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n auto cache_string =\n getConvFusionString(x_dim, conv_pad, conv_stride, conv_dilation, w_dim, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrW, devPtrY};\n int64_t uids[] = {'x', 'w', 'y'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(3, data_ptrs)\n .setUids(3, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nvoid run_dbias(int64_t* x_dim, cudnnDataType_t dataType, at::Half* devPtrX, float* devPtrY) {\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n std::stringstream log_buf;\n try {\n int convDim = 2;\n int64_t b_dim[] = {1, x_dim[1], 1, 1};\n\n int64_t stride[4];\n generateStrides(x_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto xTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, x_dim)\n .setStrides(4, stride)\n .setId('x')\n .setAlignment(16)\n .setDataType(dataType)\n .build();\n DEBUG_CUDNN_MSG(log_buf, xTensor.describe());\n\n generateStrides(b_dim, stride, 4, CUDNN_TENSOR_NHWC);\n auto yTensor = cudnn_frontend::TensorBuilder()\n .setDim(4, b_dim)\n .setStrides(4, stride)\n .setId('y')\n .setAlignment(16)\n .setDataType(CUDNN_DATA_FLOAT)\n .build();\n DEBUG_CUDNN_MSG(log_buf, yTensor.describe());\n\n // Define the bias backward operation\n auto biasDesc = cudnn_frontend::ReductionDescBuilder()\n .setMathPrecision(CUDNN_DATA_FLOAT)\n .setReductionOp(CUDNN_REDUCE_TENSOR_ADD)\n .build();\n DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());\n\n // Create bias node\n auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_REDUCTION_DESCRIPTOR)\n .setxDesc(xTensor)\n .setyDesc(yTensor)\n .setreductionDesc(biasDesc)\n .build();\n DEBUG_CUDNN_MSG(log_buf, bias_op.describe());\n\n // Create an Operation Graph. In this case it is bias only\n std::array ops = {&bias_op};\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle_).setOperationGraph(ops.size(), ops.data()).build();\n\n // Create string encoding for plan caching\n int64_t pad_dummy[] = {10, 10};\n int64_t stride_dummy[] = {10, 10};\n int64_t dilation_dummy[] = {10, 10};\n auto cache_string =\n getConvFusionString(x_dim, pad_dummy, stride_dummy, dilation_dummy, b_dim, dataType, opGraph.getTag());\n DEBUG_CUDNN_MSG(log_buf, \"[convstring] \" << cache_string);\n\n auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);\n DEBUG_CUDNN_MSG(log_buf, \"Plan tag: \" << plan.getTag());\n\n auto workspace_size = plan.getWorkspaceSize();\n DEBUG_CUDNN_MSG(log_buf, plan.describe() << \" requires workspace \" << workspace_size);\n\n void* workspace_ptr = nullptr;\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n if (workspace_size > 0) {\n workspace_ptr = workspace_tensor.data_ptr();\n }\n void* data_ptrs[] = {devPtrX, devPtrY};\n int64_t uids[] = {'x', 'y'};\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workspace_ptr)\n .setDataPointers(2, data_ptrs)\n .setUids(2, uids)\n .build();\n DEBUG_CUDNN_MSG(log_buf, \"variantPack \" << variantPack.describe());\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n checkCudnnErr(status);\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n } catch (cudnn_frontend::cudnnException e) {\n std::cout << log_buf.str() << \"[ERROR] Exception \" << e.what() << std::endl;\n }\n}\n\nstd::vector conv_bias_mask_relu_forward(std::vector inputs, int64_t padding, int64_t stride) {\n std::cout << std::fixed;\n\n // create output vector\n std::vector outputs;\n auto output_format = at::MemoryFormat::ChannelsLast;\n\n // setup dimensions\n int64_t x_dim[] = {0, 0, 0, 0};\n int64_t w_dim[] = {0, 0, 0, 0};\n\n // All dim calculation after this order of n,c,h,w\n int axis[] = {0, 1, 2, 3};\n for (int dim = 0; dim < 4; dim++) {\n x_dim[dim] = inputs[0].size(axis[dim]);\n w_dim[dim] = inputs[1].size(axis[dim]);\n }\n\n // output dim in n,c,h,w used by backend\n int64_t y_dim[] = {0, 0, 0, 0};\n\n // use these fixed values\n int64_t conv_pad[] = {padding, padding};\n int64_t conv_stride[] = {stride, stride};\n int64_t conv_dilation[] = {1, 1};\n\n // compute output from pad/stride/dilation\n y_dim[0] = x_dim[0];\n y_dim[1] = w_dim[0];\n for (int dim = 0; dim < 2; dim++) {\n y_dim[dim + 2] =\n getFwdConvOutputDim(x_dim[dim + 2], conv_pad[dim], w_dim[dim + 2], conv_stride[dim], conv_dilation[dim]);\n }\n\n // run\n at::Half* x = inputs[0].data_ptr();\n at::Half* w = inputs[1].data_ptr();\n at::Half* b = inputs[2].data_ptr();\n int8_t* m = inputs[3].data_ptr();\n auto out = at::empty(y_dim, inputs[0].type(), output_format);\n at::Half* y = out.data_ptr();\n\n run_conv_bias_mask_relu(x_dim, w_dim, y_dim, conv_pad, conv_stride, conv_dilation, CUDNN_DATA_HALF, x, w, b, m, y);\n\n DEBUG_MSG(\"[DEBUG] conv-bias-mask-relu : \" << y.to(at::kFloat).sum().item());\n\n outputs.push_back(out);\n\n return outputs;\n}\n\nat::Tensor conv_cscale_cbias_relu_forward(std::vector inputs, int64_t padding, int64_t stride) {\n std::cout << std::fixed;\n\n // setup dimensions\n int64_t x_dim[] = {0, 0, 0, 0};\n int64_t w_dim[] = {0, 0, 0, 0};\n\n // All dim calculation after this order of n,c,h,w\n int axis[] = {0, 1, 2, 3};\n for (int dim = 0; dim < 4; dim++) {\n x_dim[dim] = inputs[0].size(axis[dim]);\n w_dim[dim] = inputs[1].size(axis[dim]);\n }\n\n // output dim in n,c,h,w used by backend\n int64_t y_dim[] = {0, 0, 0, 0};\n\n // use these fixed values\n int64_t conv_pad[] = {padding, padding};\n int64_t conv_stride[] = {stride, stride};\n int64_t conv_dilation[] = {1, 1};\n\n // compute output from pad/stride/dilation\n y_dim[0] = x_dim[0];\n y_dim[1] = w_dim[0];\n for (int dim = 0; dim < 2; dim++) {\n y_dim[dim + 2] =\n getFwdConvOutputDim(x_dim[dim + 2], conv_pad[dim], w_dim[dim + 2], conv_stride[dim], conv_dilation[dim]);\n }\n\n // run\n at::Half* x = inputs[0].data_ptr();\n at::Half* w = inputs[1].data_ptr();\n at::Half* s = inputs[2].data_ptr();\n at::Half* b = inputs[3].data_ptr();\n auto out = at::empty(y_dim, inputs[0].type(), at::MemoryFormat::ChannelsLast);\n at::Half* y = out.data_ptr();\n\n run_conv_cscale_cbias_relu(x_dim, w_dim, y_dim, conv_pad, conv_stride, conv_dilation, CUDNN_DATA_HALF, x, w, s, b, y);\n\n DEBUG_MSG(\"[DEBUG] conv-cscale-cbias-relu : \" << y.to(at::kFloat).sum().item());\n\n return out;\n}\n\nstd::vector conv_cscale_cbias_relu_backward(std::vector inputs, int64_t padding,\n int64_t stride) {\n bool requires_grad = inputs[0].requires_grad();\n\n for (int i = 0; i <= 4; i++) {\n CHECK_INPUT(inputs[i]);\n }\n\n std::cout << std::fixed;\n\n // create output vector\n std::vector outputs;\n auto output_format = at::MemoryFormat::ChannelsLast;\n\n // setup dimensions\n int64_t x_dim[] = {0, 0, 0, 0};\n int64_t w_dim[] = {0, 0, 0, 0};\n int64_t y_dim[] = {0, 0, 0, 0};\n\n // All dim calculation after this order of n,c,h,w\n int axis[] = {0, 1, 2, 3};\n for (int dim = 0; dim < 4; dim++) {\n x_dim[dim] = inputs[0].size(axis[dim]);\n w_dim[dim] = inputs[1].size(axis[dim]);\n y_dim[dim] = inputs[3].size(axis[dim]);\n }\n\n int64_t b_dim[] = {1, y_dim[1], 1, 1};\n\n int64_t conv_pad[] = {padding, padding};\n int64_t conv_stride[] = {stride, stride};\n int64_t conv_dilation[] = {1, 1};\n\n // run\n // drelu-dbias\n at::Half* dy = inputs[4].data_ptr();\n at::Half* r = inputs[3].data_ptr();\n auto s = inputs[2].data_ptr();\n auto dscale = at::empty_like(inputs[4]);\n at::Half* ds = dscale.data_ptr();\n\n auto options =\n at::TensorOptions().dtype(at::kFloat).layout(inputs[0].layout()).device(inputs[0].device()).requires_grad(false);\n run_drelu_dscale(y_dim, CUDNN_DATA_HALF, dy, r, s, ds);\n\n // conv wgrad\n at::Half* x = inputs[0].data_ptr();\n auto wgrad = at::empty_like(inputs[1]);\n at::Half* dw = wgrad.data_ptr();\n run_dconv(x_dim, w_dim, y_dim, conv_pad, conv_stride, conv_dilation, CUDNN_DATA_HALF, x, dw, ds,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n\n // conv dgrad\n at::Half* w = inputs[1].data_ptr();\n auto dgrad = at::empty_like(inputs[0]);\n at::Half* dx = dgrad.data_ptr();\n run_dconv(x_dim, w_dim, y_dim, conv_pad, conv_stride, conv_dilation, CUDNN_DATA_HALF, dx, w, ds,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);\n\n outputs.push_back(dgrad);\n outputs.push_back(wgrad);\n\n return outputs;\n}\n\nstd::vector conv_bias_relu_forward(std::vector inputs, int64_t padding, int64_t stride) {\n std::cout << std::fixed;\n\n // create output vector\n std::vector outputs;\n auto output_format = at::MemoryFormat::ChannelsLast;\n\n // setup dimensions\n int64_t x_dim[] = {0, 0, 0, 0};\n int64_t w_dim[] = {0, 0, 0, 0};\n\n // All dim calculation after this order of n,c,h,w\n int axis[] = {0, 1, 2, 3};\n for (int dim = 0; dim < 4; dim++) {\n x_dim[dim] = inputs[0].size(axis[dim]);\n w_dim[dim] = inputs[1].size(axis[dim]);\n }\n\n // output dim in n,c,h,w used by backend\n int64_t y_dim[] = {0, 0, 0, 0};\n\n // use these fixed values\n int64_t conv_pad[] = {padding, padding};\n int64_t conv_stride[] = {stride, stride};\n int64_t conv_dilation[] = {1, 1};\n\n // compute output from pad/stride/dilation\n y_dim[0] = x_dim[0];\n y_dim[1] = w_dim[0];\n for (int dim = 0; dim < 2; dim++) {\n y_dim[dim + 2] =\n getFwdConvOutputDim(x_dim[dim + 2], conv_pad[dim], w_dim[dim + 2], conv_stride[dim], conv_dilation[dim]);\n }\n\n // run\n at::Half* x = inputs[0].data_ptr();\n at::Half* w = inputs[1].data_ptr();\n at::Half* b = inputs[2].data_ptr();\n auto out = at::empty(y_dim, inputs[0].type(), output_format);\n at::Half* y = out.data_ptr();\n\n run_conv_bias_relu(x_dim, w_dim, y_dim, conv_pad, conv_stride, conv_dilation, CUDNN_DATA_HALF, x, w, b, y);\n\n DEBUG_MSG(\"[DEBUG] conv-bias-relu : \" << y.to(at::kFloat).sum().item());\n\n outputs.push_back(out);\n\n return outputs;\n}\n\nstd::vector conv_bias_relu_backward(std::vector inputs, int64_t padding, int64_t stride) {\n bool requires_grad = inputs[0].requires_grad();\n\n for (int i = 0; i <= 3; i++) {\n CHECK_INPUT(inputs[i]);\n }\n\n std::cout << std::fixed;\n\n // create output vector\n std::vector outputs;\n auto output_format = at::MemoryFormat::ChannelsLast;\n\n // setup dimensions\n int64_t x_dim[] = {0, 0, 0, 0};\n int64_t w_dim[] = {0, 0, 0, 0};\n int64_t y_dim[] = {0, 0, 0, 0};\n\n // All dim calculation after this order of n,c,h,w\n int axis[] = {0, 1, 2, 3};\n for (int dim = 0; dim < 4; dim++) {\n x_dim[dim] = inputs[0].size(axis[dim]);\n w_dim[dim] = inputs[1].size(axis[dim]);\n y_dim[dim] = inputs[3].size(axis[dim]);\n }\n\n int64_t b_dim[] = {1, y_dim[1], 1, 1};\n\n int64_t conv_pad[] = {padding, padding};\n int64_t conv_stride[] = {stride, stride};\n int64_t conv_dilation[] = {1, 1};\n\n // run\n // drelu-dbias\n at::Half* dy = inputs[3].data_ptr();\n at::Half* r = inputs[2].data_ptr();\n auto drelu = at::empty_like(inputs[2]);\n at::Half* dr = drelu.data_ptr();\n auto options =\n at::TensorOptions().dtype(at::kFloat).layout(inputs[0].layout()).device(inputs[0].device()).requires_grad(false);\n auto bgrad = at::empty(b_dim, options, output_format);\n float* db = bgrad.data_ptr();\n run_drelu_dbias(y_dim, CUDNN_DATA_HALF, dy, r, dr, db);\n\n // conv wgrad\n at::Half* x = inputs[0].data_ptr();\n auto wgrad = at::empty_like(inputs[1]);\n at::Half* dw = wgrad.data_ptr();\n run_dconv(x_dim, w_dim, y_dim, conv_pad, conv_stride, conv_dilation, CUDNN_DATA_HALF, x, dw, dr,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n\n // conv dgrad\n at::Half* w = inputs[1].data_ptr();\n auto dgrad = at::empty_like(inputs[0]);\n at::Half* dx = dgrad.data_ptr();\n run_dconv(x_dim, w_dim, y_dim, conv_pad, conv_stride, conv_dilation, CUDNN_DATA_HALF, dx, w, dr,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);\n\n outputs.push_back(dgrad);\n outputs.push_back(wgrad);\n outputs.push_back(bgrad);\n\n return outputs;\n}\n\nstd::vector conv_bias_forward(std::vector inputs, int64_t padding, int64_t stride) {\n std::cout << std::fixed;\n\n // create output vector\n std::vector outputs;\n auto output_format = at::MemoryFormat::ChannelsLast;\n\n // setup dimensions\n int64_t x_dim[] = {0, 0, 0, 0};\n int64_t w_dim[] = {0, 0, 0, 0};\n\n // All dim calculation after this order of n,c,h,w\n int axis[] = {0, 1, 2, 3};\n for (int dim = 0; dim < 4; dim++) {\n x_dim[dim] = inputs[0].size(axis[dim]);\n w_dim[dim] = inputs[1].size(axis[dim]);\n }\n\n // output dim in n,c,h,w used by backend\n int64_t y_dim[] = {0, 0, 0, 0};\n\n // use these fixed values\n int64_t conv_pad[] = {padding, padding};\n int64_t conv_stride[] = {stride, stride};\n int64_t conv_dilation[] = {1, 1};\n\n // compute output from pad/stride/dilation\n y_dim[0] = x_dim[0];\n y_dim[1] = w_dim[0];\n for (int dim = 0; dim < 2; dim++) {\n y_dim[dim + 2] =\n getFwdConvOutputDim(x_dim[dim + 2], conv_pad[dim], w_dim[dim + 2], conv_stride[dim], conv_dilation[dim]);\n }\n\n // run\n at::Half* x = inputs[0].data_ptr();\n at::Half* w = inputs[1].data_ptr();\n at::Half* b = inputs[2].data_ptr();\n auto out = at::empty(y_dim, inputs[0].type(), output_format);\n at::Half* y = out.data_ptr();\n\n run_conv_bias(x_dim, w_dim, y_dim, conv_pad, conv_stride, conv_dilation, CUDNN_DATA_HALF, x, w, b, y);\n\n DEBUG_MSG(\"[DEBUG] conv-bias : \" << y.to(at::kFloat).sum().item());\n\n outputs.push_back(out);\n\n return outputs;\n}\n\nstd::vector conv_bias_backward(std::vector inputs, int64_t padding, int64_t stride) {\n bool requires_grad = inputs[0].requires_grad();\n\n for (int i = 0; i <= 2; i++) {\n CHECK_INPUT(inputs[i]);\n }\n\n std::cout << std::fixed;\n\n // create output vector\n std::vector outputs;\n auto output_format = at::MemoryFormat::ChannelsLast;\n\n // setup dimensions\n int64_t x_dim[] = {0, 0, 0, 0};\n int64_t w_dim[] = {0, 0, 0, 0};\n int64_t y_dim[] = {0, 0, 0, 0};\n\n // All dim calculation after this order of n,c,h,w\n int axis[] = {0, 1, 2, 3};\n for (int dim = 0; dim < 4; dim++) {\n x_dim[dim] = inputs[0].size(axis[dim]);\n w_dim[dim] = inputs[1].size(axis[dim]);\n y_dim[dim] = inputs[2].size(axis[dim]);\n }\n\n int64_t b_dim[] = {1, y_dim[1], 1, 1};\n\n int64_t conv_pad[] = {padding, padding};\n int64_t conv_stride[] = {stride, stride};\n int64_t conv_dilation[] = {1, 1};\n\n // run\n // dbias\n at::Half* dy = inputs[2].data_ptr();\n auto options =\n at::TensorOptions().dtype(at::kFloat).layout(inputs[0].layout()).device(inputs[0].device()).requires_grad(false);\n auto bgrad = at::empty(b_dim, options, output_format);\n float* db = bgrad.data_ptr();\n run_dbias(y_dim, CUDNN_DATA_HALF, dy, db);\n\n // conv wgrad\n at::Half* x = inputs[0].data_ptr();\n auto wgrad = at::empty_like(inputs[1]);\n at::Half* dw = wgrad.data_ptr();\n run_dconv(x_dim, w_dim, y_dim, conv_pad, conv_stride, conv_dilation, CUDNN_DATA_HALF, x, dw, dy,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);\n\n // conv dgrad\n at::Half* w = inputs[1].data_ptr();\n auto dgrad = at::empty_like(inputs[0]);\n at::Half* dx = dgrad.data_ptr();\n run_dconv(x_dim, w_dim, y_dim, conv_pad, conv_stride, conv_dilation, CUDNN_DATA_HALF, dx, w, dy,\n CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);\n\n outputs.push_back(dgrad);\n outputs.push_back(wgrad);\n outputs.push_back(bgrad);\n\n return outputs;\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"forward\", &conv_bias_relu_forward, \"Fused Conv-Bias-ReLU forward\", py::call_guard());\n m.def(\"backward\", &conv_bias_relu_backward, \"Fused Conv-Bias-ReLU backward\",\n py::call_guard());\n m.def(\"forward_no_relu\", &conv_bias_forward, \"Fused Conv-Bias forward\", py::call_guard());\n m.def(\"backward_no_relu\", &conv_bias_backward, \"Fused Conv-Bias backward\", py::call_guard());\n m.def(\"forward_mask\", &conv_bias_mask_relu_forward, \"Fused Conv-Bias-Mask-ReLU forward\",\n py::call_guard());\n m.def(\"forward_cscale_cbias_relu\", &conv_cscale_cbias_relu_forward, \"Fused Conv-(const)Scale-(const)Bias-ReLU\",\n py::call_guard());\n m.def(\"backward_cscale_cbias_relu\", &conv_cscale_cbias_relu_backward,\n \"Fused Conv-(const)Scale-(const)Bias-ReLU backward\", py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/cudnn_gbn/cudnn_gbn.cpp", "content": "#include \n#include \n#include \n\n#include \n#include \n\n#include \"norm_sample.h\"\n\n// define this enum:\nenum bn_type { BN_FWD, BN_BWD };\n\n// this is a global variable\nstatic std::map, cudnn_frontend::ExecutionPlan> gbn_plan_cache;\n\nat::Tensor gbn_forward(const at::Tensor& x, const at::Tensor& scale, const at::Tensor& bias,\n const at::Tensor& running_mean, const at::Tensor& running_var, const at::Tensor& minibatch_mean,\n const at::Tensor& minibatch_inv_var, const float momentum, const float epsilon,\n const int64_t bn_group, const int rank_id, const std::vector& peer_buffers) {\n int64_t N = x.size(0);\n int64_t C = x.size(1);\n int64_t H = x.size(2);\n int64_t W = x.size(3);\n\n int64_t tensorDims[] = {N, C, H, W};\n int64_t peerDims[] = {bn_group, 4 * C, 1, 1};\n int64_t perChannelDims[] = {1, C, 1, 1};\n int64_t epsilonDims[] = {1, 1, 1, 1};\n\n // Allocate output tensor\n at::Tensor y = at::empty_like(x);\n\n std::vector void_peer_buffers;\n for (int64_t addr : peer_buffers) {\n void_peer_buffers.push_back((void*)addr);\n }\n\n // we need the peer size for the buffer reset\n size_t peer_size = 1;\n for (size_t i = 0; i < 4; ++i) {\n peer_size *= peerDims[i];\n }\n\n // sanity check\n assert(bn_group == void_peer_buffers.size());\n\n // check if plan already exists\n std::vector fv = {(int64_t)BN_FWD, N, C, H, W, bn_group, (int64_t)CUDNN_DATA_HALF};\n if (gbn_plan_cache.find(fv) == gbn_plan_cache.end()) {\n auto plan = run_batch_norm_forward(tensorDims, perChannelDims, epsilonDims, peerDims, CUDNN_DATA_HALF);\n gbn_plan_cache.emplace(fv, std::move(plan));\n }\n\n // get plan and handle\n auto plan = gbn_plan_cache.find(fv)->second;\n\n // execute\n execute_batch_norm_forward(plan, x.data_ptr(), y.data_ptr(), scale.data_ptr(), bias.data_ptr(),\n running_mean.data_ptr(), running_var.data_ptr(), running_mean.data_ptr(),\n running_var.data_ptr(), minibatch_mean.data_ptr(), minibatch_inv_var.data_ptr(),\n void_peer_buffers, static_cast(epsilon), static_cast(momentum), peer_size,\n rank_id);\n\n return y;\n}\n\nstd::vector gbn_backward(const at::Tensor& x, const at::Tensor& dy, const at::Tensor& scale,\n const at::Tensor& minibatch_mean, const at::Tensor& minibatch_inv_var,\n const float epsilon, const int64_t bn_group, const int rank_id,\n const std::vector& peer_buffers) {\n int64_t N = x.size(0);\n int64_t C = x.size(1);\n int64_t H = x.size(2);\n int64_t W = x.size(3);\n\n int64_t tensorDims[] = {N, C, H, W};\n int64_t peerDims[] = {bn_group, 4 * C, 1, 1};\n int64_t perChannelDims[] = {1, C, 1, 1};\n int64_t epsilonDims[] = {1, 1, 1, 1};\n\n // Allocate output tensor\n // outputs\n at::Tensor x_grad, scale_grad, bias_grad;\n\n // Allocate outputs\n x_grad = at::empty_like(x);\n scale_grad = at::empty_like(scale);\n bias_grad = at::empty_like(scale);\n\n std::vector void_peer_buffers;\n for (int64_t addr : peer_buffers) {\n void_peer_buffers.push_back((void*)addr);\n }\n\n // we need the peer size for the buffer reset\n size_t peer_size = 1;\n for (size_t i = 0; i < 4; ++i) {\n peer_size *= peerDims[i];\n }\n\n assert(bn_group == void_peer_buffers.size());\n\n std::vector fv = {(int64_t)BN_BWD, N, C, H, W, bn_group, (int64_t)CUDNN_DATA_HALF};\n if (gbn_plan_cache.find(fv) == gbn_plan_cache.end()) {\n auto plan = run_batch_norm_backward(tensorDims, perChannelDims, epsilonDims, peerDims, CUDNN_DATA_HALF);\n gbn_plan_cache.emplace(fv, std::move(plan));\n }\n\n // get plan and handle\n auto plan = gbn_plan_cache.find(fv)->second;\n\n // execute\n execute_batch_norm_backward(plan, x.data_ptr(), dy.data_ptr(), scale.data_ptr(), minibatch_mean.data_ptr(),\n minibatch_inv_var.data_ptr(), void_peer_buffers, x_grad.data_ptr(), scale_grad.data_ptr(),\n bias_grad.data_ptr(), static_cast(epsilon), peer_size, rank_id);\n\n return std::vector{x_grad, scale_grad, bias_grad};\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"forward\", &gbn_forward, \"Group batch norm forward\", py::call_guard());\n m.def(\"backward\", &gbn_backward, \"Group batch backward\", py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/cudnn_gbn/norm_sample.cpp", "content": "/*\n * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n *\n * Permission is hereby granted, free of charge, to any person obtaining a\n * copy of this software and associated documentation files (the \"Software\"),\n * to deal in the Software without restriction, including without limitation\n * the rights to use, copy, modify, merge, publish, distribute, sublicense,\n * and/or sell copies of the Software, and to permit persons to whom the\n * Software is furnished to do so, subject to the following conditions:\n *\n * The above copyright notice and this permission notice shall be included in\n * all copies or substantial portions of the Software.\n *\n * THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL\n * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING\n * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER\n * DEALINGS IN THE SOFTWARE.\n */\n\n#include \"norm_sample.h\"\n\n#include // for getcudnnhandle\n#include \n#include \n#include \n\n#include \"cudnn_backend.h\"\n\n// some helpers\nint64_t checkCudaError(cudaError_t code, const char* expr, const char* file, int line) {\n if (code) {\n printf(\"CUDA error at %s:%d, code=%d (%s) in '%s'\", file, line, (int)code, cudaGetErrorString(code), expr);\n return 1;\n }\n return 0;\n}\n\nint64_t checkCudnnError(cudnnStatus_t code, const char* expr, const char* file, int line) {\n if (code) {\n printf(\"CUDNN error at %s:%d, code=%d (%s) in '%s'\\n\", file, line, (int)code, cudnnGetErrorString(code), expr);\n return 1;\n }\n return 0;\n}\n\nbool AllowAll(cudnnBackendDescriptor_t engine_config) {\n (void)engine_config;\n return false;\n}\n\nvoid generateStrides(const int64_t* dimA, int64_t* strideA, int64_t nbDims, cudnnTensorFormat_t filterFormat) {\n // For INT8x4 and INT8x32 we still compute standard strides here to input\n // into the cuDNN functions. We will manually scale by resizeFactor in the cpu ref.\n if (filterFormat == CUDNN_TENSOR_NCHW) {\n strideA[nbDims - 1] = 1;\n for (int64_t d = nbDims - 2; d >= 0; d--) {\n strideA[d] = strideA[d + 1] * dimA[d + 1];\n }\n } else {\n // Here we assume that the format is CUDNN_TENSOR_NHWC\n strideA[1] = 1;\n strideA[nbDims - 1] = strideA[1] * dimA[1];\n for (int64_t d = nbDims - 2; d >= 2; d--) {\n strideA[d] = strideA[d + 1] * dimA[d + 1];\n }\n strideA[0] = strideA[2] * dimA[2];\n }\n}\n\n// runtime\ncudnn_frontend::ExecutionPlan run_batch_norm_forward(int64_t* tensorDims, int64_t* perChannelSum, int64_t* epsilon,\n int64_t* peerDims, cudnnDataType_t data_type) {\n // get the cudnn handle\n cudnnHandle_t handle = torch::native::getCudnnHandle();\n\n // Creates the necessary tensor descriptors\n int64_t tensor_stride[4];\n int64_t stride[4];\n int64_t peer_stride[4];\n\n // NHWC format. GenerateStrides() takes care of this. Howeever, tensor dims should still be NCHW\n generateStrides(tensorDims, tensor_stride, (int64_t)4, CUDNN_TENSOR_NHWC);\n generateStrides(peerDims, peer_stride, (int64_t)4, CUDNN_TENSOR_NHWC);\n\n auto tensor_create = [&tensor_stride, &tensorDims](cudnnDataType_t type, int64_t id) {\n return cudnn_frontend::TensorBuilder()\n .setDim(4, tensorDims)\n .setStrides(4, tensor_stride)\n .setId(id)\n .setAlignment(16)\n .setDataType(type)\n .build();\n };\n\n auto peer_tensor_create = [&peer_stride, &tensorDims](cudnnDataType_t type, int64_t id) {\n return cudnn_frontend::TensorBuilder()\n .setDim(4, tensorDims)\n .setStrides(4, peer_stride)\n .setId(id)\n .setAlignment(16)\n .setDataType(type)\n .build();\n };\n\n generateStrides(perChannelSum, stride, (int64_t)4, CUDNN_TENSOR_NHWC);\n\n auto per_channel_tensor_create = [&stride, &perChannelSum](cudnnDataType_t type, int64_t id) {\n return cudnn_frontend::TensorBuilder()\n .setDim(4, perChannelSum)\n .setStrides(4, stride)\n .setId(id)\n .setAlignment(16)\n .setDataType(type)\n .build();\n };\n\n auto xTensor = tensor_create(data_type, 100);\n auto yTensor = tensor_create(data_type, 101);\n auto scaleTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 102);\n auto biasTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 103);\n auto inMeanTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 104);\n auto inVarTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 105);\n auto outMeanTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 106);\n auto outVarTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 107);\n auto savedMeanTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 108);\n auto savedInvVarTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 109);\n\n int64_t epsilon_stride[4];\n generateStrides(epsilon, epsilon_stride, (int64_t)4, CUDNN_TENSOR_NHWC);\n auto scalar_tensor_create = [&epsilon_stride, &epsilon](cudnnDataType_t type, int64_t id) {\n return cudnn_frontend::TensorBuilder()\n .setDim(4, epsilon)\n .setStrides(4, epsilon_stride)\n .setId(id)\n .setAlignment(16)\n .setDataType(type)\n .setByValue(true)\n .build();\n };\n\n auto epsilonTensor = scalar_tensor_create(CUDNN_DATA_DOUBLE, 110);\n auto expDecayTensor = scalar_tensor_create(CUDNN_DATA_DOUBLE, 111);\n\n // Create the two peer stat tensors. Jump IDs in case we need to add more tensors with UIDs\n std::vector peerStatTensors;\n for (size_t i = 112; i < 112 + peerDims[0]; ++i) {\n peerStatTensors.push_back(peer_tensor_create(CUDNN_DATA_FLOAT, i));\n }\n\n#if (CUDNN_VERSION >= 8500)\n // Batch normalization\n cudnnBackendNormMode_t normalizationMode = CUDNN_BATCH_NORM;\n\n // Forward training\n cudnnBackendNormFwdPhase_t phase = CUDNN_NORM_FWD_TRAINING;\n\n // Create a Finalize node\n auto batch_norm_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_NORM_FORWARD_DESCRIPTOR)\n .setNormalizationMode(normalizationMode)\n .setNormFwdPhase(phase)\n .setxDesc(xTensor)\n .setScaleAndBias(scaleTensor, biasTensor)\n .setPrevRunningMeanAndVar(inMeanTensor, inVarTensor)\n .setNextRunningMeanAndVar(outMeanTensor, outVarTensor)\n .setSavedMeanAndInvVar(savedMeanTensor, savedInvVarTensor)\n .setEpsilonTensor(epsilonTensor)\n .setExpDecayFactorTensor(expDecayTensor)\n .setPeerStatTensor(peerStatTensors)\n .setyDesc(yTensor)\n .build();\n\n std::array ops = {&batch_norm_op};\n#else\n std::array ops = {};\n#endif\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle).setOperationGraph(ops.size(), ops.data()).build();\n // std::cout << opGraph.describe() << std::endl;\n\n cudnn_frontend::EngineConfigList filtered_configs;\n auto statuses = cudnn_frontend::get_heuristics_list<2>({\"heuristics_instant\", \"heuristics_fallback\"}, opGraph,\n ::AllowAll, filtered_configs, true);\n\n // std::cout << \"get_heuristics_list Statuses: \";\n // for (auto i = 0u ; i < statuses.size(); i++) {\n // std::cout << cudnn_frontend::to_string(statuses[i]) << \" \";\n // }\n // std::cout << std::endl;\n // std::cout << \"Filter config list has \" << filtered_configs.size() << \" configurations \" << std::endl;\n\n // some verbose printing:\n // std::cout << \"Tensor shape: (\" << tensorDims[0] << \", \" << tensorDims[1] << \", \" << tensorDims[2] << \", \" <<\n // tensorDims[3] << \")\" << std::endl;\n\n auto plan_builder = [&filtered_configs, &opGraph, &handle]() {\n for (auto i = 0u; i < filtered_configs.size(); i++) {\n try {\n auto plan = cudnn_frontend::ExecutionPlanBuilder()\n .setHandle(handle)\n .setEngineConfig(filtered_configs[i], opGraph.getTag())\n .build();\n return plan;\n } catch (cudnn_frontend::cudnnException& e) {\n continue;\n }\n }\n return cudnn_frontend::ExecutionPlanBuilder()\n .setHandle(handle)\n .setEngineConfig(filtered_configs[0], opGraph.getTag())\n .build();\n };\n\n assert(filtered_configs.size() > 0);\n auto plan = plan_builder();\n\n return plan;\n}\n\nvoid execute_batch_norm_forward(cudnn_frontend::ExecutionPlan plan, void* xDevPtr, void* yDevPtr, void* scaledevPtr,\n void* biasdevPtr, void* in_meandevPtr, void* in_vardevPtr, void* out_meandevPtr,\n void* out_vardevPtr, void* saved_meandevPtr, void* saved_inv_vardevPtr,\n const std::vector& peer_devPtrs, double epsilon_val,\n double exponential_decay_factor, size_t peer_size, int rank_id) {\n // get handle\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n\n // get stream\n cudaStream_t stream;\n cudnnGetStream(handle_, &stream);\n\n try {\n // allocate workspace\n auto workspace_size = plan.getWorkspaceSize();\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n void* workPtr = nullptr;\n if (workspace_size > 0) {\n workPtr = workspace_tensor.data_ptr();\n }\n\n // first the data pointers\n std::vector data_ptrs{\n xDevPtr, yDevPtr, scaledevPtr, biasdevPtr, in_meandevPtr, in_vardevPtr,\n out_meandevPtr, out_vardevPtr, saved_meandevPtr, saved_inv_vardevPtr, &epsilon_val, &exponential_decay_factor};\n data_ptrs.insert(data_ptrs.end(), peer_devPtrs.begin(), peer_devPtrs.end());\n // then the uids\n std::vector uids;\n for (size_t i = 100; i < 100 + data_ptrs.size(); ++i) {\n uids.push_back(i);\n }\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workPtr)\n .setDataPointers(data_ptrs.size(), data_ptrs.data())\n .setUids(uids.size(), uids.data())\n .build();\n // std::cout << \"variantPack \" << variantPack.describe() << std::endl;\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n\n // Reset local communication buffer\n cudaMemsetAsync(peer_devPtrs[rank_id], 0, peer_size * 4, stream);\n\n } catch (cudnn_frontend::cudnnException& e) {\n struct cudaDeviceProp prop;\n checkCudaErr(cudaGetDeviceProperties(&prop, 0));\n if (prop.major == 8) {\n std::cout << \"[ERROR] Exception \" << e.what() << std::endl;\n assert(false);\n }\n }\n}\n\ncudnn_frontend::ExecutionPlan run_batch_norm_backward(int64_t* tensorDims, int64_t* perChannelSum, int64_t* epsilon,\n int64_t* peerDims, cudnnDataType_t data_type) {\n // get cudnn handle\n cudnnHandle_t handle = torch::native::getCudnnHandle();\n\n // Creates the necessary tensor descriptors\n int64_t tensor_stride[4];\n int64_t stride[4];\n int64_t peer_stride[4];\n\n // NHWC format. GenerateStrides() takes care of this. Howeever, tensor dims should still be NCHW\n generateStrides(tensorDims, tensor_stride, (int64_t)4, CUDNN_TENSOR_NHWC);\n generateStrides(peerDims, peer_stride, (int64_t)4, CUDNN_TENSOR_NHWC);\n\n auto tensor_create = [&tensor_stride, &tensorDims](cudnnDataType_t type, int64_t id) {\n return cudnn_frontend::TensorBuilder()\n .setDim(4, tensorDims)\n .setStrides(4, tensor_stride)\n .setId(id)\n .setAlignment(16)\n .setDataType(type)\n .build();\n };\n\n auto peer_tensor_create = [&peer_stride, &peerDims](cudnnDataType_t type, int64_t id) {\n return cudnn_frontend::TensorBuilder()\n .setDim(4, peerDims)\n .setStrides(4, peer_stride)\n .setId(id)\n .setAlignment(16)\n .setDataType(type)\n .build();\n };\n\n generateStrides(perChannelSum, stride, (int64_t)4, CUDNN_TENSOR_NHWC);\n\n auto per_channel_tensor_create = [&stride, &perChannelSum](cudnnDataType_t type, int64_t id) {\n return cudnn_frontend::TensorBuilder()\n .setDim(4, perChannelSum)\n .setStrides(4, stride)\n .setId(id)\n .setAlignment(16)\n .setDataType(type)\n .build();\n };\n\n auto xTensor = tensor_create(data_type, 100);\n auto dyTensor = tensor_create(data_type, 101);\n auto scaleTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 102);\n auto savedMeanTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 103);\n auto savedInvVarTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 104);\n auto dxTensor = tensor_create(data_type, 105);\n auto dScaleTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 106);\n auto dBiasTensor = per_channel_tensor_create(CUDNN_DATA_FLOAT, 107);\n\n int64_t epsilon_stride[4];\n generateStrides(epsilon, epsilon_stride, (int64_t)4, CUDNN_TENSOR_NHWC);\n auto scalar_tensor_create = [&epsilon_stride, &epsilon](cudnnDataType_t type, int64_t id) {\n return cudnn_frontend::TensorBuilder()\n .setDim(4, epsilon)\n .setStrides(4, epsilon_stride)\n .setId(id)\n .setAlignment(16)\n .setDataType(type)\n .setByValue(true)\n .build();\n };\n\n auto epsilonTensor = scalar_tensor_create(CUDNN_DATA_DOUBLE, 108);\n\n std::vector peerStatTensors;\n for (size_t i = 109; i < 109 + peerDims[0]; ++i) {\n peerStatTensors.push_back(peer_tensor_create(CUDNN_DATA_FLOAT, i));\n }\n\n#if (CUDNN_VERSION >= 8500)\n // Batch normalization\n cudnnBackendNormMode_t normalizationMode = CUDNN_BATCH_NORM;\n\n // Create a Finalize node\n auto batch_norm_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_NORM_BACKWARD_DESCRIPTOR)\n .setNormalizationMode(normalizationMode)\n .setxDesc(xTensor)\n .setSavedMeanAndInvVar(savedMeanTensor, savedInvVarTensor)\n .setdyDesc(dyTensor)\n .setScale(scaleTensor)\n .setEpsilonTensor(epsilonTensor)\n .setDScaleAndDBias(dScaleTensor, dBiasTensor)\n .setdxDesc(dxTensor)\n .setPeerStatTensor(peerStatTensors)\n .build();\n\n std::array ops = {&batch_norm_op};\n#else\n std::array ops = {};\n#endif\n\n auto opGraph =\n cudnn_frontend::OperationGraphBuilder().setHandle(handle).setOperationGraph(ops.size(), ops.data()).build();\n // std::cout << opGraph.describe() << std::endl;\n\n cudnn_frontend::EngineConfigList filtered_configs;\n auto statuses = cudnn_frontend::get_heuristics_list<2>({\"heuristics_instant\", \"heuristics_fallback\"}, opGraph,\n ::AllowAll, filtered_configs, true);\n\n auto plan_builder = [&filtered_configs, &opGraph, &handle]() {\n for (auto i = 0u; i < filtered_configs.size(); i++) {\n try {\n auto plan = cudnn_frontend::ExecutionPlanBuilder()\n .setHandle(handle)\n .setEngineConfig(filtered_configs[i], opGraph.getTag())\n .build();\n return plan;\n } catch (cudnn_frontend::cudnnException& e) {\n continue;\n }\n }\n return cudnn_frontend::ExecutionPlanBuilder()\n .setHandle(handle)\n .setEngineConfig(filtered_configs[0], opGraph.getTag())\n .build();\n };\n\n assert(filtered_configs.size() > 0);\n auto plan = plan_builder();\n\n return plan;\n}\n\nvoid execute_batch_norm_backward(cudnn_frontend::ExecutionPlan plan, void* xDevPtr, void* dyDevPtr, void* scaledevPtr,\n void* saved_meandevPtr, void* saved_inv_vardevPtr,\n const std::vector& peer_devPtrs, void* dxDevPtr, void* dscaledevPtr,\n void* dbiasdevPtr, double epsilon_val, size_t peer_size, int rank_id) {\n // get handle\n cudnnHandle_t handle_ = torch::native::getCudnnHandle();\n\n // get stream\n cudaStream_t stream;\n cudnnGetStream(handle_, &stream);\n\n try {\n // allocate workspace\n auto workspace_size = plan.getWorkspaceSize();\n auto workspace_tensor = at::empty({(workspace_size + 3) / 4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));\n void* workPtr = nullptr;\n if (workspace_size > 0) {\n workPtr = workspace_tensor.data_ptr();\n }\n\n // create helper arrays\n std::vector data_ptrs{xDevPtr, dyDevPtr, scaledevPtr, saved_meandevPtr, saved_inv_vardevPtr,\n dxDevPtr, dscaledevPtr, dbiasdevPtr, &epsilon_val};\n data_ptrs.insert(data_ptrs.end(), peer_devPtrs.begin(), peer_devPtrs.end());\n std::vector uids;\n for (size_t i = 100; i < 100 + data_ptrs.size(); ++i) {\n uids.push_back(i);\n }\n\n auto variantPack = cudnn_frontend::VariantPackBuilder()\n .setWorkspacePointer(workPtr)\n .setDataPointers(data_ptrs.size(), data_ptrs.data())\n .setUids(uids.size(), uids.data())\n .build();\n cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());\n\n cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, \"Plan execute error\", status);\n\n // Reset local communication buffer\n cudaMemsetAsync(peer_devPtrs[rank_id], 0, peer_size * 4, stream);\n\n } catch (cudnn_frontend::cudnnException& e) {\n struct cudaDeviceProp prop;\n checkCudaErr(cudaGetDeviceProperties(&prop, 0));\n if (prop.major == 8) {\n std::cout << \"[ERROR] Exception \" << e.what() << std::endl;\n assert(false);\n }\n }\n}\n" }, { "path": "apex/contrib/csrc/cudnn_gbn/norm_sample.h", "content": "#pragma once\n\n/*\n * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n *\n * Permission is hereby granted, free of charge, to any person obtaining a\n * copy of this software and associated documentation files (the \"Software\"),\n * to deal in the Software without restriction, including without limitation\n * the rights to use, copy, modify, merge, publish, distribute, sublicense,\n * and/or sell copies of the Software, and to permit persons to whom the\n * Software is furnished to do so, subject to the following conditions:\n *\n * The above copyright notice and this permission notice shall be included in\n * all copies or substantial portions of the Software.\n *\n * THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL\n * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING\n * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER\n * DEALINGS IN THE SOFTWARE.\n */\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n/* some helpers\n */\nvoid generateStrides(const int64_t* dimA, int64_t* strideA, int64_t nbDims, cudnnTensorFormat_t filterFormat);\n\nint64_t checkCudaError(cudaError_t code, const char* expr, const char* file, int line);\nint64_t checkCudnnError(cudnnStatus_t code, const char* expr, const char* file, int line);\n\n#define checkCudaErr(...) \\\n do { \\\n int64_t err = checkCudaError(__VA_ARGS__, #__VA_ARGS__, __FILE__, __LINE__); \\\n assert(err == 0); \\\n } while (0)\n\n#define checkCudnnErr(...) \\\n do { \\\n int64_t err = checkCudnnError(__VA_ARGS__, #__VA_ARGS__, __FILE__, __LINE__); \\\n assert(err == 0); \\\n } while (0)\n\n/**\n * @brief Run a Group BN forward sample with 2 peer stat tensors.\n *\n * @param tensorDims an array with shape (N, C, H, W) for input tensor dims. Stride in NHWC or NCHW will take care of\n memory format\n * @param perChannelSum an array with shape (1, C, 1, 1) to denote the sum values for each channel in the input tensor\n * @param epsilon a scalar array with shape (1, 1, 1, 1) to represent the epsilon value for the BN\n * @param peerDims an array with shape (num GPUs, 2 * C, 1, 1) to denote the tensor dimensions for peer stat tensor in\n GBN\n\n *\n */\ncudnn_frontend::ExecutionPlan run_batch_norm_forward(int64_t* tensorDims, int64_t* perChannelSum, int64_t* epsilon,\n int64_t* peerDims, cudnnDataType_t in_out_data_type);\n/**\n * @param xDevPtr input tensor device pointer\n * @param yDevPtr output tensor device pointer\n * @param scaledevPtr input scale device pointer for BN scaling\n * @param biasdevPtr input scale device pointer for BN bias\n * @param in_meandevPtr Input mean device pointer\n * @param in_vardevPtr Input variance device pointer\n * @param out_meandevPtr output mean device pointer\n * @param out_vardevPtr output variance device pointer\n * @param saved_meandevPtr saved mean device pointer for BN backward\n * @param saved_inv_vardevPtr saved inverse variance device pointer for BN backward\n * @param peer_devPtr1 peer stat tensor 1 device pointer\n * @param peer_devPtr2 peer stat tensor 2 device pointer\n * @param epsilon_val episilon value as a double\n * @param exponential_decay_factor exponential_decay_factor as a value\n *\n **/\nvoid execute_batch_norm_forward(cudnn_frontend::ExecutionPlan plan, void* xDevPtr, void* yDevPtr, void* scaledevPtr,\n void* biasdevPtr, void* in_meandevPtr, void* in_vardevPtr, void* out_meandevPtr,\n void* out_vardevPtr, void* saved_meandevPtr, void* saved_inv_vardevPtr,\n const std::vector& peer_devPtrs, double epsilon_val,\n double exponential_decay_factor, size_t peer_size, int rank_id);\n\n/**\n * @brief Run a Group BN backward sample with 2 peer stat tensors.\n *\n * @param tensorDims an array with shape (N, C, H, W) for input tensor dims. Stride in NHWC or NCHW will take care of\n * memory format\n * @param perChannelSum an array with shape (1, C, 1, 1) to denote the sum values for each channel in the input tensor\n * @param epsilon a scalar array with shape (1, 1, 1, 1) to represent the epsilon value for the BN\n * @param peerDims an array with shape (num GPUs, 2 * C, 1, 1) to denote the tensor dimensions for peer stat tensor in\n * GBN\n *\n */\ncudnn_frontend::ExecutionPlan run_batch_norm_backward(int64_t* tensorDims, int64_t* perChannelSum, int64_t* epsilon,\n int64_t* peerDims, cudnnDataType_t data_type);\n\n/**\n * @brief Run a Group BN backward sample with 2 peer stat tensors.\n *\n * @param xDevPtr input tensor device pointer\n * @param yDevPtr output tensor device pointer\n * @param scaledevPtr input scale device pointer for BN scaling\n * @param biasdevPtr input scale device pointer for BN bias\n * @param in_meandevPtr Input mean device pointer\n * @param in_vardevPtr Input variance device pointer\n * @param out_meandevPtr output mean device pointer\n * @param out_vardevPtr output variance device pointer\n * @param saved_meandevPtr saved mean device pointer for BN backward\n * @param saved_inv_vardevPtr saved inverse variance device pointer for BN backward\n * @param peer_devPtr1 peer stat tensor 1 device pointer\n * @param peer_devPtr2 peer stat tensor 2 device pointer\n * @param epsilon_val episilon value as a double\n *\n */\nvoid execute_batch_norm_backward(cudnn_frontend::ExecutionPlan plan, void* xDevPtr, void* dyDevPtr, void* scaledevPtr,\n void* saved_meandevPtr, void* saved_inv_vardevPtr,\n const std::vector& peer_devPtrs, void* dxDevPtr, void* dscaledevPtr,\n void* dbiasdevPtr, double epsilon_val, size_t peer_size, int rank_id);\n" }, { "path": "apex/contrib/csrc/fmha/fmha_api.cpp", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n\n#include \"fmha.h\"\n\nextern at::Tensor& mha_fill(at::Tensor& self, const at::Tensor& start_index);\nvoid set_params(Fused_multihead_attention_fprop_params& params,\n // sizes\n const size_t b, const size_t s, const size_t h, const size_t d,\n // device pointers\n void* qkv_packed_d, void* cu_seqlens_d, void* o_packed_d, void* s_d, float p_dropout) {\n Data_type acc_type = DATA_TYPE_FP32;\n Data_type data_type = DATA_TYPE_FP16;\n\n // Reset the parameters\n memset(¶ms, 0, sizeof(params));\n\n // Set the pointers and strides.\n params.qkv_ptr = qkv_packed_d;\n params.qkv_stride_in_bytes = get_size_in_bytes(h * 3 * d, data_type);\n params.o_ptr = o_packed_d;\n params.o_stride_in_bytes = get_size_in_bytes(h * d, data_type);\n\n params.cu_seqlens = static_cast(cu_seqlens_d);\n\n // S = softmax(P)\n params.s_ptr = s_d;\n params.s_stride_in_bytes = get_size_in_bytes(b * h * s, data_type);\n\n // Set the dimensions.\n params.b = b;\n params.h = h;\n params.s = s;\n params.d = d;\n\n // Set the different scale values.\n const float scale_bmm1 = 1.f / sqrtf(d);\n constexpr float scale_softmax = 1.f;\n constexpr float scale_bmm2 = 1.f;\n\n set_alpha(params.scale_bmm1, scale_bmm1, data_type);\n set_alpha(params.scale_softmax, scale_softmax, acc_type);\n set_alpha(params.scale_bmm2, scale_bmm2, data_type);\n\n // Set this to probability of keeping an element to simplify things.\n params.p_dropout = 1.f - p_dropout;\n params.rp_dropout = 1.f / params.p_dropout;\n TORCH_CHECK(p_dropout < 1.f);\n set_alpha(params.scale_dropout, params.rp_dropout, data_type);\n}\n\nstd::vector mha_fwd(\n const at::Tensor& qkv, // total x num_heads x 3 x head_size, total := \\sum_{i=0}^{b} s_i\n const at::Tensor& cu_seqlens, // b+1\n const float p_dropout, const int max_seq_len, const bool is_training, const bool is_nl, const bool zero_tensors,\n c10::optional gen_) {\n using namespace torch::indexing;\n auto dprops = at::cuda::getCurrentDeviceProperties();\n TORCH_CHECK((dprops->major == 8 && dprops->minor == 0) || (dprops->major == 9 && dprops->minor == 0) ||\n (dprops->major == 10 && dprops->minor == 0) || (dprops->major == 12 && dprops->minor == 0));\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n Launch_params launch_params(dprops, stream, is_training, is_nl);\n\n int seq_len = 512;\n auto launch = &run_fmha_fp16_512_64_sm80;\n if (max_seq_len <= 128) {\n seq_len = 128;\n launch = &run_fmha_fp16_128_64_sm80;\n } else if (max_seq_len <= 256) {\n seq_len = 256;\n launch = &run_fmha_fp16_256_64_sm80;\n } else if (max_seq_len <= 384) {\n seq_len = 384;\n launch = &run_fmha_fp16_384_64_sm80;\n } else if (max_seq_len <= 512) {\n seq_len = 512;\n launch = &run_fmha_fp16_512_64_sm80;\n } else {\n TORCH_CHECK(false);\n }\n\n TORCH_CHECK(qkv.is_cuda())\n TORCH_CHECK(cu_seqlens.is_cuda())\n\n TORCH_CHECK(qkv.is_contiguous())\n TORCH_CHECK(cu_seqlens.is_contiguous())\n\n TORCH_CHECK(cu_seqlens.dim() == 1);\n TORCH_CHECK(qkv.dim() == 4);\n\n const auto sizes = qkv.sizes();\n\n TORCH_CHECK(sizes[THREE_DIM] == 3);\n\n const int batch_size = cu_seqlens.numel() - 1;\n const int total = sizes[TOTAL_DIM];\n const int num_heads = sizes[H_DIM];\n const int head_size = sizes[D_DIM];\n TORCH_CHECK(batch_size > 0);\n TORCH_CHECK(head_size == 64);\n auto opts = qkv.options();\n\n auto ctx = torch::empty({total, num_heads, head_size}, opts);\n\n auto s = torch::empty({batch_size, num_heads, seq_len, seq_len}, opts);\n\n if (zero_tensors) {\n mha_fill(ctx, cu_seqlens.index({Slice(-1, None)}));\n }\n\n auto gen = at::get_generator_or_default(gen_, at::cuda::detail::getDefaultCUDAGenerator());\n\n set_params(launch_params.params, batch_size, seq_len, num_heads, head_size, qkv.data_ptr(), cu_seqlens.data_ptr(),\n ctx.data_ptr(), s.data_ptr(), p_dropout);\n\n launch(launch_params, /*configure=*/true);\n // number of times random will be generated per thread, to offset philox counter in thc random\n // state\n int64_t counter_offset = launch_params.elts_per_thread;\n at::PhiloxCudaState rng_engine_inputs;\n\n if (is_training) {\n // See Note [Acquire lock when using random generators]\n std::lock_guard lock(gen->mutex_);\n launch_params.params.philox_args = gen->philox_cuda_state(counter_offset);\n }\n\n launch(launch_params, /*configure=*/false);\n\n return {ctx, s};\n}\n\nstd::vector mha_bwd(\n const at::Tensor& dout, // total x num_heads, x head_size\n const at::Tensor& qkv, // total x num_heads x 3 x head_size, total := \\sum_{i=0}^{b} s_i\n at::Tensor& softmax, // b x h x s x s softmax and dmask - will be overwritten with dP\n const at::Tensor& cu_seqlens, // b+1\n const float p_dropout, // probability to drop\n const int max_seq_len, // max sequence length to choose the kernel\n const bool zero_tensors) {\n using namespace torch::indexing;\n auto dprops = at::cuda::getCurrentDeviceProperties();\n TORCH_CHECK((dprops->major == 8 && dprops->minor == 0) || (dprops->major == 9 && dprops->minor == 0) ||\n (dprops->major == 10 && dprops->minor == 0) || (dprops->major == 12 && dprops->minor == 0));\n int seq_len = 512;\n auto launch = &run_fmha_dgrad_fp16_512_64_sm80;\n if (max_seq_len <= 128) {\n seq_len = 128;\n launch = &run_fmha_dgrad_fp16_128_64_sm80;\n } else if (max_seq_len <= 256) {\n seq_len = 256;\n launch = &run_fmha_dgrad_fp16_256_64_sm80;\n } else if (max_seq_len <= 384) {\n seq_len = 384;\n launch = &run_fmha_dgrad_fp16_384_64_sm80;\n } else if (max_seq_len <= 512) {\n seq_len = 512;\n launch = &run_fmha_dgrad_fp16_512_64_sm80;\n } else {\n TORCH_CHECK(false);\n }\n\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n\n TORCH_CHECK(qkv.dtype() == torch::kFloat16);\n TORCH_CHECK(dout.dtype() == torch::kFloat16);\n TORCH_CHECK(softmax.dtype() == torch::kFloat16);\n TORCH_CHECK(cu_seqlens.dtype() == torch::kInt32);\n\n TORCH_CHECK(qkv.is_cuda());\n TORCH_CHECK(cu_seqlens.is_cuda());\n\n TORCH_CHECK(qkv.is_contiguous());\n TORCH_CHECK(cu_seqlens.is_contiguous());\n\n TORCH_CHECK(cu_seqlens.dim() == 1);\n TORCH_CHECK(qkv.dim() == 4);\n\n const auto sizes = qkv.sizes();\n\n TORCH_CHECK(sizes[THREE_DIM] == 3);\n\n const int batch_size = cu_seqlens.numel() - 1;\n const int num_heads = sizes[H_DIM];\n const int head_size = sizes[D_DIM];\n TORCH_CHECK(batch_size > 0);\n TORCH_CHECK(head_size == 64);\n\n auto dqkv = torch::empty_like(qkv);\n\n if (zero_tensors) {\n mha_fill(dqkv, cu_seqlens.index({Slice(-1, None)}));\n }\n\n Fused_multihead_attention_fprop_params params;\n\n set_params(params, batch_size, seq_len, num_heads, head_size, qkv.data_ptr(), cu_seqlens.data_ptr(),\n dout.data_ptr(), // we set o_ptr to dout\n softmax.data_ptr(), // softmax gets overwritten by dP!\n p_dropout);\n\n // we're re-using these scales\n Data_type acc_type = DATA_TYPE_FP32;\n set_alpha(params.scale_bmm1, 1.f, acc_type);\n set_alpha(params.scale_softmax, 1.f / sqrtf(head_size), acc_type);\n set_alpha(params.scale_bmm2, 1.f, DATA_TYPE_FP16);\n params.dqkv_ptr = dqkv.data_ptr();\n\n launch(params, stream);\n return {dqkv, softmax};\n}\n\nstd::vector mha_bwd_nl(\n const at::Tensor& dout, // total x num_heads, x head_size\n const at::Tensor& qkv, // total x num_heads x 3 x head_size, total := \\sum_{i=0}^{b} s_i\n at::Tensor& softmax, // b x h x s x s softmax and dmask - will be overwritten with dP\n const at::Tensor& cu_seqlens, // b+1\n const float p_dropout, // probability to drop\n const int max_seq_len, // max sequence length to choose the kernel\n const bool zero_tensors) {\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n\n TORCH_CHECK(qkv.is_cuda())\n TORCH_CHECK(cu_seqlens.is_cuda())\n\n TORCH_CHECK(qkv.is_contiguous())\n TORCH_CHECK(cu_seqlens.is_contiguous())\n\n TORCH_CHECK(cu_seqlens.dim() == 1);\n\n TORCH_CHECK(qkv.dim() == 4);\n\n const auto sizes = qkv.sizes();\n\n TORCH_CHECK(sizes[THREE_DIM] == 3);\n\n const int batch_size = cu_seqlens.numel() - 1;\n\n const int total = sizes[TOTAL_DIM];\n const int num_heads = sizes[H_DIM];\n const int head_size = sizes[D_DIM];\n TORCH_CHECK(batch_size > 0);\n TORCH_CHECK(head_size == 64);\n\n int seq_len = 512;\n auto launch = &run_fmha_dgrad_fp16_512_64_sm80_nl;\n\n auto opts = qkv.options();\n\n auto dqkv = torch::empty_like(qkv);\n\n if (zero_tensors) {\n dqkv.zero_();\n }\n\n int num_chunks = 2;\n if (batch_size == 1) {\n num_chunks = 4;\n } else if (batch_size == 2) {\n num_chunks = 3;\n }\n auto dkv = torch::empty({total, num_chunks, 2, num_heads, head_size}, opts);\n\n Fused_multihead_attention_fprop_params params;\n\n set_params(params, batch_size, seq_len, num_heads, head_size, qkv.data_ptr(), cu_seqlens.data_ptr(),\n dout.data_ptr(), // o_ptr = dout\n softmax.data_ptr(), // softmax gets overwritten by dP!\n p_dropout);\n\n params.dkv_ptr = dkv.data_ptr();\n\n Data_type acc_type = DATA_TYPE_FP32;\n set_alpha(params.scale_bmm1, 1.f, acc_type);\n set_alpha(params.scale_softmax, 1.f / sqrtf(head_size), acc_type);\n set_alpha(params.scale_bmm2, 1.f, DATA_TYPE_FP16);\n params.dqkv_ptr = dqkv.data_ptr();\n\n launch(params, num_chunks, stream);\n\n // SPLIT-K reduction of num_chunks dK, dV parts\n\n // The equivalent of the following Pytorch code:\n // using namespace torch::indexing;\n // at::Tensor view_out = dqkv.index({Slice(), Slice(1, None, None)});\n // torch::sum_out(view_out, dkv, 1);\n\n const int hidden_size = num_heads * head_size;\n fmha_run_noloop_reduce(dqkv.data_ptr(), dkv.data_ptr(), cu_seqlens.data_ptr(), hidden_size, batch_size, total,\n num_chunks, stream);\n\n return {dqkv, softmax, dkv};\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.doc() = \"Fused Multi-head Self-attention for BERT\";\n m.def(\"fwd\", &mha_fwd, \"Forward pass\", py::call_guard());\n m.def(\"bwd\", &mha_bwd, \"Backward pass\", py::call_guard());\n m.def(\"bwd_nl\", &mha_bwd_nl, \"Backward pass (small-batch)\", py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha/gemm.h", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\n#include \n\n#define FMHA_DIV_UP(m, n) (((m) + (n) - 1) / (n))\n\nnamespace fmha {\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Fragment_base_ {\n // The data type.\n using Data_type = Data_type_;\n // default input type\n using Input_type_ = Data_type_;\n // Does it store the array of elements.\n enum { HAS_ELTS = BITS_PER_ELT_ >= 8 };\n // The number of elements.\n enum { NUM_ELTS = NUM_ELTS_ };\n // The size of element in bits.\n enum { BITS_PER_ELT = BITS_PER_ELT_ };\n // The size of byte of a single register.\n enum { BYTES_PER_REG = 4 };\n // The size in bits.\n enum { BITS_PER_REG = BYTES_PER_REG * 8 };\n // The number of registers needed to store the fragment.\n enum { NUM_REGS = Div_up::VALUE };\n // The size in bytes (as returned by sizeof(Fragment_base<>).\n enum { SIZE_IN_BYTES = NUM_REGS * BYTES_PER_REG };\n // The alignment.\n enum { ALIGNMENT = ALIGNMENT_ > 0 ? ALIGNMENT_ : Min::VALUE };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The type of the elements.\n typename Data_type_,\n // The number of elements.\n int NUM_ELTS_,\n // The alignment if you want to force a value -- use 0 otherwise.\n int ALIGNMENT_ = 0,\n // The base class.\n typename Base_ = Fragment_base_ >\nstruct alignas(static_cast(Base_::ALIGNMENT)) Fragment : public Base_ {\n // The size of a load/store.\n enum { BYTES_PER_LOAD_STORE = Base_::NUM_REGS * sizeof(uint32_t) };\n\n // Clear the fragment. Using PTX in that code seems to produce better SASS...\n inline __device__ void clear() {\n#pragma unroll\n for (int ii = 0; ii < Base_::NUM_REGS; ++ii) {\n asm volatile(\"mov.u32 %0, 0; \\n\" : \"=r\"(this->reg(ii)) :);\n }\n }\n\n // Immutable access to a register.\n inline __device__ const uint32_t& reg(int ii) const { return this->regs_[ii]; }\n\n // Mutable access to a register.\n inline __device__ uint32_t& reg(int ii) { return this->regs_[ii]; }\n\n uint32_t regs_[Base_::NUM_REGS];\n\n // Immutable access to the elements.\n inline __device__ const Data_type_& elt(int ii) const {\n return reinterpret_cast(&this->regs_[0])[ii];\n }\n\n // Mutable access to the elements.\n inline __device__ Data_type_& elt(int ii) { return reinterpret_cast(&this->regs_[0])[ii]; }\n\n // Immutable access to the elements with a cast.\n template \n inline __device__ const Cast_type& elt_as(int ii) const {\n return reinterpret_cast(&this->regs_[0])[ii];\n }\n\n // Mutable access to the elements.\n template \n inline __device__ Cast_type& elt_as(int ii) {\n return reinterpret_cast(&this->regs_[0])[ii];\n }\n\n // Add another fragment.\n inline __device__ void add(const Fragment& other) {\n#pragma unroll\n for (int ii = 0; ii < NUM_ELTS_; ++ii) {\n this->elt(ii) += other.elt(ii);\n }\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Fragment_a : public Fragment {};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Fragment_b : public Fragment {};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Fragment_accumulator : public Fragment {\n // The base class.\n using Base = Fragment;\n\n // Add two fragments.\n template \n inline __device__ void add(const Other_fragment_& other) {\n for (int ii = 0; ii < Base::NUM_ELTS; ++ii) {\n this->elt(ii) = this->elt(ii) + other.elt(ii);\n }\n }\n\n // Do the HMMA.\n template \n inline __device__ void mma(const Fragment_a& a, const Fragment_b& b) {\n asm volatile(\n \"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 \\n\"\n \" {%0, %1, %2, %3}, \\n\"\n \" {%4, %5, %6, %7}, \\n\"\n \" {%8, %9}, \\n\"\n \" {%0, %1, %2, %3}; \\n\"\n : \"+f\"(elt(0)), \"+f\"(elt(1)), \"+f\"(elt(2)), \"+f\"(elt(3))\n : \"r\"(a.reg(0)), \"r\"(a.reg(1)), \"r\"(a.reg(2)), \"r\"(a.reg(3)), \"r\"(b.reg(0)), \"r\"(b.reg(1)));\n asm volatile(\n \"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 \\n\"\n \" {%0, %1, %2, %3}, \\n\"\n \" {%4, %5, %6, %7}, \\n\"\n \" {%8, %9}, \\n\"\n \" {%0, %1, %2, %3}; \\n\"\n : \"+f\"(elt(4)), \"+f\"(elt(5)), \"+f\"(elt(6)), \"+f\"(elt(7))\n : \"r\"(a.reg(0)), \"r\"(a.reg(1)), \"r\"(a.reg(2)), \"r\"(a.reg(3)), \"r\"(b.reg(2)), \"r\"(b.reg(3)));\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void clear(Fragment (&frag)[M][N]) {\n#pragma unroll\n for (int mi = 0; mi < M; ++mi) {\n#pragma unroll\n for (int ni = 0; ni < N; ++ni) {\n frag[mi][ni].clear();\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Clear_accumulator {};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Clear_accumulator {\n template \n static inline __device__ void apply(Acc (&acc)[M][N], bool = false) {\n fmha::clear(acc);\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void gemm(Acc (&acc)[M][N], const A (&a)[M], const B (&b)[N]) {\n#pragma unroll\n for (int mi = 0; mi < M; ++mi) {\n#pragma unroll\n for (int ni = 0; ni < N; ++ni) {\n acc[mi][ni].mma(a[mi], b[ni]);\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The number of rows in the CTA tile.\n int M_,\n // The number of cols in the CTA tile.\n int N_,\n // The number of elements in the the K dimension of the GEMM loop.\n int K_,\n // The number of rows of warps.\n int WARPS_M_,\n // The number of cols of warps.\n int WARPS_N_,\n // The number of warps in the K dimension of the GEMM loop.\n int WARPS_K_>\nstruct Cta_tile_ {\n enum { M = M_, N = N_, K = K_ };\n // The number of warps.\n enum { WARPS_M = WARPS_M_, WARPS_N = WARPS_N_, WARPS_K = WARPS_K_ };\n // The number of warps per CTA.\n enum { WARPS_PER_CTA = WARPS_M * WARPS_N * WARPS_K };\n // The number of threads per warp.\n enum { THREADS_PER_WARP = 32 };\n // The number of threads per CTA.\n enum { THREADS_PER_CTA = WARPS_PER_CTA * THREADS_PER_WARP };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Hmma_tile {\n // The number of elements computed with a single warp-MMA.\n enum { M_PER_MMA = 16, N_PER_MMA = 16, K_PER_MMA = 16 };\n\n // The number of elements computed with a single CTA-MMA.\n enum {\n M_PER_MMA_PER_CTA = M_PER_MMA * Cta_tile::WARPS_M,\n N_PER_MMA_PER_CTA = N_PER_MMA * Cta_tile::WARPS_N,\n K_PER_MMA_PER_CTA = K_PER_MMA * Cta_tile::WARPS_K\n };\n\n // The number of MMAs needed to compute the GEMM.\n enum {\n MMAS_M = Div_up::VALUE,\n MMAS_N = Div_up::VALUE,\n MMAS_K = Div_up::VALUE,\n };\n\n // The number of elements computed per warp.\n enum {\n M_PER_WARP = MMAS_M * M_PER_MMA,\n N_PER_WARP = MMAS_N * N_PER_MMA,\n K_PER_WARP = MMAS_K * K_PER_MMA,\n };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nusing A_type = uint16_t;\nusing B_type = uint16_t;\nusing C_type = uint16_t;\nusing Accumulator_type = float;\nusing Epilogue_type = float;\n\nconstexpr int BITS_PER_ELEMENT_A = sizeof(A_type) * 8;\nconstexpr int BITS_PER_ELEMENT_B = sizeof(B_type) * 8;\nconstexpr int BITS_PER_ELEMENT_C = sizeof(C_type) * 8;\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nusing Cta_tile_extd = Cta_tile_;\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nusing Cta_tile_with_k_with_padding = Cta_tile_extd::VALUE,\n Cta_tile_::WARPS_M, Cta_tile_::WARPS_N, Cta_tile_::WARPS_K>;\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace fmha\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha/gmem_tile.h", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\nnamespace fmha {\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The dimensions of the tile computed by the CTA.\n typename Cta_tile,\n // The number of bits per element.\n int BITS_PER_ELEMENT,\n // The number of rows of Q, K or V loaded by this tile.\n int ROWS,\n // The number of columns.\n int COLS,\n // The number of matrics.\n int NUM_MATS = 3>\nstruct Gmem_tile_qkv {\n // The size of each LDG.\n enum { BYTES_PER_LDG = 16 };\n // The size of a row in bytes.\n enum { BYTES_PER_ROW = COLS * BITS_PER_ELEMENT / 8 };\n\n // The number of threads to load a \"row\" of the matrix.\n enum { THREADS_PER_ROW = BYTES_PER_ROW / BYTES_PER_LDG };\n\n // The number of \"rows\" loaded per LDG.\n enum { ROWS_PER_LDG = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW };\n // The number of LDGs needed to load a chunk of the Q matrix.\n enum { LDGS = fmha::Div_up::VALUE };\n\n // Ctor.\n template \n inline __device__ Gmem_tile_qkv(const Params& params, const int qkv_offset, const BInfo& binfo, const int tidx)\n : params_qkv_stride_in_bytes_(params.qkv_stride_in_bytes),\n actual_seqlen(binfo.actual_seqlen),\n qkv_ptr_(reinterpret_cast(params.qkv_ptr)) {\n // Compute the position in the sequence (within the CTA for the moment).\n int row = tidx / THREADS_PER_ROW;\n // Compute the position of the thread in the row.\n int col = tidx % THREADS_PER_ROW;\n\n // Store the row as we need it to disable the loads.\n row_ = row;\n\n // The row offset in the batched GEMM. For each seq element, we store QKV in that order.\n int64_t row_offset = (int64_t)row * params.qkv_stride_in_bytes;\n // Add the block index.\n row_offset += (int64_t)((binfo.sum_s * NUM_MATS + qkv_offset) * binfo.h + binfo.bidh) * BYTES_PER_ROW;\n\n // Assemble the final pointer.\n qkv_ptr_ += row_offset + col * BYTES_PER_LDG;\n }\n\n // Store data to shared memory.\n template \n inline __device__ void commit(Smem_tile& smem_tile) {\n smem_tile.store(fetch_);\n }\n\n // Load data from memory.\n template \n inline __device__ void load(Smem_tile& smem_tile) {\n const void* ptrs[LDGS];\n uint32_t preds[LDGS];\n#pragma unroll\n for (int ii = 0; ii < LDGS; ++ii) {\n ptrs[ii] = qkv_ptr_ + (int64_t)ii * ROWS_PER_LDG * params_qkv_stride_in_bytes_;\n preds[ii] = ((row_ + ii * ROWS_PER_LDG) < min(ROWS, actual_seqlen));\n fetch_[ii] = make_uint4(0, 0, 0, 0);\n }\n\n // not packing predicates removes restrictions (e.g. FP16 384, 4 warps)\n Ldg_functor fct(fetch_, ptrs);\n#pragma unroll\n for (int ii = 0; ii < LDGS; ++ii) {\n fct.load(ii, preds[ii]);\n }\n }\n\n // Store data to memory.\n inline __device__ void store(const uint4 (&data)[LDGS]) {\n#pragma unroll\n for (int ii = 0; ii < LDGS; ++ii) {\n char* ptr = qkv_ptr_ + (int64_t)ii * ROWS_PER_LDG * params_qkv_stride_in_bytes_;\n if ((row_ + ii * ROWS_PER_LDG) < min(ROWS, actual_seqlen)) {\n fmha::stg(ptr, data[ii]);\n }\n }\n }\n\n // Move the pointer to the next location.\n inline __device__ void move() {\n qkv_ptr_ += (int64_t)ROWS * params_qkv_stride_in_bytes_;\n actual_seqlen -= ROWS;\n }\n\n inline __device__ void move(int steps) {\n qkv_ptr_ += (int64_t)ROWS * params_qkv_stride_in_bytes_ * steps;\n actual_seqlen -= ROWS * steps;\n }\n\n // The stride between rows for the QKV matrice.\n int64_t params_qkv_stride_in_bytes_;\n // The pointer.\n char* qkv_ptr_;\n // The fetch registers.\n uint4 fetch_[LDGS];\n // Keep track of the row the thread is processing as we move the tile.\n int row_;\n // The length of the sequence loaded by that memory tile.\n int actual_seqlen;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Gmem_tile_o {\n // The mma tile.\n using Mma_tile = fmha::Hmma_tile;\n\n // The size of each element.\n enum { BYTES_PER_ELEMENT = 2 };\n // The size of a row in bytes.\n enum { BYTES_PER_ROW = Cta_tile::N * BYTES_PER_ELEMENT };\n\n // The number of threads to store a \"row\" of the matrix.\n enum { THREADS_PER_ROW = 16 };\n // The size of each STG.\n enum { BYTES_PER_STG = BYTES_PER_ROW / THREADS_PER_ROW };\n\n // The number of \"rows\" stored per iteration of the loop. The output of 1 MMA.\n enum { ROWS = Cta_tile::M };\n // The number of \"rows\" stored per iteration of the loop. The output of 1 MMA.\n enum { ROWS_PER_LOOP = ROWS <= 64 ? ROWS : (int)Mma_tile::M_PER_MMA_PER_CTA };\n // The number of outter loop for the stores.\n enum { LOOPS = ROWS / ROWS_PER_LOOP };\n\n // The number of \"rows\" stored per STG.\n enum { ROWS_PER_STG = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW };\n // Do we have to guard against partial writes/reads.\n enum { HAS_INCOMPLETE_STG = Cta_tile::M % ROWS_PER_STG != 0 };\n // The number of STGs needed to store a chunk of the Q matrix.\n enum { STGS_PER_LOOP = fmha::Div_up::VALUE };\n // The number of STGs needed to store a chunk of the Q matrix in total.\n enum { STGS = STGS_PER_LOOP * LOOPS };\n\n // Ctor.\n template \n inline __device__ Gmem_tile_o(const Params& params, const BInfo& binfo, int tidx)\n : params_o_stride_in_bytes_(params.o_stride_in_bytes),\n actual_seqlen_(binfo.actual_seqlen),\n o_ptr_(reinterpret_cast(params.o_ptr)) {\n // Compute the position in the sequence (within the CTA for the moment).\n int row = tidx / THREADS_PER_ROW;\n // Compute the position of the thread in the row.\n int col = tidx % THREADS_PER_ROW;\n\n // Store the row as we need it to disable loads.\n row_ = row;\n\n // The row offset in the batched GEMM.\n int64_t row_offset = (int64_t)row * params.o_stride_in_bytes + binfo.bidx * BYTES_PER_ROW;\n // Assemble the final pointer.\n o_ptr_ += row_offset + col * BYTES_PER_STG;\n\n // Is that thread active on the last STG?\n if (HAS_INCOMPLETE_STG) {\n is_active_for_last_stg_ = row + (STGS - 1) * ROWS_PER_STG < Cta_tile::M;\n }\n }\n\n // Store data to global memory.\n inline __device__ void store(const uint4 (&src)[STGS_PER_LOOP], int mi) {\n#pragma unroll\n for (int ii = 0; ii < STGS_PER_LOOP; ++ii) {\n int jj = mi * STGS_PER_LOOP + ii;\n if (this->row_ + jj * ROWS_PER_STG >= this->actual_seqlen_) {\n break;\n }\n\n float x = reinterpret_cast(src[ii].x);\n float y = reinterpret_cast(src[ii].y);\n float z = reinterpret_cast(src[ii].z);\n float w = reinterpret_cast(src[ii].w);\n uint2 out = float4_to_half4(x, y, z, w);\n if (!HAS_INCOMPLETE_STG || (jj < STGS - 1 || this->is_active_for_last_stg_)) {\n fmha::stg(this->o_ptr_ + jj * ROWS_PER_STG * this->params_o_stride_in_bytes_, out);\n }\n }\n }\n\n // Move the pointer to the next location.\n inline __device__ void move() {\n row_ += ROWS;\n o_ptr_ += (int64_t)ROWS * params_o_stride_in_bytes_;\n }\n\n inline __device__ void move(const int steps) {\n row_ += ROWS * steps;\n o_ptr_ += (int64_t)ROWS * params_o_stride_in_bytes_ * steps;\n }\n\n // The stride between rows for the QKV matrice.\n int64_t params_o_stride_in_bytes_;\n // The pointer.\n char* o_ptr_;\n // Is the thread active for the last STG?\n int is_active_for_last_stg_;\n // Keep track of the row to disable loads.\n int row_;\n // The length of the sequence loaded by that memory tile.\n int actual_seqlen_;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Gmem_tile_mma_sd {\n // The mma tile.\n using Mma_tile = fmha::Hmma_tile;\n\n // Each STG stores 8 elements.\n enum { BYTES_PER_STG = BYTES_PER_ELEMENT * 8 };\n // The number of MMAs in the M dimension.\n enum { MMAS_M = Mma_tile::MMAS_M };\n // The number of MMAs in the N dimension.\n enum { MMAS_N = Mma_tile::MMAS_N };\n // The number of rows computed per MMA per thread block.\n enum { M_PER_MMA_PER_CTA = Mma_tile::M_PER_MMA_PER_CTA };\n // The number of cols computed per MMA per thread block.\n enum { N_PER_MMA_PER_CTA = Mma_tile::N_PER_MMA_PER_CTA };\n // The number of threads per block.\n enum { THREADS_PER_CTA = Cta_tile::THREADS_PER_CTA };\n // The size of each row in bytes. I.e. how many bytes are stored per STG.\n enum { BYTES_PER_ROW = THREADS_PER_CTA * BYTES_PER_STG };\n // The fixed sequence length.\n enum { SEQLEN = Cta_tile::N };\n // The distance between two blocks (in bytes).\n enum { BLOCK_STRIDE_BYTES = SEQLEN * SEQLEN * BYTES_PER_ELEMENT };\n // The distance between elements stored per loop (in bytes).\n enum { LOOP_STRIDE_BYTES = MMAS_M * MMAS_N * BYTES_PER_ROW };\n\n // The type of elements stored per STG.\n using Type = typename fmha::Uint_from_size_in_bytes::Type;\n\n // Ctor.\n template \n inline __device__ Gmem_tile_mma_sd(void* ptr, const Params& params, const int bidb, const int bidh, const int tidx)\n : ptr_(static_cast(ptr)) {\n // The block index.\n size_t bidx = bidb * params.h + bidh;\n\n // Set store location for each thread at the beginning of the loop\n ptr_ += bidx * BLOCK_STRIDE_BYTES + tidx * BYTES_PER_STG;\n }\n\n // Store to global memory.\n inline __device__ void store(const Type& data, const int mi, const int ni) {\n size_t offset = (mi * MMAS_N + ni) * BYTES_PER_ROW;\n fmha::stg(ptr_ + offset, data);\n }\n\n // Load from global memory.\n inline __device__ void load(Type& data, const int mi, const int ni) {\n size_t offset = (mi * MMAS_N + ni) * BYTES_PER_ROW;\n fmha::ldg(data, ptr_ + offset);\n }\n\n // Move to the next tile.\n inline __device__ void move() { ptr_ += LOOP_STRIDE_BYTES; }\n inline __device__ void move(const int steps) { ptr_ += LOOP_STRIDE_BYTES * steps; }\n\n // The pointer in global memory.\n char* ptr_;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate >\nstruct Gmem_tile_mma_s : public Base {\n // The number of mmas in the vertical dimension.\n enum { M = Base::MMAS_M };\n // The number of mmas in the horizontal dimension.\n enum { N = Base::MMAS_N };\n // The type of the vectors stored by each STG.\n using Type = typename Base::Type;\n\n // Ctor.\n template \n inline __device__ Gmem_tile_mma_s(const Params& params, const Block_info& binfo, const int tidx)\n : Base(params.s_ptr, params, binfo.bidb, binfo.bidh, tidx) {}\n\n // Store to global memory.\n template \n inline __device__ void store(const float (&softmax)[2 * M][4 * N], const Mask& mask) {\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n#pragma unroll\n for (int ni = 0; ni < N; ni++) {\n float tmp00 = softmax[2 * mi + 0][4 * ni + 0];\n float tmp01 = softmax[2 * mi + 0][4 * ni + 1];\n float tmp02 = softmax[2 * mi + 0][4 * ni + 2];\n float tmp03 = softmax[2 * mi + 0][4 * ni + 3];\n\n float tmp10 = softmax[2 * mi + 1][4 * ni + 0];\n float tmp11 = softmax[2 * mi + 1][4 * ni + 1];\n float tmp12 = softmax[2 * mi + 1][4 * ni + 2];\n float tmp13 = softmax[2 * mi + 1][4 * ni + 3];\n\n uint4 dst;\n dst.x = fmha::float2_to_half2(tmp00, tmp01);\n dst.y = fmha::float2_to_half2(tmp02, tmp03);\n dst.z = fmha::float2_to_half2(tmp10, tmp11);\n dst.w = fmha::float2_to_half2(tmp12, tmp13);\n if (mask.is_valid(mi, ni, 0, 0)) {\n Base::store(dst, mi, ni);\n }\n }\n }\n }\n\n // Store to global memory.\n template \n inline __device__ void store(const Fragment (&frag)[N][M], const Mask& mask) {\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n#pragma unroll\n for (int ni = 0; ni < N; ni++) {\n uint4 dst;\n dst.x = frag[ni][mi].reg(0);\n dst.y = frag[ni][mi].reg(2);\n dst.z = frag[ni][mi].reg(1);\n dst.w = frag[ni][mi].reg(3);\n if (mask.any_valid(mi, ni)) {\n Base::store(dst, mi, ni);\n }\n }\n }\n }\n\n // Load from global memory.\n template \n inline __device__ void load(uint4 (®s)[M][N], const Mask& mask) {\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n#pragma unroll\n for (int ni = 0; ni < N; ni++) {\n regs[mi][ni] = make_uint4(0, 0, 0, 0);\n if (mask.any_valid(mi, ni)) {\n Base::load(regs[mi][ni], mi, ni);\n }\n }\n }\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The dimensions of the tile computed by the CTA.\n typename Cta_tile,\n // The base class.\n typename Base = fmha::Gmem_tile_qkv >\nstruct Gmem_tile_dout : public Base {\n // Ctor.\n template \n inline __device__ Gmem_tile_dout(const Params& params, const BInfo& binfo, int tidx) : Base(params, 0, binfo, tidx) {\n this->qkv_ptr_ = reinterpret_cast(params.o_ptr);\n this->params_qkv_stride_in_bytes_ = params.o_stride_in_bytes; // needed for move\n\n // Compute the position of the thread in the row.\n int col = tidx % Base::THREADS_PER_ROW;\n\n // The row offset in the batched GEMM. For each seq element, we store O in that order.\n int64_t row_offset = (int64_t)this->row_ * params.o_stride_in_bytes + binfo.bidx * Base::BYTES_PER_ROW;\n\n // Assemble the final pointer.\n this->qkv_ptr_ += row_offset + col * Base::BYTES_PER_LDG;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate >\nstruct Gmem_tile_dq : public Base {\n // Ctor.\n template \n inline __device__ Gmem_tile_dq(const Params& params, const BInfo& binfo, int tidx) : Base(params, binfo, tidx) {\n this->o_ptr_ = reinterpret_cast(params.dqkv_ptr);\n this->params_o_stride_in_bytes_ = params.qkv_stride_in_bytes; // needed for move\n\n // Compute the position of the thread in the row.\n int col = tidx % Base::THREADS_PER_ROW;\n\n // The row offset in the batched GEMM. For each seq element, we store O in that order.\n int64_t row_offset = (int64_t)this->row_ * params.qkv_stride_in_bytes +\n (binfo.sum_s * 3 * binfo.h + binfo.bidh) * Base::BYTES_PER_ROW;\n\n // Assemble the final pointer.\n this->o_ptr_ += row_offset + col * Base::BYTES_PER_STG;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace fmha\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha/kernel_traits.h", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\n#include \"gmem_tile.h\"\n#include \"smem_tile.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct FMHA_kernel_traits {\n // The CTA description for the 1st GEMM.\n using Cta_tile_p = fmha::Cta_tile_extd;\n // The CTA description for the 2nd GEMM.\n using Cta_tile_o = fmha::Cta_tile_extd;\n\n // Do we use one buffer for K and V.\n enum { SHARE_SMEM_FOR_K_AND_V = (FLAGS & 0x08u) != 0u };\n // Do we keep K in registers.\n enum { K_IN_REGS = (FLAGS & 0x10u) == 0u };\n\n // The global memory tile to load Q.\n using Gmem_tile_q = fmha::Gmem_tile_qkv;\n\n // The shared memory tile to swizzle Q.\n using Smem_tile_q = fmha::Smem_tile_a;\n\n // The global memory tile to load K.\n using Gmem_tile_k = fmha::Gmem_tile_qkv;\n // The shared memory tile to swizzle K.\n using Smem_tile_k = fmha::Smem_tile_b;\n\n // The global memory tile to load V.\n using Gmem_tile_v = fmha::Gmem_tile_qkv;\n // The shared memory tile to swizzle V.\n using Smem_tile_v = fmha::Smem_tile_v;\n\n // The global memory tile to store O.\n using Gmem_tile_o = fmha::Gmem_tile_o;\n // The shared memory tile for O.\n using Smem_tile_o = fmha::Smem_tile_o;\n\n // The global memory tile to load/store S.\n using Gmem_tile_s = fmha::Gmem_tile_mma_s;\n\n // The shared memory tile to transpose S.\n using Smem_tile_st = fmha::Smem_tile_mma_transposed;\n\n using Gmem_tile_do = fmha::Gmem_tile_dout;\n\n // Make sure the number of threads match.\n static_assert((int)Gmem_tile_o::THREADS_PER_ROW == (int)Smem_tile_o::THREADS_PER_ROW, \"\");\n\n // The number of threads.\n enum { THREADS = Cta_tile_p::THREADS_PER_CTA };\n // Make sure the number of threads matches both CTAs.\n static_assert((int)THREADS == (int)Cta_tile_o::THREADS_PER_CTA, \"\");\n\n // The amount of shared memory needed to load Q and K.\n enum { BYTES_PER_SMEM_QK = Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE };\n // The extra amount of shared memory needed to load V.\n enum { BYTES_PER_SMEM_V = SHARE_SMEM_FOR_K_AND_V ? 0u : Smem_tile_v::BYTES_PER_TILE };\n // The amount of shared memory needed for Q, K and V..\n enum { BYTES_PER_SMEM_QKV = BYTES_PER_SMEM_QK + BYTES_PER_SMEM_V };\n // The amount of shared memory needed to load Q and store O.\n enum { BYTES_PER_SMEM_QO = Smem_tile_q::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE };\n\n // The amount of shared memory needed for Q, K, V and O.\n enum { BYTES_PER_SMEM = fmha::Max::VALUE };\n // Make sure we have enough shared memory.\n static_assert(Smem_tile_q::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE <= BYTES_PER_SMEM, \"\");\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha/mask.h", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\nnamespace fmha {\n\ntemplate \nstruct Mask {\n using Mma_tile = fmha::Hmma_tile;\n\n template \n __device__ Mask(const Params& params, const BInfo& blockInfo, int tidx) {\n actual_seqlen = blockInfo.actual_seqlen;\n\n const int warp = tidx / Cta_tile::THREADS_PER_WARP;\n const int lane = tidx % Cta_tile::THREADS_PER_WARP;\n\n static_assert(Cta_tile::WARPS_K == 1, \"\");\n\n // find the warp in the Cta tile\n const int warp_n = (warp / Cta_tile::WARPS_M);\n const int warp_m = (warp % Cta_tile::WARPS_M);\n // decompose warp into 8x4 tile\n const int quad = lane / 4;\n const int tid = (lane % 4) * 2;\n row = warp_m * 16 + quad;\n col = warp_n * 16 + tid;\n }\n\n inline __device__ bool is_valid(const int mi, const int ni, const int ii, const int jj) const {\n // ii and jj iterate over the 2x4 fragment\n const bool col_valid = (ni * Mma_tile::N_PER_MMA_PER_CTA + col + (jj & 2) * 4 + (jj & 1)) < actual_seqlen;\n //&& (row + mi * Mma_tile::M_PER_MMA_PER_CTA + ii * 8) < actual_seqlen;\n return col_valid;\n // return row_valid && col_valid;\n }\n\n // BERT Mask: if upper left is invalid, none are valid\n inline __device__ bool any_valid(int mi, int ni) const { return is_valid(mi, ni, 0, 0); }\n\n inline __device__ void load(int it) { row_offset = it * Cta_tile::M + row; }\n int row_offset;\n\n int row;\n int col;\n int actual_seqlen;\n};\n\n} // namespace fmha\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha/smem_tile.h", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\n#include \n#include \n\nnamespace fmha {\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The description of the tile computed by this CTA.\n typename Cta_tile,\n // The number of rows in the 2D shared memory buffer.\n int M_,\n // The number of cols.\n int N_,\n // The size in bits of each element.\n int BITS_PER_ELEMENT_,\n // The number of bytes per STS.\n int BYTES_PER_STS_ = 16,\n // The number of buffers. (Used in multistage and double buffer cases.)\n int BUFFERS_PER_TILE_ = 1,\n // Do we enable the fast path for LDS.128 and friends.\n int ENABLE_LDS_FAST_PATH_ = 0,\n // The number of rows that are used for the XOR swizzling to allow fast STS/LDS.\n int ROWS_PER_XOR_PATTERN_ = 8,\n // The number of cols that are used for the XOR swizzling to allow fast STS/LDS.\n int COLS_PER_XOR_PATTERN_ = 1,\n // Use or not predicates\n bool USE_PREDICATES_ = true>\nstruct Smem_tile_without_skews {\n // The size in bits of each element.\n enum { BITS_PER_ELEMENT = BITS_PER_ELEMENT_ };\n // The size in bytes of a single STS.\n enum { BYTES_PER_STS = BYTES_PER_STS_ };\n // The number of elements per STS.\n enum { ELEMENTS_PER_STS = BYTES_PER_STS * 8 / BITS_PER_ELEMENT };\n // To support arbitrary N, we pad some values to a power-of-2.\n enum { N_WITH_PADDING = Next_power_of_two::VALUE };\n // The number of bytes per row without packing of rows.\n enum { BYTES_PER_ROW_BEFORE_PACKING = N_WITH_PADDING * BITS_PER_ELEMENT / 8 };\n // The number of bytes per row -- we want at least 128B per row.\n enum { BYTES_PER_ROW = Max::VALUE };\n // The number of rows in shared memory (two rows may be packed into a single one).\n enum { ROWS = M_ * BYTES_PER_ROW_BEFORE_PACKING / BYTES_PER_ROW };\n\n // The number of threads per row.\n enum { THREADS_PER_ROW_UNBOUNDED = BYTES_PER_ROW / BYTES_PER_STS };\n // The number of threads per row.\n enum { THREADS_PER_ROW = Min::VALUE };\n\n // The number of STS per row.\n enum { STS_PER_ROW = BYTES_PER_ROW / THREADS_PER_ROW / BYTES_PER_STS };\n // It must be at least one.\n static_assert(STS_PER_ROW >= 1, \"\");\n // The number of rows written with a single STS.\n enum { ROWS_PER_STS = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW };\n // Make sure we write to at least one row per STS. Thanks Dr. Obvious ;)\n static_assert(ROWS_PER_STS >= 1, \"\");\n // The number of STS needed to store all rows.\n enum { STS_PER_COL = Div_up::VALUE };\n // The number of STS in total.\n enum { STS = STS_PER_COL * STS_PER_ROW };\n\n // The size of one buffer in bytes in shared memory.\n enum { BYTES_PER_BUFFER = STS * BYTES_PER_STS * Cta_tile::THREADS_PER_CTA };\n // The number of buffers.\n enum { BUFFERS_PER_TILE = BUFFERS_PER_TILE_ };\n // The size in bytes of total buffers.\n enum { BYTES_PER_TILE = BYTES_PER_BUFFER * BUFFERS_PER_TILE };\n // The boundary for smem_read_offset and smem_write_offset increment.\n enum { BYTES_PER_TILE_INC_BOUNDARY = BYTES_PER_TILE - BYTES_PER_BUFFER };\n\n // Do we enable the LDS.128 fast path?\n enum { ENABLE_LDS_FAST_PATH = ENABLE_LDS_FAST_PATH_ };\n static_assert(ENABLE_LDS_FAST_PATH == 0);\n // The number of rows that are used for the XOR swizzling to allow fast STS/LDS.\n enum { ROWS_PER_XOR_PATTERN = ROWS_PER_XOR_PATTERN_ };\n // The number of cols that are used for the XOR swizzling to allow fast STS/LDS.\n enum { COLS_PER_XOR_PATTERN = COLS_PER_XOR_PATTERN_ * 16 / BYTES_PER_STS };\n // Use or not predicates\n enum { USE_PREDICATES = USE_PREDICATES_ };\n\n // The type of elements that are stored in shared memory by each thread.\n using Store_type = typename Uint_from_size_in_bytes::Type;\n\n // Ctor.\n inline __device__ Smem_tile_without_skews(void* smem, int tidx) : smem_(__nvvm_get_smem_pointer(smem)) {\n // The row written by a thread. See doc/mma_smem_layout.xlsx.\n int smem_write_row = tidx / THREADS_PER_ROW;\n\n // The XOR pattern.\n int smem_write_xor = smem_write_row % ROWS_PER_XOR_PATTERN * COLS_PER_XOR_PATTERN;\n // Compute the column and apply the XOR pattern.\n int smem_write_col = (tidx % THREADS_PER_ROW) ^ smem_write_xor;\n\n // The offset.\n this->smem_write_offset_ = smem_write_row * BYTES_PER_ROW + smem_write_col * BYTES_PER_STS;\n\n // TODO: Why not merge it with the read offset?\n this->smem_read_buffer_ = __shfl_sync(0xffffffff, 0, 0);\n this->smem_write_buffer_ = __shfl_sync(0xffffffff, 0, 0);\n }\n\n // Compute the store pointers.\n template \n inline __device__ void compute_store_pointers(uint32_t (&ptrs)[N]) {\n#pragma unroll\n for (int ii = 0; ii < N; ++ii) {\n // Decompose the STS into row/col.\n int row = ii / STS_PER_ROW;\n int col = ii % STS_PER_ROW;\n\n // Assemble the offset.\n int offset = smem_write_offset_ + row * ROWS_PER_STS * BYTES_PER_ROW;\n\n // Take the column into account.\n if (STS_PER_ROW > 1) {\n offset += col * THREADS_PER_ROW * BYTES_PER_STS;\n }\n\n // Apply the XOR pattern if needed.\n if (ROWS_PER_STS < ROWS_PER_XOR_PATTERN) {\n const int m = row * ROWS_PER_STS % ROWS_PER_XOR_PATTERN;\n offset ^= m * COLS_PER_XOR_PATTERN * BYTES_PER_STS;\n }\n\n // Assemble the final pointer :)\n ptrs[ii] = smem_ + offset + smem_write_buffer_;\n }\n }\n\n inline __device__ void debug_reset() {\n for (int buffer = 0; buffer < BYTES_PER_TILE; buffer += BYTES_PER_BUFFER) {\n for (int row = 0; row < ROWS; ++row) {\n for (int col = 0; col < BYTES_PER_ROW; col += 4) {\n if (threadIdx.x == 0) {\n uint32_t val = 0x0;\n sts(val, smem_ + row * BYTES_PER_ROW + col + buffer);\n }\n }\n }\n }\n }\n\n // Print the content of the tile (only for debug ;)).\n inline __device__ void debug_print() const {\n for (int buffer = 0; buffer < BYTES_PER_TILE; buffer += BYTES_PER_BUFFER) {\n for (int row = 0; row < ROWS; ++row) {\n for (int col = 0; col < BYTES_PER_ROW; col += 4) {\n if (threadIdx.x == 0) {\n uint32_t val;\n lds(val, smem_ + row * BYTES_PER_ROW + col + buffer);\n printf(\"block=(x=%2d, y=%2d, z=%2d) (smem_=%2d, buffer=%2d, row=%2d, byte=%4d)=0x%08x\\n\", blockIdx.x,\n blockIdx.y, blockIdx.z, smem_, buffer, row, col, val);\n }\n }\n }\n }\n }\n\n // Move the read offset to next buffer.\n inline __device__ void move_to_next_read_buffer() {\n if (BUFFERS_PER_TILE > 1 && smem_read_buffer_ >= BYTES_PER_TILE_INC_BOUNDARY) {\n this->smem_read_buffer_ -= BYTES_PER_TILE_INC_BOUNDARY;\n } else if (BUFFERS_PER_TILE > 1) {\n this->smem_read_buffer_ += BYTES_PER_BUFFER;\n }\n }\n\n // Move the read offset to next buffer. TODO: Remove this member function!!!\n inline __device__ void move_next_read_buffer() { this->move_to_next_read_buffer(); }\n\n // Move the read offset to next N buffer (circular-buffer).\n inline __device__ void move_to_next_read_buffer(int N) {\n if (BUFFERS_PER_TILE > 1) {\n this->smem_read_buffer_ += N * BYTES_PER_BUFFER;\n this->smem_read_buffer_ -= smem_read_buffer_ >= BYTES_PER_TILE ? BYTES_PER_TILE : 0;\n }\n }\n\n // Move the read offset to next N buffer (circular-buffer). TODO: Remove this member function!!!\n inline __device__ void move_next_read_buffer(int N) { this->move_to_next_read_buffer(N); }\n\n // Move the write offset to next buffer.\n inline __device__ void move_to_next_write_buffer() {\n if (BUFFERS_PER_TILE > 1 && smem_write_buffer_ >= BYTES_PER_TILE_INC_BOUNDARY) {\n this->smem_write_buffer_ -= BYTES_PER_TILE_INC_BOUNDARY;\n } else if (BUFFERS_PER_TILE > 1) {\n this->smem_write_buffer_ += BYTES_PER_BUFFER;\n }\n }\n\n // Move the write offset to next buffer. TODO: Remove that member function!\n inline __device__ void move_next_write_buffer() { this->move_to_next_write_buffer(); }\n\n // Move the read offset.\n inline __device__ void move_read_offset(int delta) { this->smem_read_offset_ += delta; }\n\n // Move the write offset.\n inline __device__ void move_write_offset(int delta) { this->smem_write_offset_ += delta; }\n\n // Store to the tile in shared memory.\n template \n inline __device__ void store(const Store_type (&data)[N], uint64_t = 0) {\n uint32_t smem_ptrs[N];\n this->compute_store_pointers(smem_ptrs);\n sts(smem_ptrs, data);\n }\n\n // Store to the tile in shared memory.\n template \n inline __device__ void store(const Store_type (&data)[N], uint32_t (&preds)[M], uint64_t = 0) {\n uint32_t smem_ptrs[N];\n this->compute_store_pointers(smem_ptrs);\n sts(smem_ptrs, data, preds);\n }\n\n // Store to the tile in shared memory.\n template \n inline __device__ void store(const Store_type (&data)[N], uint32_t preds, uint64_t = 0) {\n this->store(data, preds);\n }\n\n // Store to the tile in shared memory.\n template \n inline __device__ void store(const void* (&gmem_ptrs)[N], uint32_t preds, uint64_t = 0) {\n uint32_t tmp[1] = {preds};\n this->store(gmem_ptrs, tmp);\n }\n\n // The shared memory pointer.\n uint32_t smem_;\n // The read offset. Reserve 4 offsets if needed.\n int smem_read_offset_;\n // The write offset.\n int smem_write_offset_;\n // The buffer base offset for read.\n int smem_read_buffer_;\n // The buffer base offset for write.\n int smem_write_buffer_;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The dimensions of the tile computed by the CTA.\n typename Cta_tile,\n // The layout of the tile.\n typename Layout,\n // The size of the STS.\n int BYTES_PER_STS = 16,\n // The number of buffers per tile.\n int BUFFERS_PER_TILE = 1,\n // Use or not predicates\n bool USE_PREDICATES = true>\nstruct Smem_tile_a {};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Compute_reset_mask {\n // The potential mask.\n enum { HALF = MMAS_K_WITH_PADDING / 2 };\n // The remainder.\n enum { MOD = MMAS_K % HALF };\n // The final value.\n enum { VALUE = (MMAS_K == MOD ? 0 : HALF) | Compute_reset_mask::VALUE };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Compute_reset_mask<0, MMAS_K_WITH_PADDING> {\n enum { VALUE = 0 };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Compute_reset_mask {\n enum { VALUE = MMAS_K - 1 };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Rows_per_xor_pattern_a {\n // The size in bits.\n enum { N_IN_BITS = N * fmha::BITS_PER_ELEMENT_A };\n // The number of rows.\n enum { VALUE = N_IN_BITS <= 256 ? 2 : (N_IN_BITS <= 512 ? 4 : 8) };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Rows_per_xor_pattern_row_a : public Rows_per_xor_pattern_a {};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The dimensions of the tile computed by the CTA.\n typename Cta_tile,\n // The size of the STS.\n int BYTES_PER_STS,\n // The number of buffers per tile.\n int BUFFERS_PER_TILE,\n // How many rows to use for the XOR pattern to avoid bank conflicts?\n int ROWS_PER_XOR_PATTERN_ = Rows_per_xor_pattern_row_a::VALUE>\nstruct Smem_tile_row_a : public Smem_tile_without_skews {\n // The MMA tile.\n using Mma_tile = fmha::Hmma_tile;\n // The base class.\n using Base = Smem_tile_without_skews;\n // The fragment.\n using Fragment = Fragment_a;\n\n // When we use padding to reach a power of two, special care has to be taken.\n using Cta_tile_with_padding = Cta_tile_with_k_with_padding;\n // The number of MMAs.\n using Mma_tile_with_padding = fmha::Hmma_tile;\n\n // The size of a single LDS in bytes.\n enum { BYTES_PER_LDS = 16 };\n\n // Ctor.\n inline __device__ Smem_tile_row_a(void* smem, int tidx) : Base(smem, tidx) {\n // For documentation on the layout, see doc/mma_smem_layout.xlsx.\n\n // The number of warps.\n const int WARPS_M = Cta_tile::WARPS_M;\n const int WARPS_N = Cta_tile::WARPS_N;\n const int WARPS_K = Cta_tile::WARPS_K;\n\n static_assert(WARPS_M == 1);\n static_assert(WARPS_N == 4 || WARPS_N == 8);\n static_assert(WARPS_K == 1);\n static_assert(Base::ROWS_PER_XOR_PATTERN == 8);\n\n // The row and column read by the thread.\n int smem_read_row = (tidx & 0x0f);\n int smem_read_col = (tidx & 0x07);\n smem_read_col ^= (tidx & 0x10) / 16;\n\n // The shared memory offset.\n this->smem_read_offset_ = smem_read_row * Base::BYTES_PER_ROW + smem_read_col * BYTES_PER_LDS;\n }\n\n // Rewind smem_read_offset for last LDS phase in main loop.\n inline __device__ void reverse_smem_read_offset(int ki = 0) {\n // Undo the pointer increment for the next ni.\n // Should match the load function below for ki = 0.\n if (Mma_tile_with_padding::MMAS_K >= 2) {\n this->smem_read_offset_ ^= BYTES_PER_LDS * 2;\n }\n }\n\n // Load from shared memory.\n inline __device__ void load(Fragment (&a)[Mma_tile::MMAS_M], int ki) {\n#pragma unroll\n for (int mi = 0; mi < Mma_tile::MMAS_M; ++mi) {\n // Jump by as many matrix rows as needed (a row in smem may pack multiple matrix rows).\n int offset = mi * Mma_tile::M_PER_MMA_PER_CTA * Base::BYTES_PER_ROW_BEFORE_PACKING;\n\n // Load using LDSM.M88.4.\n uint4 tmp;\n ldsm(tmp, this->smem_ + this->smem_read_offset_ + this->smem_read_buffer_ + offset);\n\n // Store the value into the fragment.\n a[mi].reg(0) = tmp.x;\n a[mi].reg(1) = tmp.y;\n a[mi].reg(2) = tmp.z;\n a[mi].reg(3) = tmp.w;\n }\n\n // Move the offset to the next possition. See doc/mma_smem_layout.xlsx.\n static_assert(Mma_tile_with_padding::MMAS_K < 64, \"Not implemented\");\n if (Mma_tile_with_padding::MMAS_K >= 32 && ki % 16 == 15) {\n this->smem_read_offset_ ^= 31 * BYTES_PER_LDS * 2;\n } else if (Mma_tile_with_padding::MMAS_K >= 16 && ki % 8 == 7) {\n this->smem_read_offset_ ^= 15 * BYTES_PER_LDS * 2;\n } else if (Mma_tile_with_padding::MMAS_K >= 8 && ki % 4 == 3) {\n this->smem_read_offset_ ^= 7 * BYTES_PER_LDS * 2;\n } else if (Mma_tile_with_padding::MMAS_K >= 4 && ki % 2 == 1) {\n this->smem_read_offset_ ^= 3 * BYTES_PER_LDS * 2;\n } else if (Mma_tile_with_padding::MMAS_K >= 2) {\n this->smem_read_offset_ ^= 1 * BYTES_PER_LDS * 2;\n }\n }\n\n // Reset the read offset.\n inline __device__ void reset_read_offset() {\n // The number of MMAs in the K dimension.\n enum { MMAS_K = Mma_tile::MMAS_K };\n // The number of MMAs in the K dimension when we include padding.\n enum { MMAS_K_WITH_PADDING = Mma_tile_with_padding::MMAS_K };\n // Assemble the mask.\n enum { MASK = Compute_reset_mask::VALUE };\n\n // Reset the read offset.\n this->smem_read_offset_ ^= MASK * BYTES_PER_LDS * 2;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The dimensions of the tile computed by the CTA.\n typename Cta_tile,\n // The size of the STS.\n int BYTES_PER_STS,\n // The number of buffers per tile.\n int BUFFERS_PER_TILE>\nstruct Smem_tile_a\n : public Smem_tile_row_a {\n // The base class.\n using Base = Smem_tile_row_a;\n\n // Ctor.\n inline __device__ Smem_tile_a(void* smem, int tidx) : Base(smem, tidx) {}\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The dimensions of the tile computed by the CTA.\n typename Cta_tile,\n // The layout of the tile.\n typename Layout,\n // The size of the STS.\n int BYTES_PER_STS = 16,\n // The number of buffers per tile.\n int BUFFERS_PER_TILE = 1,\n // Use or not predicates\n bool USE_PREDICATES = true>\nstruct Smem_tile_b {};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Rows_per_xor_pattern_b {\n // The size in bits.\n enum { N_IN_BITS = N * fmha::BITS_PER_ELEMENT_B };\n // The number of rows.\n enum { VALUE = N_IN_BITS <= 256 ? 2 : (N_IN_BITS <= 512 ? 4 : 8) };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Rows_per_xor_pattern_col_b : public Rows_per_xor_pattern_b {};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The dimensions of the tile computed by the CTA.\n typename Cta_tile,\n // The size of the STS.\n int BYTES_PER_STS,\n // The number of buffers per tile.\n int BUFFERS_PER_TILE,\n // How many rows to use for the XOR pattern to avoid bank conflicts?\n int ROWS_PER_XOR_PATTERN_ = Rows_per_xor_pattern_col_b::VALUE>\nstruct Smem_tile_col_b : public Smem_tile_without_skews {\n // The MMA tile.\n using Mma_tile = fmha::Hmma_tile;\n // The base class.\n using Base = Smem_tile_without_skews;\n // The fragment.\n using Fragment = Fragment_b;\n\n // When we use padding to reach a power of two, special care has to be taken.\n using Cta_tile_with_padding = Cta_tile_with_k_with_padding;\n // The number of MMAs.\n using Mma_tile_with_padding = fmha::Hmma_tile;\n\n // The size of a single LDS in bytes.\n enum { BYTES_PER_LDS = 16 };\n\n // The number of STS per thread\n enum { STS_PER_THREAD_ = Base::ROWS * Base::THREADS_PER_ROW / Cta_tile::THREADS_PER_CTA };\n // The number of STS per thread must be at least 1.\n enum { STS_PER_THREAD = Max<1, STS_PER_THREAD_>::VALUE };\n\n // Ctor.\n inline __device__ Smem_tile_col_b(void* smem, int tidx) : Base(smem, tidx) {\n // For documentation on the layout, see doc/mma_smem_layout.xlsx.\n\n // The number of warps.\n const int WARPS_M = Cta_tile::WARPS_M;\n const int WARPS_N = Cta_tile::WARPS_N;\n const int WARPS_K = Cta_tile::WARPS_K;\n static_assert(Base::ROWS_PER_XOR_PATTERN == 8);\n static_assert(WARPS_M == 1);\n static_assert(WARPS_N == 4 || WARPS_N == 8);\n static_assert(WARPS_K == 1);\n\n // The masks to select the warps.\n const int WARP_MASK_N = Warp_masks::N;\n\n // The divisor for the warps.\n const int WARP_DIV_N = WARPS_M * 1 * Cta_tile::THREADS_PER_WARP;\n\n // The row and column read by the thread.\n int smem_read_row = (tidx & WARP_MASK_N) / WARP_DIV_N * Mma_tile::N_PER_MMA + (tidx & 0x07) + (tidx & 0x10) / 2;\n int smem_read_col = (tidx & 0x07);\n smem_read_col ^= (tidx & 0x08) / 8;\n // The shared memory offset.\n this->smem_read_offset_ = smem_read_row * Base::BYTES_PER_ROW + smem_read_col * BYTES_PER_LDS;\n }\n\n // Rewind smem_read_offset for last LDS phase in main loop.\n inline __device__ void reverse_smem_read_offset(int ki = 0) {\n // Undo the pointer increment for the next ni.\n // Should match the load function below for ki = 0.\n if (Mma_tile_with_padding::MMAS_K >= 2) {\n this->smem_read_offset_ ^= BYTES_PER_LDS * 2;\n }\n }\n\n // Load from shared memory.\n inline __device__ void load(Fragment (&b)[Mma_tile::MMAS_N], int ki) {\n#pragma unroll\n for (int ni = 0; ni < Mma_tile::MMAS_N; ++ni) {\n // Jump by as many matrix rows as needed (a row in smem may pack multiple matrix rows).\n int offset = ni * Mma_tile::N_PER_MMA_PER_CTA * Base::BYTES_PER_ROW_BEFORE_PACKING;\n\n // Load using LDSM.M88.4.\n uint4 tmp;\n ldsm(tmp, this->smem_ + this->smem_read_offset_ + this->smem_read_buffer_ + offset);\n\n // Store the value into the fragment.\n b[ni].reg(0) = tmp.x;\n b[ni].reg(1) = tmp.y;\n b[ni].reg(2) = tmp.z;\n b[ni].reg(3) = tmp.w;\n }\n\n // Move the offset to the next possition. See doc/mma_smem_layout.xlsx.\n static_assert(Mma_tile_with_padding::MMAS_K < 64, \"Not implemented\");\n if (Mma_tile_with_padding::MMAS_K >= 32 && ki % 16 == 15) {\n this->smem_read_offset_ ^= 31 * BYTES_PER_LDS * 2;\n } else if (Mma_tile_with_padding::MMAS_K >= 16 && ki % 8 == 7) {\n this->smem_read_offset_ ^= 15 * BYTES_PER_LDS * 2;\n } else if (Mma_tile_with_padding::MMAS_K >= 8 && ki % 4 == 3) {\n this->smem_read_offset_ ^= 7 * BYTES_PER_LDS * 2;\n } else if (Mma_tile_with_padding::MMAS_K >= 4 && ki % 2 == 1) {\n this->smem_read_offset_ ^= 3 * BYTES_PER_LDS * 2;\n } else if (Mma_tile_with_padding::MMAS_K >= 2) {\n this->smem_read_offset_ ^= 1 * BYTES_PER_LDS * 2;\n }\n }\n\n // Reset the read offset.\n inline __device__ void reset_read_offset() {\n // The number of MMAs in the K dimension.\n enum { MMAS_K = Mma_tile::MMAS_K };\n // The number of MMAs in the K dimension when we include padding.\n enum { MMAS_K_WITH_PADDING = Mma_tile_with_padding::MMAS_K };\n // Assemble the mask.\n enum { MASK = Compute_reset_mask::VALUE };\n\n // Reset the read offset.\n this->smem_read_offset_ ^= MASK * BYTES_PER_LDS * 2;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The dimensions of the tile computed by the CTA.\n typename Cta_tile,\n // The size of the STS.\n int BYTES_PER_STS,\n // The number of buffers per tile.\n int BUFFERS_PER_TILE>\nstruct Smem_tile_b\n : public Smem_tile_col_b {\n // The base class.\n using Base = Smem_tile_col_b;\n\n // Ctor.\n inline __device__ Smem_tile_b(void* smem, int tidx) : Base(smem, tidx) {}\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Rows_per_xor_pattern_row_b : public Rows_per_xor_pattern_b {};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The dimensions of the tile computed by the CTA.\n typename Cta_tile,\n // The size of the STS.\n int BYTES_PER_STS,\n // The number of buffers per tile.\n int BUFFERS_PER_TILE,\n // How many rows to use for the XOR pattern to avoid bank conflicts?\n int ROWS_PER_XOR_PATTERN_ = Rows_per_xor_pattern_row_b::VALUE,\n // How many cols to use for the XOR pattern to avoid bank conflicts?\n int COLS_PER_XOR_PATTERN_ = 1>\nstruct Smem_tile_row_b\n : public Smem_tile_without_skews {\n // The MMA tile.\n using Mma_tile = fmha::Hmma_tile;\n // The base class.\n using Base = Smem_tile_without_skews;\n // The fragment.\n using Fragment = Fragment_b;\n\n // Can we use LDSM? No if the data type is 32-bit large.\n enum { USE_LDSMT = fmha::BITS_PER_ELEMENT_B == 16 };\n // The size of a single LDS in bytes.\n enum { BYTES_PER_LDS = USE_LDSMT ? 16 : 4 };\n // The number of elements per LDS.\n enum { ELEMENTS_PER_LDS = BYTES_PER_LDS * 8 / fmha::BITS_PER_ELEMENT_B };\n\n // The number of STS per thread\n enum { STS_PER_THREAD_ = Base::ROWS * Base::THREADS_PER_ROW / Cta_tile::THREADS_PER_CTA };\n // The number of STS per thread must be at least 1.\n enum { STS_PER_THREAD = Max<1, STS_PER_THREAD_>::VALUE };\n\n // Ctor.\n inline __device__ Smem_tile_row_b(void* smem, int tidx) : Base(smem, tidx) {\n // The number of warps.\n const int WARPS_M = Cta_tile::WARPS_M;\n const int WARPS_N = Cta_tile::WARPS_N;\n const int WARPS_K = Cta_tile::WARPS_K;\n static_assert(WARPS_K == 1);\n static_assert(WARPS_M == 4 || WARPS_M == 8);\n static_assert(WARPS_N == 1);\n\n // The masks to select the warps.\n const int WARP_MASK_N = Warp_masks::N;\n const int WARP_MASK_K = Warp_masks::K;\n\n // The divisor for the warps.\n const int WARP_DIV_N = WARPS_M * 1 * Cta_tile::THREADS_PER_WARP;\n const int WARP_DIV_K = WARPS_M * WARPS_N * Cta_tile::THREADS_PER_WARP;\n\n // The row/col read by the thread.\n int smem_read_row, smem_read_col;\n\n static_assert(USE_LDSMT);\n static_assert(Base::ROWS_PER_XOR_PATTERN == 8);\n\n smem_read_row = (tidx & WARP_MASK_K) / WARP_DIV_K * Mma_tile::MMAS_K * 16 + (tidx & 0x07) + (tidx & 0x08);\n smem_read_col = (tidx & 0x07);\n smem_read_col ^= (tidx & WARP_MASK_N) / WARP_DIV_N * 2 + (tidx & 0x10) / 16;\n\n // The shared memory offset.\n this->smem_read_offset_ = smem_read_row * Base::BYTES_PER_ROW + smem_read_col * BYTES_PER_LDS;\n\n // Fill zeroes for group conv\n }\n\n // Rewind smem_read_offset for last LDS phase in main loop.\n inline __device__ void reverse_smem_read_offset(int ki = 0) {\n // The size of each element in bits.\n const int BITS_PER_ELT = fmha::BITS_PER_ELEMENT_B;\n // The size in bytes of the data needed to compute an MMA per CTA.\n const int BYTES_PER_MMA_PER_CTA = Mma_tile::N_PER_MMA_PER_CTA * BITS_PER_ELT / 8;\n\n#pragma unroll\n for (int ni = 0; ni < Mma_tile::MMAS_N; ++ni) {\n // Undo the pointer increment for the next ni.\n // Should match the load function below for ki = 0.\n if (BYTES_PER_MMA_PER_CTA >= 128) {\n // Nothing to do!\n } else if (BYTES_PER_MMA_PER_CTA == 64 && Mma_tile::MMAS_N > 1) {\n this->smem_read_offset_ ^= BYTES_PER_MMA_PER_CTA;\n } else if (BYTES_PER_MMA_PER_CTA == 64) {\n // Nothing to do!\n } else if (BYTES_PER_MMA_PER_CTA == 32 && Mma_tile::MMAS_N == 4) {\n this->smem_read_offset_ ^= BYTES_PER_LDS * (ni % 2 == 0 ? 2 : 6);\n } else if (BYTES_PER_MMA_PER_CTA == 32 && Mma_tile::MMAS_N == 2) {\n this->smem_read_offset_ ^= BYTES_PER_LDS * 2;\n }\n }\n\n // Reset smem_read_offset for odd MMAS_N > 1 (npo2 kernels)\n if (BYTES_PER_MMA_PER_CTA == 64 && Mma_tile::MMAS_N > 1 && Mma_tile::MMAS_N % 2 == 1) {\n this->smem_read_offset_ ^= BYTES_PER_MMA_PER_CTA;\n }\n }\n\n // Load from shared memory.\n inline __device__ void load(Fragment (&b)[Mma_tile::MMAS_N], int ki) {\n // The size of each element in bits.\n const int BITS_PER_ELT = fmha::BITS_PER_ELEMENT_B;\n // The size in bytes of the data needed to compute an MMA per CTA.\n const int BYTES_PER_MMA_PER_CTA = Mma_tile::N_PER_MMA_PER_CTA * BITS_PER_ELT / 8;\n\n#pragma unroll\n for (int ni = 0; ni < Mma_tile::MMAS_N; ++ni) {\n // Prepare the offset.\n int offset = ki * Base::ROWS_PER_XOR_PATTERN * 2 * Base::BYTES_PER_ROW;\n if (BYTES_PER_MMA_PER_CTA == 32) {\n offset += this->smem_read_offset_;\n } else if (BYTES_PER_MMA_PER_CTA == 64) {\n offset += this->smem_read_offset_ + (ni / 2) * BYTES_PER_MMA_PER_CTA * 2;\n } else {\n offset += this->smem_read_offset_ + (ni)*BYTES_PER_MMA_PER_CTA;\n }\n\n // Load the data using LDSM.MT88.2.\n uint32_t ptr = this->smem_ + this->smem_read_buffer_ + offset;\n uint4 tmp;\n if (USE_LDSMT) {\n ldsmt(tmp, ptr);\n } else {\n lds(tmp.x, (ptr) + 0 * Base::BYTES_PER_ROW);\n lds(tmp.y, (ptr) + 4 * Base::BYTES_PER_ROW);\n lds(tmp.z, (ptr ^ 32) + 0 * Base::BYTES_PER_ROW);\n lds(tmp.w, (ptr ^ 32) + 4 * Base::BYTES_PER_ROW);\n }\n\n // Store those values in the fragment.\n b[ni].reg(0) = tmp.x;\n b[ni].reg(1) = tmp.y;\n b[ni].reg(2) = tmp.z;\n b[ni].reg(3) = tmp.w;\n\n // Move the pointer for the next ni. I expect the compiler to not recompute those.\n if (BYTES_PER_MMA_PER_CTA >= 128) {\n // Nothing to do!\n } else if (BYTES_PER_MMA_PER_CTA == 64 && Mma_tile::MMAS_N > 1) {\n this->smem_read_offset_ ^= BYTES_PER_MMA_PER_CTA;\n } else if (BYTES_PER_MMA_PER_CTA == 64) {\n // Nothing to do!\n } else if (BYTES_PER_MMA_PER_CTA == 32 && Mma_tile::MMAS_N == 4) {\n this->smem_read_offset_ ^= BYTES_PER_LDS * (ni % 2 == 0 ? 2 : 6);\n } else if (BYTES_PER_MMA_PER_CTA == 32 && Mma_tile::MMAS_N == 2) {\n this->smem_read_offset_ ^= BYTES_PER_LDS * 2;\n }\n }\n\n // Reset smem_read_offset for odd MMAS_N > 1 (npo2 kernels)\n if (BYTES_PER_MMA_PER_CTA == 64 && Mma_tile::MMAS_N > 1 && Mma_tile::MMAS_N % 2 == 1) {\n this->smem_read_offset_ ^= BYTES_PER_MMA_PER_CTA;\n }\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <\n // The dimensions of the tile computed by the CTA.\n typename Cta_tile,\n // The size of the STS.\n int BYTES_PER_STS,\n // The number of buffers per tile.\n int BUFFERS_PER_TILE>\nstruct Smem_tile_b\n : public Smem_tile_row_b {\n // The base class.\n using Base = Smem_tile_row_b;\n\n // Ctor.\n inline __device__ Smem_tile_b(void* smem, int tidx) : Base(smem, tidx) {}\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Smem_tile_v : public fmha::Smem_tile_without_skews {\n // The base class.\n using Base = Smem_tile_without_skews;\n // The MMA tile.\n using Mma_tile = fmha::Hmma_tile;\n // The fragment.\n using Fragment = Fragment_b;\n\n // The size of a single LDS in bytes.\n enum { BYTES_PER_LDS = 16 };\n\n // Ctor.\n inline __device__ Smem_tile_v(void* smem, int tidx) : Base(smem, tidx) {\n // The row/col read by the thread.\n int read_row, read_col;\n\n static_assert(Cta_tile::WARPS_M == 1 && Cta_tile::WARPS_N == 1 &&\n (Cta_tile::WARPS_K == 4 || Cta_tile::WARPS_K == 8));\n\n read_row = (tidx & 0xe0) / 2 + (tidx & 0x0f);\n read_col = (tidx & 0x07);\n read_col ^= (tidx & 0x10) / 16;\n\n // The shared memory offset.\n this->smem_read_offset_ = read_row * Base::BYTES_PER_ROW + read_col * BYTES_PER_LDS;\n }\n\n // Load from shared memory.\n inline __device__ void load(Fragment (&b)[Mma_tile::MMAS_N], int ki) {\n#pragma unroll\n for (int ni = 0; ni < Mma_tile::MMAS_N; ++ni) {\n // Jump by 16 * #warps row.\n int row = ki * 16 * Cta_tile::WARPS_K;\n\n // Load the data using LDSM.MT88.2.\n uint4 tmp;\n fmha::ldsmt(tmp, this->smem_ + this->smem_read_offset_ + row * Base::BYTES_PER_ROW);\n b[ni].reg(0) = tmp.x;\n b[ni].reg(1) = tmp.y;\n b[ni].reg(2) = tmp.z;\n b[ni].reg(3) = tmp.w;\n\n // Move the pointer for the next ni. I expect the compiler to not recompute those.\n if (Mma_tile::MMAS_N == 4) {\n this->smem_read_offset_ ^= BYTES_PER_LDS * (ni % 2 == 0 ? 2 : 6);\n } else {\n assert(false); // Not implemented!\n }\n }\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Smem_tile_o {\n // The MMA tile.\n using Mma_tile = fmha::Hmma_tile;\n // The accumulators.\n using Accumulator = fmha::Fragment_accumulator;\n // The accumulators.\n using Data_type = typename Accumulator::Data_type;\n\n // The size of each element.\n enum { BYTES_PER_ELEMENT = sizeof(Data_type) };\n // The size of each STS.\n enum { BYTES_PER_STS = 8 };\n // The size of each row in shared memory.\n enum { BYTES_PER_ROW = Cta_tile::N * Cta_tile::WARPS_K * BYTES_PER_ELEMENT };\n\n // The size of each LDS.\n enum { BYTES_PER_LDS = 16 };\n enum { THREADS_PER_ROW = 16 };\n\n // The number of rows.\n enum { ROWS = Cta_tile::M };\n // The number of \"rows\" to process per loop iteration (in the \"epilogue\").\n enum { ROWS_PER_LOOP = ROWS <= 64 ? ROWS : (int)Mma_tile::M_PER_MMA_PER_CTA };\n // The number of outer loops.\n enum { LOOPS = ROWS / ROWS_PER_LOOP };\n // Make sure it matches our expectations.\n static_assert(LOOPS == 1 || LOOPS == (int)Mma_tile::MMAS_M, \"\");\n\n // The number of rows loaded per LDS.\n enum { ROWS_PER_LDS = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW };\n // Do we have to guard against partial writes/reads.\n enum { HAS_INCOMPLETE_LDS = ROWS_PER_LOOP % ROWS_PER_LDS != 0 };\n // The total number of LDS per loop.\n enum { LDS_PER_LOOP = fmha::Div_up::VALUE };\n\n // The amount of shared memory.\n enum { BYTES_PER_TILE = ROWS_PER_LOOP * BYTES_PER_ROW };\n\n // The write pointer.\n uint32_t smem_write_, smem_read_;\n // Is the thread active for the last LDS of the series?\n int is_active_for_last_lds_;\n\n static_assert(BYTES_PER_ROW == 64 * 4 * Cta_tile::WARPS_K);\n static_assert(LOOPS == 1 || LOOPS == (int)Mma_tile::MMAS_M, \"\");\n\n // Ctor.\n inline __device__ Smem_tile_o(void* smem, int tidx) {\n // Get a 32-bit value for the shared memory address.\n uint32_t smem_ = __nvvm_get_smem_pointer(smem);\n\n static_assert(Cta_tile::WARPS_M == 1 && Cta_tile::WARPS_N == 1 &&\n (Cta_tile::WARPS_K == 4 || Cta_tile::WARPS_K == 8));\n\n int write_row = (tidx & 0x1c) / 4;\n int write_col = (tidx);\n\n // Assemble the write pointer.\n smem_write_ = smem_ + write_row * BYTES_PER_ROW + write_col * BYTES_PER_STS;\n\n // The element read by each thread.\n int read_row = tidx / THREADS_PER_ROW;\n int read_col = tidx % THREADS_PER_ROW;\n\n // Take the XOR pattern into account for the column.\n read_col ^= 2 * (read_row & 0x7);\n\n // Assemble the read pointer.\n this->smem_read_ = smem_ + read_row * BYTES_PER_ROW + read_col * BYTES_PER_LDS;\n\n // Is that thread active on the last LDS?\n if (HAS_INCOMPLETE_LDS) {\n this->is_active_for_last_lds_ = read_row + (LDS_PER_LOOP - 1) * ROWS_PER_LDS < Cta_tile::M;\n }\n }\n\n // Load the output fragments.\n inline __device__ void load(uint4 (&out)[LDS_PER_LOOP]) const {\n#pragma unroll\n for (int ii = 0; ii < LDS_PER_LOOP; ++ii) {\n // Load the elements before the reduction (split-K).\n uint4 tmp[Cta_tile::WARPS_K];\n#pragma unroll\n for (int jj = 0; jj < Cta_tile::WARPS_K; ++jj) {\n int imm = ii * ROWS_PER_LDS * BYTES_PER_ROW + jj * Cta_tile::N * BYTES_PER_ELEMENT;\n if (!HAS_INCOMPLETE_LDS || (ii < LDS_PER_LOOP - 1 || this->is_active_for_last_lds_)) {\n fmha::lds(tmp[jj], this->smem_read_ + imm);\n }\n }\n\n // Perform the reduction.\n out[ii] = tmp[0];\n#pragma unroll\n for (int jj = 1; jj < Cta_tile::WARPS_K; ++jj) {\n out[ii] = fmha::fadd4(out[ii], tmp[jj]);\n }\n }\n }\n // Store the accumulators.\n template \n inline __device__ void store(const Accumulator (&acc)[M][N], int mi) {\n enum { M_PER_MMA = Mma_tile::M_PER_MMA_PER_CTA };\n#pragma unroll\n for (int ni = 0; ni < Mma_tile::MMAS_N; ++ni) {\n // The number of MMAs that are stored per loop iteration.\n enum { MMAS_M_PER_LOOP = Mma_tile::MMAS_M / LOOPS };\n\n// Store 1st column of the different MMAs.\n#pragma unroll\n for (int mj = 0; mj < MMAS_M_PER_LOOP; ++mj) {\n // Precompute the immediates to jump between rows.\n int row_0 = (mj * M_PER_MMA + 0) * BYTES_PER_ROW;\n int row_1 = (mj * M_PER_MMA + 8) * BYTES_PER_ROW;\n uint2 tmp0, tmp1;\n tmp0.x = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(0);\n tmp0.y = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(1);\n\n tmp1.x = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(2);\n tmp1.y = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(3);\n\n // Store.\n fmha::sts(this->smem_write_ + row_0, tmp0);\n fmha::sts(this->smem_write_ + row_1, tmp1);\n }\n\n // Swizzle the write pointer using a XOR of 16B.\n this->smem_write_ ^= 32;\n\n// Store 2nd column of the different MMAs.\n#pragma unroll\n for (int mj = 0; mj < MMAS_M_PER_LOOP; ++mj) {\n // Precompute the immediates to jump between rows.\n int row_0 = (mj * M_PER_MMA + 0) * BYTES_PER_ROW;\n int row_1 = (mj * M_PER_MMA + 8) * BYTES_PER_ROW;\n\n uint2 tmp0, tmp1;\n tmp0.x = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(4);\n tmp0.y = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(5);\n\n tmp1.x = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(6);\n tmp1.y = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(7);\n // Store.\n fmha::sts(this->smem_write_ + row_0, tmp0);\n fmha::sts(this->smem_write_ + row_1, tmp1);\n }\n\n // Cancel the previous XOR of 1 + swizzle the write pointer using a XOR of 32B or 64B.\n this->smem_write_ ^= (ni & 1) ? 7 * 32 : 3 * 32;\n }\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Smem_tile_mma {\n using Mma_tile = fmha::Hmma_tile;\n using Fragment = fmha::Fragment_a;\n\n enum { COLS = Cta_tile::N };\n enum { BYTES_PER_ELT = 2 };\n enum { BYTES_PER_STS = 4 };\n enum { BYTES_PER_ROW = COLS * BYTES_PER_ELT }; // TODO\n enum { BYTES_PER_TILE = Cta_tile::M * BYTES_PER_ROW };\n\n enum { WARPS_M = Cta_tile::WARPS_M };\n enum { WARPS_N = Cta_tile::WARPS_N };\n enum { WARPS_K = Cta_tile::WARPS_K };\n\n static_assert(WARPS_K == 1);\n inline __device__ Smem_tile_mma(char* smem, int tidx) {\n smem_ = __nvvm_get_smem_pointer(smem);\n\n int write_col, write_row;\n static_assert(WARPS_M == 1 && (WARPS_N == 4 || WARPS_N == 8) || (WARPS_M == 4 || WARPS_N == 8) || WARPS_N == 1);\n if (WARPS_M == 1 && (WARPS_N == 4 || WARPS_N == 8)) {\n write_row = (tidx & 0x1c) / 4;\n write_col = (tidx & 0xe0) / 4 + (tidx & 0x03);\n } else {\n write_row = (tidx & 0xe0) / 2 + (tidx & 0x1c) / 4;\n write_col = (tidx & 0x03);\n }\n write_col ^= (write_row & 0x07) * 4;\n\n write_offset_ = write_row * BYTES_PER_ROW + write_col * BYTES_PER_STS;\n }\n\n template \n inline __device__ void store(const uint4 (®s)[M][N]) {\n static_assert(COLS == Cta_tile::N);\n for (int mi = 0; mi < M; mi++) {\n for (int ni = 0; ni < N; ni++) {\n size_t offset = write_offset_ + mi * WARPS_M * 16 * BYTES_PER_ROW + ni * WARPS_N * 16 * BYTES_PER_ELT;\n fmha::sts(smem_ + offset + 0 * BYTES_PER_ROW, regs[mi][ni].x);\n fmha::sts(smem_ + offset + 8 * BYTES_PER_ROW, regs[mi][ni].z);\n offset ^= 4 * BYTES_PER_STS;\n fmha::sts(smem_ + offset + 0 * BYTES_PER_ROW, regs[mi][ni].y);\n fmha::sts(smem_ + offset + 8 * BYTES_PER_ROW, regs[mi][ni].w);\n }\n }\n }\n\n uint32_t smem_;\n uint32_t write_offset_;\n uint32_t warp_m;\n uint32_t warp_n;\n uint32_t lane;\n};\n\ntemplate >\nstruct Smem_tile_mma_transposed : public Base {\n enum { BYTES_PER_LDS = 16 };\n enum { BYTES_PER_ROW = Base::BYTES_PER_ROW };\n enum { BYTES_PER_ELT = Base::BYTES_PER_ELT };\n enum { WARPS_M = Base::WARPS_M };\n enum { WARPS_N = Base::WARPS_N };\n static_assert(WARPS_M == 1 && (WARPS_N == 4 || WARPS_N == 8));\n using Fragment = typename Base::Fragment;\n inline __device__ Smem_tile_mma_transposed(char* smem, int tidx) : Base(smem, tidx) {\n static_assert(WARPS_M == 1 && (WARPS_N == 4 || WARPS_N == 8));\n int read_row, read_col;\n read_row = (tidx & 0x0f);\n read_col = (tidx & 0xe0) / 16 + (tidx & 0x1c) / 16;\n\n read_col ^= (read_row & 0x07);\n read_offset_ = read_row * BYTES_PER_ROW + read_col * BYTES_PER_LDS;\n }\n\n template \n inline __device__ void load(Fragment (&frag)[M][N]) {\n static_assert(Base::COLS == Cta_tile::N);\n for (int mi = 0; mi < M; mi++) {\n for (int ni = 0; ni < N; ni++) {\n size_t offset = read_offset_ + mi * WARPS_M * 16 * BYTES_PER_ROW + ni * WARPS_N * 16 * BYTES_PER_ELT;\n uint4 dst;\n fmha::ldsmt(dst, this->smem_ + offset);\n frag[mi][ni].reg(0) = dst.x;\n frag[mi][ni].reg(1) = dst.z; // Fragment A regs col major!\n frag[mi][ni].reg(2) = dst.y;\n frag[mi][ni].reg(3) = dst.w;\n }\n }\n }\n\n uint32_t read_offset_;\n};\n\ntemplate >\nstruct Smem_tile_mma_epilogue : public Base {\n enum { BYTES_PER_LDS = 16 };\n enum { BYTES_PER_ROW = Base::BYTES_PER_ROW };\n enum { BYTES_PER_ELT = Base::BYTES_PER_ELT };\n enum { THREADS_PER_ROW = BYTES_PER_ROW / BYTES_PER_LDS };\n static_assert(THREADS_PER_ROW * BYTES_PER_LDS == BYTES_PER_ROW);\n enum { ROWS_PER_LDS = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW };\n enum { NUM_LDS = Cta_tile::M / ROWS_PER_LDS };\n static_assert(NUM_LDS * ROWS_PER_LDS == Cta_tile::M);\n enum { WARPS_M = Base::WARPS_M };\n enum { WARPS_N = Base::WARPS_N };\n static_assert((WARPS_M == 4 || WARPS_N == 8) || WARPS_N == 1);\n\n using Acc = fmha::Fragment_accumulator;\n\n inline __device__ Smem_tile_mma_epilogue(char* smem, int tidx) : Base(smem, tidx) {\n const int read_row = tidx / THREADS_PER_ROW;\n int read_col = tidx % THREADS_PER_ROW;\n read_col ^= (read_row & 0x07);\n read_offset_ = read_row * BYTES_PER_ROW + read_col * BYTES_PER_LDS;\n }\n\n inline __device__ void load(uint4 (&data)[NUM_LDS]) {\n for (int ii = 0; ii < NUM_LDS; ii++) {\n size_t offset = read_offset_ + ii * ROWS_PER_LDS * BYTES_PER_ROW;\n fmha::lds(data[ii], this->smem_ + offset);\n }\n }\n\n template \n inline __device__ void store(const Acc (&acc)[M][N]) {\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n#pragma unroll\n for (int ni = 0; ni < N; ni++) {\n // 1st row - 4 elements per row.\n float tmp00 = acc[mi][ni].elt(0);\n float tmp01 = acc[mi][ni].elt(1);\n float tmp02 = acc[mi][ni].elt(4);\n float tmp03 = acc[mi][ni].elt(5);\n // 2nd row - 4 elements per row.\n float tmp10 = acc[mi][ni].elt(2);\n float tmp11 = acc[mi][ni].elt(3);\n float tmp12 = acc[mi][ni].elt(6);\n float tmp13 = acc[mi][ni].elt(7);\n\n uint32_t x = fmha::float2_to_half2(tmp00, tmp01);\n uint32_t y = fmha::float2_to_half2(tmp02, tmp03);\n uint32_t z = fmha::float2_to_half2(tmp10, tmp11);\n uint32_t w = fmha::float2_to_half2(tmp12, tmp13);\n\n size_t offset = (this->write_offset_ ^ (ni * 32)) + mi * WARPS_M * 16 * BYTES_PER_ROW;\n fmha::sts(this->smem_ + offset + 0 * BYTES_PER_ROW, x);\n fmha::sts(this->smem_ + offset + 8 * BYTES_PER_ROW, z);\n offset ^= 4 * Base::BYTES_PER_STS;\n fmha::sts(this->smem_ + offset + 0 * BYTES_PER_ROW, y);\n fmha::sts(this->smem_ + offset + 8 * BYTES_PER_ROW, w);\n }\n }\n }\n\n template \n inline __device__ void store(const uint4 (®s)[M][N]) {\n for (int mi = 0; mi < M; mi++) {\n for (int ni = 0; ni < N; ni++) {\n size_t offset = (this->write_offset_ ^ (ni * 32)) + mi * WARPS_M * 16 * BYTES_PER_ROW;\n fmha::sts(this->smem_ + offset + 0 * BYTES_PER_ROW, regs[mi][ni].x);\n fmha::sts(this->smem_ + offset + 8 * BYTES_PER_ROW, regs[mi][ni].z);\n offset ^= 4 * Base::BYTES_PER_STS;\n fmha::sts(this->smem_ + offset + 0 * BYTES_PER_ROW, regs[mi][ni].y);\n fmha::sts(this->smem_ + offset + 8 * BYTES_PER_ROW, regs[mi][ni].w);\n }\n }\n }\n\n uint32_t read_offset_;\n};\n\n} // namespace fmha\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha/softmax.h", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\nnamespace fmha {\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Sum_ {\n enum { IS_SUM = 1 };\n static inline __device__ float apply(float x, float y) { return x + y; }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Max_ {\n enum { IS_SUM = 0 };\n static inline __device__ float apply(float x, float y) { return x > y ? x : y; }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ float apply_exp_(float x, float max) { return __expf(x - max); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct ReadType {};\ntemplate <>\nstruct ReadType<4> {\n using T = float;\n};\ntemplate <>\nstruct ReadType<8> {\n using T = float2;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Smem_tile_reduce {\n // Helper class to distribute MMA tiles reduced over rows per warp over quads.\n\n // The Mma tile.\n using Mma_tile = fmha::Hmma_tile;\n\n // The number of MMAs in M/N dimensions.\n enum { MMAS_M = Mma_tile::MMAS_M };\n enum { MMAS_N = Mma_tile::MMAS_N };\n\n enum { WARPS_M = Cta_tile::WARPS_M };\n enum { WARPS_N = Cta_tile::WARPS_N };\n\n static constexpr int ROWS = WARPS_M * MMAS_M * 16;\n static constexpr int COLS = WARPS_N;\n static_assert(COLS == 4 || COLS == 8);\n static constexpr int ROWS_PER_XOR_PATTERN = (COLS == 8) ? 4 : 8;\n static constexpr int BYTES_PER_TILE = ROWS * COLS * sizeof(float);\n static constexpr int ELTS_PER_TILE = ROWS * COLS;\n\n static constexpr int THREADS_PER_GROUP = Kernel_traits::Gmem_tile_o::THREADS_PER_ROW;\n static_assert(THREADS_PER_GROUP == 16); // DEBUG\n static constexpr int ROWS_PER_WARP = 32 / THREADS_PER_GROUP;\n static constexpr int LOOPS = Kernel_traits::Gmem_tile_o::LOOPS;\n static_assert(LOOPS == 1);\n\n using read_t = typename ReadType::T;\n\n __device__ inline Smem_tile_reduce(float* smem_, const int tidx) {\n int lane = tidx % 32;\n int warp = tidx / 32;\n\n int warp_m = warp % WARPS_M;\n int warp_n = warp / WARPS_M;\n\n qid_ = lane % 4;\n int qp = lane / 4;\n\n // Swizzle the column to avoid 2-fold bank conflicts when we have 8 warps.\n // This won't affect reading as we assume commutative reduction ops.\n const int col = warp_n ^ (qp / ROWS_PER_XOR_PATTERN);\n smem_write_ = &smem_[warp_m * 16 * MMAS_M * WARPS_N + qp * WARPS_N + col];\n smem_read_ = &reinterpret_cast(smem_)[warp_m * 16 * MMAS_M * 4 + qp * 4 + qid_];\n }\n\n __device__ inline void store(float (&frag)[2 * MMAS_M]) {\n if (qid_ == 0) {\n#pragma unroll\n for (int mi = 0; mi < MMAS_M; mi++) {\n int offset = mi * 16 * WARPS_N;\n smem_write_[offset + 0 * 8 * WARPS_N] = frag[mi * 2 + 0];\n smem_write_[offset + 1 * 8 * WARPS_N] = frag[mi * 2 + 1];\n }\n }\n }\n\n __device__ inline void load(read_t (&frag)[2 * MMAS_M]) {\n#pragma unroll\n for (int mi = 0; mi < MMAS_M; mi++) {\n int offset = mi * 16 * 4;\n frag[mi * 2 + 0] = smem_read_[offset + 0 * 8 * 4];\n frag[mi * 2 + 1] = smem_read_[offset + 1 * 8 * 4];\n }\n }\n\n int qid_;\n float* smem_write_;\n read_t* smem_read_;\n};\n\ntemplate \nstruct Softmax_base {\n // The Mma tile.\n using Mma_tile = fmha::Hmma_tile;\n\n // The number of MMAs in M/N dimensions.\n enum { MMAS_M = Mma_tile::MMAS_M };\n enum { MMAS_N = Mma_tile::MMAS_N };\n\n // The number of groups of warp such that we have at most 4 warps writing consecutive elements.\n enum { GROUPS = fmha::Div_up::VALUE };\n // The number of elements that we are going to store per row.\n enum { ELEMENTS_PER_ROW = Cta_tile::WARPS_N / GROUPS };\n // The number of rows.\n enum { ROWS = Cta_tile::M * GROUPS };\n // The total number of elements.\n enum { ELEMENTS = ROWS * ELEMENTS_PER_ROW };\n\n // Ctor.\n template \n inline __device__ Softmax_base(const Params& params, void* smem, int bidb, int tidx)\n : // packed_mask_ptr_(reinterpret_cast(params.packed_mask_ptr)),\n smem_(reinterpret_cast(smem)),\n tidx_(tidx) {\n // Move to the 1st mask loaded by the thread+ tidx;\n // packed_mask_ptr_ += bidb * params.packed_mask_stride_in_bytes + tidx * sizeof(uint32_t);\n\n // Extract the position in the warp.\n int warp = tidx / Cta_tile::THREADS_PER_WARP;\n int lane = tidx % Cta_tile::THREADS_PER_WARP;\n\n // Decompose the warp index into M and N.\n int warp_m = warp % Cta_tile::WARPS_M;\n int warp_n = warp / Cta_tile::WARPS_M;\n\n // Decompose the warp-n index into group/position-inside-the-group.\n int warp_g = warp_n / ELEMENTS_PER_ROW;\n int warp_i = warp_n % ELEMENTS_PER_ROW;\n\n // The location written by the threads.\n int write_row = warp_g * (ROWS / GROUPS) + warp_m * Mma_tile::M_PER_MMA + lane / 4;\n int write_col = warp_i;\n\n // Assemble the write pointer.\n smem_write_ = &smem_[write_row * ELEMENTS_PER_ROW + write_col];\n\n // Assemble the read pointer.\n smem_read_ = &smem_[warp_m * Mma_tile::M_PER_MMA + lane / 4];\n }\n\n template \n inline __device__ void apply_mask(const Mask& mask) {\n#pragma unroll\n for (int mi = 0; mi < MMAS_M; ++mi) {\n#pragma unroll\n for (int ii = 0; ii < 2; ++ii) {\n#pragma unroll\n for (int ni = 0; ni < MMAS_N; ++ni) {\n#pragma unroll\n for (int jj = 0; jj < 4; ++jj) {\n if (!mask.is_valid(mi, ni, ii, jj)) {\n elt_[2 * mi + ii][4 * ni + jj] = -INFINITY;\n }\n }\n }\n }\n }\n }\n\n // Apply the exp to all the elements.\n inline __device__ void apply_exp(const float (&max)[MMAS_M * 2]) {\n#pragma unroll\n for (int mi = 0; mi < MMAS_M * 2; ++mi) {\n#pragma unroll\n for (int ni = 0; ni < MMAS_N * 4; ++ni) {\n elt_[mi][ni] = apply_exp_(elt_[mi][ni], max[mi]);\n }\n }\n }\n\n // Scale all the elements.\n inline __device__ void scale(const float (&sum)[MMAS_M * 2]) {\n // Precompute the inverse sum to normalize. Without -use_fast_math, it makes a huge deal.\n float inv_sum[MMAS_M * 2];\n#pragma unroll\n for (int mi = 0; mi < MMAS_M * 2; ++mi) {\n inv_sum[mi] = (sum[mi] == 0.f || sum[mi] != sum[mi]) ? 1.f : 1.f / sum[mi];\n }\n\n// Update the values.\n#pragma unroll\n for (int mi = 0; mi < MMAS_M * 2; ++mi) {\n#pragma unroll\n for (int ni = 0; ni < MMAS_N * 4; ++ni) {\n elt_[mi][ni] *= inv_sum[mi];\n }\n }\n }\n\n // The pointer to the mask.\n const char* packed_mask_ptr_;\n // Shared memory for the CTA-wide reduction.\n float *smem_, *smem_write_, *smem_read_;\n // The current thread index.\n int tidx_;\n // The elements.\n float elt_[MMAS_M * 2][MMAS_N * 4];\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Softmax : public Softmax_base {\n // The base class.\n using Base = Softmax_base;\n // The fragment.\n using Fragment_a = fmha::Fragment_a;\n\n static_assert(Fragment_a::NUM_REGS == 4);\n\n enum { WARPS_M = Cta_tile::WARPS_M };\n enum { WARPS_N = Cta_tile::WARPS_N };\n // The MMAs.\n enum { MMAS_M = Base::MMAS_M };\n enum { MMAS_N = Base::MMAS_N };\n\n // The accumulators.\n using Accumulator = fmha::Fragment_accumulator;\n using Accumulator_out = Fragment;\n static_assert(Accumulator_out::NUM_REGS == 4);\n\n static_assert(std::is_same::value);\n\n using Smem_tile_red = Smem_tile_reduce;\n static_assert(Smem_tile_red::ELTS_PER_TILE == Cta_tile::M * WARPS_N);\n // Ctor.\n template \n inline __device__ Softmax(const Params& params, void* smem, int bidb, int tidx)\n : Base(params, smem, bidb, tidx),\n params_scale_bmm1_(params.scale_bmm1),\n smem_sum_(static_cast(smem), tidx),\n smem_max_(static_cast(smem) + Smem_tile_red::ELTS_PER_TILE, tidx) {}\n\n // Pack the data to a fragment for the next GEMM.\n template \n inline __device__ void pack(Fragment_a (&dst)[K][M]) const {\n#pragma unroll\n for (int mi = 0; mi < M; ++mi) {\n#pragma unroll\n for (int ki = 0; ki < K; ++ki) {\n // 1st row - 4 elements per row.\n float tmp_00 = this->elt_[2 * mi + 0][4 * ki + 0];\n float tmp_01 = this->elt_[2 * mi + 0][4 * ki + 1];\n float tmp_02 = this->elt_[2 * mi + 0][4 * ki + 2];\n float tmp_03 = this->elt_[2 * mi + 0][4 * ki + 3];\n\n // 2nd row - 4 elements per row.\n float tmp_10 = this->elt_[2 * mi + 1][4 * ki + 0];\n float tmp_11 = this->elt_[2 * mi + 1][4 * ki + 1];\n float tmp_12 = this->elt_[2 * mi + 1][4 * ki + 2];\n float tmp_13 = this->elt_[2 * mi + 1][4 * ki + 3];\n\n // Pack to 4 registers.\n dst[ki][mi].reg(0) = fmha::float2_to_half2(tmp_00, tmp_01);\n dst[ki][mi].reg(1) = fmha::float2_to_half2(tmp_10, tmp_11);\n dst[ki][mi].reg(2) = fmha::float2_to_half2(tmp_02, tmp_03);\n dst[ki][mi].reg(3) = fmha::float2_to_half2(tmp_12, tmp_13);\n }\n }\n }\n\n // Scale FP32 fragments\n inline __device__ void unpack(const Accumulator (&acc)[MMAS_M][MMAS_N]) {\n const float scalef = reinterpret_cast(this->params_scale_bmm1_);\n\n#pragma unroll\n for (int mi = 0; mi < MMAS_M; ++mi) {\n#pragma unroll\n for (int ni = 0; ni < MMAS_N; ++ni) {\n // 1st row - 4 elements per row.\n this->elt_[2 * mi + 0][4 * ni + 0] = acc[mi][ni].elt(0) * scalef;\n this->elt_[2 * mi + 0][4 * ni + 1] = acc[mi][ni].elt(1) * scalef;\n this->elt_[2 * mi + 0][4 * ni + 2] = acc[mi][ni].elt(4) * scalef;\n this->elt_[2 * mi + 0][4 * ni + 3] = acc[mi][ni].elt(5) * scalef;\n // 2nd row - 4 elements per row.\n this->elt_[2 * mi + 1][4 * ni + 0] = acc[mi][ni].elt(2) * scalef;\n this->elt_[2 * mi + 1][4 * ni + 1] = acc[mi][ni].elt(3) * scalef;\n this->elt_[2 * mi + 1][4 * ni + 2] = acc[mi][ni].elt(6) * scalef;\n this->elt_[2 * mi + 1][4 * ni + 3] = acc[mi][ni].elt(7) * scalef;\n }\n }\n }\n // Scale FP32 fragments\n inline __device__ void unpack_noscale(const Accumulator (&acc)[MMAS_M][MMAS_N]) {\n#pragma unroll\n for (int mi = 0; mi < MMAS_M; ++mi) {\n#pragma unroll\n for (int ni = 0; ni < MMAS_N; ++ni) {\n // 1st row - 4 elements per row.\n this->elt_[2 * mi + 0][4 * ni + 0] = acc[mi][ni].elt(0);\n this->elt_[2 * mi + 0][4 * ni + 1] = acc[mi][ni].elt(1);\n this->elt_[2 * mi + 0][4 * ni + 2] = acc[mi][ni].elt(4);\n this->elt_[2 * mi + 0][4 * ni + 3] = acc[mi][ni].elt(5);\n // 2nd row - 4 elements per row.\n this->elt_[2 * mi + 1][4 * ni + 0] = acc[mi][ni].elt(2);\n this->elt_[2 * mi + 1][4 * ni + 1] = acc[mi][ni].elt(3);\n this->elt_[2 * mi + 1][4 * ni + 2] = acc[mi][ni].elt(6);\n this->elt_[2 * mi + 1][4 * ni + 3] = acc[mi][ni].elt(7);\n }\n }\n }\n\n template \n __device__ inline void reduce_(float (&frag)[2 * MMAS_M], Operator& op, Smem_tile_red& smem_red) {\n for (int mi = 0; mi < 2 * MMAS_M; mi++) {\n frag[mi] = this->elt_[mi][0];\n for (int ni = 1; ni < 4 * MMAS_N; ni++) {\n frag[mi] = op(frag[mi], this->elt_[mi][ni]);\n }\n }\n quad_reduce(frag, frag, op);\n\n smem_red.store(frag);\n __syncthreads();\n typename Smem_tile_red::read_t tmp[2 * MMAS_M];\n smem_red.load(tmp);\n\n quad_allreduce(frag, tmp, op);\n }\n\n __device__ inline void reduce_max(float (&frag)[2 * MMAS_M]) {\n MaxOp max;\n reduce_(frag, max, smem_max_);\n }\n\n __device__ inline void reduce_sum(float (&frag)[2 * MMAS_M]) {\n SumOp sum;\n reduce_(frag, sum, smem_sum_);\n }\n\n const uint32_t params_scale_bmm1_;\n Smem_tile_red smem_max_;\n Smem_tile_red smem_sum_;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace fmha\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha/utils.h", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\n#include \n#include \n#include \n\nextern \"C\" __device__ uint32_t __nvvm_get_smem_pointer(void* ptr);\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nnamespace fmha {\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Row {};\nstruct Col {};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Next_power_of_two {};\n\ntemplate \nstruct Next_power_of_two {\n enum { VALUE = M };\n};\ntemplate <>\nstruct Next_power_of_two<3, false> {\n enum { VALUE = 4 };\n};\ntemplate <>\nstruct Next_power_of_two<5, false> {\n enum { VALUE = 8 };\n};\ntemplate <>\nstruct Next_power_of_two<6, false> {\n enum { VALUE = 8 };\n};\ntemplate <>\nstruct Next_power_of_two<7, false> {\n enum { VALUE = 8 };\n};\ntemplate <>\nstruct Next_power_of_two<9, false> {\n enum { VALUE = 16 };\n};\ntemplate <>\nstruct Next_power_of_two<10, false> {\n enum { VALUE = 16 };\n};\ntemplate <>\nstruct Next_power_of_two<11, false> {\n enum { VALUE = 16 };\n};\ntemplate <>\nstruct Next_power_of_two<12, false> {\n enum { VALUE = 16 };\n};\ntemplate <>\nstruct Next_power_of_two<13, false> {\n enum { VALUE = 16 };\n};\ntemplate <>\nstruct Next_power_of_two<14, false> {\n enum { VALUE = 16 };\n};\ntemplate <>\nstruct Next_power_of_two<15, false> {\n enum { VALUE = 16 };\n};\ntemplate <>\nstruct Next_power_of_two<24, false> {\n enum { VALUE = 32 };\n};\ntemplate <>\nstruct Next_power_of_two<48, false> {\n enum { VALUE = 64 };\n};\ntemplate <>\nstruct Next_power_of_two<80, false> {\n enum { VALUE = 128 };\n};\ntemplate <>\nstruct Next_power_of_two<96, false> {\n enum { VALUE = 128 };\n};\ntemplate <>\nstruct Next_power_of_two<112, false> {\n enum { VALUE = 128 };\n};\ntemplate <>\nstruct Next_power_of_two<144, false> {\n enum { VALUE = 256 };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Prev_power_of_two {};\n\ntemplate \nstruct Prev_power_of_two {\n enum { VALUE = N };\n};\ntemplate <>\nstruct Prev_power_of_two<3, false> {\n enum { VALUE = 2 };\n};\ntemplate <>\nstruct Prev_power_of_two<5, false> {\n enum { VALUE = 4 };\n};\ntemplate <>\nstruct Prev_power_of_two<6, false> {\n enum { VALUE = 4 };\n};\ntemplate <>\nstruct Prev_power_of_two<7, false> {\n enum { VALUE = 4 };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Div_up {\n enum { VALUE = (M + N - 1) / N };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Max {\n enum { VALUE = A >= B ? A : B };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Max_3 {\n enum { VALUE = Max::VALUE, C>::VALUE };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Min {\n enum { VALUE = A <= B ? A : B };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Uint_from_size_in_bytes {};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <>\nstruct Uint_from_size_in_bytes<1> {\n using Type = uint8_t;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <>\nstruct Uint_from_size_in_bytes<2> {\n using Type = uint16_t;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <>\nstruct Uint_from_size_in_bytes<4> {\n using Type = uint32_t;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <>\nstruct Uint_from_size_in_bytes<8> {\n using Type = uint2;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <>\nstruct Uint_from_size_in_bytes<16> {\n using Type = uint4;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Warp_masks {};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <>\nstruct Warp_masks<8, 1, 1> {\n enum { M = 0xe0, N = 0x00, K = 0x00 };\n};\ntemplate <>\nstruct Warp_masks<4, 2, 1> {\n enum { M = 0x60, N = 0x80, K = 0x00 };\n};\ntemplate <>\nstruct Warp_masks<4, 1, 2> {\n enum { M = 0x60, N = 0x00, K = 0x80 };\n};\ntemplate <>\nstruct Warp_masks<4, 1, 1> {\n enum { M = 0x60, N = 0x00, K = 0x00 };\n};\ntemplate <>\nstruct Warp_masks<2, 4, 1> {\n enum { M = 0x20, N = 0xc0, K = 0x00 };\n};\ntemplate <>\nstruct Warp_masks<2, 2, 2> {\n enum { M = 0x20, N = 0x40, K = 0x80 };\n};\ntemplate <>\nstruct Warp_masks<2, 2, 1> {\n enum { M = 0x20, N = 0x40, K = 0x00 };\n};\ntemplate <>\nstruct Warp_masks<2, 1, 2> {\n enum { M = 0x20, N = 0x00, K = 0x40 };\n};\ntemplate <>\nstruct Warp_masks<2, 1, 1> {\n enum { M = 0x20, N = 0x00, K = 0x00 };\n};\ntemplate <>\nstruct Warp_masks<1, 8, 1> {\n enum { M = 0x00, N = 0xe0, K = 0x00 };\n};\ntemplate <>\nstruct Warp_masks<1, 4, 2> {\n enum { M = 0x00, N = 0x60, K = 0x80 };\n};\ntemplate <>\nstruct Warp_masks<1, 4, 1> {\n enum { M = 0x00, N = 0x60, K = 0x00 };\n};\ntemplate <>\nstruct Warp_masks<1, 2, 2> {\n enum { M = 0x00, N = 0x20, K = 0x40 };\n};\ntemplate <>\nstruct Warp_masks<1, 2, 1> {\n enum { M = 0x00, N = 0x20, K = 0x00 };\n};\ntemplate <>\nstruct Warp_masks<1, 1, 4> {\n enum { M = 0x00, N = 0x00, K = 0x60 };\n};\ntemplate <>\nstruct Warp_masks<1, 1, 2> {\n enum { M = 0x00, N = 0x00, K = 0x20 };\n};\ntemplate <>\nstruct Warp_masks<1, 1, 1> {\n enum { M = 0x00, N = 0x00, K = 0x00 };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ __host__ T div_up(T m, T n) {\n return (m + n - 1) / n;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline int clz(int x) {\n for (int i = 31; i >= 0; --i) {\n if ((1 << i) & x) {\n return 31 - i;\n }\n }\n return 32;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline int find_log_2(int x, bool round_up = false) {\n int a = 31 - clz(x);\n if (round_up) {\n a += (x & (x - 1)) ? 1 : 0;\n }\n return a;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint32_t hadd2(uint32_t a, uint32_t b) {\n uint32_t c;\n asm volatile(\"add.f16x2 %0, %1, %2;\\n\" : \"=r\"(c) : \"r\"(a), \"r\"(b));\n return c;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint32_t hmin2(uint32_t a, uint32_t b) {\n uint32_t c;\n asm volatile(\"min.f16x2 %0, %1, %2;\" : \"=r\"(c) : \"r\"(a), \"r\"(b));\n return c;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint32_t hmul2(uint32_t a, uint32_t b) {\n uint32_t c;\n asm volatile(\"mul.f16x2 %0, %1, %2;\\n\" : \"=r\"(c) : \"r\"(a), \"r\"(b));\n return c;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint2 hmul4(uint2 a, uint2 b) {\n uint2 c;\n c.x = hmul2(a.x, b.x);\n c.y = hmul2(a.y, b.y);\n return c;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint4 hmul8(uint4 a, uint4 b) {\n uint4 c;\n c.x = hmul2(a.x, b.x);\n c.y = hmul2(a.y, b.y);\n c.z = hmul2(a.z, b.z);\n c.w = hmul2(a.w, b.w);\n return c;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint4 hmul8(uint32_t a, uint4 b) {\n uint4 c;\n c.x = hmul2(a, b.x);\n c.y = hmul2(a, b.y);\n c.z = hmul2(a, b.z);\n c.w = hmul2(a, b.w);\n return c;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint32_t hrelu2(uint32_t x, uint32_t lb = 0) {\n uint32_t res;\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800\n asm volatile(\"max.f16x2 %0, %1, %2;\\n\" : \"=r\"(res) : \"r\"(x), \"r\"(lb));\n#else\n const uint32_t zero = 0u;\n asm volatile(\n \"{\\n\"\n \"\\t .reg .f16x2 sela;\\n\"\n \"\\t set.gtu.u32.f16x2 sela, %1, %2;\\n\"\n \"\\t and.b32 %0, sela, %1;\\n\"\n \"}\\n\"\n : \"=r\"(res)\n : \"r\"(x), \"r\"(zero));\n#endif\n return res;\n}\nstatic inline __device__ uint32_t habs2(uint32_t x) {\n uint32_t res;\n asm volatile(\"abs.f16x2 %0, %1;\\n\" : \"=r\"(res) : \"r\"(x));\n return res;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\ntemplate \nstatic inline __device__ T clamp(T x, T lb, T ub) {\n return x < lb ? lb : (x > ub ? ub : x);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint16_t clamp_to_zero(uint16_t x) {\n uint16_t mask;\n asm volatile(\"set.gtu %0, %1, 0;\" : \"=h\"(mask) : \"h\"(x));\n return mask & x;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint16_t float_to_half(float f) {\n uint16_t h;\n asm volatile(\"cvt.rn.f16.f32 %0, %1;\" : \"=h\"(h) : \"f\"(f));\n return h;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint32_t float2_to_half2(float a, float b) {\n uint32_t c;\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800\n asm volatile(\"cvt.rn.f16x2.f32 %0, %1, %2;\\n\" : \"=r\"(c) : \"f\"(b), \"f\"(a));\n#else\n uint16_t lo = float_to_half(a);\n uint16_t hi = float_to_half(b);\n asm volatile(\"mov.b32 %0, {%1, %2};\\n\" : \"=r\"(c) : \"h\"(lo), \"h\"(hi));\n#endif\n return c;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint32_t float_to_half2(float a) { return float2_to_half2(a, a); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint32_t float2_to_half2(const float2& f) { return float2_to_half2(f.x, f.y); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint2 float4_to_half4(float x, float y, float z, float w) {\n uint2 d;\n d.x = float2_to_half2(x, y);\n d.y = float2_to_half2(z, w);\n return d;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint32_t hfma2(uint32_t a, uint32_t b, uint32_t c) {\n uint32_t d;\n asm volatile(\"fma.rn.f16x2 %0, %1, %2, %3;\\n\" : \"=r\"(d) : \"r\"(a), \"r\"(b), \"r\"(c));\n return d;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint32_t hfma2_relu(uint32_t a, uint32_t b, uint32_t c) {\n uint32_t d;\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800\n asm volatile(\"fma.rn.f16x2.relu %0, %1, %2, %3;\" : \"=r\"(d) : \"r\"(a), \"r\"(b), \"r\"(c));\n#else\n d = hrelu2(hfma2(a, b, c));\n#endif\n return d;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint32_t h0_h0(uint32_t x) {\n uint32_t y;\n asm volatile(\"{.reg .f16 lo, hi; mov.b32 {lo, hi}, %1; mov.b32 %0, {lo, lo};}\\n\" : \"=r\"(y) : \"r\"(x));\n return y;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ float h0_to_float(uint32_t h2) {\n float f;\n asm volatile(\n \"{\\n\"\n \".reg .f16 lo, hi;\\n\"\n \"mov.b32 {lo, hi}, %1;\\n\"\n \"cvt.f32.f16 %0, lo;\\n\"\n \"}\\n\"\n : \"=f\"(f)\n : \"r\"(h2));\n return f;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint32_t h1_h1(uint32_t x) {\n uint32_t y;\n asm volatile(\"{.reg .f16 lo, hi; mov.b32 {lo, hi}, %1; mov.b32 %0, {hi, hi};}\\n\" : \"=r\"(y) : \"r\"(x));\n return y;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint16_t hadd(uint16_t a, uint16_t b) {\n uint16_t d;\n asm volatile(\"add.f16 %0, %1, %2;\" : \"=h\"(d) : \"h\"(a), \"h\"(b));\n return d;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint32_t hadd(uint32_t a, uint32_t b) { return hadd2(a, b); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint2 hadd4(uint2 a, uint2 b) {\n uint2 c;\n c.x = hadd2(a.x, b.x);\n c.y = hadd2(a.y, b.y);\n return c;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint2 hadd(uint2 a, uint2 b) { return hadd4(a, b); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint4 hadd8(uint4 a, uint4 b) {\n uint4 c;\n c.x = hadd2(a.x, b.x);\n c.y = hadd2(a.y, b.y);\n c.z = hadd2(a.z, b.z);\n c.w = hadd2(a.w, b.w);\n return c;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint4 fadd4(uint4 a, uint4 b) {\n float4 c;\n c.x = reinterpret_cast(a.x) + reinterpret_cast(b.x);\n c.y = reinterpret_cast(a.y) + reinterpret_cast(b.y);\n c.z = reinterpret_cast(a.z) + reinterpret_cast(b.z);\n c.w = reinterpret_cast(a.w) + reinterpret_cast(b.w);\n return reinterpret_cast(c);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint4 hadd(uint4 a, uint4 b) { return hadd8(a, b); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ float half_to_float(uint16_t h) {\n float f;\n asm volatile(\"cvt.f32.f16 %0, %1;\\n\" : \"=f\"(f) : \"h\"(h));\n return f;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ float2 half2_to_float2(uint32_t x) {\n uint16_t lo, hi;\n asm volatile(\"mov.b32 {%0, %1}, %2;\\n\" : \"=h\"(lo), \"=h\"(hi) : \"r\"(x));\n return make_float2(half_to_float(lo), half_to_float(hi));\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ void half2_to_float2(float& x, float& y, uint32_t h) {\n float2 tmp = half2_to_float2(h);\n x = tmp.x;\n y = tmp.y;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint16_t hfma(uint16_t a, uint16_t b, uint16_t c) {\n uint16_t d;\n asm volatile(\"fma.rn.f16 %0, %1, %2, %3;\" : \"=h\"(d) : \"h\"(a), \"h\"(b), \"h\"(c));\n return d;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ uint16_t hmul(uint16_t a, uint16_t b) {\n uint16_t d;\n asm volatile(\"mul.f16 %0, %1, %2;\" : \"=h\"(d) : \"h\"(a), \"h\"(b));\n return d;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ float sigmoid(float x) { return 1.f / (1.f + expf(-x)); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void clear(uint16_t& dst) { dst = uint16_t(0); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void clear(uint32_t& dst) { dst = 0u; }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void clear(uint2& dst) { dst = make_uint2(0u, 0u); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void clear(uint4& dst) { dst = make_uint4(0u, 0u, 0u, 0u); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// P R E D I C A T E P A C K I N G\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\nenum { BYTES_PER_REG = 4, PREDS_PER_BYTE = 4, PREDS_PER_REG = BYTES_PER_REG * PREDS_PER_BYTE };\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// G E N E R I C P R E D I C A T E D L D G S T S\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void load_(Functor& fct, const uint32_t (&preds)[M]) {\n // The number of complete bytes (where we use all the predicates in a byte).\n enum { COMPLETE = N / PREDS_PER_BYTE };\n // Make sure we did allocate enough predicates.\n static_assert(Div_up::VALUE <= M, \"\");\n // The remainder.\n enum { REMAINDER = N - COMPLETE * PREDS_PER_BYTE };\n // Make sure we got the math right and the remainder is between 0 and 3.\n static_assert(REMAINDER >= 0 && REMAINDER <= 3, \"\");\n // The mask to extract the predicates.\n enum { COMPLETE_MASK = (1 << PREDS_PER_BYTE) - 1 };\n\n// Clear the fetch registers.\n#pragma unroll\n for (int ii = 0; ii < N; ++ii) {\n fct.clear(ii);\n }\n\n // Run complete steps.\n bool p[PREDS_PER_BYTE];\n#pragma unroll\n for (int ii = 0; ii < COMPLETE; ++ii) {\n // The predicate.\n uint32_t reg = preds[ii / BYTES_PER_REG];\n\n// Extract the predicates.\n#pragma unroll\n for (int jj = 0; jj < PREDS_PER_BYTE; ++jj) {\n uint32_t mask = 1u << (ii % BYTES_PER_REG * 8 + jj);\n p[jj] = (reg & mask) != 0u;\n }\n\n// Issue the loads.\n#pragma unroll\n for (int jj = 0; jj < PREDS_PER_BYTE; ++jj) {\n fct.load(ii * PREDS_PER_BYTE + jj, p[jj]);\n }\n }\n\n // Skip the rest of the code if we do not have a remainder.\n if (REMAINDER > 0) {\n // The mask to extract the predicates.\n enum { REMAINDER_MASK = (1 << REMAINDER) - 1 };\n\n // The predicate register.\n uint32_t reg = preds[COMPLETE / BYTES_PER_REG];\n\n// Extract the predicates.\n#pragma unroll\n for (int jj = 0; jj < PREDS_PER_BYTE; ++jj) {\n uint32_t mask = 1u << (COMPLETE % BYTES_PER_REG * 8 + jj);\n p[jj] = (reg & mask) != 0u;\n }\n\n// Issue the loads.\n#pragma unroll\n for (int ii = 0; ii < REMAINDER; ++ii) {\n fct.load(COMPLETE * PREDS_PER_BYTE + ii, p[ii]);\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void load_(Functor& fct, uint32_t preds) {\n uint32_t tmp[1] = {preds};\n load_(fct, tmp);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// L D G\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void ldg(uint8_t& dst, const void* ptr) { dst = *reinterpret_cast(ptr); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void ldg(uint16_t& dst, const void* ptr) { dst = *reinterpret_cast(ptr); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void ldg(uint32_t& dst, const void* ptr) { dst = *reinterpret_cast(ptr); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void ldg(uint2& dst, const void* ptr) { dst = *reinterpret_cast(ptr); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void ldg(uint4& dst, const void* ptr) { dst = *reinterpret_cast(ptr); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Ldg_functor {\n // Ctor.\n inline __device__ Ldg_functor(Data_type (&fetch)[N], const void* (&ptrs)[N]) : fetch_(fetch), ptrs_(ptrs) {}\n\n // Clear the element.\n inline __device__ void clear(int ii) { fmha::clear(fetch_[ii]); }\n\n // Trigger the loads.\n inline __device__ void load(int ii, bool p) {\n if (p) {\n ldg(fetch_[ii], ptrs_[ii]);\n }\n }\n\n // The fetch registers.\n Data_type (&fetch_)[N];\n // The pointers.\n const void* (&ptrs_)[N];\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void ldg_(Data_type (&fetch)[N], const void* (&ptrs)[N], uint32_t (&preds)[M]) {\n Ldg_functor fct(fetch, ptrs);\n load_(fct, preds);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void ldg(uint8_t (&fetch)[N], const void* (&ptrs)[N], uint32_t (&preds)[M]) {\n ldg_(fetch, ptrs, preds);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void ldg(uint16_t (&fetch)[N], const void* (&ptrs)[N], uint32_t (&preds)[M]) {\n ldg_(fetch, ptrs, preds);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void ldg(uint32_t (&fetch)[N], const void* (&ptrs)[N], uint32_t (&preds)[M]) {\n ldg_(fetch, ptrs, preds);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void ldg(uint2 (&fetch)[N], const void* (&ptrs)[N], uint32_t (&preds)[M]) {\n ldg_(fetch, ptrs, preds);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void ldg(uint4 (&fetch)[N], const void* (&ptrs)[N], uint32_t (&preds)[M]) {\n ldg_(fetch, ptrs, preds);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// L D S\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void lds(uint16_t& dst, uint32_t ptr) {\n asm volatile(\"ld.shared.b16 %0, [%1];\\n\" : \"=h\"(dst) : \"r\"(ptr));\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void lds(uint32_t& dst, uint32_t ptr) {\n asm volatile(\"ld.shared.b32 %0, [%1];\\n\" : \"=r\"(dst) : \"r\"(ptr));\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void lds(uint2& dst, uint32_t ptr) {\n asm volatile(\"ld.shared.v2.b32 {%0, %1}, [%2];\\n\" : \"=r\"(dst.x), \"=r\"(dst.y) : \"r\"(ptr));\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void lds(uint4& dst, uint32_t ptr) {\n asm volatile(\"ld.shared.v4.b32 {%0, %1, %2, %3}, [%4];\\n\"\n : \"=r\"(dst.x), \"=r\"(dst.y), \"=r\"(dst.z), \"=r\"(dst.w)\n : \"r\"(ptr));\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// L D S M\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void ldsm(uint32_t& dst, uint32_t ptr) {\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 730\n asm volatile(\"ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\\n\" : \"=r\"(dst) : \"r\"(ptr));\n#endif\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void ldsmt(uint32_t& dst, uint32_t ptr) {\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 730\n asm volatile(\"ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16 {%0}, [%1];\\n\" : \"=r\"(dst) : \"r\"(ptr));\n#endif\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void ldsm(uint2& dst, uint32_t ptr) {\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 730\n asm volatile(\"ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\\n\" : \"=r\"(dst.x), \"=r\"(dst.y) : \"r\"(ptr));\n#endif\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void ldsmt(uint2& dst, uint32_t ptr) {\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 730\n asm volatile(\"ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0, %1}, [%2];\\n\"\n : \"=r\"(dst.x), \"=r\"(dst.y)\n : \"r\"(ptr));\n#endif\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void ldsm(uint4& dst, uint32_t ptr) {\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 730\n asm volatile(\"ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\\n\"\n : \"=r\"(dst.x), \"=r\"(dst.y), \"=r\"(dst.z), \"=r\"(dst.w)\n : \"r\"(ptr));\n#endif\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void ldsmt(uint4& dst, uint32_t ptr) {\n#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 730\n asm volatile(\"ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\\n\"\n : \"=r\"(dst.x), \"=r\"(dst.y), \"=r\"(dst.z), \"=r\"(dst.w)\n : \"r\"(ptr));\n#endif\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// S T G\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void stg(void* ptr, uint8_t val) { *reinterpret_cast(ptr) = val; }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void stg(void* ptr, uint16_t val) { *reinterpret_cast(ptr) = val; }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void stg(void* ptr, uint32_t val) { *reinterpret_cast(ptr) = val; }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void stg(void* ptr, uint2 val) { *reinterpret_cast(ptr) = val; }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void stg(void* ptr, uint4 val) { *reinterpret_cast(ptr) = val; }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// S T S\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void sts(uint32_t ptr, uint16_t val) {\n asm volatile(\"st.shared.b16 [%0], %1;\\n\" : : \"r\"(ptr), \"h\"(val));\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void sts(uint32_t ptr, uint32_t val) {\n asm volatile(\"st.shared.b32 [%0], %1;\\n\" : : \"r\"(ptr), \"r\"(val));\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void sts(uint32_t ptr, uint2 val) {\n asm volatile(\"st.shared.v2.b32 [%0], {%1, %2};\\n\" : : \"r\"(ptr), \"r\"(val.x), \"r\"(val.y));\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void sts(uint32_t ptr, uint4 val) {\n asm volatile(\"st.shared.v4.b32 [%0], {%1, %2, %3, %4};\\n\"\n :\n : \"r\"(ptr), \"r\"(val.x), \"r\"(val.y), \"r\"(val.z), \"r\"(val.w));\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void sts_(uint32_t (&ptrs)[N], const Data_type (&data)[N]) {\n#pragma unroll\n for (int ii = 0; ii < N; ++ii) {\n sts(ptrs[ii], data[ii]);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void sts(uint32_t (&ptrs)[N], const uint16_t (&data)[N]) {\n sts_(ptrs, data);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void sts(uint32_t (&ptrs)[N], const uint32_t (&data)[N]) {\n sts_(ptrs, data);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void sts(uint32_t (&ptrs)[N], const uint2 (&data)[N]) {\n sts_(ptrs, data);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void sts(uint32_t (&ptrs)[N], const uint4 (&data)[N]) {\n sts_(ptrs, data);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct MaxOp {\n __device__ inline T operator()(T const& x, T const& y) { return x > y ? x : y; }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct SumOp {\n __device__ inline T operator()(T const& x, T const& y) { return x + y; }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Allreduce {\n static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);\n template \n static __device__ inline T run(T x, Operator& op) {\n constexpr int OFFSET = THREADS / 2;\n x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));\n return Allreduce::run(x, op);\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <>\nstruct Allreduce<2> {\n template \n static __device__ inline T run(T x, Operator& op) {\n x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));\n return x;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__device__ inline void quad_reduce(float (&dst)[M], float (&src)[M], Operator& op) {\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n dst[mi] = src[mi];\n dst[mi] = op(dst[mi], __shfl_down_sync(uint32_t(-1), dst[mi], 2));\n dst[mi] = op(dst[mi], __shfl_down_sync(uint32_t(-1), dst[mi], 1));\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__device__ inline void quad_reduce(float (&dst)[M], float2 (&src)[M], Operator& op) {\n float tmp[M];\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n tmp[mi] = op(src[mi].x, src[mi].y);\n }\n quad_reduce(dst, tmp, op);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__device__ inline void quad_allreduce(float (&dst)[M], float (&src)[M], Operator& op) {\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n dst[mi] = src[mi];\n dst[mi] = Allreduce<4>::run(dst[mi], op);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__device__ inline void quad_allreduce(float (&dst)[M], float2 (&src)[M], Operator& op) {\n float tmp[M];\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n tmp[mi] = op(src[mi].x, src[mi].y);\n }\n quad_allreduce(dst, tmp, op);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace fmha\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha.h", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\n#include \n\n#include \n\n#ifdef OLD_GENERATOR_PATH\n#include \n#else\n#include \n#endif\n\n#include \n\n#include \n\nconstexpr int TOTAL_DIM = 0;\nconstexpr int THREE_DIM = 1;\nconstexpr int H_DIM = 2;\nconstexpr int D_DIM = 3;\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Qkv_params {\n // The QKV matrices.\n void* __restrict__ qkv_ptr;\n\n // The stride between rows of the Q, K and V matrices.\n size_t qkv_stride_in_bytes;\n\n // The number of heads.\n int h;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Fused_multihead_attention_fprop_params : public Qkv_params {\n // The dQKV matrices.\n void* __restrict__ dqkv_ptr;\n\n // Temporary for dKV.\n void* __restrict__ dkv_ptr;\n\n // The O matrix (output).\n void* __restrict__ o_ptr;\n\n // The stride between rows of O.\n int64_t o_stride_in_bytes;\n\n // The pointer to the S matrix, overwritten by the dP matrix (bwd).\n void* __restrict__ s_ptr;\n // The stride between rows of the S matrix.\n int64_t s_stride_in_bytes;\n\n // The dimensions.\n int b, s, d;\n\n // The scaling factors for the kernel.\n uint32_t scale_bmm1, scale_softmax, scale_bmm2;\n\n // array of length b+1 holding starting offset of each sequence.\n int* __restrict__ cu_seqlens;\n\n // The dropout probability (probability of keeping an activation).\n float p_dropout;\n\n // Scale factor of 1 / (1 - p_dropout).\n float rp_dropout;\n\n // Scale factor of 1 / (1 - p_dropout), in half2.\n uint32_t scale_dropout;\n\n // Random state.\n at::PhiloxCudaState philox_args;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Launch_params {\n Launch_params(cudaDeviceProp* props_, cudaStream_t stream_, bool is_training_, bool is_nl_)\n : elts_per_thread(0), props(props_), stream(stream_), is_training(is_training_), is_nl(is_nl_) {}\n\n size_t elts_per_thread;\n\n cudaDeviceProp* props;\n\n cudaStream_t stream;\n\n bool is_training;\n\n Kernel_params params;\n int num_full_heads;\n int num_main_groups;\n int heads_last_wave;\n int main_steps;\n int rest_steps;\n bool is_nl;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid run_fmha_fp16_128_64_sm80(Launch_params& launch_params,\n const bool configure);\nvoid run_fmha_fp16_256_64_sm80(Launch_params& launch_params,\n const bool configure);\nvoid run_fmha_fp16_384_64_sm80(Launch_params& launch_params,\n const bool configure);\nvoid run_fmha_fp16_512_64_sm80(Launch_params& launch_params,\n const bool configure);\n\nvoid run_fmha_dgrad_fp16_128_64_sm80(const Fused_multihead_attention_fprop_params& params, cudaStream_t stream);\nvoid run_fmha_dgrad_fp16_256_64_sm80(const Fused_multihead_attention_fprop_params& params, cudaStream_t stream);\nvoid run_fmha_dgrad_fp16_384_64_sm80(const Fused_multihead_attention_fprop_params& params, cudaStream_t stream);\nvoid run_fmha_dgrad_fp16_512_64_sm80(const Fused_multihead_attention_fprop_params& params, cudaStream_t stream);\n\nvoid run_fmha_fp16_512_64_sm80_nl(const Fused_multihead_attention_fprop_params& params, const bool is_training,\n const int num_chunks, cudaStream_t stream);\n\nvoid run_fmha_dgrad_fp16_512_64_sm80_nl(const Fused_multihead_attention_fprop_params& params, const int num_chunks,\n cudaStream_t stream);\n\nvoid fmha_run_noloop_reduce(void* out, const void* in, const int* cu_seqlens, const int hidden_size,\n const int batch_size, const int total, const int num_chunks, cudaStream_t stream);\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_128_64_kernel.sm80.cu", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_dgrad_kernel_1xN_reload.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<128, 64, 16, 1, 4, 0x08u>;\n\nextern \"C\" __global__ void fmha_dgrad_fp16_128_64_sm80_kernel(Fused_multihead_attention_fprop_params params) {\n fmha::compute_dv_1xN(params);\n fmha::compute_dq_dk_1xN(params);\n}\n\nvoid run_fmha_dgrad_fp16_128_64_sm80(const Fused_multihead_attention_fprop_params& params, cudaStream_t stream) {\n constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);\n constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;\n constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;\n constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;\n\n using Smem_tile_s = fmha::Smem_tile_mma_transposed;\n constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;\n static_assert(smem_size_s == 16 * 128 * 2);\n static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);\n\n constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;\n constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;\n constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(fmha_dgrad_fp16_128_64_sm80_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n dim3 grid(params.h, params.b);\n fmha_dgrad_fp16_128_64_sm80_kernel<<>>(params);\n}\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_256_64_kernel.sm80.cu", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_dgrad_kernel_1xN_reload.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<256, 64, 16, 1, 4, 0x08u>;\n\nextern \"C\" __global__ void fmha_dgrad_fp16_256_64_sm80_kernel(Fused_multihead_attention_fprop_params params) {\n fmha::compute_dv_1xN(params);\n fmha::compute_dq_dk_1xN(params);\n}\n\nvoid run_fmha_dgrad_fp16_256_64_sm80(const Fused_multihead_attention_fprop_params& params, cudaStream_t stream) {\n constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);\n constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;\n constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;\n constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;\n\n using Smem_tile_s = fmha::Smem_tile_mma_transposed;\n constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;\n static_assert(smem_size_s == 16 * 256 * 2);\n static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);\n\n constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;\n constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;\n constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(fmha_dgrad_fp16_256_64_sm80_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n dim3 grid(params.h, params.b);\n fmha_dgrad_fp16_256_64_sm80_kernel<<>>(params);\n}\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_384_64_kernel.sm80.cu", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_dgrad_kernel_1xN_reload.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<384, 64, 16, 1, 8, 0x08u>;\n\nextern \"C\" __global__ void fmha_dgrad_fp16_384_64_sm80_kernel(Fused_multihead_attention_fprop_params params) {\n fmha::compute_dv_1xN(params);\n fmha::compute_dq_dk_1xN(params);\n}\n\nvoid run_fmha_dgrad_fp16_384_64_sm80(const Fused_multihead_attention_fprop_params& params, cudaStream_t stream) {\n constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);\n constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;\n constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;\n constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;\n\n using Smem_tile_s = fmha::Smem_tile_mma_transposed;\n constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;\n static_assert(smem_size_s == 16 * 384 * 2);\n static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);\n\n constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;\n constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;\n constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(fmha_dgrad_fp16_384_64_sm80_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n dim3 grid(params.h, params.b);\n fmha_dgrad_fp16_384_64_sm80_kernel<<>>(params);\n}\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_512_64_kernel.sm80.cu", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_dgrad_kernel_1xN_reload.h\"\n#include \"fmha_dgrad_kernel_1xN_reload_nl.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<512, 64, 16, 1, 8, 0x08u>;\n\nextern \"C\" __global__ void fmha_dgrad_fp16_512_64_sm80_kernel(Fused_multihead_attention_fprop_params params) {\n fmha::compute_dv_1xN(params);\n fmha::compute_dq_dk_1xN(params);\n}\n\ntemplate \n__global__ void fmha_dgrad_fp16_512_64_sm80_nl_kernel(Fused_multihead_attention_fprop_params params) {\n fmha::compute_dv_1xN_nl(params);\n fmha::compute_dq_dk_1xN_nl(params);\n}\n\nvoid run_fmha_dgrad_fp16_512_64_sm80(const Fused_multihead_attention_fprop_params& params, cudaStream_t stream) {\n constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);\n constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;\n constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;\n constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;\n\n using Smem_tile_s = fmha::Smem_tile_mma_transposed;\n constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;\n static_assert(smem_size_s == 16 * 512 * 2);\n static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);\n\n constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;\n constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;\n constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(fmha_dgrad_fp16_512_64_sm80_kernel,\n cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n dim3 grid(params.h, params.b);\n fmha_dgrad_fp16_512_64_sm80_kernel<<>>(params);\n}\n\nvoid run_fmha_dgrad_fp16_512_64_sm80_nl(const Fused_multihead_attention_fprop_params& params, const int num_chunks,\n cudaStream_t stream) {\n constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);\n constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;\n constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;\n constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;\n\n using Smem_tile_s = fmha::Smem_tile_mma_transposed;\n constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;\n static_assert(smem_size_s == 16 * 512 * 2);\n static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);\n\n constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;\n constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;\n constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);\n\n auto kernel = fmha_dgrad_fp16_512_64_sm80_nl_kernel<2>;\n\n if (num_chunks == 2) {\n kernel = fmha_dgrad_fp16_512_64_sm80_nl_kernel<2>;\n } else if (num_chunks == 3) {\n kernel = fmha_dgrad_fp16_512_64_sm80_nl_kernel<3>;\n } else {\n assert(false && \"Unsupperted number of chunks\");\n }\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n\n dim3 grid(params.h, params.b, num_chunks);\n\n kernel<<>>(params);\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_dgrad_kernel_1xN_reload.h", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\n#include \n#include \n\n#include \"fmha_kernel.h\"\n\nnamespace fmha {\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void compute_dv_1xN(const Params& params) {\n // The description of the CTA tile for the 1st batched GEMM.\n using Cta_tile_p = typename Kernel_traits::Cta_tile_p;\n // The description of the CTA tile for the 2nd batched GEMM.\n using Cta_tile_dv =\n fmha::Cta_tile_extd;\n\n static_assert(Cta_tile_dv::M == 512 || Cta_tile_dv::M == 384 || Cta_tile_dv::M == 256 || Cta_tile_dv::M == 128);\n static_assert(Cta_tile_dv::N == 64);\n static_assert(Cta_tile_dv::K == 16);\n\n // The MMA tile for the 1st GEMM.\n using Mma_tile_p = fmha::Hmma_tile;\n // The MMA tile for the 2nd GEMM.\n using Mma_tile_dv = fmha::Hmma_tile;\n\n // The global memory tile to load Q.\n using Gmem_tile_q = typename Kernel_traits::Gmem_tile_q;\n // The shared memory tile to swizzle Q.\n // using Smem_tile_q = typename Kernel_traits::Smem_tile_q;\n using Smem_tile_q = fmha::Smem_tile_a;\n // The shared memory tile to reload Q as fragment b.\n using Smem_tile_qt = fmha::Smem_tile_b;\n\n // The global memory tile to load K.\n using Gmem_tile_k = typename Kernel_traits::Gmem_tile_k;\n // The shared memory tile to swizzle K.\n using Smem_tile_k = typename Kernel_traits::Smem_tile_k;\n\n // The global memory tile to load V.\n using Gmem_tile_v = typename Kernel_traits::Gmem_tile_v;\n // The shared memory tile to swizzle V.\n using Smem_tile_v = typename Kernel_traits::Smem_tile_v;\n\n // The global memory tile to store O.\n using Gmem_tile_o = typename Kernel_traits::Gmem_tile_o;\n // The shared memory tile to swizzle O.\n using Smem_tile_o = typename Kernel_traits::Smem_tile_o;\n\n // The global memory tile to store dV.\n using Gmem_tile_dv = typename Kernel_traits::Gmem_tile_v;\n // The shared memory tile to swizzle dV.\n using Smem_tile_dv = fmha::Smem_tile_mma_epilogue;\n static_assert(Smem_tile_dv::NUM_LDS == Gmem_tile_dv::LDGS);\n static_assert(Smem_tile_dv::THREADS_PER_ROW == Gmem_tile_dv::THREADS_PER_ROW);\n\n using Gmem_tile_s = typename Kernel_traits::Gmem_tile_s;\n using Smem_tile_st = typename Kernel_traits::Smem_tile_st;\n using Gmem_tile_do = typename Kernel_traits::Gmem_tile_do;\n\n // Shared memory.\n extern __shared__ char smem_[];\n\n // The block index for the batch.\n const int bidb = blockIdx.y;\n // The block index for the head.\n const int bidh = blockIdx.x;\n // The thread index.\n const int tidx = threadIdx.x;\n\n const BlockInfoPadded binfo(params, bidb, bidh, tidx);\n if (binfo.stop_early()) return;\n Mask mask(params, binfo, tidx);\n\n // Allocate the global memory tile loader for Q.\n Gmem_tile_do gmem_q(params, binfo, tidx); // treating dout as Q\n // Allocate the shared memory tile loader for Q.\n Smem_tile_q smem_q(&smem_[0], tidx);\n Smem_tile_qt smem_qt(&smem_[0], tidx);\n Smem_tile_st smem_s(&smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE], tidx);\n\n // Allocate the global memory tile loader for K.\n Gmem_tile_k gmem_k(params, 2, binfo, tidx); // treating V as K\n // Allocate the shared memory tile loader for K.\n Smem_tile_k smem_k(&smem_[Smem_tile_q::BYTES_PER_TILE], tidx);\n\n // Trigger the loads for Q.\n gmem_q.load(smem_q);\n // Trigger the loads for K.\n gmem_k.load(smem_k);\n\n // Commit the data for Q and K to shared memory.\n gmem_q.commit(smem_q);\n gmem_k.commit(smem_k);\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // Load the fragments for Q.\n typename Smem_tile_q::Fragment frag_q[2][Mma_tile_p::MMAS_M];\n smem_q.load(frag_q[0], 0);\n\n typename Smem_tile_qt::Fragment frag_qt[2][Mma_tile_dv::MMAS_N];\n static_assert(Smem_tile_qt::Fragment::NUM_REGS == 4);\n static_assert(Mma_tile_dv::MMAS_K == 1);\n smem_qt.load(frag_qt[0], 0);\n\n // Load the fragments for K. We keep the data in registers during the entire kernel.\n typename Smem_tile_k::Fragment frag_k[2][Mma_tile_p::MMAS_N];\n smem_k.load(frag_k[0], 0);\n\n enum { BITS_PER_ELT_S = sizeof(fmha::A_type) * 8 };\n\n Gmem_tile_s gmem_s(params, binfo, tidx);\n\n // Create the object to do the softmax.\n using Softmax = fmha::Softmax;\n Softmax softmax(params,\n &smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_st::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE],\n bidb, tidx);\n\n enum { THREADS_PER_ROW = 32 };\n enum { M = Mma_tile_p::MMAS_M };\n enum { N = Mma_tile_p::MMAS_N };\n\n // Declare the accumulators for the 2nd gemm.\n fmha::Fragment_accumulator acc_dv[Mma_tile_dv::MMAS_M][Mma_tile_dv::MMAS_N];\n fmha::Clear_accumulator::apply(acc_dv);\n\n enum { STEPS = Cta_tile_p::N / Cta_tile_p::M };\n // Load over the entire sequence length.\n for (int l = 0; l < STEPS; l++) {\n const int loop = l * Cta_tile_p::M;\n if (loop >= binfo.actual_seqlen) break;\n\n // Load S\n uint4 s_regs[M][N];\n gmem_s.load(s_regs, mask);\n fmha::Fragment_accumulator acc_p[Mma_tile_p::MMAS_M][Mma_tile_p::MMAS_N];\n fmha::Clear_accumulator::apply(acc_p);\n// Do this part of P^T = (Q * K^T)^T.\n#pragma unroll\n for (int ki = 1; ki < Mma_tile_p::MMAS_K; ++ki) {\n // Trigger the load from shared memory for the next series of Q values.\n smem_q.load(frag_q[ki & 1], ki);\n smem_k.load(frag_k[ki & 1], ki);\n // Do the math for the values already in registers.\n fmha::gemm(acc_p, frag_q[(ki - 1) & 1], frag_k[(ki - 1) & 1]);\n }\n\n // Store s * dmask to smem for transpose\n smem_s.store(s_regs);\n\n // Declare the accumulators for the 1st gemm.\n // Do the final stage of math.\n {\n int ki = Mma_tile_p::MMAS_K;\n fmha::gemm(acc_p, frag_q[(ki - 1) & 1], frag_k[(ki - 1) & 1]);\n }\n // Trigger the load for the next Q values. We're using double buffering, so reading qt is safe\n if (l < STEPS - 1) {\n smem_q.move_to_next_write_buffer();\n gmem_q.move();\n gmem_q.load(smem_q);\n }\n\n // Convert from the accumulator type to FP32 for Softmax.\n softmax.unpack(acc_p);\n\n float s_mat[2 * M][4 * N];\n\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n#pragma unroll\n for (int ni = 0; ni < N; ni++) {\n uint4& dst = s_regs[mi][ni];\n fmha::half2_to_float2(s_mat[2 * mi + 0][4 * ni + 0], s_mat[2 * mi + 0][4 * ni + 1], dst.x);\n fmha::half2_to_float2(s_mat[2 * mi + 0][4 * ni + 2], s_mat[2 * mi + 0][4 * ni + 3], dst.y);\n fmha::half2_to_float2(s_mat[2 * mi + 1][4 * ni + 0], s_mat[2 * mi + 1][4 * ni + 1], dst.z);\n fmha::half2_to_float2(s_mat[2 * mi + 1][4 * ni + 2], s_mat[2 * mi + 1][4 * ni + 3], dst.w);\n }\n }\n\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n#pragma unroll\n for (int ii = 0; ii < 2; ii++) {\n#pragma unroll\n for (int ni = 0; ni < N; ni++) {\n#pragma unroll\n for (int jj = 0; jj < 4; jj++) {\n float& s_dmask = s_mat[2 * mi + ii][4 * ni + jj];\n const bool drop = reinterpret_cast(s_dmask) & 0x80000000;\n const float d_s = drop ? 0.f : softmax.elt_[2 * mi + ii][4 * ni + jj] * params.rp_dropout;\n s_dmask = fabsf(s_dmask);\n softmax.elt_[2 * mi + ii][4 * ni + jj] = d_s * fabsf(s_dmask);\n }\n }\n }\n }\n\n float p_sum[2 * M];\n softmax.reduce_sum(p_sum);\n\n const float scalef = reinterpret_cast(params.scale_softmax);\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n#pragma unroll\n for (int ii = 0; ii < 2; ii++) {\n#pragma unroll\n for (int ni = 0; ni < N; ni++) {\n#pragma unroll\n for (int jj = 0; jj < 4; jj++) {\n softmax.elt_[2 * mi + ii][4 * ni + jj] -= p_sum[2 * mi + ii] * (s_mat[2 * mi + ii][4 * ni + jj]);\n softmax.elt_[2 * mi + ii][4 * ni + jj] *= scalef;\n }\n }\n }\n }\n typename Smem_tile_st::Fragment frag_s[Mma_tile_dv::MMAS_K][Mma_tile_dv::MMAS_M];\n smem_s.load(frag_s);\n for (int ki = 0; ki < Mma_tile_dv::MMAS_K; ki++) {\n for (int mi = 0; mi < Mma_tile_dv::MMAS_M; mi++) {\n for (int ii = 0; ii < Smem_tile_st::Fragment::NUM_REGS; ii++) {\n frag_s[ki][mi].reg(ii) = fmha::hmul2(frag_s[ki][mi].reg(ii), params.scale_dropout);\n frag_s[ki][mi].reg(ii) = fmha::hrelu2(frag_s[ki][mi].reg(ii));\n }\n }\n }\n\n gmem_s.store(softmax.elt_, mask);\n gmem_s.move();\n\n#pragma unroll\n for (int ki = 1; ki < Mma_tile_dv::MMAS_K; ++ki) {\n // Trigger the load from shared memory for the next series of Q values.\n smem_qt.load(frag_qt[ki & 1], ki);\n // Do the math for the values already in registers.\n fmha::gemm(acc_dv, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);\n }\n\n // Do the final stage of math.\n {\n int ki = Mma_tile_dv::MMAS_K;\n fmha::gemm(acc_dv, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);\n }\n // Commit the values for Q into shared memory.\n if (l < STEPS - 1) {\n gmem_q.commit(smem_q);\n }\n\n // Make sure we are reading from the correct buffer.\n smem_q.move_to_next_read_buffer();\n smem_qt.move_to_next_read_buffer();\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // Trigger the loads for the values of Q for the next iteration.\n smem_q.load(frag_q[0], 0);\n smem_k.load(frag_k[0], 0);\n smem_qt.load(frag_qt[0], 0);\n\n } // Outer loop over the sequence length.\n\n // Epilogue swizzle for dV\n Smem_tile_dv smem_dv(&smem_[Kernel_traits::Smem_tile_q::BYTES_PER_TILE], tidx);\n smem_dv.store(acc_dv);\n\n __syncthreads();\n uint4 dv_out[Smem_tile_dv::NUM_LDS];\n smem_dv.load(dv_out);\n Qkv_params dv_params;\n dv_params.qkv_ptr = params.dqkv_ptr;\n dv_params.qkv_stride_in_bytes = params.qkv_stride_in_bytes;\n dv_params.h = params.h;\n Gmem_tile_dv gmem_dv(dv_params, 2, binfo, tidx);\n gmem_dv.store(dv_out);\n}\n\ntemplate \ninline __device__ void compute_dq_dk_1xN(const Params& params) {\n // The description of the CTA tile for the 1st batched GEMM.\n using Cta_tile_p = typename Kernel_traits::Cta_tile_p;\n using Cta_tile_o = typename Kernel_traits::Cta_tile_o;\n // The description of the CTA tile for the 2nd batched GEMM.\n using Cta_tile_dk =\n fmha::Cta_tile_extd;\n static_assert(Cta_tile_dk::M == 512 || Cta_tile_dk::M == 384 || Cta_tile_dk::M == 256 || Cta_tile_dk::M == 128);\n static_assert(Cta_tile_dk::N == 64);\n static_assert(Cta_tile_dk::K == 16);\n\n // The MMA tile for the 1st GEMM.\n using Mma_tile_p = fmha::Hmma_tile;\n using Mma_tile_o = fmha::Hmma_tile;\n // The MMA tile for the 2nd GEMM.\n using Mma_tile_dk = fmha::Hmma_tile;\n\n // The global memory tile to load Q.\n using Gmem_tile_q = typename Kernel_traits::Gmem_tile_q;\n // The shared memory tile to swizzle Q.\n using Smem_tile_q = typename Kernel_traits::Smem_tile_q;\n\n // The global memory tile to load K.\n using Gmem_tile_k = typename Kernel_traits::Gmem_tile_v;\n // The shared memory tile to swizzle K.\n using Smem_tile_k = typename Kernel_traits::Smem_tile_v; // K is used like V in fprop\n\n // The global memory tile to load V.\n using Gmem_tile_v = typename Kernel_traits::Gmem_tile_v;\n // The shared memory tile to swizzle V.\n using Smem_tile_v = typename Kernel_traits::Smem_tile_v;\n\n // The global memory tile to store O.\n // using Gmem_tile_o = typename Kernel_traits::Gmem_tile_o;\n using Gmem_tile_o = fmha::Gmem_tile_dq;\n // The shared memory tile to swizzle O.\n using Smem_tile_o = typename Kernel_traits::Smem_tile_o;\n\n // The global memory tile to store dK.\n using Gmem_tile_dk = typename Kernel_traits::Gmem_tile_v;\n // The shared memory tile to swizzle dK.\n using Smem_tile_dk = fmha::Smem_tile_mma_epilogue;\n static_assert(Smem_tile_dk::NUM_LDS == Gmem_tile_dk::LDGS);\n static_assert(Smem_tile_dk::THREADS_PER_ROW == Gmem_tile_dk::THREADS_PER_ROW);\n\n // The shared memory tile to reload Q transposed.\n using Smem_tile_qt = fmha::Smem_tile_b;\n\n using Gmem_tile_s = typename Kernel_traits::Gmem_tile_s;\n\n using Smem_tile_st = typename Kernel_traits::Smem_tile_st;\n\n enum { M = Mma_tile_p::MMAS_M };\n enum { N = Mma_tile_p::MMAS_N };\n static_assert(M == Mma_tile_o::MMAS_M);\n static_assert(N == Mma_tile_o::MMAS_K);\n // Shared memory.\n extern __shared__ char smem_[];\n\n // The block index for the batch.\n const int bidb = blockIdx.y;\n // The block index for the head.\n const int bidh = blockIdx.x;\n // The thread index.\n const int tidx = threadIdx.x;\n\n const BlockInfoPadded binfo(params, bidb, bidh, tidx);\n if (binfo.stop_early()) return;\n\n Mask mask(params, binfo, tidx);\n\n // Allocate the global memory tile loader for Q.\n Gmem_tile_q gmem_q(params, 0, binfo, tidx);\n // Allocate the shared memory tile loader for Q.\n Smem_tile_q smem_q(&smem_[0], tidx);\n Smem_tile_qt smem_qt(&smem_[0], tidx);\n Smem_tile_st smem_s(&smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE],\n tidx);\n\n // Allocate the global memory tile loader for K.\n Gmem_tile_k gmem_k(params, 1, binfo, tidx);\n // Allocate the shared memory tile loader for K.\n Smem_tile_k smem_k(&smem_[Smem_tile_q::BYTES_PER_TILE], tidx);\n\n // Allocate the global memory tile loader for O.\n Gmem_tile_o gmem_o(params, binfo, tidx);\n // Allocate the shared memory tile loader for O. We use the same as K so be careful!!!\n Smem_tile_o smem_o(&smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE], tidx);\n\n // Trigger the loads for Q.\n gmem_q.load(smem_q);\n // Trigger the loads for K.\n gmem_k.load(smem_k);\n\n Gmem_tile_s gmem_s(params, binfo, tidx);\n // Load dP\n uint4 s_regs[M][N];\n gmem_s.load(s_regs, mask);\n gmem_s.move();\n\n // Commit the data for Q and K to shared memory.\n gmem_q.commit(smem_q);\n gmem_k.commit(smem_k);\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n typename Smem_tile_qt::Fragment frag_qt[2][Mma_tile_dk::MMAS_N];\n smem_qt.load(frag_qt[0], 0);\n typename Smem_tile_k::Fragment frag_k[2][Mma_tile_o::MMAS_N];\n smem_k.load(frag_k[0], 0);\n\n enum { BITS_PER_ELT_S = sizeof(fmha::A_type) * 8 };\n\n enum { THREADS_PER_ROW = 32 };\n enum { STEPS = Cta_tile_p::N / Cta_tile_p::M };\n\n // Declare the accumulators for the 2nd gemm.\n fmha::Fragment_accumulator acc_dk[Mma_tile_dk::MMAS_M][Mma_tile_dk::MMAS_N];\n fmha::Clear_accumulator::apply(acc_dk);\n\n // Load over the entire sequence length.\n for (int l = 0; l < STEPS; l++) {\n const int loop = l * Cta_tile_p::M;\n if (loop >= binfo.actual_seqlen) break;\n\n // Pack dP as Fragment_a\n fmha::Fragment_a frag_p[Mma_tile_o::MMAS_K][Mma_tile_o::MMAS_M];\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n#pragma unroll\n for (int ni = 0; ni < N; ni++) {\n uint4& dst = s_regs[mi][ni];\n frag_p[ni][mi].reg(0) = dst.x; // row 0, cols 0,1\n frag_p[ni][mi].reg(1) = dst.z; // row 8, cols 0,1\n frag_p[ni][mi].reg(2) = dst.y; // row 0, cols 8,9\n frag_p[ni][mi].reg(3) = dst.w; // row 8, cols 8,9\n }\n }\n\n // Declare the accumulators for the 1st gemm.\n fmha::Fragment_accumulator acc_o[Mma_tile_o::MMAS_M][Mma_tile_o::MMAS_N];\n fmha::Clear_accumulator::apply(acc_o);\n\n// Do this part of O = P^T * V^T. dQ = dP x dK\n#pragma unroll\n for (int ki = 1; ki < Mma_tile_o::MMAS_K; ++ki) {\n // Trigger the load from shared memory for the next series of Q values.\n smem_k.load(frag_k[ki & 1], ki);\n // Do the math for the values already in registers.\n fmha::gemm(acc_o, frag_p[ki - 1], frag_k[(ki - 1) & 1]);\n }\n\n // Do the final stage of math.\n {\n int ki = Mma_tile_o::MMAS_K;\n fmha::gemm(acc_o, frag_p[ki - 1], frag_k[(ki - 1) & 1]);\n }\n\n // Store dP to smem for transpose\n smem_s.store(s_regs);\n if (l < STEPS - 1) {\n // Load next part of S\n gmem_s.load(s_regs, mask);\n gmem_s.move();\n smem_q.move_to_next_write_buffer();\n gmem_q.move();\n gmem_q.load(smem_q);\n }\n// Loop over MMAS_M.\n#pragma unroll\n for (int ii = 0; ii < Gmem_tile_o::LOOPS; ++ii) {\n // Swizzle the elements and do the final reduction.\n smem_o.store(acc_o, ii);\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // Load from shared memory.\n uint4 out[Gmem_tile_o::STGS_PER_LOOP];\n smem_o.load(out);\n\n // Make sure the data was read from shared memory.\n if (ii < Gmem_tile_o::LOOPS - 1) {\n __syncthreads();\n }\n\n // Output the values.\n gmem_o.store(out, ii);\n }\n\n // Move to the next part of the output.\n gmem_o.move();\n\n typename Smem_tile_st::Fragment frag_s[Mma_tile_dk::MMAS_K][Mma_tile_dk::MMAS_M];\n smem_s.load(frag_s);\n\n#pragma unroll\n for (int ki = 1; ki < Mma_tile_dk::MMAS_K; ++ki) {\n // Trigger the load from shared memory for the next series of Q values.\n smem_qt.load(frag_qt[ki & 1], ki);\n // Do the math for the values already in registers.\n fmha::gemm(acc_dk, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);\n }\n\n // Do the final stage of math.\n {\n int ki = Mma_tile_dk::MMAS_K;\n fmha::gemm(acc_dk, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);\n }\n\n // Commit the values for Q into shared memory.\n if (l < STEPS - 1) {\n gmem_q.commit(smem_q);\n }\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // Trigger the loads for the values of Q for the next iteration.\n smem_qt.load(frag_qt[0], 0);\n smem_k.load(frag_k[0], 0);\n\n } // Outer loop over the sequence length.\n\n // Epilogue swizzle for dK\n Smem_tile_dk smem_dk(&smem_[0], tidx);\n smem_dk.store(acc_dk);\n __syncthreads();\n uint4 dk_out[Smem_tile_dk::NUM_LDS];\n smem_dk.load(dk_out);\n Qkv_params dk_params;\n dk_params.qkv_ptr = params.dqkv_ptr;\n dk_params.qkv_stride_in_bytes = params.qkv_stride_in_bytes;\n dk_params.h = params.h;\n Gmem_tile_dk gmem_dk(dk_params, 1, binfo, tidx);\n gmem_dk.store(dk_out);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace fmha\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_dgrad_kernel_1xN_reload_nl.h", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\n#include \n#include \n\n#include \"fmha_kernel.h\"\n\nnamespace fmha {\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void compute_dv_1xN_nl(const Params& params) {\n // The description of the CTA tile for the 1st batched GEMM.\n using Cta_tile_p = typename Kernel_traits::Cta_tile_p;\n // The description of the CTA tile for the 2nd batched GEMM.\n using Cta_tile_dv =\n fmha::Cta_tile_extd;\n\n static_assert(Cta_tile_dv::M == 512 || Cta_tile_dv::M == 384 || Cta_tile_dv::M == 256 || Cta_tile_dv::M == 128);\n static_assert(Cta_tile_dv::N == 64);\n static_assert(Cta_tile_dv::K == 16);\n\n // The MMA tile for the 1st GEMM.\n using Mma_tile_p = fmha::Hmma_tile;\n // The MMA tile for the 2nd GEMM.\n using Mma_tile_dv = fmha::Hmma_tile;\n\n // The global memory tile to load Q.\n using Gmem_tile_q = typename Kernel_traits::Gmem_tile_q;\n // The shared memory tile to swizzle Q.\n using Smem_tile_q = fmha::Smem_tile_a;\n // The shared memory tile to reload Q as fragment b.\n using Smem_tile_qt = fmha::Smem_tile_b;\n\n // The global memory tile to load K.\n using Gmem_tile_k = typename Kernel_traits::Gmem_tile_k;\n // The shared memory tile to swizzle K.\n using Smem_tile_k = typename Kernel_traits::Smem_tile_k;\n\n // The global memory tile to load V.\n using Gmem_tile_v = typename Kernel_traits::Gmem_tile_v;\n // The shared memory tile to swizzle V.\n using Smem_tile_v = typename Kernel_traits::Smem_tile_v;\n\n // The global memory tile to store dV.\n using Gmem_tile_dv = fmha::Gmem_tile_qkv;\n\n // The shared memory tile to swizzle dV.\n using Smem_tile_dv = fmha::Smem_tile_mma_epilogue;\n static_assert(Smem_tile_dv::NUM_LDS == Gmem_tile_dv::LDGS);\n static_assert(Smem_tile_dv::THREADS_PER_ROW == Gmem_tile_dv::THREADS_PER_ROW);\n\n using Gmem_tile_s = typename Kernel_traits::Gmem_tile_s;\n using Smem_tile_st = typename Kernel_traits::Smem_tile_st;\n using Gmem_tile_do = typename Kernel_traits::Gmem_tile_do;\n\n // Shared memory.\n extern __shared__ char smem_[];\n\n // The block index for the chunk.\n const int bidc = blockIdx.z;\n // The block index for the batch.\n const int bidb = blockIdx.y;\n // The block index for the head.\n const int bidh = blockIdx.x;\n // The thread index.\n const int tidx = threadIdx.x;\n\n const BlockInfoPadded binfo(params, bidb, bidh, tidx);\n if (binfo.stop_early()) return;\n fmha::Mask mask(params, binfo, tidx);\n\n // Allocate the global memory tile loader for Q.\n Gmem_tile_do gmem_q(params, binfo, tidx); // treating dout as Q\n // Allocate the shared memory tile loader for Q.\n Smem_tile_q smem_q(&smem_[0], tidx);\n Smem_tile_qt smem_qt(&smem_[0], tidx);\n Smem_tile_st smem_s(&smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE], tidx);\n\n // Allocate the global memory tile loader for K.\n Gmem_tile_k gmem_k(params, 2, binfo, tidx); // treating V as K\n // Allocate the shared memory tile loader for K.\n Smem_tile_k smem_k(&smem_[Smem_tile_q::BYTES_PER_TILE], tidx);\n\n Gmem_tile_s gmem_s(params, binfo, tidx);\n\n using Noloop = Noloop_traits;\n\n Noloop nl_traits(bidc, binfo);\n nl_traits.move_all(gmem_q, gmem_s);\n\n // Trigger the loads for Q.\n gmem_q.load(smem_q);\n // Trigger the loads for K.\n gmem_k.load(smem_k);\n\n // Commit the data for Q and K to shared memory.\n gmem_q.commit(smem_q);\n gmem_k.commit(smem_k);\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // Load the fragments for Q.\n typename Smem_tile_q::Fragment frag_q[2][Mma_tile_p::MMAS_M];\n smem_q.load(frag_q[0], 0);\n\n typename Smem_tile_qt::Fragment frag_qt[2][Mma_tile_dv::MMAS_N];\n static_assert(Smem_tile_qt::Fragment::NUM_REGS == 4);\n static_assert(Mma_tile_dv::MMAS_K == 1);\n smem_qt.load(frag_qt[0], 0);\n\n // Load the fragments for K. We keep the data in registers during the entire kernel.\n typename Smem_tile_k::Fragment frag_k[2][Mma_tile_p::MMAS_N];\n smem_k.load(frag_k[0], 0);\n\n enum { BITS_PER_ELT_S = sizeof(fmha::A_type) * 8 };\n\n // Create the object to do the softmax.\n using Softmax = fmha::Softmax;\n Softmax softmax(params,\n &smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_st::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE],\n bidb, tidx);\n\n enum { THREADS_PER_ROW = 32 };\n enum { M = Mma_tile_p::MMAS_M };\n enum { N = Mma_tile_p::MMAS_N };\n\n // Declare the accumulators for the 2nd gemm.\n fmha::Fragment_accumulator acc_dv[Mma_tile_dv::MMAS_M][Mma_tile_dv::MMAS_N];\n fmha::Clear_accumulator::apply(acc_dv);\n\n // Load over the entire sequence length.\n for (int l = 0; l < nl_traits.num_steps_; l++) {\n uint4 s_regs[M][N];\n gmem_s.load(s_regs, mask);\n fmha::Fragment_accumulator acc_p[Mma_tile_p::MMAS_M][Mma_tile_p::MMAS_N];\n fmha::Clear_accumulator::apply(acc_p);\n// Do this part of P^T = (Q * K^T)^T.\n#pragma unroll\n for (int ki = 1; ki < Mma_tile_p::MMAS_K; ++ki) {\n // Trigger the load from shared memory for the next series of Q values.\n smem_q.load(frag_q[ki & 1], ki);\n smem_k.load(frag_k[ki & 1], ki);\n // Do the math for the values already in registers.\n fmha::gemm(acc_p, frag_q[(ki - 1) & 1], frag_k[(ki - 1) & 1]);\n }\n\n smem_s.store(s_regs);\n\n // Declare the accumulators for the 1st gemm.\n // Do the final stage of math.\n {\n int ki = Mma_tile_p::MMAS_K;\n fmha::gemm(acc_p, frag_q[(ki - 1) & 1], frag_k[(ki - 1) & 1]);\n }\n // Trigger the load for the next Q values. We're using double buffering, so reading qt is safe\n if (l < nl_traits.num_steps_ - 1) {\n smem_q.move_to_next_write_buffer();\n gmem_q.move();\n gmem_q.load(smem_q);\n }\n // Convert from the accumulator type to FP32 for Softmax.\n softmax.unpack(acc_p);\n\n float s_mat[2 * M][4 * N];\n\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n#pragma unroll\n for (int ni = 0; ni < N; ni++) {\n uint4& dst = s_regs[mi][ni];\n fmha::half2_to_float2(s_mat[2 * mi + 0][4 * ni + 0], s_mat[2 * mi + 0][4 * ni + 1], dst.x);\n fmha::half2_to_float2(s_mat[2 * mi + 0][4 * ni + 2], s_mat[2 * mi + 0][4 * ni + 3], dst.y);\n fmha::half2_to_float2(s_mat[2 * mi + 1][4 * ni + 0], s_mat[2 * mi + 1][4 * ni + 1], dst.z);\n fmha::half2_to_float2(s_mat[2 * mi + 1][4 * ni + 2], s_mat[2 * mi + 1][4 * ni + 3], dst.w);\n }\n }\n\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n#pragma unroll\n for (int ii = 0; ii < 2; ii++) {\n#pragma unroll\n for (int ni = 0; ni < N; ni++) {\n#pragma unroll\n for (int jj = 0; jj < 4; jj++) {\n float& s_dmask = s_mat[2 * mi + ii][4 * ni + jj];\n const bool drop = reinterpret_cast(s_dmask) & 0x80000000;\n const float d_s = drop ? 0.f : softmax.elt_[2 * mi + ii][4 * ni + jj] * params.rp_dropout;\n s_dmask = fabsf(s_dmask);\n softmax.elt_[2 * mi + ii][4 * ni + jj] = d_s * (s_dmask);\n }\n }\n }\n }\n\n float p_sum[2 * M];\n softmax.reduce_sum(p_sum);\n\n const float scalef = reinterpret_cast(params.scale_softmax);\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n#pragma unroll\n for (int ii = 0; ii < 2; ii++) {\n#pragma unroll\n for (int ni = 0; ni < N; ni++) {\n#pragma unroll\n for (int jj = 0; jj < 4; jj++) {\n softmax.elt_[2 * mi + ii][4 * ni + jj] -= p_sum[2 * mi + ii] * (s_mat[2 * mi + ii][4 * ni + jj]);\n softmax.elt_[2 * mi + ii][4 * ni + jj] *= scalef;\n }\n }\n }\n }\n\n typename Smem_tile_st::Fragment frag_s[Mma_tile_dv::MMAS_K][Mma_tile_dv::MMAS_M];\n smem_s.load(frag_s);\n for (int ki = 0; ki < Mma_tile_dv::MMAS_K; ki++) {\n for (int mi = 0; mi < Mma_tile_dv::MMAS_M; mi++) {\n for (int ii = 0; ii < Smem_tile_st::Fragment::NUM_REGS; ii++) {\n frag_s[ki][mi].reg(ii) = fmha::hmul2(frag_s[ki][mi].reg(ii), params.scale_dropout);\n frag_s[ki][mi].reg(ii) = fmha::hrelu2(frag_s[ki][mi].reg(ii));\n }\n }\n }\n\n gmem_s.store(softmax.elt_, mask);\n gmem_s.move();\n\n static_assert(Mma_tile_dv::MMAS_K == 1); // DEBUG\n#pragma unroll\n for (int ki = 1; ki < Mma_tile_dv::MMAS_K; ++ki) {\n // Trigger the load from shared memory for the next series of Q values.\n smem_qt.load(frag_qt[ki & 1], ki);\n // Do the math for the values already in registers.\n fmha::gemm(acc_dv, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);\n }\n\n // Do the final stage of math.\n {\n int ki = Mma_tile_dv::MMAS_K;\n fmha::gemm(acc_dv, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);\n }\n // Commit the values for Q into shared memory.\n if (l < nl_traits.num_steps_ - 1) {\n gmem_q.commit(smem_q);\n }\n\n // Make sure we are reading from the correct buffer.\n smem_q.move_to_next_read_buffer();\n smem_qt.move_to_next_read_buffer();\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // Trigger the loads for the values of Q for the next iteration.\n smem_q.load(frag_q[0], 0);\n smem_k.load(frag_k[0], 0);\n smem_qt.load(frag_qt[0], 0);\n\n } // Outer loop over the sequence length.\n\n // Epilogue for dV = (S * D)' * dout'. We're fully exposed to this!\n\n // Epilogue swizzle for dV\n Smem_tile_dv smem_dv(&smem_[Kernel_traits::Smem_tile_q::BYTES_PER_TILE], tidx);\n smem_dv.store(acc_dv);\n\n __syncthreads();\n\n uint4 dv_out[Smem_tile_dv::NUM_LDS];\n smem_dv.load(dv_out);\n Qkv_params dv_params;\n dv_params.qkv_ptr = params.dkv_ptr;\n dv_params.qkv_stride_in_bytes = params.h * 2 * CHUNKS * params.d * sizeof(half);\n dv_params.h = params.h;\n Gmem_tile_dv gmem_dv(dv_params, nl_traits.get_idx_dv(), binfo, tidx);\n gmem_dv.store(dv_out);\n}\n\ntemplate \ninline __device__ void compute_dq_dk_1xN_nl(const Params& params) {\n // The description of the CTA tile for the 1st batched GEMM.\n using Cta_tile_p = typename Kernel_traits::Cta_tile_p;\n using Cta_tile_o = typename Kernel_traits::Cta_tile_o;\n // The description of the CTA tile for the 2nd batched GEMM.\n using Cta_tile_dk =\n fmha::Cta_tile_extd;\n\n static_assert(Cta_tile_dk::M == 512 || Cta_tile_dk::M == 384 || Cta_tile_dk::M == 256 || Cta_tile_dk::M == 128);\n static_assert(Cta_tile_dk::N == 64);\n static_assert(Cta_tile_dk::K == 16);\n\n // The MMA tile for the 1st GEMM.\n using Mma_tile_p = fmha::Hmma_tile;\n using Mma_tile_o = fmha::Hmma_tile;\n // The MMA tile for the 2nd GEMM.\n using Mma_tile_dk = fmha::Hmma_tile;\n\n // The global memory tile to load Q.\n using Gmem_tile_q = typename Kernel_traits::Gmem_tile_q;\n // The shared memory tile to swizzle Q.\n using Smem_tile_q = typename Kernel_traits::Smem_tile_q;\n\n // The global memory tile to load K.\n using Gmem_tile_k = typename Kernel_traits::Gmem_tile_v;\n // The shared memory tile to swizzle K.\n using Smem_tile_k = typename Kernel_traits::Smem_tile_v; // K is used like V in fprop\n\n // The global memory tile to load V.\n using Gmem_tile_v = typename Kernel_traits::Gmem_tile_v;\n // The shared memory tile to swizzle V.\n using Smem_tile_v = typename Kernel_traits::Smem_tile_v;\n\n // The global memory tile to store O.\n using Gmem_tile_o = Gmem_tile_dq;\n // The shared memory tile to swizzle O.\n using Smem_tile_o = typename Kernel_traits::Smem_tile_o;\n\n // The global memory tile to store dK.\n using Gmem_tile_dk = fmha::Gmem_tile_qkv;\n\n // The shared memory tile to swizzle dK.\n using Smem_tile_dk = fmha::Smem_tile_mma_epilogue;\n static_assert(Smem_tile_dk::NUM_LDS == Gmem_tile_dk::LDGS);\n static_assert(Smem_tile_dk::THREADS_PER_ROW == Gmem_tile_dk::THREADS_PER_ROW);\n\n // The shared memory tile to reload Q transposed.\n using Smem_tile_qt = fmha::Smem_tile_b;\n\n // The global memory tile to load dP, stored in S\n using Gmem_tile_s = Gmem_tile_mma_s;\n // The shared memory tile to transpose dP.\n using Smem_tile_st = Smem_tile_mma_transposed;\n\n using Noloop = Noloop_traits;\n\n enum { M = Mma_tile_p::MMAS_M };\n enum { N = Mma_tile_p::MMAS_N };\n static_assert(M == Mma_tile_o::MMAS_M);\n static_assert(N == Mma_tile_o::MMAS_K);\n // Shared memory.\n extern __shared__ char smem_[];\n\n const int bidc = blockIdx.z;\n // The block index for the batch.\n const int bidb = blockIdx.y;\n // The block index for the head.\n const int bidh = blockIdx.x;\n // The thread index.\n const int tidx = threadIdx.x;\n\n const BlockInfoPadded binfo(params, bidb, bidh, tidx);\n if (binfo.stop_early()) return;\n\n fmha::Mask mask(params, binfo, tidx);\n\n // Allocate the global memory tile loader for Q.\n Gmem_tile_q gmem_q(params, 0, binfo, tidx);\n // Allocate the shared memory tile loader for Q (as B).\n Smem_tile_qt smem_qt(&smem_[0], tidx);\n // Allocate the global memory tile loader for dP.\n Gmem_tile_s gmem_s(params, binfo, tidx);\n // Allocate the shared memory tile loader for dP.\n Smem_tile_st smem_s(&smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE],\n tidx);\n\n // Allocate the global memory tile loader for K.\n Gmem_tile_k gmem_k(params, 1, binfo, tidx);\n // Allocate the shared memory tile loader for K.\n Smem_tile_k smem_k(&smem_[Smem_tile_q::BYTES_PER_TILE], tidx);\n\n // Allocate the global memory tile loader for O.\n Gmem_tile_o gmem_o(params, binfo, tidx);\n // Allocate the shared memory tile loader for O. We use the same as K so be careful!!!\n Smem_tile_o smem_o(&smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE], tidx);\n\n Noloop nl_traits(bidc, binfo);\n\n nl_traits.move_all(gmem_q, gmem_o, gmem_s);\n\n // Trigger the loads for Q.\n gmem_q.load(smem_qt);\n // Trigger the loads for K.\n gmem_k.load(smem_k);\n\n uint4 s_regs[M][N];\n gmem_s.load(s_regs, mask);\n\n // Commit the data for Q and K to shared memory.\n gmem_q.commit(smem_qt);\n gmem_k.commit(smem_k);\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n typename Smem_tile_qt::Fragment frag_qt[2][Mma_tile_dk::MMAS_N];\n smem_qt.load(frag_qt[0], 0);\n typename Smem_tile_k::Fragment frag_k[2][Mma_tile_o::MMAS_N];\n smem_k.load(frag_k[0], 0);\n\n enum { BITS_PER_ELT_S = sizeof(fmha::A_type) * 8 };\n\n enum { THREADS_PER_ROW = 32 };\n\n // Declare the accumulators for the 2nd gemm.\n fmha::Fragment_accumulator acc_dk[Mma_tile_dk::MMAS_M][Mma_tile_dk::MMAS_N];\n fmha::Clear_accumulator::apply(acc_dk);\n\n // Load over the entire sequence length.\n for (int l = 0; l < nl_traits.num_steps_; l++) {\n // Pack dP as Fragment_a\n fmha::Fragment_a frag_p[Mma_tile_o::MMAS_K][Mma_tile_o::MMAS_M];\n#pragma unroll\n for (int mi = 0; mi < M; mi++) {\n#pragma unroll\n for (int ni = 0; ni < N; ni++) {\n uint4& dst = s_regs[mi][ni];\n frag_p[ni][mi].reg(0) = dst.x;\n frag_p[ni][mi].reg(1) = dst.z;\n frag_p[ni][mi].reg(2) = dst.y;\n frag_p[ni][mi].reg(3) = dst.w;\n }\n }\n smem_s.store(s_regs);\n if (l < nl_traits.num_steps_ - 1) {\n // Load next part of S\n gmem_s.move();\n gmem_s.load(s_regs, mask);\n // Trigger the load for the next Q values.\n smem_qt.move_to_next_write_buffer();\n gmem_q.move();\n gmem_q.load(smem_qt);\n }\n // Declare the accumulators for the 1st gemm.\n fmha::Fragment_accumulator acc_o[Mma_tile_o::MMAS_M][Mma_tile_o::MMAS_N];\n fmha::Clear_accumulator::apply(acc_o);\n\n// Do this part of O = P^T * V^T. dQ = dP x dK\n#pragma unroll\n for (int ki = 1; ki < Mma_tile_o::MMAS_K; ++ki) {\n // Trigger the load from shared memory for the next series of Q values.\n smem_k.load(frag_k[ki & 1], ki);\n // Do the math for the values already in registers.\n fmha::gemm(acc_o, frag_p[ki - 1], frag_k[(ki - 1) & 1]);\n }\n\n // Do the final stage of math.\n {\n int ki = Mma_tile_o::MMAS_K;\n fmha::gemm(acc_o, frag_p[ki - 1], frag_k[(ki - 1) & 1]);\n }\n\n static_assert(Gmem_tile_o::LOOPS == 1); // DEBUG\n// Loop over MMAS_M.\n#pragma unroll\n for (int ii = 0; ii < Gmem_tile_o::LOOPS; ++ii) {\n // Swizzle the elements and do the final reduction.\n smem_o.store(acc_o, ii);\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // Load from shared memory.\n uint4 out[Gmem_tile_o::STGS_PER_LOOP];\n smem_o.load(out);\n\n // Make sure the data was read from shared memory.\n if (ii < Gmem_tile_o::LOOPS - 1) {\n __syncthreads();\n }\n\n // Output the values.\n gmem_o.store(out, ii);\n }\n\n // Move to the next part of the output.\n gmem_o.move();\n\n typename Smem_tile_st::Fragment frag_s[Mma_tile_dk::MMAS_K][Mma_tile_dk::MMAS_M];\n smem_s.load(frag_s);\n\n static_assert(Mma_tile_dk::MMAS_K == 1); // DEBUG\n\n#pragma unroll\n for (int ki = 1; ki < Mma_tile_dk::MMAS_K; ++ki) {\n // Trigger the load from shared memory for the next series of Q values.\n smem_qt.load(frag_qt[ki & 1], ki);\n // Do the math for the values already in registers.\n fmha::gemm(acc_dk, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);\n }\n\n // Do the final stage of math.\n {\n int ki = Mma_tile_dk::MMAS_K;\n fmha::gemm(acc_dk, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);\n }\n\n // Commit the values for Q into shared memory.\n if (l < nl_traits.num_steps_ - 1) {\n gmem_q.commit(smem_qt);\n __syncthreads();\n // Trigger the loads for the values of Q for the next iteration.\n smem_qt.load(frag_qt[0], 0);\n smem_k.load(frag_k[0], 0);\n }\n\n } // Outer loop over the sequence length.\n\n // Epilogue for dK = dP' * dq. We're fully exposed to this!\n\n // Epilogue swizzle for dK\n Smem_tile_dk smem_dk(&smem_[0], tidx);\n smem_dk.store(acc_dk);\n\n __syncthreads();\n\n uint4 dk_out[Smem_tile_dk::NUM_LDS];\n smem_dk.load(dk_out);\n Qkv_params dk_params;\n dk_params.qkv_ptr = params.dkv_ptr;\n dk_params.qkv_stride_in_bytes = params.h * 2 * CHUNKS * params.d * sizeof(half);\n dk_params.h = params.h;\n Gmem_tile_dk gmem_dk(dk_params, nl_traits.get_idx_dk(), binfo, tidx);\n gmem_dk.store(dk_out);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace fmha\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_fill.cu", "content": "/******************************************************************************\n * Copyright (c) 2011-2023, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \n#include \n#include \n\nconstexpr int block_size = 512;\nconstexpr int ctas_per_sm = 4;\n\ntemplate \n__global__ void __launch_bounds__(block_size)\n mha_fill_kernel(scalar_t* out_tensor, const int32_t* const start_row, const size_t num_rows) {\n size_t row_stride = gridDim.y * blockDim.x;\n size_t row_index = blockIdx.x + (size_t)start_row[0];\n size_t col_index = blockIdx.y * blockDim.x + threadIdx.x;\n while (row_index < num_rows) {\n out_tensor[row_index * row_stride + col_index] = 0;\n row_index += gridDim.x;\n }\n}\n\nat::Tensor& mha_fill(at::Tensor& self, const at::Tensor& start_index) {\n auto max_tokens = self.size(0);\n auto self_2d = self.view({max_tokens, -1});\n auto fcd_size = self_2d.size(1);\n TORCH_CHECK(self.is_contiguous(), \"input not contiguous\");\n TORCH_CHECK(fcd_size % block_size == 0, \"input size not aligned to block size\");\n const int num_mp = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;\n uint64_t num_blk_y = (uint64_t)(fcd_size / block_size);\n uint64_t num_blk_x = (uint64_t)std::ceil(num_mp * ctas_per_sm / num_blk_y);\n dim3 dim_grid(num_blk_x, num_blk_y);\n dim3 dim_block(block_size);\n\n AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND2(\n at::ScalarType::Half, at::ScalarType::BFloat16, self_2d.scalar_type(), \"mha_padding_fill_\", [&]() {\n mha_fill_kernel<<>>(\n self_2d.data_ptr(), start_index.data_ptr(), max_tokens);\n C10_CUDA_KERNEL_LAUNCH_CHECK();\n });\n return self;\n}\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_fprop_fp16_128_64_kernel.sm80.cu", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_fprop_kernel_1xN.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<128, 64, 16, 1, 4, 0x08u>;\n\ntemplate \n__global__ void fmha_fprop_fp16_128_64_sm80_kernel(Fused_multihead_attention_fprop_params params,\n const int total_heads) {\n fmha::device_1xN(params, total_heads);\n}\n\nvoid run_fmha_fp16_128_64_sm80(Launch_params& launch_params,\n const bool configure) {\n auto kernel = launch_params.is_training ? &fmha_fprop_fp16_128_64_sm80_kernel\n : &fmha_fprop_fp16_128_64_sm80_kernel;\n\n constexpr int smem_size = fmha::get_dynamic_smem_size();\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n\n const int sm_count = launch_params.props->multiProcessorCount;\n int ctas_per_sm;\n FMHA_CHECK_CUDA(\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctas_per_sm, kernel, Kernel_traits::THREADS, smem_size));\n int total_ctas = sm_count * ctas_per_sm;\n\n const int heads_total = launch_params.params.b * launch_params.params.h;\n if (configure) {\n using Mma_tile_p = fmha::Hmma_tile;\n constexpr size_t STEPS = Kernel_traits::Cta_tile_p::N / Kernel_traits::Cta_tile_p::M;\n constexpr size_t MMAS_M = Mma_tile_p::MMAS_M;\n constexpr size_t MMAS_N = Mma_tile_p::MMAS_N;\n\n size_t heads_per_cta = ((heads_total + total_ctas - 1) / total_ctas);\n size_t elts_per_head = STEPS * MMAS_M * MMAS_N * 8;\n launch_params.elts_per_thread = heads_per_cta * elts_per_head;\n return;\n }\n\n dim3 grid(total_ctas);\n kernel<<>>(launch_params.params, heads_total);\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_fprop_fp16_256_64_kernel.sm80.cu", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_fprop_kernel_1xN.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<256, 64, 16, 1, 4, 0x08u>;\n\ntemplate \n__global__ void fmha_fprop_fp16_256_64_sm80_kernel(Fused_multihead_attention_fprop_params params,\n const int total_heads) {\n fmha::device_1xN(params, total_heads);\n}\n\nvoid run_fmha_fp16_256_64_sm80(Launch_params& launch_params,\n const bool configure) {\n auto kernel = launch_params.is_training ? &fmha_fprop_fp16_256_64_sm80_kernel\n : &fmha_fprop_fp16_256_64_sm80_kernel;\n\n constexpr int smem_size = fmha::get_dynamic_smem_size();\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n\n const int sm_count = launch_params.props->multiProcessorCount;\n int ctas_per_sm;\n FMHA_CHECK_CUDA(\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctas_per_sm, kernel, Kernel_traits::THREADS, smem_size));\n int total_ctas = sm_count * ctas_per_sm;\n\n const int heads_total = launch_params.params.b * launch_params.params.h;\n if (configure) {\n using Mma_tile_p = fmha::Hmma_tile;\n constexpr size_t STEPS = Kernel_traits::Cta_tile_p::N / Kernel_traits::Cta_tile_p::M;\n constexpr size_t MMAS_M = Mma_tile_p::MMAS_M;\n constexpr size_t MMAS_N = Mma_tile_p::MMAS_N;\n\n size_t heads_per_cta = ((heads_total + total_ctas - 1) / total_ctas);\n size_t elts_per_head = STEPS * MMAS_M * MMAS_N * 8;\n launch_params.elts_per_thread = heads_per_cta * elts_per_head;\n return;\n }\n\n dim3 grid(total_ctas);\n kernel<<>>(launch_params.params, heads_total);\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_fprop_fp16_384_64_kernel.sm80.cu", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_fprop_kernel_1xN.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<384, 64, 16, 1, 4, 0x18u>;\n\ntemplate \n__global__ void fmha_fprop_fp16_384_64_sm80_kernel(Fused_multihead_attention_fprop_params params,\n const int total_heads) {\n fmha::device_1xN(params, total_heads);\n}\n\nvoid run_fmha_fp16_384_64_sm80(Launch_params& launch_params,\n const bool configure) {\n auto kernel = launch_params.is_training ? &fmha_fprop_fp16_384_64_sm80_kernel\n : &fmha_fprop_fp16_384_64_sm80_kernel;\n\n constexpr int smem_size = fmha::get_dynamic_smem_size();\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n\n const int sm_count = launch_params.props->multiProcessorCount;\n int ctas_per_sm;\n FMHA_CHECK_CUDA(\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctas_per_sm, kernel, Kernel_traits::THREADS, smem_size));\n int total_ctas = sm_count * ctas_per_sm;\n\n const int heads_total = launch_params.params.b * launch_params.params.h;\n if (configure) {\n using Mma_tile_p = fmha::Hmma_tile;\n constexpr size_t STEPS = Kernel_traits::Cta_tile_p::N / Kernel_traits::Cta_tile_p::M;\n constexpr size_t MMAS_M = Mma_tile_p::MMAS_M;\n constexpr size_t MMAS_N = Mma_tile_p::MMAS_N;\n\n size_t heads_per_cta = ((heads_total + total_ctas - 1) / total_ctas);\n size_t elts_per_head = STEPS * MMAS_M * MMAS_N * 8;\n launch_params.elts_per_thread = heads_per_cta * elts_per_head;\n return;\n }\n\n dim3 grid(total_ctas);\n kernel<<>>(launch_params.params, heads_total);\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_fprop_fp16_512_64_kernel.sm80.cu", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n#include \"fmha_fprop_kernel_1xN.h\"\n\nusing Kernel_traits = FMHA_kernel_traits<512, 64, 16, 1, 8, 0x00u>;\n\ntemplate \n__global__ void fmha_fprop_fp16_512_64_sm80_kernel(Fused_multihead_attention_fprop_params params,\n const int total_heads) {\n fmha::device_1xN(params, total_heads);\n}\n\ntemplate \n__global__ void fmha_fprop_fp16_512_64_sm80_kernel_nl(Fused_multihead_attention_fprop_params params,\n const int num_full_heads, const int num_main_groups,\n const int main_group_size, const int main_steps,\n const int rest_steps) {\n fmha::device_1xN(params, num_full_heads, num_main_groups, main_group_size, main_steps,\n rest_steps);\n}\n\nvoid run_fmha_fp16_512_64_sm80_(Launch_params& launch_params,\n const bool configure) {\n auto kernel = launch_params.is_training ? &fmha_fprop_fp16_512_64_sm80_kernel\n : &fmha_fprop_fp16_512_64_sm80_kernel;\n\n constexpr int smem_size = fmha::get_dynamic_smem_size();\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n\n const int sm_count = launch_params.props->multiProcessorCount;\n int ctas_per_sm;\n FMHA_CHECK_CUDA(\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctas_per_sm, kernel, Kernel_traits::THREADS, smem_size));\n int total_ctas = sm_count * ctas_per_sm;\n\n const int heads_total = launch_params.params.b * launch_params.params.h;\n if (configure) {\n using Mma_tile_p = fmha::Hmma_tile;\n constexpr size_t STEPS = Kernel_traits::Cta_tile_p::N / Kernel_traits::Cta_tile_p::M;\n constexpr size_t MMAS_M = Mma_tile_p::MMAS_M;\n constexpr size_t MMAS_N = Mma_tile_p::MMAS_N;\n\n size_t heads_per_cta = ((heads_total + total_ctas - 1) / total_ctas);\n size_t elts_per_head = STEPS * MMAS_M * MMAS_N * 8;\n launch_params.elts_per_thread = heads_per_cta * elts_per_head;\n return;\n }\n\n dim3 grid(total_ctas);\n kernel<<>>(launch_params.params, heads_total);\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n\nvoid run_fmha_fp16_512_64_sm80_nl_(Launch_params& launch_params,\n const bool configure) {\n auto kernel = launch_params.is_training ? &fmha_fprop_fp16_512_64_sm80_kernel_nl\n : &fmha_fprop_fp16_512_64_sm80_kernel_nl;\n\n constexpr int smem_size = fmha::get_dynamic_smem_size();\n\n if (smem_size >= 48 * 1024) {\n FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));\n }\n\n const int sm_count = launch_params.props->multiProcessorCount;\n int ctas_per_sm;\n FMHA_CHECK_CUDA(\n cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctas_per_sm, kernel, Kernel_traits::THREADS, smem_size));\n int total_ctas = sm_count * ctas_per_sm;\n\n if (configure) {\n const int heads_total = launch_params.params.b * launch_params.params.h;\n std::tie(launch_params.num_full_heads, launch_params.num_main_groups, launch_params.heads_last_wave,\n launch_params.main_steps, launch_params.rest_steps, launch_params.elts_per_thread) =\n fmha::work_dist(total_ctas, heads_total);\n return;\n }\n\n dim3 grid(total_ctas);\n kernel<<>>(\n launch_params.params, launch_params.num_full_heads, launch_params.num_main_groups, launch_params.heads_last_wave,\n launch_params.main_steps, launch_params.rest_steps);\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n\nvoid run_fmha_fp16_512_64_sm80(Launch_params& launch_params,\n const bool configure) {\n if (launch_params.is_nl) {\n run_fmha_fp16_512_64_sm80_nl_(launch_params, configure);\n } else {\n run_fmha_fp16_512_64_sm80_(launch_params, configure);\n }\n}\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN.h", "content": "/***************************************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\n#include \n#include \n\n#include \"fmha_kernel.h\"\n\nnamespace fmha {\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Gemm_Q_K_base {\n using Smem_tile_o = typename Kernel_traits::Smem_tile_o;\n using Smem_tile_q = typename Kernel_traits::Smem_tile_q;\n using Smem_tile_k = typename Kernel_traits::Smem_tile_k;\n using Fragment_q = typename Smem_tile_q::Fragment;\n using Fragment_k = typename Smem_tile_k::Fragment;\n\n // The description of the CTA tile for the 1st batched GEMM.\n using Cta_tile_p = typename Kernel_traits::Cta_tile_p;\n\n // The MMA tile for the 1st GEMM.\n using Mma_tile_p = fmha::Hmma_tile;\n\n static constexpr int SMEM_BYTES_SOFTMAX = Cta_tile_p::M * Cta_tile_p::WARPS_N * sizeof(float) * 2;\n\n __device__ inline Gemm_Q_K_base(char* smem_ptr_q, char* smem_ptr_k, const int tidx)\n : smem_q(smem_ptr_q, tidx), smem_k(smem_ptr_k, tidx) {}\n\n __device__ inline void load_q() { smem_q.load(frag_q[0], 0); }\n\n __device__ inline void reload_q() { smem_q.load(frag_q[0], 0); }\n\n Fragment_q frag_q[2][Mma_tile_p::MMAS_M];\n Smem_tile_q smem_q;\n Smem_tile_k smem_k;\n};\n\ntemplate \nstruct Gemm_Q_K : public Gemm_Q_K_base {\n using Base = Gemm_Q_K_base;\n using Smem_tile_o = typename Base::Smem_tile_o;\n using Smem_tile_q = typename Base::Smem_tile_q;\n using Smem_tile_k = typename Base::Smem_tile_k;\n using Fragment_k = typename Base::Fragment_k;\n using Mma_tile_p = typename Base::Mma_tile_p;\n\n enum { SHARE_SMEM_FOR_K_AND_V = Kernel_traits::SHARE_SMEM_FOR_K_AND_V };\n\n enum { SMEM_OFFSET_O = Smem_tile_q::BYTES_PER_TILE };\n enum { SMEM_OFFSET_V = Smem_tile_q::BYTES_PER_TILE + (SHARE_SMEM_FOR_K_AND_V ? 0 : Smem_tile_k::BYTES_PER_TILE) };\n\n // Q | K / V\n // | O | SOFTMAX\n static constexpr int SMEM_BYTES =\n Smem_tile_q::BYTES_PER_TILE + std::max((SHARE_SMEM_FOR_K_AND_V ? 1 : 2) * Smem_tile_k::BYTES_PER_TILE,\n Smem_tile_o::BYTES_PER_TILE + Base::SMEM_BYTES_SOFTMAX);\n\n __device__ inline Gemm_Q_K(char* smem_, const int tidx) : Base(smem_, smem_ + Smem_tile_q::BYTES_PER_TILE, tidx) {}\n\n __device__ inline void load_k() {\n#pragma unroll\n for (int ki = 0; ki < Mma_tile_p::MMAS_K; ++ki) {\n Base::smem_k.load(frag_k[ki], ki);\n }\n }\n\n template \n __device__ inline void operator()(Acc (&acc_p)[M][N]) {\n// Do this part of P^T = (Q * K^T)^T.\n#pragma unroll\n for (int ki = 1; ki < Mma_tile_p::MMAS_K; ++ki) {\n // Trigger the load from shared memory for the next series of Q values.\n Base::smem_q.load(Base::frag_q[ki & 1], ki);\n // Do the math for the values already in registers.\n fmha::gemm(acc_p, Base::frag_q[(ki - 1) & 1], frag_k[(ki - 1)]);\n }\n // Do the final stage of math.\n {\n int ki = Mma_tile_p::MMAS_K;\n fmha::gemm(acc_p, Base::frag_q[(ki - 1) & 1], frag_k[(ki - 1)]);\n }\n }\n\n __device__ inline void reload_k() {\n // Noop.\n }\n\n Fragment_k frag_k[Mma_tile_p::MMAS_K][Mma_tile_p::MMAS_N];\n};\n\ntemplate \nstruct Gemm_Q_K : public Gemm_Q_K_base {\n using Base = Gemm_Q_K_base;\n using Smem_tile_o = typename Base::Smem_tile_o;\n using Smem_tile_q = typename Base::Smem_tile_q;\n using Smem_tile_k = typename Base::Smem_tile_k;\n using Smem_tile_v = typename Kernel_traits::Smem_tile_v;\n using Fragment_k = typename Base::Fragment_k;\n using Mma_tile_p = typename Base::Mma_tile_p;\n Fragment_k frag_k[2][Mma_tile_p::MMAS_N];\n\n enum { SHARE_SMEM_FOR_K_AND_V = Kernel_traits::SHARE_SMEM_FOR_K_AND_V };\n\n enum { SMEM_OFFSET_V = Smem_tile_q::BYTES_PER_TILE + (SHARE_SMEM_FOR_K_AND_V ? 0 : Smem_tile_k::BYTES_PER_TILE) };\n static_assert(Smem_tile_v::BYTES_PER_TILE == (int)Smem_tile_k::BYTES_PER_TILE);\n enum { SMEM_OFFSET_O = SMEM_OFFSET_V + Smem_tile_v::BYTES_PER_TILE };\n\n // Q | K/V + O + SOFTMAX\n static constexpr int SMEM_BYTES = Smem_tile_q::BYTES_PER_TILE +\n (SHARE_SMEM_FOR_K_AND_V ? 1 : 2) * Smem_tile_k::BYTES_PER_TILE +\n Smem_tile_o::BYTES_PER_TILE + Base::SMEM_BYTES_SOFTMAX;\n\n __device__ inline Gemm_Q_K(char* smem_, const int tidx) : Base(smem_, smem_ + Smem_tile_q::BYTES_PER_TILE, tidx) {}\n\n __device__ inline void load_k() { Base::smem_k.load(frag_k[0], 0); }\n\n template \n __device__ inline void operator()(Acc (&acc_p)[M][N]) {\n// Do this part of P^T = (Q * K^T)^T.\n#pragma unroll\n for (int ki = 1; ki < Mma_tile_p::MMAS_K; ++ki) {\n // Trigger the load from shared memory for the next series of Q values.\n Base::smem_q.load(Base::frag_q[ki & 1], ki);\n Base::smem_k.load(frag_k[ki & 1], ki);\n // Do the math for the values already in registers.\n fmha::gemm(acc_p, Base::frag_q[(ki - 1) & 1], frag_k[(ki - 1) & 1]);\n }\n // Do the final stage of math.\n {\n int ki = Mma_tile_p::MMAS_K;\n fmha::gemm(acc_p, Base::frag_q[(ki - 1) & 1], frag_k[(ki - 1) & 1]);\n }\n }\n\n __device__ inline void reload_k() { Base::smem_k.load(frag_k[0], 0); }\n};\n\ntemplate \nconstexpr size_t get_dynamic_smem_size() {\n return Gemm_Q_K::SMEM_BYTES;\n}\n\ntemplate \ninline __device__ void device_1xN_(const Params& params, const int bidb, const int bidh, const int begin,\n const int steps, Prng& ph) {\n // The description of the CTA tile for the 1st batched GEMM.\n using Cta_tile_p = typename Kernel_traits::Cta_tile_p;\n // The description of the CTA tile for the 2nd batched GEMM.\n using Cta_tile_o = typename Kernel_traits::Cta_tile_o;\n\n // The MMA tile for the 1st GEMM.\n using Mma_tile_p = fmha::Hmma_tile;\n // The MMA tile for the 2nd GEMM.\n using Mma_tile_o = fmha::Hmma_tile;\n\n // The global memory tile to load Q.\n using Gmem_tile_q = typename Kernel_traits::Gmem_tile_q;\n\n // The global memory tile to load K.\n using Gmem_tile_k = typename Kernel_traits::Gmem_tile_k;\n\n // The global memory tile to load V.\n using Gmem_tile_v = typename Kernel_traits::Gmem_tile_v;\n // The shared memory tile to swizzle V.\n using Smem_tile_v = typename Kernel_traits::Smem_tile_v;\n\n // The global memory tile to store O.\n using Gmem_tile_o = typename Kernel_traits::Gmem_tile_o;\n // The shared memory tile to swizzle O.\n using Smem_tile_o = typename Kernel_traits::Smem_tile_o;\n\n using Gmem_tile_s = typename Kernel_traits::Gmem_tile_s;\n\n using Gemm1 = Gemm_Q_K;\n\n using Softmax = fmha::Softmax;\n\n // The number of threads per row.\n enum { THREADS_PER_ROW = 32 };\n\n enum { BITS_PER_ELT_S = sizeof(fmha::A_type) * 8 };\n\n // Shared memory.\n extern __shared__ char smem_[];\n\n // The thread index.\n const int tidx = threadIdx.x;\n\n const BlockInfoPadded binfo(params, bidb, bidh, tidx);\n if (binfo.stop_early()) return;\n\n Gemm1 gemm_q_k(smem_, tidx);\n // Allocate the global memory tile loader for Q.\n Gmem_tile_q gmem_q(params, 0, binfo, tidx);\n // Allocate the global memory tile loader for O.\n Gmem_tile_o gmem_o(params, binfo, tidx);\n // Allocate the global memory tile loader for S.\n Gmem_tile_s gmem_s(params, binfo, tidx);\n // Wind gmem tiles to the correct position.\n for (int it = 0; it < begin; it++) {\n gmem_q.move();\n gmem_s.move();\n gmem_o.move();\n }\n\n fmha::Mask mask(params, binfo, tidx);\n\n // Allocate the global memory tile loader for K.\n Gmem_tile_k gmem_k(params, 1, binfo, tidx);\n // Allocate the global memory tile loader for V.\n Gmem_tile_v gmem_v(params, 2, binfo, tidx);\n // The base pointer of smem_v;\n char* smem_v_ = &smem_[Gemm1::SMEM_OFFSET_V];\n\n // Allocate the shared memory tile loader for V. We use the same as K so be careful!!!\n Smem_tile_v smem_v(smem_v_, tidx);\n\n // Allocate the shared memory tile loader for O. We use the same as K so be careful!!!\n Smem_tile_o smem_o(&smem_[Gemm1::SMEM_OFFSET_O], tidx);\n\n // Trigger the loads for K.\n gmem_k.load(gemm_q_k.smem_k);\n // Trigger the loads for Q.\n gmem_q.load(gemm_q_k.smem_q);\n // Trigger the loads for V.\n gmem_v.load(smem_v);\n\n const uint32_t scale_bmm1 = reinterpret_cast(params.scale_bmm1);\n#pragma unroll\n for (int it = 0; it < Gmem_tile_k::LDGS; it++) {\n gmem_k.fetch_[it] = fmha::hmul8(scale_bmm1, gmem_k.fetch_[it]);\n }\n\n // Commit the data for Q and V to shared memory.\n gmem_q.commit(gemm_q_k.smem_q);\n gmem_v.commit(smem_v);\n\n // Commit the data for K to shared memory.\n if (!Kernel_traits::SHARE_SMEM_FOR_K_AND_V) {\n gmem_k.commit(gemm_q_k.smem_k);\n }\n\n __syncthreads();\n\n // Load the fragments for Q.\n gemm_q_k.load_q();\n\n // Load the fragments for V. We keep the data in registers during the entire kernel.\n typename Smem_tile_v::Fragment frag_v[Mma_tile_o::MMAS_K][Mma_tile_o::MMAS_N];\n#pragma unroll\n for (int ki = 0; ki < Mma_tile_o::MMAS_K; ++ki) {\n smem_v.load(frag_v[ki], ki);\n }\n\n // Commit the data for V to shared memory if it has not been done already.\n if (Kernel_traits::SHARE_SMEM_FOR_K_AND_V) {\n // Make sure we are done loading the fragments for K.\n __syncthreads();\n\n // Commit the data to shared memory for V.\n gmem_k.commit(gemm_q_k.smem_k);\n\n // Make sure the data is in shared memory.\n __syncthreads();\n }\n\n // Load the fragments for K.\n gemm_q_k.load_k();\n uint32_t p_scaled = (uint32_t)256.0 * params.p_dropout;\n\n // Create the object to do the softmax.\n Softmax softmax(params, &smem_[Gemm1::SMEM_OFFSET_O + Smem_tile_o::BYTES_PER_TILE], bidb, tidx);\n\n // Load over the entire sequence length.\n for (int l = 0; l < steps; l++) {\n if (begin + l * Cta_tile_p::M >= binfo.actual_seqlen) break;\n\n // Declare the accumulators for the 1st gemm.\n fmha::Fragment_accumulator acc_p[Mma_tile_p::MMAS_M][Mma_tile_p::MMAS_N];\n fmha::Clear_accumulator::apply(acc_p);\n\n // Do this part of P^T = (Q * K^T)^T.\n gemm_q_k(acc_p);\n\n // Trigger the load for the next Q values.\n if (l < steps - 1) {\n gemm_q_k.smem_q.move_to_next_write_buffer();\n gmem_q.move();\n gmem_q.load(gemm_q_k.smem_q);\n }\n\n // Load the mask for that iteration.\n mask.load(begin + l);\n\n // Convert from the accumulator type to FP32 for Softmax.\n softmax.unpack_noscale(acc_p);\n\n // Apply the mask.\n softmax.apply_mask(mask);\n\n if (Kernel_traits::SHARE_SMEM_FOR_K_AND_V && l == 0) {\n // if we share K and V, it could be that V was not fully read yet but we write into smem for reduction\n __syncthreads();\n }\n // Compute the max.\n float p_max[Mma_tile_p::MMAS_M * 2];\n // softmax.template reduce(p_max);\n softmax.reduce_max(p_max);\n\n // Compute the exponential value.\n softmax.apply_exp(p_max);\n\n // Compute the sum.\n float p_sum[Mma_tile_p::MMAS_M * 2];\n softmax.reduce_sum(p_sum);\n\n // Finalize softmax on the accumulators of P^T.\n softmax.scale(p_sum);\n\n using Frag_p = fmha::Fragment_a;\n Frag_p frag_p[Mma_tile_o::MMAS_K][Mma_tile_o::MMAS_M];\n if (Is_training) {\n auto encode_dropout = [](bool keep, float val) { return keep ? val : -val; };\n#pragma unroll\n for (int mi = 0; mi < Mma_tile_p::MMAS_M; mi++) {\n#pragma unroll\n for (int ii = 0; ii < 2; ii++) {\n#pragma unroll\n for (int ni = 0; ni < Mma_tile_p::MMAS_N / 4; ni++) {\n uint8_t* rand_arr = (uint8_t*)&ph();\n // We encode the dropout pattern in the sign bit of the non-negative softmax to distinguish from\n // pre-existing zeros\n for (int ind = 0; ind < 16; ind++) {\n softmax.elt_[2 * mi + ii][16 * ni + ind] =\n encode_dropout(rand_arr[ind] <= p_scaled, softmax.elt_[2 * mi + ii][16 * ni + ind]);\n }\n }\n }\n }\n softmax.pack(frag_p);\n gmem_s.store(frag_p, mask);\n gmem_s.move();\n } else {\n softmax.pack(frag_p);\n }\n\n // Commit the values for Q into shared memory.\n if (l < steps - 1) {\n gmem_q.commit(gemm_q_k.smem_q);\n }\n\n if (Is_training) {\n#pragma unroll\n for (int ki = 0; ki < Mma_tile_o::MMAS_K; ki++) {\n#pragma unroll\n for (int mi = 0; mi < Mma_tile_o::MMAS_M; mi++) {\n#pragma unroll\n for (int ii = 0; ii < Frag_p::NUM_REGS; ii++) {\n //\"Apply\" the dropout.\n frag_p[ki][mi].reg(ii) = fmha::hmul2(frag_p[ki][mi].reg(ii), params.scale_dropout);\n frag_p[ki][mi].reg(ii) = fmha::hrelu2(frag_p[ki][mi].reg(ii));\n }\n }\n }\n }\n\n // Declare the accumulators for the 1st gemm.\n fmha::Fragment_accumulator acc_o[Mma_tile_o::MMAS_M][Mma_tile_o::MMAS_N];\n fmha::Clear_accumulator::apply(acc_o);\n\n// Do this part of O = P^T * V^T.\n#pragma unroll\n for (int ki = 0; ki < Mma_tile_o::MMAS_K; ++ki) {\n fmha::gemm(acc_o, frag_p[ki], frag_v[ki]);\n }\n\n// Loop over MMAS_M.\n#pragma unroll\n for (int ii = 0; ii < Gmem_tile_o::LOOPS; ++ii) {\n // Swizzle the elements and do the final reduction.\n smem_o.store(acc_o, ii);\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // Load from shared memory.\n uint4 out[Gmem_tile_o::STGS_PER_LOOP];\n smem_o.load(out);\n\n // Make sure the data was read from shared memory.\n if (ii < Gmem_tile_o::LOOPS - 1) {\n __syncthreads();\n }\n\n // Output the values.\n gmem_o.store(out, ii);\n }\n\n // Move to the next part of the output.\n gmem_o.move();\n gemm_q_k.reload_k();\n\n // Commit the values for Q into shared memory.\n if (l < steps - 1) {\n gemm_q_k.reload_q();\n }\n\n } // Outer loop over the sequence length.\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void device_1xN(const Params& params, const int num_full_heads, const int num_main_groups,\n const int main_group_size, const int main_steps, const int rest_steps) {\n constexpr int STEPS = Kernel_traits::Cta_tile_p::N / Kernel_traits::Cta_tile_p::M;\n const int tidx_global = blockIdx.x * gridDim.x + threadIdx.x;\n auto seeds = at::cuda::philox::unpack(params.philox_args);\n Philox ph(std::get<0>(seeds), tidx_global, std::get<1>(seeds));\n for (int it = 0; it < num_full_heads; it++) {\n const int bidx = it * gridDim.x + blockIdx.x;\n const int bidh = bidx % params.h;\n const int bidb = bidx / params.h;\n fmha::device_1xN_(params, bidb, bidh, 0, STEPS, ph);\n __syncthreads();\n }\n if (main_group_size == 0) return;\n const int head_offset = num_full_heads * gridDim.x;\n\n if (blockIdx.x < main_group_size * num_main_groups) {\n // process within heads\n const int group = blockIdx.x % num_main_groups;\n const int bidx = blockIdx.x / num_main_groups;\n const int bidh = (head_offset + bidx) % params.h;\n const int bidb = (head_offset + bidx) / params.h;\n const int offset = group * main_steps;\n fmha::device_1xN_(params, bidb, bidh, offset, main_steps, ph);\n } else {\n if (rest_steps == 0) return;\n // process across heads\n const int bidx = blockIdx.x - main_group_size * num_main_groups;\n const int offset = num_main_groups * main_steps;\n const int total_heads = params.b * params.h;\n const int rest_ctas = gridDim.x - main_group_size * num_main_groups;\n for (int it = head_offset + bidx; it < total_heads; it += rest_ctas) {\n const int bidh = it % params.h;\n const int bidb = it / params.h;\n fmha::device_1xN_(params, bidb, bidh, offset, rest_steps, ph);\n __syncthreads();\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void device_1xN(const Params& params, const int total_heads) {\n const int tidx_global = blockIdx.x * gridDim.x + threadIdx.x;\n auto seeds = at::cuda::philox::unpack(params.philox_args);\n Philox ph(std::get<0>(seeds), tidx_global, std::get<1>(seeds));\n constexpr int STEPS = Kernel_traits::Cta_tile_p::N / Kernel_traits::Cta_tile_p::M;\n\n for (int bidx = blockIdx.x; bidx < total_heads; bidx += gridDim.x) {\n const int bidh = bidx % params.h;\n const int bidb = bidx / params.h;\n fmha::device_1xN_(params, bidb, bidh, 0, STEPS, ph);\n __syncthreads();\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace fmha\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_kernel.h", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\nnamespace fmha {\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct BlockInfoPadded {\n template \n __device__ BlockInfoPadded(const Params& params, const int bidb, const int bidh, const int tidx)\n : bidb(bidb), bidh(bidh), h(params.h) {\n // The block index.\n sum_s = params.cu_seqlens[bidb];\n actual_seqlen = params.cu_seqlens[bidb + 1] - sum_s;\n bidx = sum_s * params.h + bidh;\n\n tidx_global = (bidb * params.h + bidh) * THREADS_PER_CTA + tidx;\n }\n\n __device__ bool stop_early() const { return actual_seqlen == 0; }\n\n int actual_seqlen;\n int bidx;\n int sum_s;\n int bidh;\n int bidb;\n int tidx_global;\n int h;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Noloop_traits {\n // Interpretation of Cta_tile dims, i.e. Cta_tile_p:\n enum { STEP = Cta_tile::M };\n enum { SEQLEN = Cta_tile::N };\n\n template \n inline __device__ Noloop_traits(const int bidc, const Block_info& binfo) : bidc_(bidc) {\n const int seqlen = binfo.actual_seqlen;\n const int steps = (seqlen + STEP - 1) / STEP;\n const int steps_per_chunk = (steps + CHUNKS - 1) / CHUNKS;\n\n const int step_begin = bidc_ * steps_per_chunk;\n const int step_end = min(steps, (bidc_ + 1) * steps_per_chunk);\n const int actual_steps = max(0, step_end - step_begin);\n loop_offset_ = step_begin;\n num_steps_ = actual_steps;\n }\n\n template \n inline __device__ void move_all(Tiles&... tiles) const {\n using expand_type = int[];\n for (int s = 0; s < loop_offset_; s++) {\n expand_type{(tiles.move(), 0)...};\n }\n }\n\n inline __device__ int get_idx_dk() const {\n // return bidc_;\n return bidc_ * 2 + 0;\n }\n\n inline __device__ int get_idx_dv() const {\n // return CHUNKS + bidc_;\n return bidc_ * 2 + 1;\n }\n\n inline __device__ int offset_loop_count(const int l) {\n // convert loop counter to position in the outer sequence\n return (loop_offset_ + l) * STEP;\n }\n\n const uint32_t bidc_;\n int loop_offset_;\n int num_steps_;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstd::tuple work_dist(const int total_ctas, const int heads_total) {\n constexpr int STEPS_PER_HEAD = Kernel_traits::Cta_tile_p::N / Kernel_traits::Cta_tile_p::M;\n\n const int num_full_heads = heads_total / total_ctas;\n const int heads_last_wave = heads_total % total_ctas;\n\n int num_main_groups = 0;\n int main_steps = 0;\n int rest_steps = 0;\n if (heads_last_wave > 0) {\n // Number of CTA groups that process within heads.\n num_main_groups = total_ctas / heads_last_wave;\n // Remaining CTAs that process between heads.\n const int rest_ctas = total_ctas - (heads_last_wave * num_main_groups);\n if (rest_ctas == 0) {\n // We have exactly \"num_main_groups\" CTAs to process each of the remaining heads.\n main_steps = (STEPS_PER_HEAD + num_main_groups - 1) / num_main_groups;\n num_main_groups = STEPS_PER_HEAD / main_steps; // Here: main_step > 0\n rest_steps = STEPS_PER_HEAD % main_steps;\n\n } else {\n // Ideal number of steps if we could load-balance as evenly as possible.\n const int steps_ideal = (heads_last_wave * STEPS_PER_HEAD + total_ctas - 1) / total_ctas;\n // Iterations that a \"rest\" CTA has to do at most.\n const int max_rest_iters = (heads_last_wave + rest_ctas - 1) / rest_ctas;\n // Find the first step distribution, s.t. the maximum work of the \"rest\" CTAs is less than the work of the main\n // CTAs.\n main_steps = steps_ideal;\n rest_steps = STEPS_PER_HEAD - main_steps * num_main_groups;\n for (; main_steps * num_main_groups < STEPS_PER_HEAD; main_steps++) {\n rest_steps = STEPS_PER_HEAD - main_steps * num_main_groups;\n const int max_rest_total_steps = rest_steps * max_rest_iters;\n if (max_rest_total_steps < main_steps) break;\n }\n rest_steps = STEPS_PER_HEAD - main_steps * num_main_groups;\n }\n }\n\n using Cta_tile_p = typename Kernel_traits::Cta_tile_p;\n using Mma_tile_p = fmha::Hmma_tile;\n\n const int max_steps = STEPS_PER_HEAD * num_full_heads + std::max(main_steps, rest_steps);\n const int elts_per_thread_per_step = Mma_tile_p::MMAS_M * Mma_tile_p::MMAS_N * 8;\n const int elts_per_thread = max_steps * elts_per_thread_per_step;\n\n return {num_full_heads, num_main_groups, heads_last_wave, main_steps, rest_steps, elts_per_thread};\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace fmha\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_noloop_reduce.cu", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#include \"fmha.h\"\n\ninline __device__ float4 ldg128(const void* ptr) { return *static_cast(ptr); }\n\ninline __device__ void stg128(void* ptr, const float4& data) { *static_cast(ptr) = data; }\n\ntemplate \n__global__ __launch_bounds__(THREADS) void fmha_noloop_reduce_kernel(void* __restrict__ out,\n const void* __restrict__ in,\n const int* __restrict__ cu_seqlens,\n const int batch_size) {\n enum { BYTES_PER_LDG = 16 };\n enum { NUM_ELTS = BYTES_PER_LDG / sizeof(T) };\n\n // One CTA hidden vector for K and V\n enum { BYTES_PER_ROW = HIDDEN_SIZE * sizeof(T) * 2 };\n // The stride in bytes in dQKV\n enum { OUT_STRIDE_BYTES = 3 * HIDDEN_SIZE * sizeof(T) };\n // The offset in bytes in dQKV to the dKV part for non-interleaved heads\n enum { OUT_OFFSET_KV_BYTES = HIDDEN_SIZE * sizeof(T) };\n\n static_assert(BYTES_PER_ROW == HIDDEN_SIZE * 2 * sizeof(T));\n\n // Size in bytes of the input tile\n enum { BYTES_PER_TILE = CHUNKS * BYTES_PER_ROW };\n\n enum { BYTES_PER_CTA = THREADS * BYTES_PER_LDG };\n\n enum { LDGS = BYTES_PER_ROW / BYTES_PER_CTA };\n static_assert(BYTES_PER_CTA * LDGS == BYTES_PER_ROW);\n\n union Vec_t {\n float4 raw;\n T elt[NUM_ELTS];\n };\n\n // ZERO-OUT invalid positions in dQKV\n const int total = cu_seqlens[batch_size];\n if (blockIdx.x >= total) {\n enum { BYTES_PER_QKV_ROW = 3 * HIDDEN_SIZE * sizeof(T) };\n enum { STGS = BYTES_PER_QKV_ROW / BYTES_PER_LDG };\n\n const float4 zeros = make_float4(0.f, 0.f, 0.f, 0.f);\n\n char* base_ptr = static_cast(out) + blockIdx.x * OUT_STRIDE_BYTES;\n\n for (int tidx = threadIdx.x; tidx < STGS; tidx += THREADS) {\n stg128(base_ptr + tidx * BYTES_PER_LDG, zeros);\n }\n\n return;\n }\n\n // SETUP\n const int offset_in = blockIdx.x * BYTES_PER_TILE + threadIdx.x * BYTES_PER_LDG;\n const char* ptr_in = static_cast(in) + offset_in;\n\n const int offset_out = blockIdx.x * OUT_STRIDE_BYTES + threadIdx.x * BYTES_PER_LDG;\n char* ptr_out = static_cast(out) + OUT_OFFSET_KV_BYTES + offset_out;\n\n // LOAD\n\n Vec_t local_in[CHUNKS][LDGS];\n\n#pragma unroll\n for (int c = 0; c < CHUNKS; c++) {\n#pragma unroll\n for (int l = 0; l < LDGS; l++) {\n int offset = c * BYTES_PER_ROW + l * BYTES_PER_CTA;\n local_in[c][l].raw = ldg128(ptr_in + offset);\n }\n }\n\n // UNPACK\n float acc[LDGS][NUM_ELTS];\n\n#pragma unroll\n for (int l = 0; l < LDGS; l++) {\n#pragma unroll\n for (int e = 0; e < NUM_ELTS; e++) {\n acc[l][e] = float(local_in[0][l].elt[e]);\n }\n }\n\n// COMPUTE\n#pragma unroll\n for (int c = 1; c < CHUNKS; c++) {\n#pragma unroll\n for (int l = 0; l < LDGS; l++) {\n#pragma unroll\n for (int e = 0; e < NUM_ELTS; e++) {\n acc[l][e] += float(local_in[c][l].elt[e]);\n }\n }\n }\n\n // PACK\n Vec_t local_out[LDGS];\n\n#pragma unroll\n for (int l = 0; l < LDGS; l++) {\n#pragma unroll\n for (int e = 0; e < NUM_ELTS; e++) {\n local_out[l].elt[e] = T(acc[l][e]);\n }\n }\n\n// STORE\n#pragma unroll\n for (int l = 0; l < LDGS; l++) {\n const int offset = l * BYTES_PER_CTA;\n stg128(ptr_out + offset, local_out[l].raw);\n }\n}\n\nvoid fmha_run_noloop_reduce(void* out, const void* in, const int* cu_seqlens, const int hidden_size,\n const int batch_size, const int total, const int num_chunks, cudaStream_t stream) {\n const int blocks = total;\n\n if (hidden_size == 1024) {\n constexpr int HIDDEN_SIZE = 1024;\n constexpr int THREADS = 256;\n\n if (num_chunks == 2) {\n fmha_noloop_reduce_kernel\n <<>>(out, in, cu_seqlens, batch_size);\n } else if (num_chunks == 3) {\n fmha_noloop_reduce_kernel\n <<>>(out, in, cu_seqlens, batch_size);\n } else {\n assert(false && \"Unsupported num_chunks\");\n }\n\n } else {\n assert(false && \"Unsupported hidden_size\");\n }\n\n FMHA_CHECK_CUDA(cudaPeekAtLastError());\n}\n" }, { "path": "apex/contrib/csrc/fmha/src/fmha_utils.h", "content": "/******************************************************************************\n * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n * * Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * * Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n * * Neither the name of the NVIDIA CORPORATION nor the\n * names of its contributors may be used to endorse or promote products\n * derived from this software without specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY\n * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n ******************************************************************************/\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define FMHA_CHECK_CUDA(call) \\\n do { \\\n cudaError_t status_ = call; \\\n if (status_ != cudaSuccess) { \\\n fprintf(stderr, \"CUDA error (%s:%d): %s\\n\", __FILE__, __LINE__, cudaGetErrorString(status_)); \\\n exit(1); \\\n } \\\n } while (0)\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nenum Data_type { DATA_TYPE_FP16, DATA_TYPE_FP32, DATA_TYPE_INT32, DATA_TYPE_INT8 };\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline void set_alpha(uint32_t& alpha, float norm, Data_type dtype) {\n if (dtype == DATA_TYPE_FP16) {\n half x = __float2half_rn(norm);\n uint16_t h = reinterpret_cast(x);\n ushort2 h2 = {h, h};\n alpha = reinterpret_cast(h2);\n } else if (dtype == DATA_TYPE_FP32) {\n alpha = reinterpret_cast(norm);\n } else if (dtype == DATA_TYPE_INT32) {\n int32_t inorm = static_cast(norm);\n alpha = reinterpret_cast(inorm);\n } else {\n assert(false);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline size_t get_size_in_bytes(size_t n, Data_type dtype) {\n switch (dtype) {\n case DATA_TYPE_FP32:\n return n * 4;\n case DATA_TYPE_FP16:\n return n * 2;\n case DATA_TYPE_INT32:\n return n * 4;\n case DATA_TYPE_INT8:\n return n;\n default:\n assert(false);\n return 0;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n" }, { "path": "apex/contrib/csrc/focal_loss/focal_loss_cuda.cpp", "content": "#include \n\n#include \n#include \n\n// CUDA forward declarations\n\nstd::vector focal_loss_forward_cuda(const at::Tensor& cls_output, const at::Tensor& cls_targets_at_level,\n const at::Tensor& num_positives_sum, const int64_t num_real_classes,\n const float alpha, const float gamma, const float smoothing_factor);\n\nat::Tensor focal_loss_backward_cuda(const at::Tensor& grad_output, const at::Tensor& partial_grad,\n const at::Tensor& num_positives_sum);\n\n// C++ interface\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n#define CHECK_INPUT(x) \\\n CHECK_CUDA(x); \\\n CHECK_CONTIGUOUS(x)\n\nstd::vector focal_loss_forward(const at::Tensor& cls_output, const at::Tensor& cls_targets_at_level,\n const at::Tensor& num_positives_sum, const int64_t num_real_classes,\n const float alpha, const float gamma, const float smoothing_factor) {\n CHECK_INPUT(cls_output);\n CHECK_INPUT(cls_targets_at_level);\n CHECK_INPUT(num_positives_sum);\n\n return focal_loss_forward_cuda(cls_output, cls_targets_at_level, num_positives_sum, num_real_classes, alpha, gamma,\n smoothing_factor);\n}\n\nat::Tensor focal_loss_backward(const at::Tensor& grad_output, const at::Tensor& partial_grad,\n const at::Tensor& num_positives_sum) {\n CHECK_INPUT(grad_output);\n CHECK_INPUT(partial_grad);\n\n return focal_loss_backward_cuda(grad_output, partial_grad, num_positives_sum);\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"forward\", &focal_loss_forward, \"Focal loss calculation forward (CUDA)\",\n py::call_guard());\n m.def(\"backward\", &focal_loss_backward, \"Focal loss calculation backward (CUDA)\",\n py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/focal_loss/focal_loss_cuda_kernel.cu", "content": "#include \n#include \n#include \n\n// Use 128-bit vectorization\ntypedef uint4 vector_t;\n\n#define ASSERT_ALIGNED(DTYPE, PTR) \\\n TORCH_INTERNAL_ASSERT(is_aligned(PTR), \"Tensor \" #PTR \" is not \" #DTYPE \" aligned\")\n\ntemplate \nbool is_aligned(const void* ptr) noexcept {\n auto iptr = reinterpret_cast(ptr);\n return !(iptr % alignof(T));\n}\n\ntemplate \n__global__ void focal_loss_forward_cuda_kernel(outscalar_t* loss, scalar_t* partial_grad,\n const scalar_t* __restrict__ cls_output,\n const labelscalar_t* __restrict__ cls_targets_at_level,\n const float* __restrict__ num_positives_sum, const int64_t num_examples,\n const int64_t num_classes, const int64_t num_real_classes,\n const float alpha, const float gamma, const float smoothing_factor) {\n extern __shared__ unsigned char shm[];\n accscalar_t* loss_shm = reinterpret_cast(shm);\n loss_shm[threadIdx.x] = 0;\n accscalar_t loss_acc = 0;\n\n accscalar_t one = accscalar_t(1.0);\n accscalar_t K = accscalar_t(2.0);\n accscalar_t normalizer = one / static_cast(num_positives_sum[0]);\n accscalar_t nn_norm, np_norm, pn_norm, pp_norm;\n\n // *_norm is used for label smoothing only\n if (SMOOTHING) {\n nn_norm = one - smoothing_factor / K;\n np_norm = smoothing_factor / K;\n pn_norm = smoothing_factor - smoothing_factor / K;\n pp_norm = one - smoothing_factor + smoothing_factor / K;\n }\n\n vector_t p_vec, grad_vec;\n\n // Accumulate loss on each thread\n for (int64_t i = (blockIdx.x * blockDim.x + threadIdx.x) * ILP; i < num_examples * num_classes;\n i += gridDim.x * blockDim.x * ILP) {\n int64_t idy = i / num_classes;\n labelscalar_t y = cls_targets_at_level[idy];\n int64_t base_yid = i % num_classes;\n\n int64_t pos_idx = idy * num_classes + y;\n p_vec = *(vector_t*)&cls_output[i]; // Vectorized load\n\n // Skip ignored matches\n if (y == -2) {\n#pragma unroll\n for (int j = 0; j < ILP; j++) {\n *((scalar_t*)(&grad_vec) + j) = 0;\n }\n *(vector_t*)&partial_grad[i] = grad_vec;\n continue;\n }\n\n#pragma unroll\n for (int j = 0; j < ILP; j++) {\n // Skip the pad classes\n if (base_yid + j >= num_real_classes) {\n *((scalar_t*)(&grad_vec) + j) = 0;\n continue;\n }\n\n accscalar_t p = static_cast(*((scalar_t*)(&p_vec) + j));\n accscalar_t exp_np = ::exp(-p);\n accscalar_t exp_pp = ::exp(p);\n accscalar_t sigma = one / (one + exp_np);\n accscalar_t logee = (p >= 0) ? exp_np : exp_pp;\n accscalar_t addee = (p >= 0) ? 0 : -p;\n accscalar_t off_a = addee + ::log(one + logee);\n\n // Negative matches\n accscalar_t base = SMOOTHING ? nn_norm * p : p;\n accscalar_t off_b = (SMOOTHING ? np_norm : 0) - sigma;\n accscalar_t coeff_f1 = one - alpha;\n accscalar_t coeff_f2 = sigma;\n accscalar_t coeff_b1 = gamma;\n accscalar_t coeff_b2 = one - sigma;\n\n // Positive matches\n if (y >= 0 && (i + j == pos_idx)) {\n base = SMOOTHING ? pn_norm * p : 0;\n off_b = (SMOOTHING ? pp_norm : one) - sigma;\n coeff_f1 = alpha;\n coeff_f2 = one - sigma;\n coeff_b1 = -gamma;\n coeff_b2 = sigma;\n }\n\n accscalar_t coeff_f = coeff_f1 * ::pow(coeff_f2, gamma);\n accscalar_t coeff_b = coeff_b1 * coeff_b2;\n\n accscalar_t loss_t = coeff_f * (base + off_a);\n accscalar_t grad = coeff_f * (coeff_b * (base + off_a) - off_b);\n\n // Delay the normalize of partial gradient by num_positives_sum to back\n // propagation because scalar_t reduces precision. Focal loss is very\n // sensitive to the small gradient. No worry on overflow here since\n // gradient has relative smaller range than input.\n loss_acc += loss_t;\n *((scalar_t*)(&grad_vec) + j) = static_cast(grad);\n }\n\n // This may generate two vectorized stores instead of one\n *(vector_t*)&partial_grad[i] = grad_vec;\n }\n loss_shm[threadIdx.x] = loss_acc;\n\n // Intra-CTA reduction\n __syncthreads();\n for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {\n if (threadIdx.x < s) {\n loss_shm[threadIdx.x] += loss_shm[threadIdx.x + s];\n }\n __syncthreads();\n }\n\n // Inter-CTA reduction\n if (threadIdx.x == 0) {\n loss_acc = loss_shm[0] * normalizer;\n atomicAdd(loss, loss_acc);\n }\n}\n\ntemplate \n__global__ void focal_loss_backward_cuda_kernel(scalar_t* partial_grad, const outscalar_t* __restrict__ grad_output,\n const float* __restrict__ num_positives_sum, const uint64_t numel) {\n int64_t idx = (blockIdx.x * blockDim.x + threadIdx.x) * ILP;\n\n accscalar_t normalizer = static_cast(grad_output[0]) / static_cast(num_positives_sum[0]);\n\n // The input is enforced to pad to use vector load, thus there's no need to\n // check whether the last element of ILP can out of bound.\n if (idx >= numel) return;\n\n vector_t grad_vec;\n grad_vec = *(vector_t*)&partial_grad[idx];\n#pragma unroll(ILP)\n for (int i = 0; i < ILP; i++) {\n auto grad = static_cast(*((scalar_t*)(&grad_vec) + i));\n grad *= normalizer;\n *((scalar_t*)(&grad_vec) + i) = static_cast(grad);\n }\n *(vector_t*)&partial_grad[idx] = grad_vec;\n}\n\nstd::vector focal_loss_forward_cuda(const at::Tensor& cls_output, const at::Tensor& cls_targets_at_level,\n const at::Tensor& num_positives_sum, const int64_t num_real_classes,\n const float alpha, const float gamma, const float smoothing_factor) {\n // Checks required for correctness\n TORCH_INTERNAL_ASSERT(cls_output.size(-1) >= num_real_classes, \"Incorrect number of real classes.\");\n TORCH_INTERNAL_ASSERT(cls_targets_at_level.scalar_type() == at::kLong, \"Invalid label type.\");\n TORCH_INTERNAL_ASSERT((num_positives_sum.numel() == 1) && (num_positives_sum.scalar_type() == at::kFloat),\n \"Expect num_positives_sum to be a float32 tensor with only one element.\");\n TORCH_INTERNAL_ASSERT(cls_output.dim() == cls_targets_at_level.dim() + 1,\n \"Mis-matched dimensions between class output and label.\");\n for (int64_t i = 0; i < cls_targets_at_level.dim(); i++)\n TORCH_INTERNAL_ASSERT(cls_output.size(i) == cls_targets_at_level.size(i),\n \"Mis-matched shape between class output and label.\");\n\n // Checks required for better performance\n const int ILP = sizeof(vector_t) / cls_output.element_size();\n ASSERT_ALIGNED(vector_t, cls_output.data_ptr());\n TORCH_INTERNAL_ASSERT(cls_output.size(-1) % ILP == 0,\n \"Pad number of classes first to take advantage of vectorized load.\");\n TORCH_INTERNAL_ASSERT(num_real_classes >= ILP, \"Too few classes.\");\n\n int64_t num_classes = cls_output.size(-1);\n int64_t num_examples = cls_output.numel() / num_classes;\n at::Tensor loss = at::zeros({}, cls_output.options().dtype(at::kFloat));\n\n // Compute the incompelete gradient during fprop since most of the heavy\n // functions of bprop are the same as fprop, thus trade memory for compute\n // helps with focal loss.\n at::Tensor partial_grad = at::empty_like(cls_output);\n\n // Set the number of CTAs per SM according to the compute capability.\n // Each CTA loops on input with stride till the last item.\n cudaDeviceProp props;\n cudaGetDeviceProperties(&props, at::cuda::current_device());\n int cta_per_sm = 2;\n if (props.major >= 10) {\n cta_per_sm = 8;\n }\n dim3 block(512);\n dim3 grid(cta_per_sm * props.multiProcessorCount);\n\n // Specialize on label smoothing or not to reduce redundant operations\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if (smoothing_factor == 0.0f) {\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(cls_output.scalar_type(), \"focal_loss_fprop\", [&] {\n using accscalar_t = at::acc_type;\n using labelscalar_t = int64_t;\n using outscalar_t = float;\n const int ILP = sizeof(vector_t) / sizeof(scalar_t);\n focal_loss_forward_cuda_kernel\n <<>>(\n loss.data_ptr(), partial_grad.data_ptr(), cls_output.data_ptr(),\n cls_targets_at_level.data_ptr(), num_positives_sum.data_ptr(), num_examples,\n num_classes, num_real_classes, alpha, gamma, smoothing_factor);\n });\n } else {\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(cls_output.scalar_type(), \"focal_loss_fprop\", [&] {\n using accscalar_t = at::acc_type;\n using labelscalar_t = int64_t;\n using outscalar_t = float;\n const int ILP = sizeof(vector_t) / sizeof(scalar_t);\n focal_loss_forward_cuda_kernel\n <<>>(\n loss.data_ptr(), partial_grad.data_ptr(), cls_output.data_ptr(),\n cls_targets_at_level.data_ptr(), num_positives_sum.data_ptr(), num_examples,\n num_classes, num_real_classes, alpha, gamma, smoothing_factor);\n });\n }\n\n AT_CUDA_CHECK(cudaGetLastError());\n return {loss, partial_grad};\n}\n\nat::Tensor focal_loss_backward_cuda(const at::Tensor& grad_output, const at::Tensor& partial_grad,\n const at::Tensor& num_positives_sum) {\n // Each thread process ILP elements\n const int ILP = sizeof(vector_t) / partial_grad.element_size();\n dim3 block(512);\n dim3 grid((partial_grad.numel() + block.x * ILP - 1) / (block.x * ILP));\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(partial_grad.scalar_type(), \"focal_loss_bprop\", [&] {\n using accscalar_t = at::acc_type;\n using outscalar_t = float;\n const int ILP = sizeof(vector_t) / sizeof(scalar_t);\n focal_loss_backward_cuda_kernel\n <<>>(partial_grad.data_ptr(), grad_output.data_ptr(),\n num_positives_sum.data_ptr(), partial_grad.numel());\n });\n\n AT_CUDA_CHECK(cudaGetLastError());\n return partial_grad;\n}\n" }, { "path": "apex/contrib/csrc/gpu_direct_storage/gds.cpp", "content": "// Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n\n#include \n\n// torch\n#include \n#include \n\n// cuda\n#include \n#include \n\n// file io\n#include \n\nnamespace apex::contrib::gds {\n\n// POSIX\ntemplate ::value, std::nullptr_t>::type = nullptr>\nstd::string cuFileGetErrorString(T status) {\n status = std::abs(status);\n return IS_CUFILE_ERR(status) ? std::string(CUFILE_ERRSTR(status)) : std::string(std::strerror(errno));\n}\n\n// CUfileError_t\ntemplate ::value, std::nullptr_t>::type = nullptr>\nstd::string cuFileGetErrorString(T status) {\n std::string errStr = cuFileGetErrorString(static_cast(status.err));\n if (IS_CUDA_ERR(status)) errStr.append(\".\").append(cudaGetErrorString(static_cast(status.cu_err)));\n return errStr;\n}\n\nFile::File() : is_open(false) {};\n\nFile::File(const std::string& filename, const std::string& mode) : filename(filename), mode(mode), is_open(false) {\n open(filename, mode);\n}\n\nFile::~File() {\n if (is_open) {\n close();\n }\n}\n\nvoid File::open(const std::string& other_filename, const std::string& other_mode) {\n TORCH_CHECK(is_open == false, \"file\", filename, \"is already open\");\n if (!filename.empty()) {\n TORCH_CHECK(other_filename == filename, \"file\", filename, \"is already open with mode\", mode);\n }\n if (!mode.empty()) {\n TORCH_CHECK(other_mode == mode, \"file\", filename, \"is already open with mode\", mode);\n }\n\n maybe_register = true;\n // Open the binary file\n if (mode == \"r\") {\n // for reading\n fd = ::open(filename.c_str(), O_RDONLY | O_DIRECT);\n } else if (mode == \"w\") {\n // for writing\n fd = ::open(filename.c_str(), O_CREAT | O_WRONLY | O_DIRECT, 0664);\n } else if (mode == \"rn\") {\n // for reading\n fd = ::open(filename.c_str(), O_RDONLY);\n maybe_register = false;\n } else if (mode == \"wn\") {\n // for writing\n fd = ::open(filename.c_str(), O_CREAT | O_WRONLY, 0664);\n maybe_register = false;\n } else {\n TORCH_CHECK(false, \"only r and w modes are currently supported, but got:\", mode);\n }\n TORCH_CHECK(fd >= 0, \"fcntl cannot open file: \", filename);\n\n // Register cuFile handle\n if (maybe_register) {\n memset((void*)&cf_descr, 0, sizeof(CUfileDescr_t));\n cf_descr.handle.fd = fd;\n cf_descr.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD;\n status = cuFileHandleRegister(&cf_handle, &cf_descr);\n if (status.err != CU_FILE_SUCCESS) {\n TORCH_CHECK(false, \"cuFileHandleRegister failed: \", cuFileGetErrorString(status));\n }\n }\n is_open = true;\n}\n\nvoid File::close() {\n // Deregister cuFile handle and close the file\n if (is_open) {\n if (maybe_register) {\n cuFileHandleDeregister(cf_handle);\n }\n ::close(fd);\n fd = -1;\n }\n is_open = false;\n}\n\nvoid File::load_data(const torch::Tensor& tensor) {\n TORCH_CHECK(mode == \"r\", filename, \" was opened for read only\");\n c10::cuda::CUDAGuard gpuGuard(tensor.device());\n\n void* dataPtr = tensor.data_ptr();\n const size_t nbytes = tensor.nbytes();\n\n // Read the binary file\n ssize_t ret = cuFileRead(cf_handle, (void*)dataPtr, nbytes, 0, 0);\n TORCH_CHECK(ret >= 0, \"cuFileWrite failed: \", cuFileGetErrorString(ret));\n}\n\nvoid File::save_data(const torch::Tensor& tensor) {\n TORCH_CHECK(mode == \"w\", filename, \" was opened for write only\");\n c10::cuda::CUDAGuard gpuGuard(tensor.device());\n\n void* dataPtr = tensor.data_ptr();\n const size_t nbytes = tensor.nbytes();\n\n // Register device memory\n status = cuFileBufRegister(dataPtr, nbytes, 0);\n TORCH_CHECK(status.err == CU_FILE_SUCCESS, \"cuFileBufRegister failed: \", cuFileGetErrorString(status));\n\n // Write device memory contents to the file\n ssize_t ret = cuFileWrite(cf_handle, dataPtr, nbytes, 0, 0);\n status = cuFileBufDeregister(dataPtr);\n\n TORCH_CHECK(ret >= 0, \"cuFileWrite failed: \", cuFileGetErrorString(ret));\n TORCH_CHECK(status.err == CU_FILE_SUCCESS, \"cuFileBufDeregister failed:\", cuFileGetErrorString(status));\n}\n\n// Just for benchmarking purposes\n\nvoid File::load_data_no_gds(const torch::Tensor& tensor) {\n TORCH_CHECK(mode == \"rn\", filename, \" was opened for read only\");\n c10::cuda::CUDAGuard gpuGuard(tensor.device());\n\n void* dataPtrCPU = nullptr;\n void* dataPtr = tensor.data_ptr();\n const size_t nbytes = tensor.nbytes();\n dataPtrCPU = malloc(nbytes);\n TORCH_CHECK(dataPtrCPU != nullptr, \"malloc failed\");\n\n const ssize_t nbytes_read = pread(fd, dataPtrCPU, nbytes, 0);\n TORCH_CHECK(nbytes_read == nbytes || nbytes_read == 0, \"fcntl pread failed\");\n C10_CUDA_CHECK(cudaMemcpy(dataPtr, dataPtrCPU, nbytes, cudaMemcpyHostToDevice));\n free(dataPtrCPU);\n}\n\nvoid File::save_data_no_gds(const torch::Tensor& tensor) {\n TORCH_CHECK(mode == \"wn\", filename, \" was opened for write only\");\n c10::cuda::CUDAGuard gpuGuard(tensor.device());\n\n void* dataPtrCPU = nullptr;\n void* dataPtr = tensor.data_ptr();\n const size_t nbytes = tensor.nbytes();\n dataPtrCPU = malloc(nbytes);\n TORCH_CHECK(dataPtrCPU != nullptr, \"malloc failed\");\n C10_CUDA_CHECK(cudaMemcpy(dataPtrCPU, dataPtr, nbytes, cudaMemcpyDeviceToHost));\n\n const ssize_t nbytes_written = pwrite(fd, dataPtrCPU, nbytes, 0);\n TORCH_CHECK(nbytes_written == nbytes, \"fcntl pwrite failed\");\n free(dataPtrCPU);\n}\n\n} // namespace apex::contrib::gds\n" }, { "path": "apex/contrib/csrc/gpu_direct_storage/gds.h", "content": "// Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n\n#pragma once\n\n#include \n#include \n\n#include \n\nnamespace apex::contrib::gds {\nclass File {\n public:\n File();\n File(const std::string& filename, const std::string& mode);\n ~File();\n\n void open(const std::string& filename, const std::string& mode);\n void close();\n\n void load_data(const torch::Tensor& tensor);\n void save_data(const torch::Tensor& tensor);\n void load_data_no_gds(const torch::Tensor& tensor);\n void save_data_no_gds(const torch::Tensor& tensor);\n\n private:\n std::string filename;\n std::string mode;\n\n CUfileDescr_t cf_descr;\n CUfileHandle_t cf_handle;\n CUfileError_t status;\n\n int fd = -1;\n bool is_open = false;\n bool maybe_register = true;\n};\n} // namespace apex::contrib::gds\n" }, { "path": "apex/contrib/csrc/gpu_direct_storage/gds_pybind.cpp", "content": "// Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n\n#include \n#include \n#include \n\n#include \n\n// python bindings\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n py::class_>(m, \"_GDSFile\")\n .def(py::init<>())\n .def(py::init())\n .def(\"open\", &apex::contrib::gds::File::open)\n .def(\"close\", &apex::contrib::gds::File::close)\n .def(\"load_data\", &apex::contrib::gds::File::load_data)\n .def(\"save_data\", &apex::contrib::gds::File::save_data)\n .def(\"load_data_no_gds\", &apex::contrib::gds::File::load_data_no_gds)\n .def(\"save_data_no_gds\", &apex::contrib::gds::File::save_data_no_gds);\n}\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc.cpp", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\ntemplate \nfloat inline unpack(const T& x) {\n return {};\n}\n\ntemplate <>\nfloat inline unpack(const __half& x) {\n return __half2float(x);\n}\n\ntemplate <>\nfloat inline unpack(const __nv_bfloat16& x) {\n return __bfloat162float(x);\n}\n\ntemplate <>\nfloat inline unpack(const float& x) {\n return x;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nvoid check_results(const char* name, const T* out, const T* ref, size_t elts, float tol) {\n // The number of errors.\n int failed = 0;\n // The number of infinite value.\n int infs = 0;\n // The min/max values.\n float min_val = FLT_MAX, max_val = -FLT_MAX, max_err = 0.f;\n // The total sum of error.\n double sum_err = 0.0;\n\n // The case we are checking.\n printf(\"\\e[1;34mchecking.....................: %s\\e[0m\\n\", name);\n fflush(stdout);\n\n // Iterate over the different values.\n for (size_t ii = 0; ii < elts; ++ii) {\n float a = unpack(out[ii]);\n float b = unpack(ref[ii]);\n\n // Compute the absolute norms.\n float abs_a = fabsf(a);\n float abs_b = fabsf(b);\n\n // Compute the error.\n float den = abs_a + abs_b;\n // Is one of the quantities very small?\n bool is_small = abs_a <= tol || abs_b <= tol || den <= tol;\n // The error.\n float err = is_small ? fabsf(a - b) : fabsf(a - b) / den;\n // Is the result ok?\n bool ok = !isnan(a) && !isnan(b) && err <= tol;\n\n // Print the error.\n if (!ok && (failed < 10 || err > max_err)) {\n fprintf(stderr, \">> invalid result for ii=%lu:\\n\", ii);\n if (std::is_same::value || std::is_same::value) {\n // The data.\n fprintf(stderr, \">> found...: 0x%04x (%10.6f)\\n\", reinterpret_cast(out[ii]), a);\n fprintf(stderr, \">> expected: 0x%04x (%10.6f)\\n\", reinterpret_cast(ref[ii]), b);\n } else if (std::is_same::value) {\n fprintf(stderr, \">> found...: 0x%08x (%10.6f)\\n\", reinterpret_cast(a), a);\n fprintf(stderr, \">> expected: 0x%08x (%10.6f)\\n\", reinterpret_cast(b), b);\n } else {\n fprintf(stderr, \"\\e[1;34mUnknown type of check_results\\e[0m\\n\");\n exit(1);\n }\n fprintf(stderr, \">> error...: %.6f\\n\", err);\n }\n\n // Update the number of failures.\n failed += ok ? 0 : 1;\n\n // Measure min/max errors.\n min_val = fminf(min_val, a);\n max_val = fmaxf(max_val, a);\n max_err = fmaxf(max_err, err);\n\n // Accumulate the sum.\n sum_err = sum_err + (double)err;\n\n infs += !isfinite(a);\n infs += !isfinite(b);\n }\n\n if (!failed && infs < 10) {\n printf(\"\\e[1;32mcheck........................: OK\\e[0m\\n\");\n } else {\n printf(\"\\e[1;31mcheck........................: FAILED\\e[0m\\n\");\n }\n\n printf(\"tested.......................: %lu\\n\", elts);\n printf(\"failures.....................: %d\\n\", failed);\n printf(\"failure rate.................: %.2lf%%\\n\", (double)failed * 100.0 / (double)elts);\n printf(\"infs.........................: %d\\n\", infs);\n printf(\"tolerance....................: %.8f\\n\", tol);\n printf(\"\\n\");\n\n printf(\"min. value...................: %.6f\\n\", min_val);\n printf(\"max. value...................: %.6f\\n\", max_val);\n printf(\"max. error...................: %.6f\\n\", max_err);\n printf(\"sum. error...................: %.6lf\\n\", sum_err);\n printf(\"avg. error...................: %.6lf\\n\", sum_err / (double)elts);\n printf(\"\\n\");\n}\n\ntemplate void check_results(const char* name, const __half* out, const __half* ref, size_t elts, float tol);\n\ntemplate void check_results(const char* name, const __nv_bfloat16* out, const __nv_bfloat16* ref, size_t elts,\n float tol);\n\ntemplate void check_results(const char* name, const float* out, const float* ref, size_t elts, float tol);\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic void group_norm_nhwc_bwd_(void* dx_h, float* dgamma_h, float* dbeta_h, const void* dy_h, const void* x_h,\n const float* gamma_h, const float* beta_h, const float2* sums_h, float epsilon, int n,\n int h, int w, int c, int groups, bool with_swish, bool use_fp32, bool use_bf16) {\n // The number of channels in each group.\n int channels_per_group = c / groups;\n // The normalization term to compute the means.\n float rcp_hwc_per_group = 1.f / (float)(h * w * channels_per_group);\n\n // The array to compute gamma.\n float* dgamma = (float*)malloc(c * sizeof(float));\n // The array to compute beta.\n float* dbeta = (float*)malloc(c * sizeof(float));\n\n // Set gamma/beta to 0.\n memset(dgamma, 0, c * sizeof(float));\n memset(dbeta, 0, c * sizeof(float));\n\n // Normalize the activations.\n for (int ni = 0; ni < n; ++ni) {\n for (int gi = 0; gi < groups; ++gi) {\n // The sums from the fwd pass.\n float2 sums = sums_h[ni * groups + gi];\n // The mean of X (computed during the fwd pass -- one value per batch*group).\n float x_mean = sums.x;\n // The mean of squares of X (computed during the fwd pass -- one value per batch*group).\n float x_sq_mean = sums.y;\n // The variance.\n float x_var = x_sq_mean - x_mean * x_mean;\n // The reciprocal of the standard deviation (i.e. 1.f / sqrt(var + epsilon)).\n float rcp_x_stddev = x_var <= 0.f ? 1.f : 1.f / sqrtf(x_var + epsilon);\n\n // TODO: We should store rcp_x_stddev instead of the sums of squares.\n\n // The following nested loops compute 2 means.\n float mean_1 = 0.f, mean_2 = 0.f;\n\n // Iterate over the activations in the group.\n for (int hi = 0; hi < h; ++hi) {\n for (int wi = 0; wi < w; ++wi) {\n for (int ii = 0; ii < channels_per_group; ++ii) {\n // The channel.\n int ci = gi * channels_per_group + ii;\n // Compute the src/dst offset.\n size_t offset = (size_t)ni * h * w * c + (size_t)hi * w * c + (size_t)wi * c + (size_t)ci;\n // Convert the element at that position to float.\n float x;\n if (use_fp32) {\n x = reinterpret_cast(x_h)[offset];\n } else if (use_bf16) {\n x = __bfloat162float(reinterpret_cast(x_h)[offset]);\n } else {\n x = __half2float(reinterpret_cast(x_h)[offset]);\n }\n // The output.\n float dy;\n if (use_fp32) {\n dy = reinterpret_cast(dy_h)[offset];\n } else if (use_bf16) {\n dy = __bfloat162float(reinterpret_cast(dy_h)[offset]);\n } else {\n dy = __half2float(reinterpret_cast(dy_h)[offset]);\n }\n\n // Gamma.\n float gamma = gamma_h[ci];\n\n // X - X_mean.\n float x_minus_x_mean = x - x_mean;\n // Normalize X.\n float x_norm = x_minus_x_mean * rcp_x_stddev;\n\n if (with_swish) {\n // Beta\n float beta = beta_h[ci];\n\n float x_gn = x_norm * gamma + beta;\n float s = sigmoid(x_gn);\n dy = dy * s * (1.f + x_gn * (1.f - s));\n }\n\n // Compute the gradient for beta.\n dbeta[ci] += dy;\n\n // Compute the gradient for gamma.\n dgamma[ci] += dy * x_norm;\n\n // The gradient that enters the x_norm node.\n float dx_norm = dy * gamma;\n\n // Accumulators over 2 means\n mean_1 += x_norm * dx_norm;\n mean_2 += dx_norm;\n\n } // ii\n } // wi\n } // hi\n\n mean_1 *= rcp_hwc_per_group;\n mean_2 *= rcp_hwc_per_group;\n\n // Iterate over the activations in the group.\n for (int hi = 0; hi < h; ++hi) {\n for (int wi = 0; wi < w; ++wi) {\n for (int ii = 0; ii < channels_per_group; ++ii) {\n // The channel.\n int ci = gi * channels_per_group + ii;\n // Compute the src/dst offset.\n size_t offset = (size_t)ni * h * w * c + (size_t)hi * w * c + (size_t)wi * c + (size_t)ci;\n float x;\n if (use_fp32) {\n x = reinterpret_cast(x_h)[offset];\n } else if (use_bf16) {\n x = __bfloat162float(reinterpret_cast(x_h)[offset]);\n } else {\n x = __half2float(reinterpret_cast(x_h)[offset]);\n }\n // The output.\n float dy;\n if (use_fp32) {\n dy = reinterpret_cast(dy_h)[offset];\n } else if (use_bf16) {\n dy = __bfloat162float(reinterpret_cast(dy_h)[offset]);\n } else {\n dy = __half2float(reinterpret_cast(dy_h)[offset]);\n }\n\n // Gamma.\n float gamma = gamma_h[ci];\n\n // X - X_mean.\n float x_minus_x_mean = x - x_mean;\n // Normalize X.\n float x_norm = x_minus_x_mean * rcp_x_stddev;\n\n if (with_swish) {\n // Beta\n float beta = beta_h[ci];\n\n float x_gn = x_norm * gamma + beta;\n float s = sigmoid(x_gn);\n dy = dy * s * (1.f + x_gn * (1.f - s));\n }\n\n // The gradient that enters the x_norm node.\n float dx_norm = dy * gamma;\n\n // Input gradient\n float dx = (dx_norm - (x_norm * mean_1 + mean_2)) * rcp_x_stddev;\n\n // Set the output gradient.\n if (use_fp32) {\n reinterpret_cast(dx_h)[offset] = dx;\n } else if (use_bf16) {\n reinterpret_cast<__nv_bfloat16*>(dx_h)[offset] = __float2bfloat16_rn(dx);\n } else {\n reinterpret_cast<__half*>(dx_h)[offset] = __float2half_rn(dx);\n }\n\n } // ii\n } // wi\n } // hi\n\n } // gi\n } // ni\n\n // Store gamma/beta.\n for (int ci = 0; ci < c; ++ci) {\n dgamma_h[ci] = dgamma[ci];\n dbeta_h[ci] = dbeta[ci];\n }\n\n // Release temporary memory.\n free(dgamma);\n free(dbeta);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic void group_norm_nhwc_fwd_(void* y_h, const void* x_h, const float* gamma_h, const float* beta_h, float epsilon,\n int n, int h, int w, int c, int groups, bool with_swish, bool use_fp32,\n bool use_bf16) {\n // The number of channels in each group.\n int channels_per_group = c / groups;\n\n // The normalization term to compute the means.\n float inv_hwcg = 1.f / (float)(h * w * channels_per_group);\n\n // Normalize the activations.\n for (int ni = 0; ni < n; ++ni) {\n for (int gi = 0; gi < groups; ++gi) {\n // The sums to compute the mean/variance for that group.\n float sum = 0.f, sum_sq = 0.f;\n\n // Iterate over the activations in the group.\n for (int hi = 0; hi < h; ++hi) {\n for (int wi = 0; wi < w; ++wi) {\n for (int ii = 0; ii < channels_per_group; ++ii) {\n // The channel.\n int ci = gi * channels_per_group + ii;\n // Compute the src/dst offset.\n size_t offset = (size_t)ni * h * w * c + (size_t)hi * w * c + (size_t)wi * c + (size_t)ci;\n // Convert the element at that position to float.\n float x;\n if (use_fp32) {\n x = reinterpret_cast(x_h)[offset];\n } else if (use_bf16) {\n x = __bfloat162float(reinterpret_cast(x_h)[offset]);\n } else {\n x = __half2float(reinterpret_cast(x_h)[offset]);\n }\n\n // Update the sums.\n sum += x;\n sum_sq += x * x;\n\n } // ii\n } // wi\n } // hi\n\n // Compute the mean.\n float mean = sum * inv_hwcg;\n // Compute the average value for the squares.\n float mean_sq = sum_sq * inv_hwcg;\n // Compute the variance.\n float var = mean_sq - (mean * mean);\n // Invert the variance.\n float inv_stddev = var <= 0.f ? 1.f : (1.f / sqrtf(var + epsilon));\n\n // Iterate over the data to normalize the output.\n for (int hi = 0; hi < h; ++hi) {\n for (int wi = 0; wi < w; ++wi) {\n for (int ii = 0; ii < channels_per_group; ++ii) {\n // The channel.\n int ci = gi * channels_per_group + ii;\n // Compute the src/dst offset.\n size_t offset = (size_t)ni * h * w * c + (size_t)hi * w * c + (size_t)wi * c + (size_t)ci;\n // Normalize.\n float x;\n if (use_fp32) {\n x = reinterpret_cast(x_h)[offset];\n } else if (use_bf16) {\n x = __bfloat162float(reinterpret_cast(x_h)[offset]);\n } else {\n x = __half2float(reinterpret_cast(x_h)[offset]);\n }\n float y = (x - mean) * inv_stddev;\n // Scale with gamma and add beta.\n y = y * gamma_h[ci] + beta_h[ci];\n // Apply swish (if needed).\n if (with_swish) {\n y = y * sigmoid(y);\n }\n // Store the result.\n if (use_fp32) {\n reinterpret_cast(y_h)[offset] = y;\n } else if (use_bf16) {\n reinterpret_cast<__nv_bfloat16*>(y_h)[offset] = __float2bfloat16_rn(y);\n } else {\n reinterpret_cast<__half*>(y_h)[offset] = __float2half_rn(y);\n }\n\n } // ii\n } // wi\n } // hi\n } // gi\n } // ni\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nvoid random_data(T* dst_h, size_t n, bool use_1s, int range = 3) {\n for (size_t ii = 0; ii < n; ++ii) {\n float x = 1.f;\n if (!use_1s) {\n x = (float)(rand() % range - (range / 2));\n }\n if (std::is_same::value) {\n dst_h[ii] = __float2half_rn(x);\n } else if (std::is_same::value) {\n dst_h[ii] = x;\n } else if (std::is_same::value) {\n dst_h[ii] = __float2bfloat16_rn(x);\n } else {\n fprintf(stderr, \"\\e[1;34mUnknown type of random_data\\e[0m\\n\");\n exit(1);\n }\n }\n}\n\ntemplate void random_data(float* dst_h, size_t n, bool use_1s, int range);\n\ntemplate void random_data(__half* dst_h, size_t n, bool use_1s, int range);\n\ntemplate void random_data(__nv_bfloat16* dst_h, size_t n, bool use_1s, int range);\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nenum class Mode { FWD_INFERENCE, FWD_TRAINING, BWD };\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nint main(int argc, char** argv) {\n // The tensor size.\n int n = 2, h = 64, w = 64, c = 320, groups = 32;\n // The default mode is inference.\n Mode mode = Mode::FWD_INFERENCE;\n // The constant epsilon for sqrt(var + epsilon).\n float epsilon = 1.e-5f;\n // Do we fuse with the Swish activation function?\n bool with_swish = false;\n // Do we use the one-pass kernel?\n bool use_one_pass = false;\n // The number of runs to time the code.\n int runs = 1;\n // Do we use 1s for the input data.\n bool use_1s = false;\n // The tolerance to check the results.\n float tol = 1.e-3f;\n // Do we skip the checks?\n bool skip_checks = false;\n // Do we output csv format only\n bool csv_output = false;\n // Use fp32 IO\n bool use_fp32 = false;\n // Use bf16 IO\n bool use_bf16 = false;\n\n // Parse the parameters.\n for (int ii = 1; ii < argc; ++ii) {\n if (!strcmp(argv[ii], \"-1s\")) {\n use_1s = true;\n } else if (!strcmp(argv[ii], \"-bwd\")) {\n mode = Mode::BWD;\n } else if (!strcmp(argv[ii], \"-c\") && ++ii < argc) {\n c = strtol(argv[ii], nullptr, 10);\n } else if (!strcmp(argv[ii], \"-epsilon\") && ++ii < argc) {\n epsilon = (float)strtod(argv[ii], nullptr);\n } else if (!strcmp(argv[ii], \"-fwd\")) {\n mode = Mode::FWD_INFERENCE;\n } else if (!strcmp(argv[ii], \"-fwd-tr\")) {\n mode = Mode::FWD_TRAINING;\n } else if (!strcmp(argv[ii], \"-groups\") && ++ii < argc) {\n groups = strtol(argv[ii], nullptr, 10);\n } else if (!strcmp(argv[ii], \"-h\") && ++ii < argc) {\n h = strtol(argv[ii], nullptr, 10);\n } else if (!strcmp(argv[ii], \"-n\") && ++ii < argc) {\n n = strtol(argv[ii], nullptr, 10);\n } else if (!strcmp(argv[ii], \"-one-pass\")) {\n use_one_pass = true;\n } else if (!strcmp(argv[ii], \"-runs\") && ++ii < argc) {\n runs = strtol(argv[ii], nullptr, 10);\n } else if (!strcmp(argv[ii], \"-skip-checks\")) {\n skip_checks = true;\n } else if (!strcmp(argv[ii], \"-tol\") && ++ii < argc) {\n tol = (float)strtod(argv[ii], nullptr);\n } else if (!strcmp(argv[ii], \"-w\") && ++ii < argc) {\n w = strtol(argv[ii], nullptr, 10);\n } else if (!strcmp(argv[ii], \"-with-swish\")) {\n with_swish = true;\n } else if (!strcmp(argv[ii], \"-csv\")) {\n csv_output = true;\n } else if (!strcmp(argv[ii], \"-fp32\")) {\n use_fp32 = true;\n } else if (!strcmp(argv[ii], \"-bf16\")) {\n use_bf16 = true;\n } else if (ii < argc) {\n fprintf(stderr, \"Unknown argument: %s\\n\", argv[ii]);\n return 1;\n } else {\n fprintf(stderr, \"Argument %s requires a value\\n\", argv[ii - 1]);\n return 1;\n }\n }\n\n if (use_bf16 && use_fp32) {\n fprintf(stderr, \"Can't use fp32 and bf16 IO at the same time\\n\");\n return 1;\n }\n\n // Header.\n if (!csv_output) {\n printf(\"\\n\");\n printf(\"#######################################################################\\n\");\n printf(\"# Group Norm NHWC + Swish kernel\\n\");\n printf(\"# --\\n\");\n printf(\"# Compiled on %s\\n\", __DATE__);\n printf(\"#######################################################################\\n\");\n printf(\"\\n\");\n }\n\n // GPU info.\n cudaDeviceProp props;\n CHECK_CUDA(cudaGetDeviceProperties(&props, 0));\n if (!csv_output) {\n printf(\"device.......................: %s\\n\", props.name);\n printf(\"cc...........................: %d.%d\\n\", props.major, props.minor);\n printf(\"# of sms.....................: %d\\n\", props.multiProcessorCount);\n }\n\n // Dram peak bandwidth.\n float dram_clock = props.memoryClockRate / 1.e6f;\n float dram_peak = 2.f * dram_clock * props.memoryBusWidth / 8.f;\n if (!csv_output) {\n printf(\"dram clock...................: %.3f GHz\\n\", dram_clock);\n printf(\"dram peak....................: %.3f TB/s\\n\", dram_peak * 1.e-3f);\n printf(\"\\n\");\n }\n\n // Output the problem size.\n if (!csv_output) {\n printf(\"n............................: %d\\n\", n);\n printf(\"h............................: %d\\n\", h);\n printf(\"w............................: %d\\n\", w);\n printf(\"c............................: %d\\n\", c);\n printf(\"groups.......................: %d\\n\", groups);\n printf(\"epsilon......................: %f\\n\", epsilon);\n printf(\"with swish...................: %s\\n\", with_swish ? \"true\" : \"false\");\n printf(\"channels per group...........: %d\\n\", c / groups);\n if (mode == Mode::BWD) {\n printf(\"mode.........................: bwd\\n\");\n } else if (mode == Mode::FWD_INFERENCE) {\n printf(\"mode.........................: fwd inference\\n\");\n } else if (mode == Mode::FWD_TRAINING) {\n printf(\"mode.........................: fwd training\\n\");\n } else {\n assert(false);\n }\n printf(\"\\n\");\n }\n\n // Compute the SOL.\n double bytes = 0;\n int32_t io_bytes = use_fp32 ? sizeof(float) : sizeof(__half);\n if (mode != Mode::BWD) {\n bytes = (double)n * h * w * c * io_bytes + // src\n (double)c * 4 + // gamma\n (double)c * 4 + // beta\n (double)n * h * w * c * io_bytes; // out\n } else {\n bytes = (double)n * h * w * c * io_bytes * 2 + // src, dsrc\n (double)c * 4 * 2 + // gamma, dgamma\n (double)c * 4 * 2 + // beta, dbeta\n (double)n * h * w * c * io_bytes * 1; // dout\n }\n double gbytes = bytes * 1.e-9;\n double dram_sol = gbytes / dram_peak * 1.e3;\n if (!csv_output) {\n printf(\"bytes........................: %.3lfGB\\n\", gbytes);\n printf(\"dram sol.....................: %.6lfms\\n\", dram_sol);\n\n // The number of runs to measure performance.\n printf(\"runs.........................: %d\\n\", runs);\n printf(\"\\n\");\n }\n\n // The number of elements in the x tensor. The layout is N x H x W x C.\n size_t x_elts = (size_t)n * h * w * c;\n // The size of the src in bytes.\n size_t x_sz = x_elts * io_bytes;\n\n // Allocate the src/dst on the host.\n void* x_h = malloc(x_sz);\n void* y_h = malloc(x_sz);\n\n // Allocate src/dst on the device.\n void *x_d, *y_d;\n CHECK_CUDA(cudaMalloc((void**)&x_d, x_sz));\n CHECK_CUDA(cudaMalloc((void**)&y_d, x_sz));\n\n // The number of elements in the gamma/beta array.\n size_t gamma_elts = (size_t)c;\n // The size of the gamma/beta array in bytes.\n size_t gamma_sz = gamma_elts * sizeof(float);\n // Allocate gamma/beta on the host.\n float* gamma_h = (float*)malloc(gamma_sz);\n // Allocate gamma/beta on the device.\n float* gamma_d;\n CHECK_CUDA(cudaMalloc((void**)&gamma_d, gamma_sz));\n\n // Allocate gamma/beta on the host.\n float* beta_h = (float*)malloc(gamma_sz);\n // Allocate gamma/beta on the device.\n float* beta_d;\n CHECK_CUDA(cudaMalloc((void**)&beta_d, gamma_sz));\n\n // Allocate the reference on the host (to be computed on the host).\n void* y_ref_h = nullptr;\n if (!skip_checks) {\n y_ref_h = malloc(x_sz);\n }\n\n // Allocate the src/dst on the host for the gradients (bwd).\n void *dx_h = nullptr, *dy_h = nullptr;\n if (mode == Mode::BWD) {\n dx_h = malloc(x_sz);\n dy_h = malloc(x_sz);\n }\n\n // Allocate src/dst on the device.\n void *dx_d = nullptr, *dy_d = nullptr;\n if (mode == Mode::BWD) {\n CHECK_CUDA(cudaMalloc((void**)&dx_d, x_sz));\n CHECK_CUDA(cudaMalloc((void**)&dy_d, x_sz));\n }\n\n // The gradients for gamma and beta on the host.\n float *dgamma_h = nullptr, *dbeta_h = nullptr;\n if (mode == Mode::BWD) {\n dgamma_h = (float*)malloc(gamma_sz);\n dbeta_h = (float*)malloc(gamma_sz);\n }\n\n // The gradients for gamma and beta on the device.\n float *dgamma_d = nullptr, *dbeta_d = nullptr;\n if (mode == Mode::BWD) {\n CHECK_CUDA(cudaMalloc((void**)&dgamma_d, gamma_sz));\n CHECK_CUDA(cudaMalloc((void**)&dbeta_d, gamma_sz));\n }\n\n // The number of sums for the bwd pass.\n size_t sums_elts = mode == Mode::FWD_INFERENCE ? 0 : n * groups;\n // The size needed to store that array.\n size_t sums_sz = sums_elts * sizeof(float2);\n\n // The sums for the bwd pass on the host.\n float2* sums_h = nullptr;\n if (sums_sz > 0) {\n sums_h = (float2*)malloc(sums_sz);\n }\n\n // The sums for the bwd pass on the device.\n float2* sums_d = nullptr;\n if (sums_sz > 0) {\n CHECK_CUDA(cudaMalloc((void**)&sums_d, sums_sz));\n }\n\n // Allocate the reference on the host (to be computed on the host).\n void* dx_ref_h = nullptr;\n if (mode == Mode::BWD && !skip_checks) {\n dx_ref_h = malloc(x_sz);\n }\n\n // Allocate the reference on the host (to be computed on the host).\n float *dgamma_ref_h = nullptr, *dbeta_ref_h = nullptr;\n if (mode == Mode::BWD && !skip_checks) {\n dgamma_ref_h = (float*)malloc(gamma_sz);\n dbeta_ref_h = (float*)malloc(gamma_sz);\n }\n\n // Generate random input data for the forward pass.\n if (use_fp32) {\n random_data(reinterpret_cast(x_h), x_elts, use_1s);\n } else if (use_bf16) {\n random_data<__nv_bfloat16>(reinterpret_cast<__nv_bfloat16*>(x_h), x_elts, use_1s);\n } else {\n random_data<__half>(reinterpret_cast<__half*>(x_h), x_elts, use_1s);\n }\n random_data(gamma_h, gamma_elts, use_1s);\n random_data(beta_h, gamma_elts, use_1s);\n\n // Generate the gradients for the bwd pass.\n if (mode == Mode::BWD) {\n if (use_fp32) {\n random_data(reinterpret_cast(dy_h), x_elts, use_1s);\n } else if (use_bf16) {\n random_data<__nv_bfloat16>(reinterpret_cast<__nv_bfloat16*>(dy_h), x_elts, use_1s);\n } else {\n random_data<__half>(reinterpret_cast<__half*>(dy_h), x_elts, use_1s);\n }\n }\n\n // Precompute the sums (from the fwd pass) for bwd.\n if (mode == Mode::BWD) {\n // Clear the array of sums (all the elements are set to 0.f).\n memset(sums_h, 0, sums_sz);\n\n // The number of channels in each group.\n int channels_per_group = c / groups;\n // Iterate over the different groups.\n for (int ni = 0; ni < n; ++ni) {\n for (int gi = 0; gi < groups; ++gi) {\n for (int hi = 0; hi < h; ++hi) {\n for (int wi = 0; wi < w; ++wi) {\n for (int ii = 0; ii < channels_per_group; ++ii) {\n // The position of the channel.\n int ci = gi * channels_per_group + ii;\n // The offset to the element.\n int64_t offset = (int64_t)ni * h * w * c + hi * w * c + wi * c + ci;\n // The element in float.\n float x;\n if (use_fp32) {\n x = reinterpret_cast(x_h)[offset];\n } else if (use_bf16) {\n x = __bfloat162float(reinterpret_cast<__nv_bfloat16*>(x_h)[offset]);\n } else {\n x = __half2float(reinterpret_cast<__half*>(x_h)[offset]);\n }\n\n // Update the sums (sum of X and sum of squares).\n sums_h[ni * groups + gi].x += x;\n sums_h[ni * groups + gi].y += x * x;\n }\n }\n }\n }\n }\n\n // The normalization term to compute the means.\n float rcp_hwc_per_group = 1.f / (float)(h * w * channels_per_group);\n // Normalize the sums.\n for (int ngi = 0; ngi < n * groups; ++ngi) {\n sums_h[ngi].x *= rcp_hwc_per_group;\n sums_h[ngi].y *= rcp_hwc_per_group;\n }\n }\n\n // Compute the golden reference on the host.\n if (!skip_checks) {\n if (mode == Mode::BWD) {\n group_norm_nhwc_bwd_(dx_ref_h, dgamma_ref_h, dbeta_ref_h, dy_h, x_h, gamma_h, beta_h, sums_h, epsilon, n, h, w, c,\n groups, with_swish, use_fp32, use_bf16);\n } else {\n group_norm_nhwc_fwd_(y_ref_h, x_h, gamma_h, beta_h, epsilon, n, h, w, c, groups, with_swish, use_fp32, use_bf16);\n }\n }\n\n // Copy to the device.\n CHECK_CUDA(cudaMemcpyAsync(x_d, x_h, x_sz, cudaMemcpyHostToDevice, cudaStreamDefault));\n CHECK_CUDA(cudaMemcpyAsync(gamma_d, gamma_h, gamma_sz, cudaMemcpyHostToDevice, cudaStreamDefault));\n CHECK_CUDA(cudaMemcpyAsync(beta_d, beta_h, gamma_sz, cudaMemcpyHostToDevice, cudaStreamDefault));\n\n if (mode == Mode::BWD) {\n CHECK_CUDA(cudaMemcpyAsync(dy_d, dy_h, x_sz, cudaMemcpyHostToDevice, cudaStreamDefault));\n\n // // DEBUG.\n // printf(\"sums_h[0] = %8.3f, %8.3f\\n\", sums_h[0].x, sums_h[0].y);\n // // END OF DEBUG.\n\n CHECK_CUDA(cudaMemcpyAsync(sums_d, sums_h, sums_sz, cudaMemcpyHostToDevice, cudaStreamDefault));\n }\n\n // Reset the output buffer with garbage to detect invalid results.\n if (mode == Mode::BWD) {\n CHECK_CUDA(cudaMemsetAsync(dx_d, 0xdc, x_sz, cudaStreamDefault));\n CHECK_CUDA(cudaMemsetAsync(dgamma_d, 0xdc, gamma_sz, cudaStreamDefault));\n CHECK_CUDA(cudaMemsetAsync(dbeta_d, 0xdc, gamma_sz, cudaStreamDefault));\n } else {\n CHECK_CUDA(cudaMemsetAsync(y_d, 0xdc, x_sz, cudaStreamDefault));\n }\n\n // Declare the parameters.\n Group_norm_nhwc_fwd_params params_fwd;\n memset(¶ms_fwd, 0, sizeof(params_fwd));\n Group_norm_nhwc_bwd_params params_bwd;\n memset(¶ms_bwd, 0, sizeof(params_bwd));\n\n const auto precision = [&]() -> PrecisionMode {\n if (use_fp32) {\n return PrecisionMode::FP32IOFP32W;\n } else if (use_bf16) {\n return PrecisionMode::BF16IOFP32W;\n } else {\n return PrecisionMode::FP16IOFP32W;\n }\n }();\n\n // Initialize the parameters.\n if (mode == Mode::BWD) {\n params_bwd.dx = dx_d;\n params_bwd.dgamma = dgamma_d;\n params_bwd.dbeta = dbeta_d;\n params_bwd.sums = sums_d;\n params_bwd.dy = dy_d;\n params_bwd.x = x_d;\n params_bwd.gamma = gamma_d;\n params_bwd.beta = beta_d;\n params_bwd.epsilon = epsilon;\n params_bwd.n = n;\n params_bwd.h = h;\n params_bwd.w = w;\n params_bwd.c = c;\n params_bwd.groups = groups;\n params_bwd.with_swish = with_swish;\n params_bwd.precision = precision;\n } else {\n params_fwd.y = y_d;\n params_fwd.sums = sums_d;\n params_fwd.x = x_d;\n params_fwd.gamma = gamma_d;\n params_fwd.beta = beta_d;\n params_fwd.epsilon = epsilon;\n params_fwd.n = n;\n params_fwd.h = h;\n params_fwd.w = w;\n params_fwd.c = c;\n params_fwd.groups = groups;\n params_fwd.with_swish = with_swish;\n params_fwd.precision = precision;\n }\n\n // The number of barriers.\n size_t barriers_elts = 0;\n // The number of elements in the reduction buffer.\n size_t red_buffer_elts = 0;\n // The number of elements in the reduction buffer that must be zeroed.\n size_t zeroed_red_buffer_elts = 0;\n\n // Finalize the parameters.\n dim3 grid;\n if (mode == Mode::BWD && use_one_pass) {\n group_norm_nhwc_bwd_one_pass_setup(params_bwd, barriers_elts, red_buffer_elts, zeroed_red_buffer_elts, grid, props);\n } else if (mode == Mode::BWD) {\n group_norm_nhwc_bwd_two_passes_setup(params_bwd, zeroed_red_buffer_elts);\n } else if (use_one_pass) {\n group_norm_nhwc_fwd_one_pass_setup(params_fwd, barriers_elts, red_buffer_elts, grid, props);\n } else {\n group_norm_nhwc_fwd_two_passes_setup(params_fwd, zeroed_red_buffer_elts);\n }\n\n // The size in bytes for the reduction buffer.\n size_t red_buffer_sz = red_buffer_elts * sizeof(float);\n // Allocate on the device.\n if (red_buffer_sz > 0) {\n float** ptr = mode == Mode::BWD ? ¶ms_bwd.red_buffer : ¶ms_fwd.red_buffer;\n CHECK_CUDA(cudaMalloc((void**)ptr, red_buffer_sz));\n }\n\n // The size of the array of barriers.\n size_t barriers_sz = barriers_elts * sizeof(int);\n // The size in bytes for the reduction buffer that must be zeroed.\n size_t zeroed_red_buffer_sz = barriers_sz + zeroed_red_buffer_elts * sizeof(float);\n\n // Allocate the buffer if needed.\n void* zeroed_red_buffer_d_ = nullptr;\n if (zeroed_red_buffer_sz > 0) {\n CHECK_CUDA(cudaMalloc((void**)&zeroed_red_buffer_d_, zeroed_red_buffer_sz));\n }\n\n // The buffer of barriers. DO NOT CALL cudaFree on it!!!\n int* barriers_d = reinterpret_cast(zeroed_red_buffer_d_);\n // The zeroed red buffer. DO NOT CALL cudaFree on it!!!\n float* zeroed_red_buffer_d = reinterpret_cast(&barriers_d[barriers_elts]);\n // Must be aligned on 4B for floats. It obviously is (unless someone changes the code ;)).\n assert(reinterpret_cast(zeroed_red_buffer_d) % sizeof(float) == 0);\n\n // Set the barriers if needed.\n if (mode == Mode::BWD) {\n params_bwd.barriers = barriers_d;\n params_bwd.zeroed_red_buffer = zeroed_red_buffer_d;\n } else {\n params_fwd.barriers = barriers_d;\n params_fwd.zeroed_red_buffer = zeroed_red_buffer_d;\n }\n\n // Create events to time the reference code.\n cudaEvent_t start, stop;\n CHECK_CUDA(cudaEventCreate(&start));\n CHECK_CUDA(cudaEventCreate(&stop));\n\n // Time the reference code.\n CHECK_CUDA(cudaEventRecord(start));\n for (int ii = 0; ii < runs; ++ii) {\n // Clear the zeroed buffer if needed.\n if (zeroed_red_buffer_sz > 0) {\n CHECK_CUDA(cudaMemsetAsync(zeroed_red_buffer_d_, 0, zeroed_red_buffer_sz, cudaStreamDefault));\n }\n if (use_one_pass && mode == Mode::BWD) {\n group_norm_nhwc_bwd_one_pass_run(params_bwd, grid, cudaStreamDefault);\n } else if (use_one_pass) {\n group_norm_nhwc_fwd_one_pass_run(params_fwd, grid, cudaStreamDefault);\n } else if (mode == Mode::BWD) {\n group_norm_nhwc_bwd_two_passes_sum(params_bwd, cudaStreamDefault);\n group_norm_nhwc_bwd_two_passes_scale(params_bwd, cudaStreamDefault);\n } else {\n group_norm_nhwc_fwd_two_passes_sum(params_fwd, cudaStreamDefault);\n group_norm_nhwc_fwd_two_passes_scale(params_fwd, cudaStreamDefault);\n }\n }\n CHECK_CUDA(cudaEventRecord(stop));\n CHECK_CUDA(cudaDeviceSynchronize());\n\n // Print the runtime.\n float elapsed = 0.f;\n CHECK_CUDA(cudaEventElapsedTime(&elapsed, start, stop));\n if (!csv_output) {\n printf(\"elapsed......................: %.3fms\\n\", elapsed);\n printf(\"elapsed per run..............: %.3fms\\n\", elapsed / (float)runs);\n printf(\"efficiency...................: %.3lf%%\\n\", dram_sol * runs / elapsed * 100.0);\n printf(\"\\n\");\n }\n\n // Copy the results to the host.\n if (mode == Mode::BWD) {\n CHECK_CUDA(cudaMemcpyAsync(dx_h, dx_d, x_sz, cudaMemcpyDeviceToHost, cudaStreamDefault));\n CHECK_CUDA(cudaMemcpyAsync(dgamma_h, dgamma_d, gamma_sz, cudaMemcpyDeviceToHost, cudaStreamDefault));\n CHECK_CUDA(cudaMemcpyAsync(dbeta_h, dbeta_d, gamma_sz, cudaMemcpyDeviceToHost, cudaStreamDefault));\n } else {\n CHECK_CUDA(cudaMemcpyAsync(y_h, y_d, x_sz, cudaMemcpyDeviceToHost, cudaStreamDefault));\n }\n\n // Make sure the data has been transferred.\n CHECK_CUDA(cudaStreamSynchronize(cudaStreamDefault));\n\n // Check the results.\n if (!csv_output) {\n if (mode == Mode::BWD && !skip_checks) {\n if (use_fp32) {\n check_results(\"dx\", reinterpret_cast(dx_h), reinterpret_cast(dx_ref_h), x_elts, tol);\n } else if (use_bf16) {\n check_results<__nv_bfloat16>(\"dx\", reinterpret_cast<__nv_bfloat16*>(dx_h),\n reinterpret_cast<__nv_bfloat16*>(dx_ref_h), x_elts, tol);\n } else {\n check_results<__half>(\"dx\", reinterpret_cast<__half*>(dx_h), reinterpret_cast<__half*>(dx_ref_h), x_elts, tol);\n }\n check_results(\"dgamma\", dgamma_h, dgamma_ref_h, gamma_elts, tol);\n check_results(\"dbeta\", dbeta_h, dbeta_ref_h, gamma_elts, tol);\n } else if (!skip_checks) {\n if (use_fp32) {\n check_results(\"y\", reinterpret_cast(y_h), reinterpret_cast(y_ref_h), x_elts, tol);\n } else if (use_bf16) {\n check_results<__nv_bfloat16>(\"y\", reinterpret_cast<__nv_bfloat16*>(y_h),\n reinterpret_cast<__nv_bfloat16*>(y_ref_h), x_elts, tol);\n } else {\n check_results<__half>(\"y\", reinterpret_cast<__half*>(y_h), reinterpret_cast<__half*>(y_ref_h), x_elts, tol);\n }\n }\n } else {\n printf(\"%d,%d,%d,%d,%d,%d,%d,%f\\n\", n, h, w, c, groups, (uint32_t)use_one_pass, (uint32_t)mode,\n elapsed / (float)runs);\n }\n\n // Destroy the cuda events.\n CHECK_CUDA(cudaEventDestroy(start));\n CHECK_CUDA(cudaEventDestroy(stop));\n\n // Release device memory.\n CHECK_CUDA(cudaFree(x_d));\n CHECK_CUDA(cudaFree(y_d));\n CHECK_CUDA(cudaFree(gamma_d));\n CHECK_CUDA(cudaFree(beta_d));\n CHECK_CUDA(cudaFree(dx_d));\n CHECK_CUDA(cudaFree(dy_d));\n CHECK_CUDA(cudaFree(dgamma_d));\n CHECK_CUDA(cudaFree(dbeta_d));\n CHECK_CUDA(cudaFree(sums_d));\n CHECK_CUDA(cudaFree(zeroed_red_buffer_d_));\n CHECK_CUDA(cudaFree(params_bwd.red_buffer));\n CHECK_CUDA(cudaFree(params_fwd.red_buffer));\n\n // Release host memory.\n free(x_h);\n free(y_h);\n free(gamma_h);\n free(beta_h);\n free(dx_h);\n free(dy_h);\n free(dgamma_h);\n free(dbeta_h);\n free(sums_h);\n free(y_ref_h);\n free(dx_ref_h);\n free(dgamma_ref_h);\n free(dbeta_ref_h);\n\n // Release the GPU.\n CHECK_CUDA(cudaDeviceReset());\n return 0;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc.h", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define CHECK_CUDA(call) \\\n do { \\\n cudaError_t status_ = call; \\\n if (status_ != cudaSuccess) { \\\n fprintf(stderr, \"CUDA error (%s:%d): %s\\n\", __FILE__, __LINE__, cudaGetErrorString(status_)); \\\n exit(1); \\\n } \\\n } while (0)\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ __host__ int div_up(int m, int n) { return (m + n - 1) / n; }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ __host__ float sigmoid(float x) { return 1.f / (1.f + expf(-x)); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline __device__ void spin_wait_(int* barrier, int step, int expected) {\n // THE FOLLOWING CODE MUST BE EXECUTED BY A SINGLE THREAD IN THE CTA.\n\n // Update the global counter. Make sure prior writes are visible.\n asm volatile(\"red.release.gpu.global.add.s32 [%0], %1;\" ::\"l\"(barrier), \"r\"(step));\n\n // Busy wait. We could use found = old + step with old = atomicAdd(...) but it's not faster.\n for (volatile int found = -1; found != expected;) {\n asm volatile(\"ld.global.acquire.gpu.b32 %0, [%1];\" : \"=r\"(found) : \"l\"(barrier));\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n// Input type followed by parameter type\nenum PrecisionMode {\n FP32IOFP16W,\n FP32IOBF16W,\n FP32IOFP32W,\n FP16IOFP16W,\n FP16IOBF16W,\n FP16IOFP32W,\n BF16IOFP16W,\n BF16IOBF16W,\n BF16IOFP32W,\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Group_sums {\n // Is it the 1st element of the group?\n int flag;\n // The sum.\n float sum;\n // The sum of squares.\n float sum_sq;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Group_sums_op {\n inline __device__ Group_sums operator()(const Group_sums& a, const Group_sums& b) {\n Group_sums dst;\n dst.sum = b.flag ? b.sum : (a.sum + b.sum);\n dst.sum_sq = b.flag ? b.sum_sq : (a.sum_sq + b.sum_sq);\n dst.flag = a.flag + b.flag;\n return dst;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Group_norm_nhwc_fwd_params {\n // The output buffer. Layout NHWC.\n void* y;\n // The sums for the bwd pass. Not written if it is a nullptr.\n float2* sums;\n // The input buffer. Layout NHWC.\n const void* x;\n // The gamma scaling factor.\n const void* gamma;\n // The beta term to add in GN.\n const void* beta;\n // The constant epsilon for sqrt(var + epsilon).\n float epsilon;\n // The barriers for the persistent kernel.\n int* barriers;\n // The extra storage for multi-CTA reductions as well as to pass data to the bwd.\n float *red_buffer, *zeroed_red_buffer;\n\n // The number of instances in the batch.\n int n;\n // The height and width of each activation map. The number of channels.\n int64_t h, w, c, hw, hwc;\n // The number of groups.\n int groups;\n // Do we apply the Swish activation function?\n bool with_swish;\n\n // Precomputed values and parameters to control the execution of the kernels.\n\n // The number of batch instances per block.\n int instances_per_block;\n // The number of activations computed per block.\n int acts_per_block;\n // The number of groups in each block.\n int groups_per_block;\n // The number of channels per group = c / groups.\n int channels_per_group;\n // The number of channels per block = groups_per_block * channels_per_group.\n int channels_per_block;\n // The inverse of hwc in floats (to compute mean/var).\n float inv_hwc_per_group;\n // IO precision\n PrecisionMode precision;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_fwd_two_passes_setup(Group_norm_nhwc_fwd_params&, size_t& red_buffer_elts);\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_fwd_two_passes_sum(const Group_norm_nhwc_fwd_params&, cudaStream_t);\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_fwd_two_passes_scale(const Group_norm_nhwc_fwd_params&, cudaStream_t);\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Group_norm_nhwc_bwd_params {\n // The output buffer. Layout NHWC.\n void* dx;\n // The output buffer. Layout NHWC.\n void* dgamma;\n // The output buffer. Layout NHWC.\n void* dbeta;\n // The input buffer. Layout NHWC.\n const void* dy;\n // The input buffer. Layout NHWC.\n const void* x;\n // The gamma scaling factor.\n const void* gamma;\n // The beta term to add in GN.\n const void* beta;\n // The sums from the fwd pass.\n const float2* sums;\n // The constant epsilon for sqrt(var + epsilon).\n float epsilon;\n // The barriers for the persistent kernel.\n int* barriers;\n // The extra storage for multi-CTA reductions as well as to pass data to the bwd.\n float *red_buffer, *zeroed_red_buffer;\n\n // The number of instances in the batch.\n int n;\n // The height and width of each activation map. The number of channels.\n int64_t h, w, c, hw, hwc;\n // The number of groups.\n int groups;\n // Do we apply the Swish activation function?\n bool with_swish;\n\n // Precomputed values and parameters to control the execution of the kernels.\n\n // The number of batch instances per block.\n int instances_per_block;\n // The number of activations computed per block.\n int acts_per_block;\n // The number of groups in each block.\n int groups_per_block;\n // The number of channels per group = c / groups.\n int channels_per_group;\n // The number of channels per block = groups_per_block * channels_per_group.\n int channels_per_block;\n // The inverse of hwc in floats (to compute mean/var).\n float inv_hwc_per_group;\n // IO precision\n PrecisionMode precision;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_two_passes_setup(Group_norm_nhwc_bwd_params&, size_t& red_buffer_elts);\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_two_passes_sum(const Group_norm_nhwc_bwd_params&, cudaStream_t);\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_two_passes_scale(const Group_norm_nhwc_bwd_params&, cudaStream_t);\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_one_pass.h", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#include \n\n#include \n\n#include \"group_norm_nhwc.h\"\n#include \"macros.h\"\n#include \"traits.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// B A C K W A R D\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define GN_BWD_SELECT(FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(4, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(8, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(10, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(12, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(14, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(16, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(20, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(26, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(24, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(28, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(30, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(32, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(40, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(42, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(48, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(56, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(60, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(64, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(70, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(80, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(84, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(96, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(98, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(112, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(120, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(128, FUNC_POSTFIX, function) \\\n GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(160, FUNC_POSTFIX, function) { \\\n assert(false && \"Not implemented\"); \\\n }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define GN_BWD_RUNNER_SELECT(function) GN_BWD_SELECT(_run, function)\n\n#define GN_BWD_BLOCKS_PER_SM_SELECT(function) GN_BWD_SELECT(_blocks_per_sm, function)\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 4)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 8)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 10)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 12)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 14)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 16)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 20)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 26)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 24)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 28)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 30)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 32)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 40)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 42)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 48)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 56)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 60)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 64)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 70)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 80)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 84)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 96)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 98)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 112)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 120)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 128)\nGN_BWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 160)\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_one_pass_setup(Group_norm_nhwc_bwd_params& params, size_t& barriers_elts,\n size_t& red_buffer_elts, size_t& zeroed_red_buffer_elts, dim3& grid,\n const cudaDeviceProp& props) {\n // The pre-computed dimensions.\n params.hw = params.h * params.w;\n params.hwc = params.c * params.hw;\n\n // The number of channels per group.\n params.channels_per_group = params.c / params.groups;\n // The inverse to compute the mean/variance.\n params.inv_hwc_per_group = 1.f / (float)(params.hw * params.channels_per_group);\n\n // Define how many activations are computed per block.\n if ((params.hw >= 1024 && params.channels_per_group >= 80) ||\n (params.hw >= 256 && params.channels_per_group >= 160)) {\n params.acts_per_block = 8 * 16;\n } else if (params.hw >= 512) {\n params.acts_per_block = 32 * 16;\n } else if (params.hw >= 256) {\n params.acts_per_block = 16 * 16;\n } else if (params.hw >= 128) {\n params.acts_per_block = 8 * 16;\n } else if (params.hw > 0) {\n params.acts_per_block = 8 * 8;\n } else {\n // We should never be here if params are set correctly.\n assert(false);\n }\n\n // Define the number of blocks per activation map. TODO: Make sure it matches the kernel sizes.\n int blocks_per_slice = div_up(params.hw, params.acts_per_block);\n\n // Select the kernel.\n using Function_t = int (*)();\n\n Function_t blocks_per_sm_function;\n GN_BWD_BLOCKS_PER_SM_SELECT(blocks_per_sm_function);\n // The number of blocks that can be run per SM.\n int blocks_per_sm = blocks_per_sm_function();\n\n // The number of blocks per grid.\n int max_blocks_per_grid = blocks_per_sm * props.multiProcessorCount;\n\n // Make sure we are safe to run that many blocks\n assert(blocks_per_slice <= max_blocks_per_grid);\n\n // The number of blocks per slice is the X dimension of the grid.\n grid.x = blocks_per_slice;\n // The number of groups * is the X dimension of the grid.\n grid.y = std::min(max_blocks_per_grid / blocks_per_slice, params.groups * params.n);\n\n // The number of barriers.\n barriers_elts = blocks_per_slice > 1 ? grid.y * 2 : 0;\n\n // Add 1 for the final conversion for dgamma/dbeta.\n barriers_elts += 1;\n\n // The number of elements in the reduction buffer (for the sums and sums of squared).\n if (blocks_per_slice == 1) {\n red_buffer_elts = 0;\n } else {\n // The first 2 is for double-buffering. The 2nd one is for the fact that we have two floats.\n red_buffer_elts = 2 * grid.x * grid.y * 2;\n }\n\n // The number of elements in the buffer that has to be zeroed.\n zeroed_red_buffer_elts = params.c * 2;\n\n // Make sure a group does not span multiple blocks.\n assert(params.channels_per_block % params.channels_per_group == 0);\n}\n\ninline void group_norm_nhwc_bwd_one_pass_run(const Group_norm_nhwc_bwd_params& params, const dim3& grid,\n cudaStream_t stream) {\n using Function_t = void (*)(const Group_norm_nhwc_bwd_params&, const dim3&, cudaStream_t);\n\n Function_t runner;\n GN_BWD_RUNNER_SELECT(runner);\n\n runner(params, grid, stream);\n}\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_one_pass_kernel.cuh", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#include \n\n#include \n\n#include \"group_norm_nhwc.h\"\n#include \"traits.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// B A C K W A R D\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ __launch_bounds__(THREADS_PER_BLOCK_) void group_norm_nhwc_bwd_one_pass_kernel(\n Group_norm_nhwc_bwd_params params) {\n // The IO traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n // The Weights traits.\n using WTraits = typename Traits::WTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // Weights type\n using WType = typename WTraits::Type;\n // Weights doubled type\n using WType2 = typename WTraits::Type2;\n\n // The number of activations per block.\n constexpr int ACTS_PER_BLOCK = ACTS_PER_BLOCK_;\n // The number of channels per group.\n constexpr int CHANNELS_PER_GROUP = CHANNELS_PER_GROUP_;\n // The number of threads per block.\n constexpr int THREADS_PER_BLOCK = THREADS_PER_BLOCK_;\n // The number of channels per thread (load fp16x2 numbers).\n constexpr int CHANNELS_PER_THREAD = 2;\n\n // The number of threads needed per activation.\n constexpr int THREADS_PER_ACT = CHANNELS_PER_GROUP / CHANNELS_PER_THREAD;\n // The number of activations that are loaded per loop.\n constexpr int ACTS_PER_LOOP = THREADS_PER_BLOCK / THREADS_PER_ACT;\n // The number of rows per thread.\n constexpr int ACTS_PER_THREAD = (ACTS_PER_BLOCK + ACTS_PER_LOOP - 1) / ACTS_PER_LOOP;\n\n // The number of active threads.\n constexpr int ACTIVE_THREADS = THREADS_PER_BLOCK / THREADS_PER_ACT * THREADS_PER_ACT;\n\n // The object in charge of doing the sums for the block.\n typedef cub::BlockReduce Block_reduce;\n // Allocate shared memory for Block_reduce.\n __shared__ typename Block_reduce::TempStorage temp_storage;\n // Allocate shared memory to store the sums.\n __shared__ float2 smem_sums;\n // Allocate shared memory to store the gamma/beta gradients.\n __shared__ float4 smem_dgamma_dbeta[THREADS_PER_BLOCK];\n\n // Shared memory to store the gradients for gamma and beta.\n\n // The first activation loaded by that thread.\n int hwi = blockIdx.x * params.acts_per_block + threadIdx.x / THREADS_PER_ACT;\n // The first channel loaded by that thread.\n int ci = threadIdx.x % THREADS_PER_ACT * CHANNELS_PER_THREAD;\n\n // Is it an active thread?\n const bool is_active = threadIdx.x < ACTIVE_THREADS;\n\n // Iterate over the iterms in the batch.\n for (int ngi = blockIdx.y, step = 0; ngi < params.n * params.groups; ngi += gridDim.y, ++step) {\n // The instance and the group. TODO: Use fast divmod?\n int ni = ngi / params.groups;\n int gi = ngi % params.groups;\n\n // The sums from the fwd pass.\n float2 fwd = params.sums[ngi];\n // The mean of X (computed during the fwd pass -- one value per batch*group).\n float x_mean = fwd.x;\n // The mean of squares of X (computed during the fwd pass -- one value per batch*group).\n float x_sq_mean = fwd.y;\n // The variance.\n float x_var = x_sq_mean - x_mean * x_mean;\n // The reciprocal of the standard deviation (i.e. 1.f / sqrt(var + epsilon)).\n float rcp_x_stddev = x_var <= 0.f ? 1.f : 1.f / sqrtf(x_var + params.epsilon);\n\n // The offset to the first activation loaded by that thread.\n const int64_t offset = (int64_t)ni * params.hwc + gi * CHANNELS_PER_GROUP + ci;\n // The pointer to the first activation loaded by that thread.\n const IOType* x_ptr = &reinterpret_cast(params.x)[offset];\n // The pointer to the first gradient loaded by that thread.\n const IOType* dy_ptr = &reinterpret_cast(params.dy)[offset];\n\n // Load the X and dY into registers.\n IOType2 x[ACTS_PER_THREAD], dy[ACTS_PER_THREAD];\n#pragma unroll\n for (int ii = 0; ii < ACTS_PER_THREAD; ++ii) {\n int hwj = hwi + ii * ACTS_PER_LOOP;\n x[ii] = IOTraits::zero();\n dy[ii] = IOTraits::zero();\n if (is_active && hwj < params.hw) {\n x[ii] = *reinterpret_cast(&x_ptr[hwj * params.c]);\n dy[ii] = *reinterpret_cast(&dy_ptr[hwj * params.c]);\n }\n }\n\n // Load gamma as well.\n float2 gamma_f2 = make_float2(0.f, 0.f);\n float2 beta_f2 = make_float2(0.f, 0.f);\n if (is_active) {\n gamma_f2 = WTraits::unpack(*reinterpret_cast(\n &reinterpret_cast(params.gamma)[gi * CHANNELS_PER_GROUP + ci]));\n if (params.with_swish) {\n beta_f2 = WTraits::unpack(*reinterpret_cast(\n &reinterpret_cast(params.beta)[gi * CHANNELS_PER_GROUP + ci]));\n }\n }\n\n // Gradients for gamma and beta (for this particular group).\n float4 dgamma_dbeta = make_float4(0.f, 0.f, 0.f, 0.f);\n // Accumulated gradients for dgrad calculation.\n float mean_1 = 0.f, mean_2 = 0.f;\n\n// Compute the sum and the sum of squares for each thread.\n#pragma unroll\n for (int ii = 0; ii < ACTS_PER_THREAD; ++ii) {\n // Convert x to float.\n float2 x_f2 = IOTraits::unpack(x[ii]);\n // Convert dY to float.\n float2 dy_f2 = IOTraits::unpack(dy[ii]);\n\n // X - X_mean.\n float x_minus_x_mean_x = x_f2.x - x_mean;\n float x_minus_x_mean_y = x_f2.y - x_mean;\n\n // Normalize X.\n float x_norm_x = x_minus_x_mean_x * rcp_x_stddev;\n float x_norm_y = x_minus_x_mean_y * rcp_x_stddev;\n\n if (params.with_swish) {\n float x_gn_x = x_norm_x * gamma_f2.x + beta_f2.x;\n float x_gn_y = x_norm_y * gamma_f2.y + beta_f2.y;\n float s_x = sigmoid(x_gn_x);\n float s_y = sigmoid(x_gn_y);\n dy_f2.x = dy_f2.x * s_x * (1.f + x_gn_x * (1.f - s_x));\n dy_f2.y = dy_f2.y * s_y * (1.f + x_gn_y * (1.f - s_y));\n }\n\n // Update beta.\n dgamma_dbeta.z += dy_f2.x;\n dgamma_dbeta.w += dy_f2.y;\n\n // Update dgamma.\n dgamma_dbeta.x += dy_f2.x * x_norm_x;\n dgamma_dbeta.y += dy_f2.y * x_norm_y;\n\n // The gradient that enters the x_norm node.\n float dx_norm_x = dy_f2.x * gamma_f2.x;\n float dx_norm_y = dy_f2.y * gamma_f2.y;\n\n // Add to the 1st mean.\n mean_1 += dx_norm_x * x_norm_x;\n mean_1 += dx_norm_y * x_norm_y;\n\n // Add to the 2nd mean.\n mean_2 += dx_norm_x;\n mean_2 += dx_norm_y;\n }\n\n // Pack valid gradients.\n float2 sums = make_float2(0.f, 0.f);\n if (ACTIVE_THREADS == THREADS_PER_BLOCK || is_active) {\n sums = make_float2(mean_1, mean_2);\n }\n\n // Store dgamma and dbeta to shared memory.\n smem_dgamma_dbeta[threadIdx.x] = dgamma_dbeta;\n\n // Compute the sums for the block.\n sums = Block_reduce(temp_storage).Reduce(sums, [](const float2& a, const float2& b) {\n return make_float2(a.x + b.x, a.y + b.y);\n });\n\n // Make sure we can read gamma/beta from smemory. Block_reduce uses one syncthread already.\n __syncthreads();\n\n // Compute gamma/beta for the block.\n if (threadIdx.x < THREADS_PER_ACT) {\n for (int ii = 1; ii < ACTS_PER_LOOP; ++ii) {\n float4 other = smem_dgamma_dbeta[threadIdx.x + ii * THREADS_PER_ACT];\n dgamma_dbeta.x += other.x;\n dgamma_dbeta.y += other.y;\n dgamma_dbeta.z += other.z;\n dgamma_dbeta.w += other.w;\n }\n }\n\n // The position in the channel dimension - 2 channels per thread.\n int cj = gi * THREADS_PER_ACT + threadIdx.x;\n // The reduction buffer dfor gamma/dbeta.\n float* red_buffer_dgamma_dbeta = ¶ms.zeroed_red_buffer[cj];\n\n // The first threads store their gradients for gamma/beta.\n if (threadIdx.x < THREADS_PER_ACT) {\n atomicAdd(&red_buffer_dgamma_dbeta[0 * params.c / 2], dgamma_dbeta.x);\n atomicAdd(&red_buffer_dgamma_dbeta[1 * params.c / 2], dgamma_dbeta.y);\n atomicAdd(&red_buffer_dgamma_dbeta[2 * params.c / 2], dgamma_dbeta.z);\n atomicAdd(&red_buffer_dgamma_dbeta[3 * params.c / 2], dgamma_dbeta.w);\n }\n\n // The block leader stores to global memory, if needed.\n if (gridDim.x > 1) {\n // The index of the buffer.\n int red_buffer_idx = step & 1;\n // The barrier.\n int* barrier = ¶ms.barriers[red_buffer_idx * gridDim.y + blockIdx.y];\n // The offset to the reduction buffer.\n int red_buffer_offset = red_buffer_idx * gridDim.x * gridDim.y * 2;\n // The reduction buffer.\n float2* red_buffer = reinterpret_cast(¶ms.red_buffer[red_buffer_offset]);\n\n // The offset to the reduction buffer for dgamma/dbeta.\n\n // The first thread stores its sums.\n if (threadIdx.x == 0) {\n red_buffer[blockIdx.x * gridDim.y + blockIdx.y] = sums;\n }\n\n // Make sure the data is in memory.\n if (threadIdx.x == 0) {\n spin_wait_(barrier, (step & 2) ? -1 : 1, (step & 2) ? 0 : gridDim.x);\n }\n __syncthreads();\n\n // Update the sums.\n for (int ii = 0; ii < gridDim.x; ++ii) {\n if (ii != blockIdx.x && threadIdx.x == 0) {\n float2 other_sums = red_buffer[ii * gridDim.y + blockIdx.y];\n sums.x += other_sums.x;\n sums.y += other_sums.y;\n }\n }\n }\n\n // Store the result for other threads.\n if (threadIdx.x == 0) {\n smem_sums = sums;\n }\n\n // Make sure the sums are in shared memory.\n __syncthreads();\n\n // Read the 1st mean from shared memory.\n mean_1 = smem_sums.x;\n // Read the 2nd mean from shared memory.\n mean_2 = smem_sums.y;\n\n mean_1 *= params.inv_hwc_per_group;\n mean_2 *= params.inv_hwc_per_group;\n\n // The pointer to the first activation stored by that thread.\n IOType* dx_ptr = &reinterpret_cast(params.dx)[offset];\n\n // Iterate over the activations to normalize the activations and store the results.\n for (int ii = 0; ii < ACTS_PER_THREAD; ++ii) {\n // Convert x to float.\n float2 x_f2 = IOTraits::unpack(x[ii]);\n // Convert dY to float.\n float2 dy_f2 = IOTraits::unpack(dy[ii]);\n\n // X - X_mean.\n float2 x_minus_x_mean_f2;\n x_minus_x_mean_f2.x = x_f2.x - x_mean;\n x_minus_x_mean_f2.y = x_f2.y - x_mean;\n // Normalize X.\n float2 x_norm_f2;\n x_norm_f2.x = x_minus_x_mean_f2.x * rcp_x_stddev;\n x_norm_f2.y = x_minus_x_mean_f2.y * rcp_x_stddev;\n\n if (params.with_swish) {\n float x_gn_x = x_norm_f2.x * gamma_f2.x + beta_f2.x;\n float x_gn_y = x_norm_f2.y * gamma_f2.y + beta_f2.y;\n float s_x = sigmoid(x_gn_x);\n float s_y = sigmoid(x_gn_y);\n dy_f2.x = dy_f2.x * s_x * (1.f + x_gn_x * (1.f - s_x));\n dy_f2.y = dy_f2.y * s_y * (1.f + x_gn_y * (1.f - s_y));\n }\n\n // The gradient that enters the x_norm node.\n float2 dx_norm;\n dx_norm.x = dy_f2.x * gamma_f2.x;\n dx_norm.y = dy_f2.y * gamma_f2.y;\n\n // The gradient along the input path.\n float2 dx;\n dx.x = (dx_norm.x - (x_norm_f2.x * mean_1 + mean_2)) * rcp_x_stddev;\n dx.y = (dx_norm.y - (x_norm_f2.y * mean_1 + mean_2)) * rcp_x_stddev;\n\n // Store the scaled values.\n int hwj = hwi + ii * ACTS_PER_LOOP;\n if (is_active && hwj < params.hw) {\n *reinterpret_cast(&dx_ptr[hwj * params.c]) = IOTraits::pack(dx);\n }\n }\n }\n\n // The completion barrier.\n int* barrier = ¶ms.barriers[gridDim.x == 1 ? 0 : gridDim.y * 2];\n\n // Mark the completion of the threadblock.\n if (threadIdx.x == 0) {\n asm volatile(\"red.release.gpu.global.add.s32 [%0], 1;\" ::\"l\"(barrier));\n }\n\n // Exit if that's not the last thread block.\n if (blockIdx.x != gridDim.x - 1 || blockIdx.y != gridDim.y - 1) {\n return;\n }\n\n // Busy wait. We could use found = old + step with old = atomicAdd(...) but it's not faster.\n if (threadIdx.x == 0) {\n for (int found = -1; found != gridDim.x * gridDim.y;) {\n asm volatile(\"ld.global.acquire.gpu.b32 %0, [%1];\" : \"=r\"(found) : \"l\"(barrier));\n }\n }\n __syncthreads();\n\n // The last block converts dgamma and dbeta to half.\n for (int idx = threadIdx.x; idx < params.c / 2; idx += THREADS_PER_BLOCK) {\n // Load dgamma.\n float2 dgamma;\n dgamma.x = params.zeroed_red_buffer[idx + 0 * params.c / 2];\n dgamma.y = params.zeroed_red_buffer[idx + 1 * params.c / 2];\n\n // Load dbeta.\n float2 dbeta;\n dbeta.x = params.zeroed_red_buffer[idx + 2 * params.c / 2];\n dbeta.y = params.zeroed_red_buffer[idx + 3 * params.c / 2];\n\n // Store to global memory.\n *reinterpret_cast(&reinterpret_cast(params.dgamma)[idx * 2]) = WTraits::pack(dgamma);\n *reinterpret_cast(&reinterpret_cast(params.dbeta)[idx * 2]) = WTraits::pack(dbeta);\n }\n}\n\n//////////////////////////////////////////////////////////////////////////////////////////////////\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_bwd_two_pass.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#include \n\n#include \n\n#include \"group_norm_nhwc.h\"\n#include \"macros.h\"\n#include \"traits.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// B A C K W A R D\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void group_norm_nhwc_bwd_sum_kernel(Group_norm_nhwc_bwd_params params) {\n // The IO traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n // The Weights traits.\n using WTraits = typename Traits::WTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // Weights type\n using WType = typename WTraits::Type;\n // Weights doubled type\n using WType2 = typename WTraits::Type2;\n\n // The object in charge of doing the sums for the different blocks.\n typedef cub::BlockScan Block_scan;\n\n // Allocate shared memory for Block_scan.\n __shared__ typename Block_scan::TempStorage temp_storage;\n // Allocate shared memory for the groups. We could reduce the amount of shared memory reserved.\n __shared__ float2 smem[THREADS_PER_BLOCK];\n\n // The instance in the batch.\n int ni = blockIdx.z;\n // The channel loaded by that thread (2 channels per thread for F16x2).\n int ci = blockIdx.x * params.channels_per_block + threadIdx.x * 2;\n // The group that thread works on and the channel in the group (modulus).\n int gi = ci / params.channels_per_group;\n\n // The sums from the fwd pass.\n float2 fwd = params.sums[ni * params.groups + gi];\n // The mean of X (computed during the fwd pass -- one value per batch*group).\n float x_mean = fwd.x;\n // The mean of squares of X (computed during the fwd pass -- one value per batch*group).\n float x_sq_mean = fwd.y;\n // The variance.\n float x_var = x_sq_mean - x_mean * x_mean;\n // The reciprocal of the standard deviation (i.e. 1.f / sqrt(var + epsilon)).\n float rcp_x_stddev = x_var <= 0.f ? 1.f : 1.f / sqrtf(x_var + params.epsilon);\n\n // Load gamma.\n float2 gamma_f2 = make_float2(0.f, 0.f);\n float2 beta_f2 = make_float2(0.f, 0.f);\n if (ci < params.c) {\n gamma_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.gamma)[ci]));\n if (params.with_swish) {\n beta_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.beta)[ci]));\n }\n }\n\n // The group that thread works on and the channel in the group (modulus).\n int gj = threadIdx.x * 2 / params.channels_per_group;\n int cj = threadIdx.x * 2 - params.channels_per_group * gj;\n\n // The first activation loaded by that block.\n int hw_begin = blockIdx.y * params.acts_per_block;\n // The last activation loaded by that block.\n int hw_end = min((int64_t)hw_begin + params.acts_per_block, params.hw);\n\n // The gradients for gamma/beta.\n float2 dgamma = make_float2(0.f, 0.f), dbeta = make_float2(0.f, 0.f);\n // Accumulated gradients for dgrad calculation\n float mean_1 = 0.f, mean_2 = 0.f;\n\n // Iterate over the activations to compute the sums.\n for (int hwi = hw_begin; hwi < hw_end; ++hwi) {\n // The offset.\n int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;\n\n // Fetch two channels per thread.\n IOType2 x_v2 = IOTraits::zero();\n IOType2 dy_v2 = IOTraits::zero();\n if (ci < params.c) {\n x_v2 = *reinterpret_cast(&reinterpret_cast(params.x)[offset]);\n dy_v2 = *reinterpret_cast(&reinterpret_cast(params.dy)[offset]);\n }\n\n // Extract the two half values.\n float2 x_f2 = IOTraits::unpack(x_v2);\n float2 dy_f2 = IOTraits::unpack(dy_v2);\n\n // X - X_mean.\n float x_minus_x_mean_x = x_f2.x - x_mean;\n float x_minus_x_mean_y = x_f2.y - x_mean;\n\n // Normalize X.\n float x_norm_x = x_minus_x_mean_x * rcp_x_stddev;\n float x_norm_y = x_minus_x_mean_y * rcp_x_stddev;\n\n if (params.with_swish) {\n float x_gn_x = x_norm_x * gamma_f2.x + beta_f2.x;\n float x_gn_y = x_norm_y * gamma_f2.y + beta_f2.y;\n float s_x = sigmoid(x_gn_x);\n float s_y = sigmoid(x_gn_y);\n dy_f2.x = dy_f2.x * s_x * (1.f + x_gn_x * (1.f - s_x));\n dy_f2.y = dy_f2.y * s_y * (1.f + x_gn_y * (1.f - s_y));\n }\n\n // Update beta.\n dbeta.x += dy_f2.x;\n dbeta.y += dy_f2.y;\n\n // Update dgamma.\n dgamma.x += dy_f2.x * x_norm_x;\n dgamma.y += dy_f2.y * x_norm_y;\n\n // The gradient that enters the x_norm node.\n float dx_norm_x = dy_f2.x * gamma_f2.x;\n float dx_norm_y = dy_f2.y * gamma_f2.y;\n\n // Add to the 1st mean.\n mean_1 += dx_norm_x * x_norm_x;\n mean_1 += dx_norm_y * x_norm_y;\n\n // Add to the 2nd mean.\n mean_2 += dx_norm_x;\n mean_2 += dx_norm_y;\n }\n\n // The data for the summations.\n Group_sums inp{cj == 0 ? 1 : 0, mean_1, mean_2};\n\n // Do the segmented scan.\n Group_sums out;\n Block_scan(temp_storage).InclusiveScan(inp, out, Group_sums_op());\n\n // Store the results for the groups in shared memory (to produce coalesced stores later).\n if (cj == params.channels_per_group - 2 /* 2 channels per thread */) {\n smem[gj] = make_float2(out.sum, out.sum_sq);\n }\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // The global group index.\n int gk = blockIdx.x * params.groups_per_block + threadIdx.x;\n\n // The first threads (those storing to global memory, load the values).\n float2 sums = smem[threadIdx.x];\n\n // Store to global memory.\n if (threadIdx.x < params.groups_per_block && gk < params.groups) {\n atomicAdd(¶ms.zeroed_red_buffer[(2 * ni + 0) * params.groups + gk], sums.x);\n atomicAdd(¶ms.zeroed_red_buffer[(2 * ni + 1) * params.groups + gk], sums.y);\n }\n\n // The base pointer for the gradients for gamma and beta.\n float* dgamma_beta_ptr = ¶ms.zeroed_red_buffer[params.n * params.groups * 2];\n\n // The 1st channel in the output tensor. NOTE: Two channels per thread store interleaved.\n int ck = blockIdx.x * params.channels_per_block + threadIdx.x;\n\n // Store dgamma and dbeta as well.\n if (ck < params.c) {\n atomicAdd(&dgamma_beta_ptr[0 * params.c + 0 * blockDim.x + ck], dgamma.x);\n atomicAdd(&dgamma_beta_ptr[0 * params.c + 1 * blockDim.x + ck], dgamma.y);\n atomicAdd(&dgamma_beta_ptr[1 * params.c + 0 * blockDim.x + ck], dbeta.x);\n atomicAdd(&dgamma_beta_ptr[1 * params.c + 1 * blockDim.x + ck], dbeta.y);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_two_passes_setup(Group_norm_nhwc_bwd_params& params, size_t& zeroed_red_buffer_elts) {\n // The pre-computed dimensions.\n params.hw = params.h * params.w;\n params.hwc = params.c * params.hw;\n\n // The number of channels per group.\n params.channels_per_group = params.c / params.groups;\n // The inverse to compute the mean/variance.\n params.inv_hwc_per_group = 1.f / (float)(params.hw * params.channels_per_group);\n\n // Define the number of blocks per activation map. That's a simple heuristic.\n int blocks_per_act_slice = 0;\n if (params.c >= 1280) {\n blocks_per_act_slice = 128 / params.n;\n } else if (params.c >= 640) {\n blocks_per_act_slice = 256 / params.n;\n } else {\n blocks_per_act_slice = 512 / params.n;\n }\n\n // Clamp to at least 1 to avoid divide-by-zero when batch size is large.\n blocks_per_act_slice = max(blocks_per_act_slice, 1);\n\n // Make sure we launch blocks per activation is no less than activations\n blocks_per_act_slice = min(blocks_per_act_slice, div_up(params.hw, params.n));\n\n // Define how many activations are computed per block.\n params.acts_per_block = div_up(params.hw, blocks_per_act_slice);\n\n // The number of channels per block.\n params.channels_per_block = 320;\n // Special case to deal with 30 channels per group.\n if (params.channels_per_block % params.channels_per_group != 0) {\n params.channels_per_block = 240;\n }\n\n // Special case to deal with 70 channels per group.\n if (params.c == 2240) {\n params.channels_per_block = 280;\n } else if (params.c == 832) {\n params.channels_per_block = 208;\n }\n\n if (params.c % params.channels_per_block != 0) {\n if (params.c % 512 == 0 && params.c != 1536 && params.c != 3072 && params.c % 448 != 0) {\n params.channels_per_block = 512;\n } else if (params.c % 42 == 0) {\n params.channels_per_block = 336;\n } else if (params.c % 384 == 0) {\n params.channels_per_block = 384;\n } else if (params.c % 256 == 0 && params.c % 448 != 0 && params.c % 392 != 0) {\n params.channels_per_block = 256;\n } else if (params.c % 128 == 0 && params.c % 448 != 0 && params.c % 392 != 0) {\n params.channels_per_block = 128;\n } else if (params.c % 448 == 0 && params.c % 392 != 0) {\n params.channels_per_block = 448;\n } else if (params.c % 392 == 0) {\n params.channels_per_block = 392;\n }\n }\n\n // The number of groups per block.\n params.groups_per_block = params.channels_per_block / params.channels_per_group;\n\n // Make sure the number of channels is a multiple of the number of channels per block.\n assert(params.c % params.channels_per_block == 0);\n // Make sure a group does not span multiple blocks.\n assert(params.channels_per_block % params.channels_per_group == 0);\n\n // The number of elements in the reduction buffer (for the sums and sums of squared).\n zeroed_red_buffer_elts = params.n * params.groups * 2 + params.c * 2;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_two_passes_sum(const Group_norm_nhwc_bwd_params& params, cudaStream_t stream) {\n // The dimension of the grid.\n dim3 grid;\n\n // The number of blocks to compute all the channels.\n grid.x = params.c / params.channels_per_block;\n // The number of blocks to compute all the activations in a given instance.\n grid.y = div_up(params.hw, params.acts_per_block);\n // The number of instances.\n grid.z = params.n;\n\n if (params.precision == PrecisionMode::FP16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp16IOFp16W)\n } else if (params.precision == PrecisionMode::FP16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp16IOBf16W)\n } else if (params.precision == PrecisionMode::FP16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp16IOFp32W)\n } else if (params.precision == PrecisionMode::BF16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Bf16IOFp16W)\n } else if (params.precision == PrecisionMode::BF16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Bf16IOBf16W)\n } else if (params.precision == PrecisionMode::BF16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Bf16IOFp32W)\n } else if (params.precision == PrecisionMode::FP32IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp32IOFp16W)\n } else if (params.precision == PrecisionMode::FP32IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp32IOBf16W)\n } else {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_sum_kernel, Fp32IOFp32W)\n }\n\n // Make sure it launched ok.\n CHECK_CUDA(cudaGetLastError());\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void group_norm_nhwc_bwd_scale_kernel(Group_norm_nhwc_bwd_params params) {\n // The IO traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n // The Weights traits.\n using WTraits = typename Traits::WTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // Weights type\n using WType = typename WTraits::Type;\n // Weights doubled type\n using WType2 = typename WTraits::Type2;\n\n // The instance in the batch.\n int ni = blockIdx.z;\n // The channel loaded by that thread (2 channels per thread for F16x2).\n int ci = blockIdx.x * params.channels_per_block + threadIdx.x * 2;\n // The group that thread works on and the channel in the group (modulus).\n int gi = ci / params.channels_per_group;\n\n // Load the gradients for the group.\n float mean_1 = 0.f, mean_2 = 0.f;\n if (gi < params.groups) {\n mean_1 = params.zeroed_red_buffer[(2 * ni + 0) * params.groups + gi];\n mean_2 = params.zeroed_red_buffer[(2 * ni + 1) * params.groups + gi];\n }\n\n // The sums from the fwd pass.\n float2 fwd = params.sums[ni * params.groups + gi];\n // The mean of X (computed during the fwd pass -- one value per batch*group).\n float x_mean = fwd.x;\n // The mean of squares of X (computed during the fwd pass -- one value per batch*group).\n float x_sq_mean = fwd.y;\n // The variance.\n float x_var = x_sq_mean - x_mean * x_mean;\n // The reciprocal of the standard deviation (i.e. 1.f / sqrt(var + epsilon)).\n float rcp_x_stddev = x_var <= 0.f ? 1.f : 1.f / sqrtf(x_var + params.epsilon);\n\n // Mutiply by 1/(HWC) to get real mean\n mean_1 *= params.inv_hwc_per_group;\n mean_2 *= params.inv_hwc_per_group;\n\n // Load gamma.\n float2 gamma_f2 = make_float2(0.f, 0.f);\n float2 beta_f2 = make_float2(0.f, 0.f);\n if (ci < params.c) {\n gamma_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.gamma)[ci]));\n if (params.with_swish) {\n beta_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.beta)[ci]));\n }\n }\n\n // The first activation loaded by that block.\n int hw_begin = blockIdx.y * params.acts_per_block;\n // The last activation loaded by that block.\n int hw_end = min((int64_t)hw_begin + params.acts_per_block, params.hw);\n\n // Iterate over the activations to compute the sums.\n for (int hwi = hw_begin; hwi < hw_end; ++hwi) {\n // The src/dst offset.\n int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;\n\n // Fetch two channels per thread.\n IOType2 x_v2 = IOTraits::zero();\n IOType2 dy_v2 = IOTraits::zero();\n if (ci < params.c) {\n x_v2 = *reinterpret_cast(&reinterpret_cast(params.x)[offset]);\n dy_v2 = *reinterpret_cast(&reinterpret_cast(params.dy)[offset]);\n }\n\n // Extract the two half values.\n float2 x_f2 = IOTraits::unpack(x_v2);\n float2 dy_f2 = IOTraits::unpack(dy_v2);\n\n // X - X_mean.\n float2 x_minus_x_mean_f2;\n x_minus_x_mean_f2.x = x_f2.x - x_mean;\n x_minus_x_mean_f2.y = x_f2.y - x_mean;\n\n // Normalize X.\n float2 x_norm_f2;\n x_norm_f2.x = x_minus_x_mean_f2.x * rcp_x_stddev;\n x_norm_f2.y = x_minus_x_mean_f2.y * rcp_x_stddev;\n\n if (params.with_swish) {\n float x_gn_x = x_norm_f2.x * gamma_f2.x + beta_f2.x;\n float x_gn_y = x_norm_f2.y * gamma_f2.y + beta_f2.y;\n float s_x = sigmoid(x_gn_x);\n float s_y = sigmoid(x_gn_y);\n dy_f2.x = dy_f2.x * s_x * (1.f + x_gn_x * (1.f - s_x));\n dy_f2.y = dy_f2.y * s_y * (1.f + x_gn_y * (1.f - s_y));\n }\n\n // The gradient that enters the x_norm node.\n float2 dx_norm;\n dx_norm.x = dy_f2.x * gamma_f2.x;\n dx_norm.y = dy_f2.y * gamma_f2.y;\n\n // The gradient along the input path.\n float2 dx;\n dx.x = (dx_norm.x - (x_norm_f2.x * mean_1 + mean_2)) * rcp_x_stddev;\n dx.y = (dx_norm.y - (x_norm_f2.y * mean_1 + mean_2)) * rcp_x_stddev;\n\n // Store the scaled values.\n if (ci < params.c) {\n *reinterpret_cast(&reinterpret_cast(params.dx)[offset]) = IOTraits::pack(dx);\n }\n }\n\n // Load gamma/beta and convert to half.\n if (blockIdx.y > 0 || blockIdx.z > 0 || ci >= params.c) {\n return;\n }\n\n // The base pointer for the gradients for gamma and beta.\n float* dgamma_beta_ptr = ¶ms.zeroed_red_buffer[params.n * params.groups * 2];\n\n // The 1st channel in the output tensor. NOTE: Two channels per thread store interleaved.\n int ck = blockIdx.x * params.channels_per_block + threadIdx.x;\n\n // Load the FP32 version of dgamma and dbeta.\n float2 dgamma, dbeta;\n if (ck < params.c) {\n dgamma.x = dgamma_beta_ptr[0 * params.c + 0 * blockDim.x + ck];\n dgamma.y = dgamma_beta_ptr[0 * params.c + 1 * blockDim.x + ck];\n dbeta.x = dgamma_beta_ptr[1 * params.c + 0 * blockDim.x + ck];\n dbeta.y = dgamma_beta_ptr[1 * params.c + 1 * blockDim.x + ck];\n\n // Convert to half2 and store to memory.\n *reinterpret_cast(&reinterpret_cast(params.dgamma)[ci]) = WTraits::pack(dgamma);\n *reinterpret_cast(&reinterpret_cast(params.dbeta)[ci]) = WTraits::pack(dbeta);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_bwd_two_passes_scale(const Group_norm_nhwc_bwd_params& params, cudaStream_t stream) {\n // The dimension of the grid.\n dim3 grid;\n\n // The number of blocks to compute all the channels.\n grid.x = params.c / params.channels_per_block;\n // The number of blocks to compute all the activations in a given instance.\n grid.y = div_up(params.hw, params.acts_per_block);\n // The number of instances.\n grid.z = params.n;\n\n if (params.precision == PrecisionMode::FP16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp16IOFp16W)\n } else if (params.precision == PrecisionMode::FP16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp16IOBf16W)\n } else if (params.precision == PrecisionMode::FP16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp16IOFp32W)\n } else if (params.precision == PrecisionMode::BF16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Bf16IOFp16W)\n } else if (params.precision == PrecisionMode::BF16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Bf16IOBf16W)\n } else if (params.precision == PrecisionMode::BF16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Bf16IOFp32W)\n } else if (params.precision == PrecisionMode::FP32IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp32IOFp16W)\n } else if (params.precision == PrecisionMode::FP32IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp32IOBf16W)\n } else {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_bwd_scale_kernel, Fp32IOFp32W)\n }\n\n // Make sure it launched ok.\n CHECK_CUDA(cudaGetLastError());\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_fwd_one_pass.h", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#include \n\n#include \n\n#include \"group_norm_nhwc.h\"\n#include \"macros.h\"\n#include \"traits.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// F O R W A R D\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define GN_FWD_SELECT(FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(4, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(8, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(10, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(12, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(14, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(16, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(20, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(26, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(24, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(28, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(30, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(32, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(40, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(42, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(48, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(56, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(60, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(64, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(70, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(80, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(84, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(96, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(98, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(112, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(120, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(128, FUNC_POSTFIX, function) \\\n GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(160, FUNC_POSTFIX, function) { \\\n assert(false && \"Not implemented\"); \\\n }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define GN_FWD_RUNNER_SELECT(function) GN_FWD_SELECT(_run, function)\n\n#define GN_FWD_BLOCKS_PER_SM_SELECT(function) GN_FWD_SELECT(_blocks_per_sm, function)\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 4)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 8)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 10)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 12)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 14)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 16)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 20)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 26)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 24)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 28)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 30)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 32)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 40)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 42)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 48)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 56)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 60)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 64)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 70)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 80)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 84)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 96)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 98)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 112)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 120)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 128)\nGN_FWD_ONE_PASS_DECLARATION(/* CHANNELS_PER_GROUP */ 160)\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline void group_norm_nhwc_fwd_one_pass_setup(Group_norm_nhwc_fwd_params& params, size_t& barriers_elts,\n size_t& red_buffer_elts, dim3& grid, const cudaDeviceProp& props) {\n // The pre-computed dimensions.\n params.hw = params.h * params.w;\n params.hwc = params.c * params.hw;\n\n // The number of channels per group.\n params.channels_per_group = params.c / params.groups;\n // The inverse to compute the mean/variance.\n params.inv_hwc_per_group = 1.f / (float)(params.hw * params.channels_per_group);\n\n // Select the kernel.\n using Function_t = int (*)();\n\n Function_t blocks_per_sm_function;\n GN_FWD_BLOCKS_PER_SM_SELECT(blocks_per_sm_function);\n\n // Define how many activations are computed per block.\n if (params.hw >= 1024 && params.channels_per_group >= 80 || (params.hw >= 256 && params.channels_per_group >= 160)) {\n params.acts_per_block = 8 * 16;\n } else if (params.hw >= 512) {\n params.acts_per_block = 16 * 32;\n } else if (params.hw >= 256) {\n params.acts_per_block = 16 * 16;\n } else if (params.hw >= 128) {\n params.acts_per_block = 8 * 16;\n } else if (params.hw > 0) {\n params.acts_per_block = 8 * 8;\n } else {\n // We should never be here if params are set correctly.\n assert(false);\n }\n\n // Define the number of blocks per activation map. TODO: Make sure it matches the kernel sizes.\n int blocks_per_slice = div_up(params.hw, params.acts_per_block);\n\n // The number of blocks that can be run per SM.\n int blocks_per_sm = blocks_per_sm_function();\n\n // The number of blocks per grid.\n int max_blocks_per_grid = blocks_per_sm * props.multiProcessorCount;\n\n // Make sure we are safe to run that many blocks\n assert(blocks_per_slice <= max_blocks_per_grid);\n\n // The number of blocks per slice is the X dimension of the grid.\n grid.x = blocks_per_slice;\n // The number of groups * is the X dimension of the grid.\n grid.y = std::min(max_blocks_per_grid / blocks_per_slice, params.groups * params.n);\n\n // The number of barriers.\n barriers_elts = blocks_per_slice > 1 ? grid.y * 2 : 0;\n\n // The number of elements in the reduction buffer (for the sums and sums of squared).\n if (blocks_per_slice == 1) {\n red_buffer_elts = 0;\n } else {\n // The first 2 is for double-buffering. The 2nd one is for the fact that we have two floats.\n red_buffer_elts = 2 * grid.x * grid.y * 2;\n }\n}\n\ninline void group_norm_nhwc_fwd_one_pass_run(const Group_norm_nhwc_fwd_params& params, const dim3& grid,\n cudaStream_t stream) {\n using Function_t = void (*)(const Group_norm_nhwc_fwd_params&, const dim3&, cudaStream_t);\n\n Function_t runner;\n GN_FWD_RUNNER_SELECT(runner);\n\n runner(params, grid, stream);\n}\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_fwd_one_pass_kernel.cuh", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#include \n\n#include \n\n#include \"group_norm_nhwc.h\"\n#include \"traits.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// F O R W A R D\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ __launch_bounds__(THREADS_PER_BLOCK_) void group_norm_nhwc_fwd_one_pass_kernel(\n Group_norm_nhwc_fwd_params params) {\n // The traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n // The Weights traits.\n using WTraits = typename Traits::WTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // Weights type\n using WType = typename WTraits::Type;\n // Weights doubled type\n using WType2 = typename WTraits::Type2;\n\n // The number of activations per block.\n constexpr int ACTS_PER_BLOCK = ACTS_PER_BLOCK_;\n // The number of channels per group.\n constexpr int CHANNELS_PER_GROUP = CHANNELS_PER_GROUP_;\n // The number of threads per block.\n constexpr int THREADS_PER_BLOCK = THREADS_PER_BLOCK_;\n // The number of channels per thread (load fp16x2 numbers).\n constexpr int CHANNELS_PER_THREAD = 2;\n\n // The number of threads needed per activation.\n constexpr int THREADS_PER_ACT = CHANNELS_PER_GROUP / CHANNELS_PER_THREAD;\n // The number of activations that are loaded per loop.\n constexpr int ACTS_PER_LOOP = THREADS_PER_BLOCK / THREADS_PER_ACT;\n // The number of rows per thread.\n constexpr int ACTS_PER_THREAD = (ACTS_PER_BLOCK + ACTS_PER_LOOP - 1) / ACTS_PER_LOOP;\n\n // The number of active threads.\n constexpr int ACTIVE_THREADS = THREADS_PER_BLOCK / THREADS_PER_ACT * THREADS_PER_ACT;\n\n // The object in charge of doing the sums for the block.\n typedef cub::BlockReduce Block_reduce;\n // Allocate shared memory for Block_reduce.\n __shared__ typename Block_reduce::TempStorage temp_storage;\n // Allocate shared memory to store the sums.\n __shared__ float2 smem_sums;\n\n // The first activation loaded by that thread.\n int hwi = blockIdx.x * params.acts_per_block + threadIdx.x / THREADS_PER_ACT;\n // The first channel loaded by that thread.\n int ci = threadIdx.x % THREADS_PER_ACT * CHANNELS_PER_THREAD;\n\n // Is it an active thread?\n const bool is_active = threadIdx.x < ACTIVE_THREADS;\n\n // Iterate over the iterms in the batch.\n for (int ngi = blockIdx.y, step = 0; ngi < params.n * params.groups; ngi += gridDim.y, ++step) {\n // The instance and the group. TODO: Use fast divmod?\n int ni = ngi / params.groups;\n int gi = ngi % params.groups;\n\n // The offset to the first activation loaded by that thread.\n const int64_t offset = (int64_t)ni * params.hwc + gi * CHANNELS_PER_GROUP + ci;\n // The pointer to the first activation loaded by that thread.\n const IOType* x_ptr = &reinterpret_cast(params.x)[offset];\n\n // Load the activations into registers.\n IOType2 x[ACTS_PER_THREAD];\n#pragma unroll\n for (int ii = 0; ii < ACTS_PER_THREAD; ++ii) {\n int hwj = hwi + ii * ACTS_PER_LOOP;\n x[ii] = IOTraits::zero();\n if (is_active && hwj < params.hw) {\n x[ii] = *reinterpret_cast(&x_ptr[hwj * params.c]);\n }\n }\n\n // Compute the sum and the sum of squares for each thread.\n float2 sums = make_float2(0.f, 0.f);\n#pragma unroll\n for (int ii = 0; ii < ACTS_PER_THREAD; ++ii) {\n float2 f2 = IOTraits::unpack(x[ii]);\n sums.x += f2.x + f2.y;\n sums.y += f2.x * f2.x + f2.y * f2.y;\n }\n\n // Clear invalid threads.\n if (ACTIVE_THREADS < THREADS_PER_BLOCK && !is_active) {\n sums = make_float2(0.f, 0.f);\n }\n\n // Compute the sums for the block.\n sums = Block_reduce(temp_storage).Reduce(sums, [](const float2& a, const float2& b) {\n return make_float2(a.x + b.x, a.y + b.y);\n });\n\n // The block leader stores to global memory, if needed.\n if (gridDim.x > 1) {\n // The index of the buffer (double-buffering).\n int red_buffer_idx = step & 1;\n // The barrier.\n int* barrier = ¶ms.barriers[red_buffer_idx * gridDim.y + blockIdx.y];\n // The offset to the reduction buffer.\n int red_buffer_offset = red_buffer_idx * gridDim.x * gridDim.y * 2;\n // The reduction buffer.\n float2* red_buffer = reinterpret_cast(¶ms.red_buffer[red_buffer_offset]);\n\n // The first thread stores its sums.\n if (threadIdx.x == 0) {\n red_buffer[blockIdx.x * gridDim.y + blockIdx.y] = sums;\n }\n\n // Make sure the data is in memory.\n if (threadIdx.x == 0) {\n spin_wait_(barrier, (step & 2) ? -1 : 1, (step & 2) ? 0 : gridDim.x);\n }\n __syncthreads();\n\n // Update the sums.\n for (int ii = 0; ii < gridDim.x; ++ii) {\n if (ii != blockIdx.x && threadIdx.x == 0) {\n float2 other_sums = red_buffer[ii * gridDim.y + blockIdx.y];\n sums.x += other_sums.x;\n sums.y += other_sums.y;\n }\n }\n }\n\n // Store the result for other threads.\n if (threadIdx.x == 0) {\n smem_sums = sums;\n }\n\n // Store the results to global memory as well (for training).\n if (params.sums != nullptr && blockIdx.x == 0 && threadIdx.x == 0) {\n sums.x *= params.inv_hwc_per_group;\n sums.y *= params.inv_hwc_per_group;\n params.sums[ngi] = sums;\n }\n\n // Make sure the sums are in shared memory.\n __syncthreads();\n\n // Load gamma/beta.\n float2 gamma_f2 = WTraits::unpack(\n *reinterpret_cast(&reinterpret_cast(params.gamma)[gi * CHANNELS_PER_GROUP + ci]));\n float2 beta_f2 = WTraits::unpack(\n *reinterpret_cast(&reinterpret_cast(params.beta)[gi * CHANNELS_PER_GROUP + ci]));\n\n // Compute the mean.\n float mean = smem_sums.x * params.inv_hwc_per_group;\n // Compute the variance.\n float var = smem_sums.y * params.inv_hwc_per_group - (mean * mean);\n // Compute the inverse of the stddev.\n float inv_stddev = var <= 0.f ? 1.f : rsqrtf(var + params.epsilon);\n\n // The pointer to the first activation stored by that thread.\n IOType* y_ptr = &reinterpret_cast(params.y)[offset];\n\n // Iterate over the activations to normalize the activations and store the results.\n for (int ii = 0; ii < ACTS_PER_THREAD; ++ii) {\n // Extract the two half values.\n float2 f2 = IOTraits::unpack(x[ii]);\n\n // Normalize the channels.\n f2.x = (f2.x - mean) * inv_stddev;\n f2.y = (f2.y - mean) * inv_stddev;\n\n // Scale by gamma and add beta.\n f2.x = gamma_f2.x * f2.x + beta_f2.x;\n f2.y = gamma_f2.y * f2.y + beta_f2.y;\n\n // Apply Swish if needed.\n if (params.with_swish) {\n f2.x = f2.x * sigmoid(f2.x);\n f2.y = f2.y * sigmoid(f2.y);\n }\n\n // Store the scaled values.\n int hwj = hwi + ii * ACTS_PER_LOOP;\n if (is_active && hwj < params.hw) {\n *reinterpret_cast(&y_ptr[hwj * params.c]) = IOTraits::pack(f2);\n }\n }\n }\n}\n\n//////////////////////////////////////////////////////////////////////////////////////////////////\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_fwd_two_pass.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#include \n\n#include \n\n#include \"group_norm_nhwc.h\"\n#include \"macros.h\"\n#include \"traits.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n//\n// F O R W A R D\n//\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void group_norm_nhwc_fwd_sum_kernel(Group_norm_nhwc_fwd_params params) {\n // The traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // The object in charge of doing the sums for the different blocks.\n typedef cub::BlockScan Block_scan;\n\n // Allocate shared memory for Block_scan.\n __shared__ typename Block_scan::TempStorage temp_storage;\n // Allocate shared memory for the groups. We could reduce the amount of shared memory reserved.\n __shared__ float2 smem[THREADS_PER_BLOCK];\n\n // The instance in the batch.\n int ni = blockIdx.z;\n // The channel loaded by that thread (2 channels per thread for F16x2).\n int ci = blockIdx.x * params.channels_per_block + threadIdx.x * 2;\n\n // The first activation loaded by that block.\n int hw_begin = blockIdx.y * params.acts_per_block;\n // The last activation loaded by that block.\n int hw_end = min((int64_t)hw_begin + params.acts_per_block, params.hw);\n\n // The sums.\n float sum = 0.f, sum_sq = 0.f;\n\n // Iterate over the activations to compute the sums.\n for (int hwi = hw_begin; hwi < hw_end; ++hwi) {\n // The offset.\n int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;\n\n // Fetch two channels per thread.\n IOType2 v2 = IOTraits::zero();\n if (ci < params.c) {\n v2 = *reinterpret_cast(&reinterpret_cast(params.x)[offset]);\n }\n\n // Extract the two values.\n float2 f2 = IOTraits::unpack(v2);\n\n // Update the sum.\n sum += f2.x + f2.y;\n // Update the sum of squares.\n sum_sq += f2.x * f2.x + f2.y * f2.y;\n }\n\n // The group that thread works on and the channel in the group (modulus).\n int gj = threadIdx.x * 2 / params.channels_per_group;\n int cj = threadIdx.x * 2 - params.channels_per_group * gj;\n\n // The data for the summations.\n Group_sums inp{cj == 0 ? 1 : 0, sum, sum_sq};\n\n // Do the segmented scan.\n Group_sums out;\n Block_scan(temp_storage).InclusiveScan(inp, out, Group_sums_op());\n\n // Store the results for the groups in shared memory (to produce coalesced stores later).\n if (cj == params.channels_per_group - 2 /* 2 channels per thread */) {\n smem[gj] = make_float2(out.sum, out.sum_sq);\n }\n\n // Make sure the data is in shared memory.\n __syncthreads();\n\n // The global group index.\n int gk = blockIdx.x * params.groups_per_block + threadIdx.x;\n\n // Threads that have nothing left to do, exit.\n if (threadIdx.x >= params.groups_per_block || gk >= params.groups) {\n return;\n }\n\n // The first threads (those storing to global memory, load the values).\n float2 sums = smem[threadIdx.x];\n\n // Store to global memory.\n atomicAdd(¶ms.zeroed_red_buffer[(2 * ni + 0) * params.groups + gk], sums.x);\n atomicAdd(¶ms.zeroed_red_buffer[(2 * ni + 1) * params.groups + gk], sums.y);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_fwd_two_passes_setup(Group_norm_nhwc_fwd_params& params, size_t& zeroed_red_buffer_elts) {\n // The pre-computed dimensions.\n params.hw = params.h * params.w;\n params.hwc = params.c * params.hw;\n\n // The number of channels per group.\n params.channels_per_group = params.c / params.groups;\n // The inverse to compute the mean/variance.\n params.inv_hwc_per_group = 1.f / (float)(params.hw * params.channels_per_group);\n\n // Define the number of blocks per activation map. That's a simple heuristic.\n int blocks_per_act_slice = 0;\n if (params.c >= 1280) {\n blocks_per_act_slice = 128 / params.n;\n } else if (params.c >= 640) {\n blocks_per_act_slice = 256 / params.n;\n } else {\n blocks_per_act_slice = 512 / params.n;\n }\n\n // Clamp to at least 1 to avoid divide-by-zero when batch size is large.\n blocks_per_act_slice = max(blocks_per_act_slice, 1);\n\n // Make sure we launch blocks per activation is no less than activations\n blocks_per_act_slice = min(blocks_per_act_slice, div_up(params.hw, params.n));\n\n // Define how many activations are computed per block.\n params.acts_per_block = div_up(params.hw, blocks_per_act_slice);\n // The number of channels per block.\n params.channels_per_block = 320;\n // Special case to deal with 30 channels per group.\n if (params.channels_per_block % params.channels_per_group != 0) {\n params.channels_per_block = 240;\n }\n\n // Special case to deal with 70 channels per group.\n if (params.c == 2240) {\n params.channels_per_block = 280;\n } else if (params.c == 832) {\n params.channels_per_block = 208;\n }\n\n if (params.c % params.channels_per_block != 0) {\n if (params.c % 512 == 0 && params.c != 1536 && params.c != 3072 && params.c % 448 != 0) {\n params.channels_per_block = 512;\n } else if (params.c % 42 == 0) {\n params.channels_per_block = 336;\n } else if (params.c % 384 == 0) {\n params.channels_per_block = 384;\n } else if (params.c % 256 == 0 && params.c % 448 != 0 && params.c % 392 != 0) {\n params.channels_per_block = 256;\n } else if (params.c % 128 == 0 && params.c % 448 != 0 && params.c % 392 != 0) {\n params.channels_per_block = 128;\n } else if (params.c % 448 == 0 && params.c % 392 != 0) {\n params.channels_per_block = 448;\n } else if (params.c % 392 == 0) {\n params.channels_per_block = 392;\n }\n }\n\n // The number of groups per block.\n params.groups_per_block = params.channels_per_block / params.channels_per_group;\n\n // Make sure the number of channels is a multiple of the number of channels per block.\n assert(params.c % params.channels_per_block == 0);\n // Make sure a group does not span multiple blocks.\n assert(params.channels_per_block % params.channels_per_group == 0);\n\n // The number of elements in the reduction buffer (for the sums and sums of squared).\n zeroed_red_buffer_elts = params.n * params.groups * 2;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_fwd_two_passes_sum(const Group_norm_nhwc_fwd_params& params, cudaStream_t stream) {\n // The dimension of the grid.\n dim3 grid;\n\n // The number of blocks to compute all the channels.\n grid.x = params.c / params.channels_per_block;\n // The number of blocks to compute all the activations in a given instance.\n grid.y = div_up(params.hw, params.acts_per_block);\n // The number of instances.\n grid.z = params.n;\n\n // Launch the kernel.\n if (params.precision == PrecisionMode::FP16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp16IOFp16W)\n } else if (params.precision == PrecisionMode::FP16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp16IOBf16W)\n } else if (params.precision == PrecisionMode::FP16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp16IOFp32W)\n } else if (params.precision == PrecisionMode::BF16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Bf16IOFp16W)\n } else if (params.precision == PrecisionMode::BF16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Bf16IOBf16W)\n } else if (params.precision == PrecisionMode::BF16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Bf16IOFp32W)\n } else if (params.precision == PrecisionMode::FP32IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp32IOFp16W)\n } else if (params.precision == PrecisionMode::FP32IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp32IOBf16W)\n } else {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_sum_kernel, Fp32IOFp32W)\n }\n\n // Make sure it launched ok.\n CHECK_CUDA(cudaGetLastError());\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void group_norm_nhwc_fwd_scale_kernel(Group_norm_nhwc_fwd_params params) {\n // The traits.\n using Traits = Traits_;\n // The IO traits.\n using IOTraits = typename Traits::IOTraits;\n // The Weights traits.\n using WTraits = typename Traits::WTraits;\n\n // The IO type\n using IOType = typename IOTraits::Type;\n // The IO doubled type\n using IOType2 = typename IOTraits::Type2;\n\n // Weights type\n using WType = typename WTraits::Type;\n // Weights doubled type\n using WType2 = typename WTraits::Type2;\n\n // The instance in the batch.\n int ni = blockIdx.z;\n // The channel loaded by that thread (2 channels per thread for F16x2).\n int ci = blockIdx.x * params.channels_per_block + threadIdx.x * 2;\n // The group that thread works on and the channel in the group (modulus).\n int gi = ci / params.channels_per_group;\n\n // Load the sum and sum of squares for the group.\n float sum = 0.f, sum_sq = 0.f;\n if (gi < params.groups) {\n sum = params.zeroed_red_buffer[(2 * ni + 0) * params.groups + gi];\n sum_sq = params.zeroed_red_buffer[(2 * ni + 1) * params.groups + gi];\n }\n\n // Load gamma/beta.\n float2 gamma_f2, beta_f2;\n if (ci < params.c) {\n gamma_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.gamma)[ci]));\n beta_f2 = WTraits::unpack(*reinterpret_cast(&reinterpret_cast(params.beta)[ci]));\n }\n\n // Compute the mean.\n float mean = sum * params.inv_hwc_per_group;\n // Compute the variance.\n float var = sum_sq * params.inv_hwc_per_group - (mean * mean);\n // Compute the inverse of the stddev.\n float inv_stddev = var <= 0.f ? 1.f : rsqrtf(var + params.epsilon);\n\n // The first activation loaded by that block.\n int hw_begin = blockIdx.y * params.acts_per_block;\n // The last activation loaded by that block.\n int hw_end = min((int64_t)hw_begin + params.acts_per_block, params.hw);\n\n // Iterate over the activations to compute the sums.\n for (int hwi = hw_begin; hwi < hw_end; ++hwi) {\n // The src/dst offset.\n int64_t offset = (int64_t)ni * params.hwc + hwi * params.c + ci;\n\n // Fetch two channels per thread.\n IOType2 v2 = IOTraits::zero();\n if (ci < params.c) {\n v2 = *reinterpret_cast(&reinterpret_cast(params.x)[offset]);\n }\n\n // Extract the two values.\n float2 f2 = IOTraits::unpack(v2);\n\n // Normalize the channels.\n f2.x = (f2.x - mean) * inv_stddev;\n f2.y = (f2.y - mean) * inv_stddev;\n\n // Scale by gamma and add beta.\n f2.x = gamma_f2.x * f2.x + beta_f2.x;\n f2.y = gamma_f2.y * f2.y + beta_f2.y;\n\n // Apply Swish if needed.\n if (params.with_swish) {\n f2.x = f2.x * sigmoid(f2.x);\n f2.y = f2.y * sigmoid(f2.y);\n }\n\n // Store the scaled values.\n if (ci < params.c) {\n *reinterpret_cast(&reinterpret_cast(params.y)[offset]) = IOTraits::pack(f2);\n }\n }\n\n // Write the sums if needed.\n if (params.sums != nullptr && gi < params.groups) {\n float2 sums;\n sums.x = sum * params.inv_hwc_per_group;\n sums.y = sum_sq * params.inv_hwc_per_group;\n params.sums[ni * params.groups + gi] = sums;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nvoid group_norm_nhwc_fwd_two_passes_scale(const Group_norm_nhwc_fwd_params& params, cudaStream_t stream) {\n // The dimension of the grid.\n dim3 grid;\n\n // The number of blocks to compute all the channels.\n grid.x = params.c / params.channels_per_block;\n // The number of blocks to compute all the activations in a given instance.\n grid.y = div_up(params.hw, params.acts_per_block);\n // The number of instances.\n grid.z = params.n;\n\n if (params.precision == PrecisionMode::FP16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp16IOFp16W)\n } else if (params.precision == PrecisionMode::FP16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp16IOBf16W)\n } else if (params.precision == PrecisionMode::FP16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp16IOFp32W)\n } else if (params.precision == PrecisionMode::BF16IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Bf16IOFp16W)\n } else if (params.precision == PrecisionMode::BF16IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Bf16IOBf16W)\n } else if (params.precision == PrecisionMode::BF16IOFP32W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Bf16IOFp32W)\n } else if (params.precision == PrecisionMode::FP32IOFP16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp32IOFp16W)\n } else if (params.precision == PrecisionMode::FP32IOBF16W) {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp32IOBf16W)\n } else {\n CALL_TWO_PASS_KERNEL(group_norm_nhwc_fwd_scale_kernel, Fp32IOFp32W)\n }\n\n // Make sure it launched ok.\n CHECK_CUDA(cudaGetLastError());\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_10.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 10, /* THREADS_PER_BLOCK */ 640)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_112.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 112, /* THREADS_PER_BLOCK */ 448)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_12.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 12, /* THREADS_PER_BLOCK */ 384)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_120.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 120, /* THREADS_PER_BLOCK */ 480)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_128.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 128, /* THREADS_PER_BLOCK */ 512)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_14.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 14, /* THREADS_PER_BLOCK */ 224)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_16.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 16, /* THREADS_PER_BLOCK */ 256)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_160.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 160, /* THREADS_PER_BLOCK */ 640)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_20.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 20, /* THREADS_PER_BLOCK */ 640)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_24.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 24, /* THREADS_PER_BLOCK */ 384)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_26.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 26, /* THREADS_PER_BLOCK */ 416)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_28.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 28, /* THREADS_PER_BLOCK */ 448)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_30.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 30, /* THREADS_PER_BLOCK */ 480)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_32.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 32, /* THREADS_PER_BLOCK */ 512)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_4.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 4, /* THREADS_PER_BLOCK */ 128)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_40.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 40, /* THREADS_PER_BLOCK */ 640)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_42.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 42, /* THREADS_PER_BLOCK */ 672)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_48.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 48, /* THREADS_PER_BLOCK */ 384)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_56.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 56, /* THREADS_PER_BLOCK */ 448)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_60.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 60, /* THREADS_PER_BLOCK */ 480)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_64.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 64, /* THREADS_PER_BLOCK */ 512)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_70.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 70, /* THREADS_PER_BLOCK */ 560)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_8.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 8, /* THREADS_PER_BLOCK */ 128)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_80.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 80, /* THREADS_PER_BLOCK */ 640)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_84.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 84, /* THREADS_PER_BLOCK */ 672)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_96.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 96, /* THREADS_PER_BLOCK */ 768)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_one_pass_98.cu", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \"group_norm_nhwc_bwd_one_pass_kernel.cuh\"\n#include \"group_norm_nhwc_fwd_one_pass_kernel.cuh\"\n#include \"macros.h\"\n\nGN_FWD_BWD_ONE_PASS_DEFINITION(/* CHANNELS_PER_GROUP */ 98, /* THREADS_PER_BLOCK */ 392)\n" }, { "path": "apex/contrib/csrc/group_norm/group_norm_nhwc_op.cpp", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#include \n#include \n\n#include \"group_norm_nhwc.h\"\n#include \"group_norm_nhwc_bwd_one_pass.h\"\n#include \"group_norm_nhwc_fwd_one_pass.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define CHECK_CUDA_STATUS(call) \\\n do { \\\n cudaError_t status_ = call; \\\n if (status_ != cudaSuccess) { \\\n fprintf(stderr, \"CUDA error (%s:%d): %s\\n\", __FILE__, __LINE__, cudaGetErrorString(status_)); \\\n exit(1); \\\n } \\\n } while (0)\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n#define CHECK_CHANNELS_LAST(x) TORCH_CHECK(x.is_contiguous(at::MemoryFormat::ChannelsLast), #x \" must be channels last\")\n#define CHECK_INPUT(x) \\\n CHECK_CUDA(x); \\\n CHECK_CONTIGUOUS(x)\n#define CHECK_NHWC_INPUT(x) \\\n CHECK_CUDA(x); \\\n CHECK_CHANNELS_LAST(x)\n\nstatic bool initialized = false;\nstatic cudaDeviceProp props;\n\nconst std::unordered_set supported_c_values = {128, 256, 320, 384, 448, 512, 640, 768,\n 896, 960, 1024, 1280, 1344, 1536, 1792, 1920,\n 2048, 2240, 2560, 2688, 3072, 3136, 3584, 4096};\nconst std::unordered_set supported_groups_values = {16, 32};\n\nstd::vector group_norm_fwd(torch::Tensor input, int groups, torch::Tensor weight, torch::Tensor bias,\n float eps, int passes, bool with_swish = false) {\n if (!initialized) {\n CHECK_CUDA_STATUS(cudaGetDeviceProperties(&props, 0));\n initialized = true;\n }\n CHECK_NHWC_INPUT(input);\n auto stream = at::cuda::getCurrentCUDAStream();\n\n // Achieve group norm arguments\n int n = input.size(0);\n int c = input.size(1);\n int h = input.size(2);\n int w = input.size(3);\n\n // Check kernel constraints\n TORCH_CHECK(supported_groups_values.count(groups), \"`groups` of {16, 32} are only supported but \", groups,\n \" is passed\");\n TORCH_CHECK(supported_c_values.count(c), \"`c` of \", c, \" is not included in supported_c_values\");\n\n // Allocate tensors\n auto options = at::TensorOptions(at::kCUDA);\n auto output = at::empty_like(input, at::MemoryFormat::Preserve);\n auto sums_d = at::empty({2 * n * groups}, options.dtype(at::kFloat));\n\n // Declare the parameters.\n Group_norm_nhwc_fwd_params params_fwd;\n memset(¶ms_fwd, 0, sizeof(params_fwd));\n\n // Initialize the parameters.\n params_fwd.y = reinterpret_cast(output.data_ptr());\n params_fwd.sums = reinterpret_cast(sums_d.data_ptr());\n params_fwd.x = const_cast(reinterpret_cast(input.data_ptr()));\n params_fwd.gamma = const_cast(reinterpret_cast(weight.data_ptr()));\n params_fwd.beta = const_cast(reinterpret_cast(bias.data_ptr()));\n params_fwd.epsilon = eps;\n params_fwd.n = n;\n params_fwd.h = h;\n params_fwd.w = w;\n params_fwd.c = c;\n params_fwd.groups = groups;\n params_fwd.with_swish = with_swish;\n\n PrecisionMode mode;\n if (input.dtype() == torch::kFloat32) {\n if (weight.dtype() == torch::kFloat16) {\n mode = PrecisionMode::FP32IOFP16W;\n } else if (weight.dtype() == torch::kBFloat16) {\n mode = PrecisionMode::FP32IOBF16W;\n } else {\n mode = PrecisionMode::FP32IOFP32W;\n }\n } else if (input.dtype() == torch::kBFloat16) {\n if (weight.dtype() == torch::kFloat16) {\n mode = PrecisionMode::BF16IOFP16W;\n } else if (weight.dtype() == torch::kBFloat16) {\n mode = PrecisionMode::BF16IOBF16W;\n } else {\n mode = PrecisionMode::BF16IOFP32W;\n }\n } else {\n if (weight.dtype() == torch::kFloat16) {\n mode = PrecisionMode::FP16IOFP16W;\n } else if (weight.dtype() == torch::kBFloat16) {\n mode = PrecisionMode::FP16IOBF16W;\n } else {\n mode = PrecisionMode::FP16IOFP32W;\n }\n }\n params_fwd.precision = mode;\n\n // The number of barriers.\n size_t barriers_elts = 0;\n // The number of elements in the reduction buffer.\n size_t red_buffer_elts = 0;\n // The number of elements in the reduction buffer that must be zeroed.\n size_t zeroed_red_buffer_elts = 0;\n\n // Finalize the parameters.\n dim3 grid;\n if (passes == 1) {\n group_norm_nhwc_fwd_one_pass_setup(params_fwd, barriers_elts, red_buffer_elts, grid, props);\n } else {\n group_norm_nhwc_fwd_two_passes_setup(params_fwd, zeroed_red_buffer_elts);\n }\n\n // Allocate on the device.\n auto red_buffer = at::empty({red_buffer_elts}, options.dtype(at::kFloat));\n params_fwd.red_buffer = red_buffer.data_ptr();\n\n // Allocate the buffer if needed.\n auto barriers = at::zeros({barriers_elts}, options.dtype(at::kInt));\n params_fwd.barriers = barriers.data_ptr();\n auto zeroed_red_buffer = at::zeros({zeroed_red_buffer_elts}, options.dtype(at::kFloat));\n params_fwd.zeroed_red_buffer = zeroed_red_buffer.data_ptr();\n\n if (passes == 1) {\n group_norm_nhwc_fwd_one_pass_run(params_fwd, grid, stream);\n } else {\n group_norm_nhwc_fwd_two_passes_sum(params_fwd, stream);\n group_norm_nhwc_fwd_two_passes_scale(params_fwd, stream);\n }\n\n return {output, sums_d};\n}\n\nstd::vector group_norm_bwd(torch::Tensor grad_output, torch::Tensor sums, torch::Tensor input,\n int groups, torch::Tensor weight, torch::Tensor bias, float eps, int passes,\n bool with_swish = false) {\n if (!initialized) {\n CHECK_CUDA_STATUS(cudaGetDeviceProperties(&props, 0));\n initialized = true;\n }\n CHECK_NHWC_INPUT(grad_output);\n auto stream = at::cuda::getCurrentCUDAStream();\n\n // Achieve group norm arguments\n int n = input.size(0);\n int c = input.size(1);\n int h = input.size(2);\n int w = input.size(3);\n\n // Check kernel constraints\n TORCH_CHECK(supported_groups_values.count(groups), \"`groups` of {16, 32} are only supported but \", groups,\n \" is passed\");\n TORCH_CHECK(supported_c_values.count(c), \"`c` of \", c, \" is not included in supported_c_values\");\n\n // Allocate tensors\n auto options = at::TensorOptions(at::kCUDA);\n auto grad_input = at::empty_like(input, at::MemoryFormat::Preserve);\n auto grad_weight = at::empty_like(weight, at::MemoryFormat::Preserve);\n auto grad_bias = at::empty_like(bias, at::MemoryFormat::Preserve);\n auto sums_d = at::empty({2 * n * groups}, options.dtype(at::kFloat));\n\n // Declare the parameters.\n Group_norm_nhwc_bwd_params params_bwd;\n memset(¶ms_bwd, 0, sizeof(params_bwd));\n\n // Initialize the parameters.\n params_bwd.dx = reinterpret_cast(grad_input.data_ptr());\n params_bwd.dgamma = reinterpret_cast(grad_weight.data_ptr());\n params_bwd.dbeta = reinterpret_cast(grad_bias.data_ptr());\n params_bwd.sums = const_cast(reinterpret_cast(sums.data_ptr()));\n params_bwd.dy = const_cast(reinterpret_cast(grad_output.data_ptr()));\n params_bwd.x = const_cast(reinterpret_cast(input.data_ptr()));\n ;\n params_bwd.gamma = const_cast(reinterpret_cast(weight.data_ptr()));\n params_bwd.beta = const_cast(reinterpret_cast(bias.data_ptr()));\n ;\n params_bwd.epsilon = eps;\n params_bwd.n = n;\n params_bwd.h = h;\n params_bwd.w = w;\n params_bwd.c = c;\n params_bwd.groups = groups;\n params_bwd.with_swish = with_swish;\n\n PrecisionMode mode;\n if (input.dtype() == torch::kFloat32) {\n if (weight.dtype() == torch::kFloat16) {\n mode = PrecisionMode::FP32IOFP16W;\n } else if (weight.dtype() == torch::kBFloat16) {\n mode = PrecisionMode::FP32IOBF16W;\n } else {\n mode = PrecisionMode::FP32IOFP32W;\n }\n } else if (input.dtype() == torch::kBFloat16) {\n if (weight.dtype() == torch::kFloat16) {\n mode = PrecisionMode::BF16IOFP16W;\n } else if (weight.dtype() == torch::kBFloat16) {\n mode = PrecisionMode::BF16IOBF16W;\n } else {\n mode = PrecisionMode::BF16IOFP32W;\n }\n } else {\n if (weight.dtype() == torch::kFloat16) {\n mode = PrecisionMode::FP16IOFP16W;\n } else if (weight.dtype() == torch::kBFloat16) {\n mode = PrecisionMode::FP16IOBF16W;\n } else {\n mode = PrecisionMode::FP16IOFP32W;\n }\n }\n params_bwd.precision = mode;\n\n // The number of barriers.\n size_t barriers_elts = 0;\n // The number of elements in the reduction buffer.\n size_t red_buffer_elts = 0;\n // The number of elements in the reduction buffer that must be zeroed.\n size_t zeroed_red_buffer_elts = 0;\n\n // Finalize the parameters.\n dim3 grid;\n if (passes == 1) {\n group_norm_nhwc_bwd_one_pass_setup(params_bwd, barriers_elts, red_buffer_elts, zeroed_red_buffer_elts, grid, props);\n } else {\n group_norm_nhwc_bwd_two_passes_setup(params_bwd, zeroed_red_buffer_elts);\n }\n\n // Allocate on the device.\n auto red_buffer = at::empty({red_buffer_elts}, options.dtype(at::kFloat));\n params_bwd.red_buffer = red_buffer.data_ptr();\n\n // Allocate the buffer if needed.\n auto barriers = at::zeros({barriers_elts}, options.dtype(at::kInt));\n params_bwd.barriers = barriers.data_ptr();\n auto zeroed_red_buffer = at::zeros({zeroed_red_buffer_elts}, options.dtype(at::kFloat));\n params_bwd.zeroed_red_buffer = zeroed_red_buffer.data_ptr();\n\n if (passes == 1) {\n group_norm_nhwc_bwd_one_pass_run(params_bwd, grid, stream);\n } else {\n group_norm_nhwc_bwd_two_passes_sum(params_bwd, stream);\n group_norm_nhwc_bwd_two_passes_scale(params_bwd, stream);\n }\n\n return {grad_input, grad_weight, grad_bias};\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"forward\", &group_norm_fwd, \"NHWC group norm forward\", py::call_guard());\n m.def(\"backward\", &group_norm_bwd, \"NHWC group norm backward\", py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/group_norm/macros.h", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n\n#define GN_ONE_PASS_RUN_FUNCTION_NAME(Traits, ACTS_PER_BLOCK, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME) \\\n void group_norm_nhwc_##PASS_NAME##_one_pass_##CHANNELS_PER_GROUP##_##ACTS_PER_BLOCK##_##Traits##_run( \\\n const Group_norm_nhwc_##PASS_NAME##_params& params, const dim3& grid, cudaStream_t stream)\n\n#define GN_ONE_PASS_RUN_FUNCTION(Traits, ACTS_PER_BLOCK, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME) \\\n GN_ONE_PASS_RUN_FUNCTION_NAME(Traits, ACTS_PER_BLOCK, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME) { \\\n auto kernel = \\\n group_norm_nhwc_##PASS_NAME##_one_pass_kernel; \\\n \\\n const Group_norm_nhwc_##PASS_NAME##_params* params_ = ¶ms; \\\n if (grid.x > 1) { \\\n CHECK_CUDA(cudaLaunchCooperativeKernel((const void*)kernel, grid, dim3(THREADS_PER_BLOCK), (void**)¶ms_, 0, \\\n stream)); \\\n \\\n } else { \\\n CHECK_CUDA(cudaLaunchKernel((const void*)kernel, grid, dim3(THREADS_PER_BLOCK), (void**)¶ms_, 0, stream)); \\\n } \\\n \\\n CHECK_CUDA(cudaGetLastError()); \\\n }\n\n//////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION_NAME(Traits, ACTS_PER_BLOCK, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, \\\n PASS_NAME) \\\n int group_norm_nhwc_##PASS_NAME##_one_pass_##CHANNELS_PER_GROUP##_##ACTS_PER_BLOCK##_##Traits##_blocks_per_sm()\n\n#define GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION(Traits, ACTS_PER_BLOCK, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME) \\\n GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION_NAME(Traits, ACTS_PER_BLOCK, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME) { \\\n auto kernel = \\\n group_norm_nhwc_##PASS_NAME##_one_pass_kernel; \\\n \\\n int blocks_per_sm = 0; \\\n CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks_per_sm, kernel, THREADS_PER_BLOCK, 0)); \\\n \\\n CHECK_CUDA(cudaGetLastError()); \\\n return blocks_per_sm; \\\n }\n\n//////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define GN_ONE_PASS_(FUNCTION, Traits, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME) \\\n FUNCTION(Traits, 512, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n FUNCTION(Traits, 256, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n FUNCTION(Traits, 128, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n FUNCTION(Traits, 64, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME);\n\n#define GN_ONE_PASS_RUN_DEFINITION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME) \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION, Fp32IOFp16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION, Fp32IOBf16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION, Fp32IOFp32W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION, Fp16IOFp16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION, Fp16IOBf16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION, Fp16IOFp32W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION, Bf16IOFp16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION, Bf16IOBf16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION, Bf16IOFp32W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME);\n\n#define GN_ONE_PASS_RUN_DECLARATION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME) \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION_NAME, Fp32IOFp16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION_NAME, Fp32IOBf16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION_NAME, Fp32IOFp32W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION_NAME, Fp16IOFp16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION_NAME, Fp16IOBf16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION_NAME, Fp16IOFp32W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION_NAME, Bf16IOFp16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION_NAME, Bf16IOBf16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_RUN_FUNCTION_NAME, Bf16IOFp32W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME);\n\n#define GN_ONE_PASS_BLOCKS_PER_SM_DEFINITION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME) \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION, Fp32IOFp16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION, Fp32IOBf16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION, Fp32IOFp32W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION, Fp16IOFp16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION, Fp16IOBf16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION, Fp16IOFp32W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION, Bf16IOFp16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION, Bf16IOBf16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION, Bf16IOFp32W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME);\n\n#define GN_ONE_PASS_BLOCKS_PER_SM_DECLARATION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME) \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION_NAME, Fp32IOFp16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, \\\n PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION_NAME, Fp32IOBf16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, \\\n PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION_NAME, Fp32IOFp32W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, \\\n PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION_NAME, Fp16IOFp16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, \\\n PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION_NAME, Fp16IOBf16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, \\\n PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION_NAME, Fp16IOFp32W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, \\\n PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION_NAME, Bf16IOFp16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, \\\n PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION_NAME, Bf16IOBf16W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, \\\n PASS_NAME); \\\n GN_ONE_PASS_(GN_ONE_PASS_BLOCKS_PER_SM_FUNCTION_NAME, Bf16IOFp32W, CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME);\n\n#define GN_ONE_PASS_DEFINITION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME) \\\n GN_ONE_PASS_RUN_DEFINITION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME) \\\n GN_ONE_PASS_BLOCKS_PER_SM_DEFINITION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK, PASS_NAME)\n\n#define GN_FWD_ONE_PASS_DEFINITION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK) \\\n GN_ONE_PASS_DEFINITION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK, fwd)\n\n#define GN_BWD_ONE_PASS_DEFINITION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK) \\\n GN_ONE_PASS_DEFINITION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK, bwd)\n\n#define GN_FWD_BWD_ONE_PASS_DEFINITION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK) \\\n GN_FWD_ONE_PASS_DEFINITION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK) \\\n GN_BWD_ONE_PASS_DEFINITION(CHANNELS_PER_GROUP, THREADS_PER_BLOCK)\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define GN_SELECTION_STATEMENT(function, Traits, PRECISION, FUNC_POSTFIX, HW_THRESHOLD, ACTS_PER_BLOCK, \\\n CHANNELS_PER_GROUP, PASS_NAME) \\\n if (params.hw >= HW_THRESHOLD && params.channels_per_group == CHANNELS_PER_GROUP && \\\n params.precision == PrecisionMode::PRECISION) { \\\n function = \\\n group_norm_nhwc_##PASS_NAME##_one_pass_##CHANNELS_PER_GROUP##_##ACTS_PER_BLOCK##_##Traits##FUNC_POSTFIX; \\\n } else\n\n#define GN_SELECTION_STATEMENT_WITH_CPG_LIMIT(function, Traits, PRECISION, FUNC_POSTFIX, HW_THRESHOLD, ACTS_PER_BLOCK, \\\n CHANNELS_PER_GROUP, PASS_NAME, LIMIT_CPG) \\\n if (params.hw >= HW_THRESHOLD && params.channels_per_group == CHANNELS_PER_GROUP && \\\n params.precision == PrecisionMode::PRECISION && CHANNELS_PER_GROUP >= LIMIT_CPG) { \\\n function = \\\n group_norm_nhwc_##PASS_NAME##_one_pass_##CHANNELS_PER_GROUP##_##ACTS_PER_BLOCK##_##Traits##FUNC_POSTFIX; \\\n } else\n\n#define GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Traits, PRECISION, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, PASS_NAME) \\\n GN_SELECTION_STATEMENT_WITH_CPG_LIMIT(function, Traits, PRECISION, FUNC_POSTFIX, 1024, 128, CHANNELS_PER_GROUP, \\\n PASS_NAME, 80) \\\n GN_SELECTION_STATEMENT_WITH_CPG_LIMIT(function, Traits, PRECISION, FUNC_POSTFIX, 256, 128, CHANNELS_PER_GROUP, \\\n PASS_NAME, 160) \\\n GN_SELECTION_STATEMENT(function, Traits, PRECISION, FUNC_POSTFIX, 512, 512, CHANNELS_PER_GROUP, PASS_NAME) \\\n GN_SELECTION_STATEMENT(function, Traits, PRECISION, FUNC_POSTFIX, 256, 256, CHANNELS_PER_GROUP, PASS_NAME) \\\n GN_SELECTION_STATEMENT(function, Traits, PRECISION, FUNC_POSTFIX, 128, 128, CHANNELS_PER_GROUP, PASS_NAME) \\\n GN_SELECTION_STATEMENT(function, Traits, PRECISION, FUNC_POSTFIX, 0, 64, CHANNELS_PER_GROUP, PASS_NAME)\n\n#define GN_FWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(CHANNELS_PER_GROUP, FUNC_POSTFIX, function) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Fp32IOFp16W, FP32IOFP16W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, fwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Fp32IOBf16W, FP32IOBF16W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, fwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Fp32IOFp32W, FP32IOFP32W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, fwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Fp16IOFp16W, FP16IOFP16W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, fwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Fp16IOBf16W, FP16IOBF16W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, fwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Fp16IOFp32W, FP16IOFP32W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, fwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Bf16IOFp16W, BF16IOFP16W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, fwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Bf16IOBf16W, BF16IOBF16W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, fwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Bf16IOFp32W, BF16IOFP32W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, fwd)\n\n#define GN_BWD_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(CHANNELS_PER_GROUP, FUNC_POSTFIX, function) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Fp32IOFp16W, FP32IOFP16W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, bwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Fp32IOBf16W, FP32IOBF16W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, bwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Fp32IOFp32W, FP32IOFP32W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, bwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Fp16IOFp16W, FP16IOFP16W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, bwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Fp16IOBf16W, FP16IOBF16W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, bwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Fp16IOFp32W, FP16IOFP32W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, bwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Bf16IOFp16W, BF16IOFP16W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, bwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Bf16IOBf16W, BF16IOBF16W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, bwd) \\\n GN_SELECTION_STATEMENT_HW_THRESHOLD_ACTS_PER_BLOCK_DISPATCH(Bf16IOFp32W, BF16IOFP32W, CHANNELS_PER_GROUP, \\\n FUNC_POSTFIX, function, bwd)\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define GN_ONE_PASS_DECLARATION(CHANNELS_PER_GROUP, PASS_NAME) \\\n GN_ONE_PASS_RUN_DECLARATION(CHANNELS_PER_GROUP, /* dummy value */ 0, PASS_NAME) \\\n GN_ONE_PASS_BLOCKS_PER_SM_DECLARATION(CHANNELS_PER_GROUP, /* dummy value */ 0, PASS_NAME)\n\n#define GN_FWD_ONE_PASS_DECLARATION(CHANNELS_PER_GROUP) GN_ONE_PASS_DECLARATION(CHANNELS_PER_GROUP, fwd)\n\n#define GN_BWD_ONE_PASS_DECLARATION(CHANNELS_PER_GROUP) GN_ONE_PASS_DECLARATION(CHANNELS_PER_GROUP, bwd)\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define CALL_TWO_PASS_KERNEL(Kernel, Precision) \\\n if (params.channels_per_block == 320) { \\\n Kernel<<>>(params); \\\n } else if (params.channels_per_block == 280) { \\\n Kernel<<>>(params); \\\n } else if (params.channels_per_block == 208) { \\\n Kernel<<>>(params); \\\n } else if (params.channels_per_block == 240) { \\\n Kernel<<>>(params); \\\n } else if (params.channels_per_block == 512) { \\\n Kernel<<>>(params); \\\n } else if (params.channels_per_block == 448) { \\\n Kernel<<>>(params); \\\n } else if (params.channels_per_block == 384) { \\\n Kernel<<>>(params); \\\n } else if (params.channels_per_block == 256) { \\\n Kernel<<>>(params); \\\n } else if (params.channels_per_block == 128) { \\\n Kernel<<>>(params); \\\n } else if (params.channels_per_block == 336) { \\\n Kernel<<>>(params); \\\n } else if (params.channels_per_block == 392) { \\\n Kernel<<>>(params); \\\n } else { \\\n assert(false); \\\n }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n" }, { "path": "apex/contrib/csrc/group_norm/traits.h", "content": "/*\n * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n * SPDX-License-Identifier: BSD-3-Clause\n */\n#pragma once\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Fp32 {\n // Type is float32_t\n using Type = float;\n // Doubled type\n using Type2 = float2;\n\n // Unpack input to accumulators type\n static inline __device__ float2 unpack(const float2& f2) { return f2; }\n\n // Pack the accumulators into outputs.\n static inline __device__ float2 pack(const float2& f2) { return f2; }\n\n static inline __device__ float2 zero() { return {0.f, 0.f}; }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Fp16 {\n // Type is __half\n using Type = __half;\n // Doubled type\n using Type2 = __half2;\n\n // Unpack input to accumulators type\n static inline __device__ float2 unpack(const __half2& h2) {\n // FIXME(nkorobov): __half22float2 makes compilation error in container\n return {__half2float(h2.x), __half2float(h2.y)};\n }\n\n // Pack the accumulators into outputs.\n static inline __device__ __half2 pack(const float2& f2) {\n // FIXME(nkorobov): __float22half2_rn makes compilation error in container\n return {__float2half_rn(f2.x), __float2half_rn(f2.y)};\n }\n\n static inline __device__ __half2 zero() {\n uint32_t zero = 0;\n return *reinterpret_cast<__half2*>(&zero);\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Bf16 {\n // Type is __nv_bfloat16\n using Type = __nv_bfloat16;\n // Doubled type\n using Type2 = __nv_bfloat162;\n\n // Unpack input to accumulators type\n static inline __device__ float2 unpack(const __nv_bfloat162& h2) {\n // FIXME(nkorobov): __half22float2 makes compilation error in container\n return {__bfloat162float(h2.x), __bfloat162float(h2.y)};\n }\n\n // Pack the accumulators into outputs.\n static inline __device__ __nv_bfloat162 pack(const float2& f2) {\n // FIXME(nkorobov): __float22bfloat162_rn makes compilation error in container\n return {__float2bfloat16_rn(f2.x), __float2bfloat16_rn(f2.y)};\n }\n\n static inline __device__ __nv_bfloat162 zero() {\n uint32_t zero = 0;\n return *reinterpret_cast<__nv_bfloat162*>(&zero);\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\nstruct Fp32IOFp16W {\n // IO traits\n using IOTraits = Fp32;\n // Weigths traits\n using WTraits = Fp16;\n};\n\nstruct Fp32IOBf16W {\n // IO traits\n using IOTraits = Fp32;\n // Weigths traits\n using WTraits = Bf16;\n};\n\nstruct Fp32IOFp32W {\n // IO traits\n using IOTraits = Fp32;\n // Weigths traits\n using WTraits = Fp32;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct Fp16IOFp16W {\n // IO traits\n using IOTraits = Fp16;\n // Weigths traits\n using WTraits = Fp16;\n};\n\nstruct Fp16IOBf16W {\n // IO traits\n using IOTraits = Fp16;\n // Weigths traits\n using WTraits = Bf16;\n};\n\nstruct Fp16IOFp32W {\n // IO traits\n using IOTraits = Fp16;\n // Weigths traits\n using WTraits = Fp32;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\nstruct Bf16IOFp16W {\n // IO traits\n using IOTraits = Bf16;\n // Weigths traits\n using WTraits = Fp16;\n};\n\nstruct Bf16IOBf16W {\n // IO traits\n using IOTraits = Bf16;\n // Weigths traits\n using WTraits = Bf16;\n};\n\nstruct Bf16IOFp32W {\n // IO traits\n using IOTraits = Bf16;\n // Weigths traits\n using WTraits = Fp32;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n" }, { "path": "apex/contrib/csrc/group_norm_v2/generate_gn_cuda_inst.py", "content": "import pathlib\n\n\nhw_c_list = [\n (8 * 8, 1280),\n (8 * 8, 2560),\n (16 * 16, 640),\n (16 * 16, 1280),\n (16 * 16, 1920),\n (16 * 16, 2560),\n (32 * 32, 320),\n (32 * 32, 640),\n (32 * 32, 960),\n (32 * 32, 1280),\n (32 * 32, 1920),\n (64 * 64, 320),\n (64 * 64, 640),\n (64 * 64, 960),\n]\n\n\ndef run():\n src_path = pathlib.Path(__file__).parent.absolute()\n\n for f in src_path.glob(\"gn_cuda_inst_*.cu\"):\n f.unlink()\n\n for hw, c in hw_c_list:\n print(f\"GN_CUDA_INST_DEFINE({hw}, {c})\")\n with open(src_path / f\"gn_cuda_inst_{hw}_{c}.cu\", \"w\") as f:\n f.write('#include \"gn_cuda_host_template.cuh\"\\n')\n f.write(\"\\n\")\n f.write(\"\\n\")\n f.write(\"namespace group_norm_v2 {\\n\")\n f.write(\"\\n\")\n f.write(f\"GN_CUDA_INST_DEFINE({hw}, {c})\\n\")\n f.write(\"\\n\")\n f.write(\"} // namespace group_norm_v2\\n\")\n\n with open(src_path / \"gn_dispatch_hw_c.hpp\", \"w\") as f:\n f.write(\"#pragma once\\n\")\n f.write(\"\\n\")\n f.write(\"#define DISPATCH_HW_C(hw, c, HW, C, ...) [&] { \\\\\\n\")\n for hw, c in hw_c_list:\n f.write(\n f\" if (hw == {hw} && c == {c}) {{ constexpr int HW = {hw}, C = {c}; return __VA_ARGS__(); }} \\\\\\n\"\n )\n f.write(\n ' throw std::invalid_argument(\"DISPATCH_HW_C \" + std::to_string(hw) + \" \" + std::to_string(c)); \\\\\\n'\n )\n f.write(\" }()\\n\")\n\n\nif __name__ == \"__main__\":\n run()\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn.cpp", "content": "#include \"gn.hpp\"\n\n#include \n#include \n\nnamespace group_norm_v2 {\n\ntorch::Tensor gn(torch::Tensor x, torch::Tensor w, torch::Tensor b, float eps, bool silu, int num_groups,\n std::optional mean_var_out, int sm_margin) {\n if (w.dtype() != b.dtype() || (mean_var_out.has_value() && mean_var_out->dtype() != torch::kFloat32)) {\n throw std::invalid_argument(\"gn dtype mismatch\");\n }\n torch::Tensor out = torch::empty_like(x);\n float* ptr_mean_var_out = mean_var_out.has_value() ? mean_var_out->data_ptr() : nullptr;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n int device_id = at::cuda::getCurrentCUDAStream().device().index();\n group_norm_v2::Meta meta;\n if (x.dtype() == torch::kHalf && w.dtype() == torch::kHalf) {\n group_norm_v2::gn_cuda((half*)out.data_ptr(), (half*)x.data_ptr(), (half*)w.data_ptr(), (half*)b.data_ptr(), eps,\n silu, x.size(0), x.size(2) * x.size(3), num_groups, x.size(1) / num_groups, ptr_mean_var_out,\n nullptr, nullptr, sm_margin, stream, device_id, &meta, true);\n } else if (x.dtype() == torch::kBFloat16 && w.dtype() == torch::kBFloat16) {\n group_norm_v2::gn_cuda((__nv_bfloat16*)out.data_ptr(), (__nv_bfloat16*)x.data_ptr(), (__nv_bfloat16*)w.data_ptr(),\n (__nv_bfloat16*)b.data_ptr(), eps, silu, x.size(0), x.size(2) * x.size(3), num_groups,\n x.size(1) / num_groups, ptr_mean_var_out, nullptr, nullptr, sm_margin, stream, device_id,\n &meta, true);\n } else {\n throw std::invalid_argument(\"gn only supports half or bfloat16 input and weight\");\n }\n torch::Tensor red_buffer =\n torch::empty({meta.red_buffer_size}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA));\n thread_local torch::Tensor barrier;\n if (barrier.size(0) < meta.barrier_size) {\n barrier = torch::zeros({meta.barrier_size}, torch::TensorOptions().dtype(torch::kUInt32).device(torch::kCUDA));\n }\n if (x.dtype() == torch::kHalf && w.dtype() == torch::kHalf) {\n group_norm_v2::gn_cuda((half*)out.data_ptr(), (half*)x.data_ptr(), (half*)w.data_ptr(), (half*)b.data_ptr(), eps,\n silu, x.size(0), x.size(2) * x.size(3), num_groups, x.size(1) / num_groups, ptr_mean_var_out,\n red_buffer.data_ptr(), barrier.data_ptr(), sm_margin, stream, device_id,\n nullptr, false);\n } else if (x.dtype() == torch::kBFloat16 && w.dtype() == torch::kBFloat16) {\n group_norm_v2::gn_cuda((__nv_bfloat16*)out.data_ptr(), (__nv_bfloat16*)x.data_ptr(), (__nv_bfloat16*)w.data_ptr(),\n (__nv_bfloat16*)b.data_ptr(), eps, silu, x.size(0), x.size(2) * x.size(3), num_groups,\n x.size(1) / num_groups, ptr_mean_var_out, red_buffer.data_ptr(),\n barrier.data_ptr(), sm_margin, stream, device_id, nullptr, false);\n } else {\n throw std::invalid_argument(\"gn only supports half or bfloat16 input and weight\");\n }\n return out;\n}\n\nauto gn_bwd(torch::Tensor grad_output, torch::Tensor x, torch::Tensor w, torch::Tensor b, torch::Tensor mean_var,\n float eps, bool silu, int num_groups, int sm_margin) {\n if (w.dtype() != b.dtype() || x.dtype() != grad_output.dtype() || mean_var.dtype() != torch::kFloat32) {\n throw std::invalid_argument(\"gn_bwd dtype mismatch\");\n }\n torch::Tensor grad_input = torch::empty_like(x);\n torch::Tensor grad_weight = torch::empty_like(w);\n torch::Tensor grad_bias = torch::empty_like(w);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n int device_id = at::cuda::getCurrentCUDAStream().device().index();\n group_norm_v2::Meta meta;\n if (x.dtype() == torch::kHalf && w.dtype() == torch::kHalf) {\n group_norm_v2::gn_bwd_cuda((half*)grad_input.data_ptr(), (half*)grad_weight.data_ptr(), (half*)grad_bias.data_ptr(),\n (half*)grad_output.data_ptr(), (half*)x.data_ptr(), (half*)w.data_ptr(),\n (half*)b.data_ptr(), mean_var.data_ptr(), eps, silu, x.size(0),\n x.size(2) * x.size(3), num_groups, x.size(1) / num_groups, nullptr, nullptr, sm_margin,\n stream, device_id, &meta, true);\n } else if (x.dtype() == torch::kBFloat16 && w.dtype() == torch::kBFloat16) {\n group_norm_v2::gn_bwd_cuda((__nv_bfloat16*)grad_input.data_ptr(), (__nv_bfloat16*)grad_weight.data_ptr(),\n (__nv_bfloat16*)grad_bias.data_ptr(), (__nv_bfloat16*)grad_output.data_ptr(),\n (__nv_bfloat16*)x.data_ptr(), (__nv_bfloat16*)w.data_ptr(), (__nv_bfloat16*)b.data_ptr(),\n mean_var.data_ptr(), eps, silu, x.size(0), x.size(2) * x.size(3), num_groups,\n x.size(1) / num_groups, nullptr, nullptr, sm_margin, stream, device_id, &meta, true);\n } else {\n throw std::invalid_argument(\"gn only supports half or bfloat16 input and weight\");\n }\n torch::Tensor red_buffer =\n torch::empty({meta.red_buffer_size}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA));\n thread_local torch::Tensor barrier;\n if (barrier.size(0) < meta.barrier_size) {\n barrier = torch::zeros({meta.barrier_size}, torch::TensorOptions().dtype(torch::kUInt32).device(torch::kCUDA));\n }\n if (x.dtype() == torch::kHalf && w.dtype() == torch::kHalf) {\n group_norm_v2::gn_bwd_cuda((half*)grad_input.data_ptr(), (half*)grad_weight.data_ptr(), (half*)grad_bias.data_ptr(),\n (half*)grad_output.data_ptr(), (half*)x.data_ptr(), (half*)w.data_ptr(),\n (half*)b.data_ptr(), mean_var.data_ptr(), eps, silu, x.size(0),\n x.size(2) * x.size(3), num_groups, x.size(1) / num_groups, red_buffer.data_ptr(),\n barrier.data_ptr(), sm_margin, stream, device_id, nullptr, false);\n } else if (x.dtype() == torch::kBFloat16 && w.dtype() == torch::kBFloat16) {\n group_norm_v2::gn_bwd_cuda((__nv_bfloat16*)grad_input.data_ptr(), (__nv_bfloat16*)grad_weight.data_ptr(),\n (__nv_bfloat16*)grad_bias.data_ptr(), (__nv_bfloat16*)grad_output.data_ptr(),\n (__nv_bfloat16*)x.data_ptr(), (__nv_bfloat16*)w.data_ptr(), (__nv_bfloat16*)b.data_ptr(),\n mean_var.data_ptr(), eps, silu, x.size(0), x.size(2) * x.size(3), num_groups,\n x.size(1) / num_groups, red_buffer.data_ptr(), barrier.data_ptr(),\n sm_margin, stream, device_id, nullptr, false);\n } else {\n throw std::invalid_argument(\"gn only supports half or bfloat16 input and weight\");\n }\n return std::make_tuple(grad_input, grad_weight, grad_bias);\n}\n\n} // namespace group_norm_v2\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"gn\", &group_norm_v2::gn, py::arg(\"x\"), py::arg(\"w\"), py::arg(\"b\"), py::arg(\"eps\"), py::arg(\"silu\"),\n py::arg(\"num_groups\"), py::arg(\"mean_var_out\") = py::none(), py::arg(\"sm_margin\") = 0, \"\");\n m.def(\"gn_bwd\", &group_norm_v2::gn_bwd, py::arg(\"grad_output\"), py::arg(\"x\"), py::arg(\"w\"), py::arg(\"b\"),\n py::arg(\"mean_var\"), py::arg(\"eps\"), py::arg(\"silu\"), py::arg(\"num_groups\"), py::arg(\"sm_margin\") = 0, \"\");\n}\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn.hpp", "content": "#pragma once\n\n#include \n\n#include \n\nnamespace group_norm_v2 {\n\nstruct Meta {\n int64_t red_buffer_size;\n int64_t barrier_size;\n int BLOCK_DIM_X;\n int C_PER_BLOCK;\n int ROWS_PER_BLOCK;\n int VEC_ELEMS;\n bool LOAD_TWICE;\n int BLOCKS_PER_SM;\n bool HARDWARE_CLUSTER;\n int wgrad_sync_method;\n};\n\ntemplate \nvoid gn_cuda(T* out, T* x, T* w, T* b, float eps, bool silu, int64_t n, int64_t hw, int num_groups,\n int channels_per_group, float* mean_var_out, float* red_buffer, unsigned* barrier, int sm_margin,\n cudaStream_t stream, int device_id, Meta* meta_ptr, bool meta_only);\n\ntemplate \nvoid gn_bwd_cuda(T* grad_input, T* grad_weight, T* grad_bias, T* grad_output, T* x, T* w, T* b, float* mean_var,\n float eps, bool silu, int64_t n, int64_t hw, int num_groups, int channels_per_group, float* red_buffer,\n unsigned* barrier, int sm_margin, cudaStream_t stream, int device_id, Meta* meta_ptr, bool meta_only);\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda.cu", "content": "#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"gn.hpp\"\n#include \"gn_dispatch_hw_c.hpp\"\n#include \"gn_utils.hpp\"\n\n#define DISPATCH_NUM_GROUPS_AND_SILU(num_groups, silu, NUM_GROUPS, SILU, ...) \\\n [&] { \\\n if (num_groups == 16 && silu == true) { \\\n constexpr int NUM_GROUPS = 16; \\\n constexpr bool SILU = true; \\\n return __VA_ARGS__(); \\\n } \\\n if (num_groups == 32 && silu == false) { \\\n constexpr int NUM_GROUPS = 32; \\\n constexpr bool SILU = false; \\\n return __VA_ARGS__(); \\\n } \\\n throw std::invalid_argument(\"DISPATCH_NUM_GROUPS_AND_SILU \" + std::to_string(num_groups) + \" \" + \\\n std::to_string(silu)); \\\n }()\n\nnamespace group_norm_v2 {\n\ntemplate \nvoid gn_cuda_single_shape(GN_CUDA_HOST_PARAMS(T));\n\ntemplate \nvoid gn_bwd_cuda_single_shape(GN_BWD_CUDA_HOST_PARAMS(T));\n\ntemplate \nvoid gn_cuda(GN_CUDA_HOST_PARAMS(T)) {\n DISPATCH_HW_C(hw, num_groups * channels_per_group, HW, C, [&] {\n DISPATCH_NUM_GROUPS_AND_SILU(num_groups, silu, G, SILU,\n [&] { return gn_cuda_single_shape(GN_CUDA_HOST_ARGS); });\n });\n}\n\ntemplate \nvoid gn_bwd_cuda(GN_BWD_CUDA_HOST_PARAMS(T)) {\n DISPATCH_HW_C(hw, num_groups * channels_per_group, HW, C, [&] {\n DISPATCH_NUM_GROUPS_AND_SILU(num_groups, silu, G, SILU,\n [&] { return gn_bwd_cuda_single_shape(GN_BWD_CUDA_HOST_ARGS); });\n });\n}\n\ntemplate void gn_cuda(GN_CUDA_HOST_PARAMS(half));\ntemplate void gn_cuda(GN_CUDA_HOST_PARAMS(__nv_bfloat16));\n\ntemplate void gn_bwd_cuda(GN_BWD_CUDA_HOST_PARAMS(half));\ntemplate void gn_bwd_cuda(GN_BWD_CUDA_HOST_PARAMS(__nv_bfloat16));\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_host_template.cuh", "content": "#pragma once\n\n#include \n#include \n#include \n\n#include \n#include \n\n#include \"gn_cuda_kernel.cuh\"\n#include \"gn_utils.hpp\"\n\nnamespace group_norm_v2 {\n\n#define DISPATCH_LOWER_BOUND_N(VALUE, CONST_NAME, ...) \\\n [&] { \\\n if (VALUE >= 16) { \\\n constexpr int CONST_NAME = 16; \\\n return __VA_ARGS__(); \\\n } \\\n if (VALUE >= 8) { \\\n constexpr int CONST_NAME = 8; \\\n return __VA_ARGS__(); \\\n } \\\n if (VALUE >= 4) { \\\n constexpr int CONST_NAME = 4; \\\n return __VA_ARGS__(); \\\n } \\\n if (VALUE >= 2) { \\\n constexpr int CONST_NAME = 2; \\\n return __VA_ARGS__(); \\\n } \\\n if (VALUE >= 1) { \\\n constexpr int CONST_NAME = 1; \\\n return __VA_ARGS__(); \\\n } \\\n throw std::invalid_argument(\"DISPATCH_LOWER_BOUND_N \" + std::to_string(VALUE)); \\\n }()\n\n#define DISPATCH_CUDA_ARCH_AND_LOWER_BOUND_SM_COUNT(runtime_cuda_arch, sm_count, RUNTIME_CUDA_ARCH, LB_SM_COUNT, ...) \\\n [&] { \\\n if (runtime_cuda_arch == 1000 && sm_count >= 148) { \\\n constexpr int RUNTIME_CUDA_ARCH = 1000, LB_SM_COUNT = 148; \\\n return __VA_ARGS__(); \\\n } \\\n throw std::invalid_argument(\"DISPATCH_CUDA_ARCH_AND_LOWER_BOUND_SM_COUNT \" + std::to_string(runtime_cuda_arch) + \\\n \" \" + std::to_string(sm_count)); \\\n }()\n\n#define DISPATCH_SM_MARGIN(VALUE, CONST_NAME, ...) \\\n [&] { \\\n if (VALUE == 0) { \\\n constexpr int CONST_NAME = 0; \\\n return __VA_ARGS__(); \\\n } \\\n if (VALUE == 32) { \\\n constexpr int CONST_NAME = 32; \\\n return __VA_ARGS__(); \\\n } \\\n throw std::invalid_argument(\"DISPATCH_SM_MARGIN \" + std::to_string(VALUE)); \\\n }()\n\ninline constexpr int get_max_cuda_arch() {\n int cuda_arch_list[] = {__CUDA_ARCH_LIST__};\n int max_cuda_arch = -1;\n for (int cuda_arch_item : cuda_arch_list) {\n if (cuda_arch_item > max_cuda_arch) {\n max_cuda_arch = cuda_arch_item;\n }\n }\n return max_cuda_arch;\n}\n\ntemplate \nconstexpr auto compute_gn_params() {\n constexpr int C = G * CPG;\n\n // Initialize each variable to comply with C++17\n int BLOCK_DIM_X = 0;\n int C_PER_BLOCK = 0;\n int ROWS_PER_BLOCK = 0;\n bool LOAD_TWICE = false;\n int BLOCKS_PER_SM = 0;\n WgradSyncMethod wgrad_sync_method = WGRAD_SYNC_UNSPECIFIED;\n\n // There are two tiling strategies:\n // - block sync: each block handles a whole group, i.e., a multiple of (G * HW) elements\n // - virtual cluster sync: each virtual cluster handles a group\n // Block sync can avoid cross-block synchronization latency, but it may cause low occupancy.\n // Use block sync if the IO size is small, when latency rather than occupancy dominates the kernel running time.\n\n // Elements to load for forward pass is `x`, elements to load for backward pass are `x` and `grad_output`, hence there\n // is a factor of (1 + BWD)\n if (HW * CPG * (1 + BWD) * sizeof(T) <= 20480) {\n // Strategy 1: block sync\n C_PER_BLOCK = CPG;\n ROWS_PER_BLOCK = HW;\n BLOCK_DIM_X = lcm(32, C_PER_BLOCK);\n while (BLOCK_DIM_X < 256) {\n BLOCK_DIM_X *= 2;\n }\n BLOCKS_PER_SM = 1;\n // The size of registers is 65536 registers * 4 bytes per register.\n // We have to leave some room for other variables and compiler optimizations,\n // so we use 36000 as the threshold.\n LOAD_TWICE = BLOCKS_PER_SM * ROWS_PER_BLOCK * C_PER_BLOCK * (1 + BWD) * sizeof(T) > 36000 * 4;\n } else {\n // Strategy 2: virtual cluster sync\n // A virtual cluster is a group of blocks that are synchronized with each other.\n // Each group, i.e., a multiple of (G * HW) elements, should be handled on the same virtual cluster.\n // If the virtual cluster size is supported by the hardware, HARDWARE_CLUSTER is preferred;\n // otherwise, cooperative groups are used (i.e., PERSISTENT kernels).\n int c_per_cluster = lcm(128 / (int)sizeof(T), CPG);\n\n C_PER_BLOCK = c_per_cluster;\n BLOCK_DIM_X = C_PER_BLOCK == 320 ? 320 : 480;\n\n // Maximum number of rows that should reside in registers\n int register_max_rows = 36000 * 4 / (C_PER_BLOCK * (1 + BWD) * sizeof(T));\n\n std::tuple best_candidate{};\n BLOCKS_PER_SM = 0;\n ROWS_PER_BLOCK = 0;\n for (int blocks_per_sm = 1; blocks_per_sm <= 3; blocks_per_sm++) {\n for (int rows_per_block = HW; rows_per_block >= 1; rows_per_block /= 2) {\n int virtual_cluster_size = (HW / rows_per_block) * (c_per_cluster / C_PER_BLOCK);\n if (virtual_cluster_size > blocks_per_sm * (LB_SM_COUNT - SM_MARGIN)) {\n continue;\n }\n int num_clusters = blocks_per_sm * (LB_SM_COUNT - SM_MARGIN) / virtual_cluster_size;\n int num_tasks = LB_N * (C / c_per_cluster);\n int num_waves = up_div(num_tasks, num_clusters);\n bool load_twice = rows_per_block > register_max_rows / blocks_per_sm;\n\n // Wave utilization: the percent of SMs that are used for each wave\n // For example, SM_COUNT=100 and VIRTUAL_CLUSTER_SIZE=64,\n // if BLOCKS_PER_SM=1, num_clusters=1, wave_util=64%;\n // if BLOCKS_PER_SM=2, num_clusters=3, wave_util=96%.\n // This helps select a good number of BLOCKS_PER_SM\n int wave_util = 10000 * std::min(num_tasks, num_clusters) * virtual_cluster_size /\n (blocks_per_sm * (LB_SM_COUNT - SM_MARGIN));\n\n decltype(best_candidate) candidate = {\n true,\n !load_twice, // Prefer no load twice\n !(num_waves >= 2 &&\n blocks_per_sm ==\n 1), // When there are multiple waves, prefer multiple blocks per SM to ensure overlapping\n -num_waves, // Prefer fewer waves\n std::min(9000, wave_util), // Prefer high wave utilization\n -blocks_per_sm, // Prefer fewer blocks per SM in order to reduce threads overhead\n };\n if (candidate > best_candidate) {\n // Assign each element respectively to comply with C++17\n std::get<0>(best_candidate) = std::get<0>(candidate);\n std::get<1>(best_candidate) = std::get<1>(candidate);\n std::get<2>(best_candidate) = std::get<2>(candidate);\n std::get<3>(best_candidate) = std::get<3>(candidate);\n std::get<4>(best_candidate) = std::get<4>(candidate);\n std::get<5>(best_candidate) = std::get<5>(candidate);\n static_assert(std::tuple_size::value == 6, \"missing assignments\");\n\n BLOCKS_PER_SM = blocks_per_sm;\n ROWS_PER_BLOCK = rows_per_block;\n }\n }\n }\n\n LOAD_TWICE = ROWS_PER_BLOCK > register_max_rows / BLOCKS_PER_SM;\n }\n\n int c_per_cluster = lcm(CPG, C_PER_BLOCK);\n int virtual_cluster_size = (c_per_cluster / C_PER_BLOCK) * (HW / ROWS_PER_BLOCK);\n\n // The occupancy is affected if cluster size is large.\n // For example, on H100, when gridDim=128 and each block occupies the whole SM,\n // if cluster is not used, all blocks can be active simultaneously.\n // if cluster size is 16, not all blocks can be active simultaneously (which can be queried by\n // cudaOccupancyMaxActiveClusters),\n // so there will be two waves which impacts efficiency.\n // When SM_MARGIN is set, no cluster should be used because other kernels may occupy a part of the cluster.\n bool HARDWARE_CLUSTER = virtual_cluster_size <= 2 && virtual_cluster_size != 1 && SM_MARGIN == 0;\n\n int MAX_VEC_BYTES =\n 8; // Sometimes 4 or 16 is better, but there is no trivial way to select the best vectorization size.\n int VEC_ELEMS = std::min(gcd(MAX_VEC_BYTES / (int)sizeof(T), C_PER_BLOCK),\n gcd(MAX_VEC_BYTES / (int)sizeof(T), ROWS_PER_BLOCK * C_PER_BLOCK / BLOCK_DIM_X));\n\n return std::make_tuple(BLOCK_DIM_X, C_PER_BLOCK, ROWS_PER_BLOCK, VEC_ELEMS, LOAD_TWICE, BLOCKS_PER_SM,\n HARDWARE_CLUSTER, wgrad_sync_method);\n}\n\n// Save compilation time for unused CUDA_ARCHs\n// For each template argument from DISPATCH_CUDA_ARCH_AND_LOWER_BOUND_SM_COUNT, the kernel is only compiled for the\n// corresponding CUDA_ARCH\ntemplate \nclass CompileCondition {\n public:\n __host__ __device__ static constexpr bool matches() {\n#if defined(__CUDA_ARCH__)\n return __CUDA_ARCH__ == EFFECTIVE_CUDA_ARCH;\n#else\n return false;\n#endif\n }\n};\n\ntemplate \nvoid gn_cuda_single_shape(GN_CUDA_HOST_PARAMS(T)) {\n if (out == x) {\n throw std::invalid_argument(\"not __restrict__\");\n }\n\n cudaDeviceProp const& deviceProp = get_device_prop(device_id);\n int runtime_cuda_arch = deviceProp.major * 100 + deviceProp.minor * 10;\n int sm_count = deviceProp.multiProcessorCount;\n\n DISPATCH_LOWER_BOUND_N(n, LB_N, [&] {\n DISPATCH_CUDA_ARCH_AND_LOWER_BOUND_SM_COUNT(runtime_cuda_arch, sm_count, RUNTIME_CUDA_ARCH, LB_SM_COUNT, [&] {\n DISPATCH_SM_MARGIN(sm_margin, SM_MARGIN, [&] {\n if (hw != HW) {\n throw std::invalid_argument(\"wrong HW\");\n }\n if (num_groups * channels_per_group != C) {\n throw std::invalid_argument(\"wrong C\");\n }\n if (num_groups != G) {\n throw std::invalid_argument(\"wrong G\");\n }\n if (silu != SILU) {\n throw std::invalid_argument(\"wrong SILU\");\n }\n if (n < LB_N) {\n throw std::invalid_argument(\"wrong LB_N\");\n }\n if (runtime_cuda_arch != RUNTIME_CUDA_ARCH) {\n throw std::invalid_argument(\"wrong RUNTIME_CUDA_ARCH\");\n }\n if (sm_count < LB_SM_COUNT) {\n throw std::invalid_argument(\"wrong LB_SM_COUNT\");\n }\n if (sm_margin != SM_MARGIN) {\n throw std::invalid_argument(\"wrong SM_MARGIN\");\n }\n constexpr int EFFECTIVE_CUDA_ARCH =\n std::min(RUNTIME_CUDA_ARCH, get_max_cuda_arch()); // Assume the max CUDA_ARCH is used to generate PTX\n\n constexpr int CPG = C / G;\n\n constexpr auto params = compute_gn_params();\n constexpr int BLOCK_DIM_X = std::get<0>(params);\n constexpr int C_PER_BLOCK = std::get<1>(params);\n constexpr int ROWS_PER_BLOCK = std::get<2>(params);\n constexpr int VEC_ELEMS = std::get<3>(params);\n constexpr bool LOAD_TWICE = std::get<4>(params);\n constexpr int BLOCKS_PER_SM = std::get<5>(params);\n constexpr bool HARDWARE_CLUSTER = std::get<6>(params);\n\n constexpr int C_PER_CLUSTER = lcm(CPG, C_PER_BLOCK);\n constexpr int VIRTUAL_CLUSTER_SIZE = (C_PER_CLUSTER / C_PER_BLOCK) * (HW / ROWS_PER_BLOCK);\n constexpr int NUM_VIRTUAL_CLUSTERS = ((LB_SM_COUNT - SM_MARGIN) * BLOCKS_PER_SM) / VIRTUAL_CLUSTER_SIZE;\n constexpr bool PERSISTENT =\n !HARDWARE_CLUSTER &&\n VIRTUAL_CLUSTER_SIZE >=\n 2; // Only virtual cluster sync (not include hardware cluster sync) requires PERSISTENT kernels\n\n if (meta_ptr) {\n constexpr int MAX_NUM_GROUPS_PER_BLOCK =\n C_PER_BLOCK % CPG == 0 ? C_PER_BLOCK / CPG : up_div(C_PER_BLOCK - gcd(C_PER_BLOCK, CPG), CPG) + 1;\n meta_ptr->red_buffer_size = 2 * NUM_VIRTUAL_CLUSTERS * VIRTUAL_CLUSTER_SIZE * MAX_NUM_GROUPS_PER_BLOCK * 2;\n meta_ptr->barrier_size = NUM_VIRTUAL_CLUSTERS;\n meta_ptr->BLOCK_DIM_X = BLOCK_DIM_X;\n meta_ptr->C_PER_BLOCK = C_PER_BLOCK;\n meta_ptr->ROWS_PER_BLOCK = ROWS_PER_BLOCK;\n meta_ptr->VEC_ELEMS = VEC_ELEMS;\n meta_ptr->LOAD_TWICE = LOAD_TWICE;\n meta_ptr->BLOCKS_PER_SM = BLOCKS_PER_SM;\n meta_ptr->HARDWARE_CLUSTER = HARDWARE_CLUSTER;\n meta_ptr->wgrad_sync_method = (int)WGRAD_SYNC_UNSPECIFIED;\n }\n if (meta_only) {\n return;\n }\n\n cudaLaunchConfig_t config = {0};\n config.gridDim = dim3(\n VIRTUAL_CLUSTER_SIZE,\n PERSISTENT ? std::min((int)n * (C / C_PER_CLUSTER), NUM_VIRTUAL_CLUSTERS) : n * (C / C_PER_CLUSTER), 1);\n config.blockDim = BLOCK_DIM_X;\n config.stream = stream;\n\n cudaLaunchAttribute attribute[2];\n if constexpr (HARDWARE_CLUSTER) {\n attribute[0].id = cudaLaunchAttributeClusterDimension;\n attribute[0].val.clusterDim.x = VIRTUAL_CLUSTER_SIZE; // Cluster size in X-dimension\n attribute[0].val.clusterDim.y = 1;\n attribute[0].val.clusterDim.z = 1;\n config.attrs = attribute;\n config.numAttrs++;\n }\n if constexpr (PERSISTENT) {\n attribute[config.numAttrs].id = cudaLaunchAttributeCooperative;\n attribute[config.numAttrs].val.cooperative = 1;\n config.attrs = attribute;\n config.numAttrs++;\n }\n\n auto kernel = &gn_cuda_kernel >;\n if constexpr (HARDWARE_CLUSTER) {\n if constexpr (VIRTUAL_CLUSTER_SIZE > 8) {\n CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1));\n }\n int max_cluster_size;\n int active_clusters;\n CUDA_CHECK(cudaOccupancyMaxPotentialClusterSize(&max_cluster_size, (void*)kernel, &config));\n if (VIRTUAL_CLUSTER_SIZE <= max_cluster_size && PERSISTENT) {\n attribute[0].val.clusterDim.x = VIRTUAL_CLUSTER_SIZE;\n CUDA_CHECK(cudaOccupancyMaxActiveClusters(&active_clusters, (void*)kernel, &config));\n }\n if (VIRTUAL_CLUSTER_SIZE <= max_cluster_size &&\n (!PERSISTENT || PERSISTENT && NUM_VIRTUAL_CLUSTERS <= active_clusters)) {\n attribute[0].val.clusterDim.x = VIRTUAL_CLUSTER_SIZE;\n } else {\n // Fallback to cooperative groups because hardware cluster cannot be active simultaneously\n constexpr bool HARDWARE_CLUSTER_NEW = false;\n constexpr bool PERSISTENT_NEW = !HARDWARE_CLUSTER_NEW && VIRTUAL_CLUSTER_SIZE >= 2;\n config.gridDim = dim3(\n VIRTUAL_CLUSTER_SIZE,\n PERSISTENT_NEW ? std::min((int)n * (C / C_PER_CLUSTER), NUM_VIRTUAL_CLUSTERS) : n * (C / C_PER_CLUSTER),\n 1);\n config.attrs = nullptr;\n config.numAttrs = 0;\n if constexpr (PERSISTENT_NEW) {\n attribute[config.numAttrs].id = cudaLaunchAttributeCooperative;\n attribute[config.numAttrs].val.cooperative = 1;\n config.attrs = attribute;\n config.numAttrs++;\n }\n kernel = &gn_cuda_kernel >;\n }\n }\n CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, out, x, w, b, eps, n, mean_var_out, red_buffer, barrier));\n });\n });\n });\n}\n\ntemplate \nvoid gn_bwd_cuda_single_shape(GN_BWD_CUDA_HOST_PARAMS(T)) {\n if (grad_input == grad_output || grad_input == x) {\n throw std::invalid_argument(\"not __restrict__\");\n }\n\n cudaDeviceProp const& deviceProp = get_device_prop(device_id);\n int runtime_cuda_arch = deviceProp.major * 100 + deviceProp.minor * 10;\n int sm_count = deviceProp.multiProcessorCount;\n\n DISPATCH_LOWER_BOUND_N(n, LB_N, [&] {\n DISPATCH_CUDA_ARCH_AND_LOWER_BOUND_SM_COUNT(runtime_cuda_arch, sm_count, RUNTIME_CUDA_ARCH, LB_SM_COUNT, [&] {\n DISPATCH_SM_MARGIN(sm_margin, SM_MARGIN, [&] {\n if (hw != HW) {\n throw std::invalid_argument(\"wrong HW\");\n }\n if (num_groups * channels_per_group != C) {\n throw std::invalid_argument(\"wrong C\");\n }\n if (num_groups != G) {\n throw std::invalid_argument(\"wrong G\");\n }\n if (silu != SILU) {\n throw std::invalid_argument(\"wrong SILU\");\n }\n if (n < LB_N) {\n throw std::invalid_argument(\"wrong LB_N\");\n }\n if (runtime_cuda_arch != RUNTIME_CUDA_ARCH) {\n throw std::invalid_argument(\"wrong RUNTIME_CUDA_ARCH\");\n }\n if (sm_count < LB_SM_COUNT) {\n throw std::invalid_argument(\"wrong LB_SM_COUNT\");\n }\n if (sm_margin != SM_MARGIN) {\n throw std::invalid_argument(\"wrong SM_MARGIN\");\n }\n constexpr int EFFECTIVE_CUDA_ARCH =\n std::min(RUNTIME_CUDA_ARCH, get_max_cuda_arch()); // Assume the max CUDA_ARCH is used to generate PTX\n\n constexpr bool REQUIRES_WGRAD = true;\n constexpr int CPG = C / G;\n\n constexpr auto params = compute_gn_params();\n constexpr int BLOCK_DIM_X = std::get<0>(params);\n constexpr int C_PER_BLOCK = std::get<1>(params);\n constexpr int ROWS_PER_BLOCK = std::get<2>(params);\n constexpr int VEC_ELEMS = std::get<3>(params);\n constexpr bool LOAD_TWICE = std::get<4>(params);\n constexpr int BLOCKS_PER_SM = std::get<5>(params);\n constexpr bool HARDWARE_CLUSTER = std::get<6>(params);\n constexpr WgradSyncMethod wgrad_sync_method_hint = std::get<7>(params);\n\n constexpr int C_PER_CLUSTER = lcm(CPG, C_PER_BLOCK);\n constexpr int VIRTUAL_CLUSTER_SIZE = (C_PER_CLUSTER / C_PER_BLOCK) * (HW / ROWS_PER_BLOCK);\n constexpr int NUM_VIRTUAL_CLUSTERS_NOT_ALIGNED =\n ((LB_SM_COUNT - SM_MARGIN) * BLOCKS_PER_SM) / VIRTUAL_CLUSTER_SIZE;\n\n // PERSISTENT is required because wgrad reduction requires synchronization.\n // TODO: specilize for the case that REQUIRES_WGRAD == false\n constexpr bool PERSISTENT = true;\n\n // Determine whether to align each virtual cluster to a fixed range of channels\n // If aligned, WGRAD_REUSE_SUM_SYNC_GROUP can be used, then less local wgrad memory is used (leave more room\n // for compiler\n // optimizations), and wgrad reduction is more efficient.\n // However, aligning can cause low occupancy.\n // There is a trade-off, and the condition to align is `NUM_VIRTUAL_CLUSTERS_NOT_ALIGNED > 2 * (C /\n // C_PER_CLUSTER)`\n constexpr WgradSyncMethod wgrad_sync_method =\n wgrad_sync_method_hint == WGRAD_SYNC_UNSPECIFIED\n ? NUM_VIRTUAL_CLUSTERS_NOT_ALIGNED > 2 * (C / C_PER_CLUSTER) ||\n NUM_VIRTUAL_CLUSTERS_NOT_ALIGNED % (C / C_PER_CLUSTER) == 0\n ? (HARDWARE_CLUSTER ? WGRAD_ARRIVE_AND_WAIT_GROUP : WGRAD_REUSE_SUM_SYNC_GROUP)\n : WGRAD_REUSE_SUM_SYNC_GRID\n : wgrad_sync_method_hint;\n constexpr int NUM_VIRTUAL_CLUSTERS =\n wgrad_sync_method == WGRAD_ARRIVE_AND_WAIT_GROUP || wgrad_sync_method == WGRAD_REUSE_SUM_SYNC_GROUP\n ? NUM_VIRTUAL_CLUSTERS_NOT_ALIGNED / (C / C_PER_CLUSTER) * (C / C_PER_CLUSTER)\n : NUM_VIRTUAL_CLUSTERS_NOT_ALIGNED;\n\n if (meta_ptr) {\n constexpr int MAX_NUM_GROUPS_PER_BLOCK =\n C_PER_BLOCK % CPG == 0 ? C_PER_BLOCK / CPG : up_div(C_PER_BLOCK - gcd(C_PER_BLOCK, CPG), CPG) + 1;\n meta_ptr->red_buffer_size =\n 2 * NUM_VIRTUAL_CLUSTERS * VIRTUAL_CLUSTER_SIZE * MAX_NUM_GROUPS_PER_BLOCK * 2 +\n std::max(n, (int64_t)NUM_VIRTUAL_CLUSTERS / (C / C_PER_CLUSTER)) * (HW / ROWS_PER_BLOCK) * C * 2;\n meta_ptr->barrier_size = NUM_VIRTUAL_CLUSTERS + C / C_PER_CLUSTER;\n meta_ptr->BLOCK_DIM_X = BLOCK_DIM_X;\n meta_ptr->C_PER_BLOCK = C_PER_BLOCK;\n meta_ptr->ROWS_PER_BLOCK = ROWS_PER_BLOCK;\n meta_ptr->VEC_ELEMS = VEC_ELEMS;\n meta_ptr->LOAD_TWICE = LOAD_TWICE;\n meta_ptr->BLOCKS_PER_SM = BLOCKS_PER_SM;\n meta_ptr->HARDWARE_CLUSTER = HARDWARE_CLUSTER;\n meta_ptr->wgrad_sync_method = (int)wgrad_sync_method;\n }\n if (meta_only) {\n return;\n }\n\n cudaLaunchConfig_t config = {0};\n config.gridDim = dim3(VIRTUAL_CLUSTER_SIZE, PERSISTENT ? NUM_VIRTUAL_CLUSTERS : n * (C / C_PER_CLUSTER), 1);\n config.blockDim = BLOCK_DIM_X;\n config.stream = stream;\n\n cudaLaunchAttribute attribute[2];\n if constexpr (HARDWARE_CLUSTER) {\n attribute[0].id = cudaLaunchAttributeClusterDimension;\n attribute[0].val.clusterDim.x = 1; // Cluster size in X-dimension\n attribute[0].val.clusterDim.y = 1;\n attribute[0].val.clusterDim.z = 1;\n config.attrs = attribute;\n config.numAttrs++;\n }\n if constexpr (PERSISTENT) {\n attribute[config.numAttrs].id = cudaLaunchAttributeCooperative;\n attribute[config.numAttrs].val.cooperative = 1;\n config.attrs = attribute;\n config.numAttrs++;\n }\n\n auto kernel =\n &gn_bwd_cuda_kernel >;\n if constexpr (HARDWARE_CLUSTER) {\n if constexpr (VIRTUAL_CLUSTER_SIZE > 8) {\n CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1));\n }\n int max_cluster_size;\n int active_clusters;\n CUDA_CHECK(cudaOccupancyMaxPotentialClusterSize(&max_cluster_size, (void*)kernel, &config));\n if (VIRTUAL_CLUSTER_SIZE <= max_cluster_size && PERSISTENT) {\n attribute[0].val.clusterDim.x = VIRTUAL_CLUSTER_SIZE;\n CUDA_CHECK(cudaOccupancyMaxActiveClusters(&active_clusters, (void*)kernel, &config));\n }\n if (VIRTUAL_CLUSTER_SIZE <= max_cluster_size &&\n (!PERSISTENT || PERSISTENT && NUM_VIRTUAL_CLUSTERS <= active_clusters)) {\n attribute[0].val.clusterDim.x = VIRTUAL_CLUSTER_SIZE;\n } else {\n // Fallback to cooperative groups for dgrad computation because hardware cluster cannot be active\n // simultaneously\n attribute[0].val.clusterDim.x = 1;\n kernel =\n &gn_bwd_cuda_kernel >;\n }\n }\n CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, grad_input, grad_weight, grad_bias, grad_output, x, w, b,\n mean_var, eps, n, red_buffer, barrier));\n });\n });\n });\n}\n\n#define GN_CUDA_INST_DEFINE(HW, C) \\\n template void gn_cuda_single_shape(GN_CUDA_HOST_PARAMS(half)); \\\n template void gn_cuda_single_shape(GN_CUDA_HOST_PARAMS(half)); \\\n template void gn_bwd_cuda_single_shape(GN_BWD_CUDA_HOST_PARAMS(half)); \\\n template void gn_bwd_cuda_single_shape(GN_BWD_CUDA_HOST_PARAMS(half)); \\\n template void gn_cuda_single_shape<__nv_bfloat16, HW, C, 16, true>(GN_CUDA_HOST_PARAMS(__nv_bfloat16)); \\\n template void gn_cuda_single_shape<__nv_bfloat16, HW, C, 32, false>(GN_CUDA_HOST_PARAMS(__nv_bfloat16)); \\\n template void gn_bwd_cuda_single_shape<__nv_bfloat16, HW, C, 16, true>(GN_BWD_CUDA_HOST_PARAMS(__nv_bfloat16)); \\\n template void gn_bwd_cuda_single_shape<__nv_bfloat16, HW, C, 32, false>(GN_BWD_CUDA_HOST_PARAMS(__nv_bfloat16));\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_1024_1280.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(1024, 1280)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_1024_1920.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(1024, 1920)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_1024_320.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(1024, 320)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_1024_640.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(1024, 640)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_1024_960.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(1024, 960)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_256_1280.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(256, 1280)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_256_1920.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(256, 1920)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_256_2560.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(256, 2560)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_256_640.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(256, 640)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_4096_320.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(4096, 320)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_4096_640.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(4096, 640)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_4096_960.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(4096, 960)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_64_1280.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(64, 1280)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_inst_64_2560.cu", "content": "#include \"gn_cuda_host_template.cuh\"\n\nnamespace group_norm_v2 {\n\nGN_CUDA_INST_DEFINE(64, 2560)\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_cuda_kernel.cuh", "content": "#pragma once\n\n#include \n\n#include \"gn_utils.hpp\"\n\nnamespace group_norm_v2 {\n\nnamespace cg = cooperative_groups;\n\ntemplate \ninline constexpr T up_div(T a, T b) {\n return (a + b - 1) / b;\n}\n\ntemplate \ninline constexpr T round_up(T a, T b) {\n return up_div(a, b) * b;\n}\n\ninline constexpr unsigned round_up_pow2(unsigned x) {\n int log = 0;\n x--;\n while (x) {\n x /= 2;\n log++;\n }\n return 1U << log;\n}\n\ninline constexpr unsigned round_down_pow2(unsigned x) { return round_up_pow2(x + 1) / 2; }\n\ntemplate \ninline constexpr T gcd(T a, T b) {\n while (b != 0) {\n int t = b;\n b = a % b;\n a = t;\n }\n return a;\n}\n\ntemplate \ninline constexpr T lcm(T a, T b) {\n return (a * b) / gcd(a, b);\n}\n\ntemplate \ninline constexpr T relative_prime(T x, T min) {\n int p = min;\n while (gcd(p, x) != 1) {\n p++;\n }\n return p;\n}\n\ntemplate \ninline constexpr T max_divisor(T x, T max) {\n int p = max;\n while (x % p != 0) {\n p--;\n }\n return p;\n}\n\nconstexpr unsigned FINAL_MASK = 0xffffffff;\n\ntemplate \n__device__ void virtual_cluster_sync(unsigned int* barrier) {\n if constexpr (VIRTUAL_CLUSTER_SIZE == 1) {\n __syncthreads();\n } else if constexpr (HARDWARE_CLUSTER) {\n cg::this_cluster().sync();\n } else {\n static_assert(PERSISTENT, \"potential deadlock\");\n volatile unsigned int* arrived = &barrier[blockIdx.y];\n __syncthreads();\n if (threadIdx.x == 0) {\n unsigned int expected = VIRTUAL_CLUSTER_SIZE;\n bool gpu_master = blockIdx.x == 0;\n unsigned int nb = 1;\n if (gpu_master) {\n nb = 0x80000000 - (expected - 1);\n }\n unsigned int oldArrive;\n asm volatile(\"atom.add.release.gpu.u32 %0,[%1],%2;\"\n : \"=r\"(oldArrive)\n : _CG_ASM_PTR_CONSTRAINT((unsigned int*)arrived), \"r\"(nb)\n : \"memory\");\n unsigned int current_arrive;\n do {\n asm volatile(\"ld.acquire.gpu.u32 %0,[%1];\"\n : \"=r\"(current_arrive)\n : _CG_ASM_PTR_CONSTRAINT((unsigned int*)arrived)\n : \"memory\");\n } while (!cooperative_groups::details::bar_has_flipped(oldArrive, current_arrive));\n }\n __syncthreads();\n }\n}\n\ntemplate \n__device__ unsigned int group_barrier_arrive(unsigned int* barrier, bool gpu_master) {\n static_assert(PERSISTENT, \"potential deadlock\");\n volatile unsigned int* arrived = &barrier[0];\n __syncthreads();\n if (threadIdx.x == 0) {\n unsigned int expected = NUM_BLOCKS;\n unsigned int nb = 1;\n if (gpu_master) {\n nb = 0x80000000 - (expected - 1);\n }\n unsigned int oldArrive;\n asm volatile(\"atom.add.release.gpu.u32 %0,[%1],%2;\"\n : \"=r\"(oldArrive)\n : _CG_ASM_PTR_CONSTRAINT((unsigned int*)arrived), \"r\"(nb)\n : \"memory\");\n return oldArrive;\n } else {\n return 0;\n }\n}\n\n__device__ inline void group_barrier_wait(unsigned int* barrier, unsigned int oldArrive) {\n volatile unsigned int* arrived = &barrier[0];\n if (threadIdx.x == 0) {\n unsigned int current_arrive;\n do {\n asm volatile(\"ld.acquire.gpu.u32 %0,[%1];\"\n : \"=r\"(current_arrive)\n : _CG_ASM_PTR_CONSTRAINT((unsigned int*)arrived)\n : \"memory\");\n } while (!cooperative_groups::details::bar_has_flipped(oldArrive, current_arrive));\n }\n __syncthreads();\n}\n\n// Calculate `n` (batch id) and `c` (channel range id) for each loop\ntemplate \nclass NCScheduler;\n\ntemplate \nclass NCScheduler {\n public:\n __device__ NCScheduler(int64_t n) {\n nc_loop_ = blockIdx.y;\n at_end_ = nc_loop_ >= n * (C / C_PER_CLUSTER);\n }\n __device__ auto get_nc() {\n int64_t n_loop = nc_loop_ / (C / C_PER_CLUSTER);\n int c_loop = nc_loop_ % (C / C_PER_CLUSTER);\n return std::make_tuple(n_loop, c_loop);\n }\n __device__ void next(int64_t n) {\n if constexpr (PERSISTENT) {\n nc_loop_ += NUM_VIRTUAL_CLUSTERS;\n at_end_ = nc_loop_ >= n * (C / C_PER_CLUSTER);\n }\n }\n __device__ bool at_end(int64_t n) { return !PERSISTENT || at_end_; }\n\n private:\n int64_t nc_loop_;\n bool at_end_;\n};\n\ntemplate \nclass NCScheduler {\n public:\n __device__ NCScheduler(int64_t n) {\n n_loop_ = blockIdx.y / (C / C_PER_CLUSTER);\n c_loop_ = blockIdx.y % (C / C_PER_CLUSTER);\n }\n __device__ auto get_nc() { return std::make_tuple(n_loop_, c_loop_); }\n __device__ void next(int64_t n) {\n if constexpr (PERSISTENT) {\n n_loop_ += NUM_VIRTUAL_CLUSTERS / (C / C_PER_CLUSTER);\n }\n }\n __device__ bool at_end(int64_t n) { return !PERSISTENT || n_loop_ >= n; }\n\n private:\n int64_t n_loop_;\n int c_loop_;\n};\n\nclass CompileConditionAlwaysTrue {\n public:\n __device__ static constexpr bool matches() { return true; }\n};\n\ntemplate \n__global__ __launch_bounds__(BLOCK_DIM_X, BLOCKS_PER_SM) void gn_cuda_kernel(\n T* __restrict__ out, T const* __restrict__ x, T const* __restrict__ w, T const* __restrict__ b, float eps,\n int64_t n, float* __restrict__ mean_var_out, float* __restrict__ red_buffer, unsigned* __restrict__ barrier) {\n // Procedure Overview\n // 1. Thread sum: read from gmem, write partial sum to smem, store input in registers (if no LOAD_TWICE)\n // 2. Block sum: read from smem, write partial sum to gmem (or distributed shared memory if HARDWARE_CLUSTER is\n // used)\n // 3. Group sum: read from gmem, write mean&var to smem\n // 4. Scale: read mean&var from smem, read input from gmem (if LOAD_TWICE), write output to gmem\n\n static_assert(BLOCK_DIM_X % 32 == 0, \"warp shuffle error\");\n\n constexpr int C = G * CPG;\n static_assert(C % C_PER_CLUSTER == 0, \"cannot divide channels into clusters\");\n static_assert(C_PER_CLUSTER % C_PER_BLOCK == 0, \"cannot divide a cluster into blocks\");\n static_assert(C_PER_CLUSTER % CPG == 0, \"no reduce between clusters, would produce incorrect results\");\n static_assert(!(C_PER_BLOCK % CPG == 0 && C_PER_CLUSTER != C_PER_BLOCK),\n \"inefficient configuration, please reduce C_PER_CLUSTER\");\n\n static_assert(ROWS_PER_BLOCK * C_PER_BLOCK % BLOCK_DIM_X == 0, \"cannot divide tile into threads\");\n struct alignas(VEC_ELEMS * sizeof(T)) U {\n T data[VEC_ELEMS];\n };\n\n auto compute_mean_var = [&](float2 sum) {\n float mean = sum.x / (HW * CPG);\n float var = std::max(0.f, sum.y / (HW * CPG) - mean * mean);\n return float2{mean, var};\n };\n\n static_assert(HW % ROWS_PER_BLOCK == 0,\n \"HW must be divisible by ROWS_PER_BLOCK to determine the number of blocks on the HW axis\");\n constexpr int MAX_NUM_GROUPS_PER_BLOCK =\n C_PER_BLOCK % CPG == 0 ? C_PER_BLOCK / CPG : up_div(C_PER_BLOCK - gcd(C_PER_BLOCK, CPG), CPG) + 1;\n constexpr int VIRTUAL_CLUSTER_SIZE = (C_PER_CLUSTER / C_PER_BLOCK) * (HW / ROWS_PER_BLOCK);\n constexpr int virtual_cluster_dim_x = C_PER_CLUSTER / C_PER_BLOCK;\n constexpr int virtual_cluster_dim_y = HW / ROWS_PER_BLOCK;\n int virtual_block_idx_x = (blockIdx.x % VIRTUAL_CLUSTER_SIZE) % virtual_cluster_dim_x;\n int virtual_block_idx_y = (blockIdx.x % VIRTUAL_CLUSTER_SIZE) / virtual_cluster_dim_x;\n\n if constexpr (CompileCondition::matches()) {\n int step = 0;\n constexpr bool CONSTANT_C_LOOP = PERSISTENT && NUM_VIRTUAL_CLUSTERS % (C / C_PER_CLUSTER) == 0;\n NCScheduler nc_scheduler(n);\n while (true) { // TODO: unroll the loop\n if constexpr (PERSISTENT) {\n if (nc_scheduler.at_end(n)) {\n break;\n }\n }\n auto [n_loop, c_loop] = nc_scheduler.get_nc();\n if constexpr (PERSISTENT) {\n nc_scheduler.next(n);\n }\n static_assert(C_PER_BLOCK % VEC_ELEMS == 0, \"cannot vectorize\");\n static_assert((BLOCK_DIM_X * VEC_ELEMS) % C_PER_BLOCK == 0,\n \"each block should load one or more C_PER_BLOCK at once\");\n constexpr int ROWS_PER_IO = BLOCK_DIM_X * VEC_ELEMS / C_PER_BLOCK;\n static_assert(ROWS_PER_BLOCK % ROWS_PER_IO == 0, \"cannot determine the IO times per batch\");\n int block_channel_start = virtual_block_idx_x * C_PER_BLOCK + c_loop * C_PER_CLUSTER;\n int block_group_start = block_channel_start / CPG;\n int thread_channel_start = block_channel_start + threadIdx.x % (C_PER_BLOCK / VEC_ELEMS) * VEC_ELEMS;\n U frag[ROWS_PER_BLOCK / ROWS_PER_IO];\n\n // GCD_VEC_CPG is an important constant that determines how many channels can be merged in reduction computation\n // For example, VEC_ELEMS=4 and CPG=10, then GCD_VEC_CPG=2,\n // so we need to store only 2 sums on each thread, and compute only 2 mean&var for each thread.\n constexpr int GCD_VEC_CPG = gcd(VEC_ELEMS, CPG);\n\n // If each block handles only one group, run warpReduce and store the sum to `sum_per_channel_single_group`;\n // otherwise store (VEC_ELEMS / GCD_VEC_CPG) sums to `sum_per_channel_multi_group`, where `relative_prime` is used\n // for swizzle.\n constexpr bool SINGLE_GROUP_PER_BLOCK = CPG % C_PER_BLOCK == 0;\n [[maybe_unused]] __shared__ float2 sum_per_channel_single_group[BLOCK_DIM_X / 32];\n [[maybe_unused]] __shared__ float2 sum_per_channel_multi_group[C_PER_BLOCK / GCD_VEC_CPG][relative_prime(\n 128 / (int)sizeof(float2), ROWS_PER_IO)];\n\n if constexpr (LOAD_TWICE) {\n float2 frag_sum_per_channel[VEC_ELEMS / GCD_VEC_CPG]{};\n for (int j = 0; j < ROWS_PER_BLOCK / ROWS_PER_IO; j++) {\n int64_t input_idx =\n n_loop * HW * C +\n (virtual_block_idx_y * ROWS_PER_BLOCK + j * ROWS_PER_IO + threadIdx.x / (C_PER_BLOCK / VEC_ELEMS)) * C +\n thread_channel_start;\n U val = *reinterpret_cast(&x[input_idx]);\n for (int i = 0; i < VEC_ELEMS / GCD_VEC_CPG; i++) {\n float2 sum = frag_sum_per_channel[i];\n for (int k = 0; k < GCD_VEC_CPG; k++) {\n sum.x += (float)val.data[i * GCD_VEC_CPG + k];\n sum.y += (float)val.data[i * GCD_VEC_CPG + k] * (float)val.data[i * GCD_VEC_CPG + k];\n }\n frag_sum_per_channel[i] = sum;\n }\n }\n for (int i = 0; i < VEC_ELEMS / GCD_VEC_CPG; i++) {\n if constexpr (SINGLE_GROUP_PER_BLOCK) {\n for (int mask = 16; mask > 0; mask >>= 1) {\n frag_sum_per_channel[i].x += __shfl_xor_sync(FINAL_MASK, frag_sum_per_channel[i].x, mask, 32);\n frag_sum_per_channel[i].y += __shfl_xor_sync(FINAL_MASK, frag_sum_per_channel[i].y, mask, 32);\n }\n static_assert(VEC_ELEMS / GCD_VEC_CPG == 1, \"process only one element for each warp\");\n if (threadIdx.x % 32 == 0) {\n sum_per_channel_single_group[threadIdx.x / 32] = frag_sum_per_channel[i];\n }\n } else {\n sum_per_channel_multi_group[i * (C_PER_BLOCK / VEC_ELEMS) + threadIdx.x % (C_PER_BLOCK / VEC_ELEMS)]\n [threadIdx.x / (C_PER_BLOCK / VEC_ELEMS)] = frag_sum_per_channel[i];\n }\n }\n __syncthreads();\n } else {\n for (int j = 0; j < ROWS_PER_BLOCK / ROWS_PER_IO; j++) {\n int64_t input_idx =\n n_loop * HW * C +\n (virtual_block_idx_y * ROWS_PER_BLOCK + j * ROWS_PER_IO + threadIdx.x / (C_PER_BLOCK / VEC_ELEMS)) * C +\n thread_channel_start;\n frag[j] = *reinterpret_cast(&x[input_idx]);\n }\n\n for (int i = 0; i < VEC_ELEMS / GCD_VEC_CPG; i++) {\n float2 sum = {0.f, 0.f};\n for (int j = 0; j < ROWS_PER_BLOCK / ROWS_PER_IO; j++) {\n for (int k = 0; k < GCD_VEC_CPG; k++) {\n sum.x += (float)frag[j].data[i * GCD_VEC_CPG + k];\n sum.y += (float)frag[j].data[i * GCD_VEC_CPG + k] * (float)frag[j].data[i * GCD_VEC_CPG + k];\n }\n }\n if constexpr (SINGLE_GROUP_PER_BLOCK) {\n for (int mask = 16; mask > 0; mask >>= 1) {\n sum.x += __shfl_xor_sync(FINAL_MASK, sum.x, mask, 32);\n sum.y += __shfl_xor_sync(FINAL_MASK, sum.y, mask, 32);\n }\n static_assert(VEC_ELEMS / GCD_VEC_CPG == 1, \"process only one element for each warp\");\n if (threadIdx.x % 32 == 0) {\n sum_per_channel_single_group[threadIdx.x / 32] = sum;\n }\n } else {\n sum_per_channel_multi_group[i * (C_PER_BLOCK / VEC_ELEMS) + threadIdx.x % (C_PER_BLOCK / VEC_ELEMS)]\n [threadIdx.x / (C_PER_BLOCK / VEC_ELEMS)] = sum;\n }\n }\n __syncthreads();\n }\n\n U uw = *reinterpret_cast(&w[thread_channel_start]);\n U ub = *reinterpret_cast(&b[thread_channel_start]);\n\n // Three cases for the red_buffer:\n // - Block sync (VIRTUAL_CLUSTER_SIZE=1): use shared memory\n // - Virtual cluster sync with HARDWARE_CLUSTER: use distributed shared memory\n // - Virtual cluster sync without HARDWARE_CLUSTER: use global memory, i.e., `red_buffer`\n constexpr bool USE_SHARED_RED_BUFFER = HARDWARE_CLUSTER || VIRTUAL_CLUSTER_SIZE == 1;\n\n // Specialize for the case that each group is handled by only one block\n // For common cases, blockSum produces partial sum and stores it to the red_buffer, and groupSum produces\n // mean&var For the special case, blockSum produces mean&var directly\n constexpr bool STORE_MEAN_VAR_IN_SHARED_RED_BUFFER =\n VIRTUAL_CLUSTER_SIZE == 1 &&\n MAX_NUM_GROUPS_PER_BLOCK == 1; // MAX_NUM_GROUPS_PER_BLOCK > 1 is possible but not implemented\n\n [[maybe_unused]] __align__(16)\n __shared__ float2 shared_red_buffer[MAX_NUM_GROUPS_PER_BLOCK * (STORE_MEAN_VAR_IN_SHARED_RED_BUFFER ? 1 : 2)];\n\n // Block sum\n if constexpr (SINGLE_GROUP_PER_BLOCK) {\n // block reduce\n if (threadIdx.x < 32) {\n float2 sum_local_group =\n threadIdx.x < BLOCK_DIM_X / 32 ? sum_per_channel_single_group[threadIdx.x] : float2{0.f, 0.f};\n constexpr int warp_num_pow2 = round_up_pow2(BLOCK_DIM_X / 32);\n for (int mask = warp_num_pow2 / 2; mask > 0; mask >>= 1) {\n sum_local_group.x += __shfl_xor_sync(FINAL_MASK, sum_local_group.x, mask, 32);\n sum_local_group.y += __shfl_xor_sync(FINAL_MASK, sum_local_group.y, mask, 32);\n }\n if (threadIdx.x == 0) {\n if constexpr (USE_SHARED_RED_BUFFER) {\n if constexpr (STORE_MEAN_VAR_IN_SHARED_RED_BUFFER) {\n shared_red_buffer[0] = compute_mean_var(sum_local_group);\n } else {\n shared_red_buffer[step * MAX_NUM_GROUPS_PER_BLOCK + 0] = sum_local_group;\n }\n } else {\n *reinterpret_cast(\n &red_buffer[((step * gridDim.y + blockIdx.y) * VIRTUAL_CLUSTER_SIZE * MAX_NUM_GROUPS_PER_BLOCK +\n virtual_block_idx_x * virtual_cluster_dim_y * MAX_NUM_GROUPS_PER_BLOCK +\n // (threadIdx.x / THREADS_PER_GROUP) * virtual_cluster_dim_y +\n virtual_block_idx_y) *\n 2]) = sum_local_group;\n }\n }\n }\n } else {\n // The number of threads to calculate the sum of each group (should be a power of 2 for warp reduce)\n constexpr int THREADS_PER_GROUP = std::min(std::min(32U, round_up_pow2(ROWS_PER_IO)),\n round_up_pow2(BLOCK_DIM_X / MAX_NUM_GROUPS_PER_BLOCK / 2 + 1));\n static_assert(BLOCK_DIM_X >= MAX_NUM_GROUPS_PER_BLOCK * THREADS_PER_GROUP, \"not enough threads\");\n float2 sum_local_group = {0.f, 0.f};\n if (threadIdx.x / THREADS_PER_GROUP < MAX_NUM_GROUPS_PER_BLOCK) {\n int local_group_idx = block_group_start + threadIdx.x / THREADS_PER_GROUP;\n // TODO: map threads to both the CPG loop and the ROWS loop\n for (int local_c_loop = 0; local_c_loop < CPG; local_c_loop += GCD_VEC_CPG) {\n int c = local_group_idx * CPG + local_c_loop;\n if (C_PER_BLOCK % CPG == 0 || (c >= block_channel_start && c < block_channel_start + C_PER_BLOCK)) {\n for (int src_thread_tile_y = threadIdx.x % THREADS_PER_GROUP; src_thread_tile_y < ROWS_PER_IO;\n src_thread_tile_y += THREADS_PER_GROUP) {\n int channel_idx = (c - block_channel_start) / GCD_VEC_CPG;\n channel_idx = channel_idx % (VEC_ELEMS / GCD_VEC_CPG) * (C_PER_BLOCK / VEC_ELEMS) +\n channel_idx / (VEC_ELEMS / GCD_VEC_CPG);\n sum_local_group.x += sum_per_channel_multi_group[channel_idx][src_thread_tile_y].x;\n sum_local_group.y += sum_per_channel_multi_group[channel_idx][src_thread_tile_y].y;\n }\n }\n }\n }\n static_assert(32 % THREADS_PER_GROUP == 0, \"cannot shuffle\");\n for (int mask = THREADS_PER_GROUP / 2; mask > 0; mask >>= 1) {\n sum_local_group.x += __shfl_xor_sync(FINAL_MASK, sum_local_group.x, mask, 32);\n sum_local_group.y += __shfl_xor_sync(FINAL_MASK, sum_local_group.y, mask, 32);\n }\n if (threadIdx.x % THREADS_PER_GROUP == 0 && threadIdx.x / THREADS_PER_GROUP < MAX_NUM_GROUPS_PER_BLOCK) {\n if constexpr (USE_SHARED_RED_BUFFER) {\n static_assert(HARDWARE_CLUSTER || VIRTUAL_CLUSTER_SIZE == 1, \"no distributed shared memory\");\n if constexpr (STORE_MEAN_VAR_IN_SHARED_RED_BUFFER) {\n shared_red_buffer[threadIdx.x / THREADS_PER_GROUP] = compute_mean_var(sum_local_group);\n } else {\n shared_red_buffer[step * MAX_NUM_GROUPS_PER_BLOCK + threadIdx.x / THREADS_PER_GROUP] = sum_local_group;\n }\n } else {\n *reinterpret_cast(\n &red_buffer[((step * gridDim.y + blockIdx.y) * VIRTUAL_CLUSTER_SIZE * MAX_NUM_GROUPS_PER_BLOCK +\n virtual_block_idx_x * virtual_cluster_dim_y * MAX_NUM_GROUPS_PER_BLOCK +\n (threadIdx.x / THREADS_PER_GROUP) * virtual_cluster_dim_y + virtual_block_idx_y) *\n 2]) = sum_local_group;\n }\n }\n }\n\n virtual_cluster_sync(barrier);\n\n // Group sum\n __shared__ float2 mean_var[MAX_NUM_GROUPS_PER_BLOCK];\n if constexpr (!STORE_MEAN_VAR_IN_SHARED_RED_BUFFER) {\n // The number of threads to calculate the sum of each group (should be a power of 2 for warp reduce)\n constexpr int THREADS_PER_GROUP = std::min(std::min(32U, round_up_pow2(virtual_cluster_dim_y)),\n round_up_pow2(BLOCK_DIM_X / MAX_NUM_GROUPS_PER_BLOCK / 2 + 1));\n static_assert(BLOCK_DIM_X >= MAX_NUM_GROUPS_PER_BLOCK * THREADS_PER_GROUP, \"not enough threads\");\n float2 sum_global_group = {0.f, 0.f};\n if (threadIdx.x / THREADS_PER_GROUP < MAX_NUM_GROUPS_PER_BLOCK) {\n if constexpr (C_PER_BLOCK % CPG == 0) {\n // Special case: no cross-virtual_cluster_dim_x reduction\n float2 buffer[up_div(virtual_cluster_dim_y, THREADS_PER_GROUP)];\n for (int i = threadIdx.x % THREADS_PER_GROUP; i < virtual_cluster_dim_y; i += THREADS_PER_GROUP) {\n float2 val;\n if constexpr (USE_SHARED_RED_BUFFER) {\n if constexpr (VIRTUAL_CLUSTER_SIZE == 1) {\n val = shared_red_buffer[step * MAX_NUM_GROUPS_PER_BLOCK + threadIdx.x / THREADS_PER_GROUP];\n } else {\n static_assert(HARDWARE_CLUSTER, \"no distributed shared memory\");\n float2 const* src_shared_red_buffer = cg::this_cluster().map_shared_rank(\n shared_red_buffer, i * virtual_cluster_dim_x + virtual_block_idx_x);\n val = src_shared_red_buffer[step * MAX_NUM_GROUPS_PER_BLOCK + threadIdx.x / THREADS_PER_GROUP];\n }\n } else {\n val = *reinterpret_cast(\n &red_buffer[((step * gridDim.y + blockIdx.y) * VIRTUAL_CLUSTER_SIZE * MAX_NUM_GROUPS_PER_BLOCK +\n virtual_block_idx_x * virtual_cluster_dim_y * MAX_NUM_GROUPS_PER_BLOCK +\n (threadIdx.x / THREADS_PER_GROUP) * virtual_cluster_dim_y + i) *\n 2]);\n }\n buffer[i / THREADS_PER_GROUP] = val;\n }\n for (int i = threadIdx.x % THREADS_PER_GROUP; i < virtual_cluster_dim_y; i += THREADS_PER_GROUP) {\n float2 val = buffer[i / THREADS_PER_GROUP];\n sum_global_group.x += val.x;\n sum_global_group.y += val.y;\n }\n } else {\n // Common case: cross-virtual_cluster_dim_x reduction\n int local_group_idx = block_group_start + threadIdx.x / THREADS_PER_GROUP;\n for (int i = threadIdx.x % THREADS_PER_GROUP; i < VIRTUAL_CLUSTER_SIZE; i += THREADS_PER_GROUP) {\n int src_virtual_block_idx_x = i % virtual_cluster_dim_x;\n int src_block_channel_start = src_virtual_block_idx_x * C_PER_BLOCK + c_loop * C_PER_CLUSTER;\n int src_block_group_start = src_block_channel_start / CPG;\n int relative_group_idx = local_group_idx - src_block_group_start;\n if (0 <= relative_group_idx && relative_group_idx < MAX_NUM_GROUPS_PER_BLOCK) {\n float2 val;\n if constexpr (USE_SHARED_RED_BUFFER) {\n static_assert(HARDWARE_CLUSTER, \"no distributed shared memory\");\n static_assert(VIRTUAL_CLUSTER_SIZE != 1,\n \"layout error: should not add (step * MAX_NUM_GROUPS_PER_BLOCK)\");\n float2 const* src_shared_red_buffer = cg::this_cluster().map_shared_rank(shared_red_buffer, i);\n val = src_shared_red_buffer[step * MAX_NUM_GROUPS_PER_BLOCK + relative_group_idx];\n } else {\n val = *reinterpret_cast(\n &red_buffer[((step * gridDim.y + blockIdx.y) * VIRTUAL_CLUSTER_SIZE * MAX_NUM_GROUPS_PER_BLOCK +\n src_virtual_block_idx_x * virtual_cluster_dim_y * MAX_NUM_GROUPS_PER_BLOCK +\n relative_group_idx * virtual_cluster_dim_y + i / virtual_cluster_dim_x) *\n 2]);\n }\n sum_global_group.x += val.x;\n sum_global_group.y += val.y;\n }\n }\n }\n }\n if constexpr (USE_SHARED_RED_BUFFER && VIRTUAL_CLUSTER_SIZE > 1) {\n // Need cluster sync after distributed shared memory access, otherwise behavior is undefined\n if constexpr (PERSISTENT) {\n if (nc_scheduler.at_end(n)) {\n cg::this_cluster().barrier_arrive();\n }\n } else {\n cg::this_cluster().barrier_arrive();\n }\n }\n static_assert(32 % THREADS_PER_GROUP == 0, \"cannot shuffle\");\n for (int mask = THREADS_PER_GROUP / 2; mask > 0; mask >>= 1) {\n sum_global_group.x += __shfl_xor_sync(FINAL_MASK, sum_global_group.x, mask, 32);\n sum_global_group.y += __shfl_xor_sync(FINAL_MASK, sum_global_group.y, mask, 32);\n }\n if (threadIdx.x % THREADS_PER_GROUP == 0 && threadIdx.x / THREADS_PER_GROUP < MAX_NUM_GROUPS_PER_BLOCK) {\n mean_var[threadIdx.x / THREADS_PER_GROUP] = compute_mean_var(sum_global_group);\n }\n __syncthreads();\n }\n\n auto get_mean_var = [&](int relative_group_idx) {\n return STORE_MEAN_VAR_IN_SHARED_RED_BUFFER ? shared_red_buffer[relative_group_idx]\n : mean_var[relative_group_idx];\n };\n\n if (mean_var_out) {\n static_assert(MAX_NUM_GROUPS_PER_BLOCK <= BLOCK_DIM_X, \"need loop\");\n if (virtual_block_idx_y == 0 && threadIdx.x < MAX_NUM_GROUPS_PER_BLOCK) {\n int g = block_group_start + threadIdx.x;\n if (C_PER_BLOCK % CPG == 0 || g < G) {\n *reinterpret_cast(&mean_var_out[(n_loop * G + g) * 2]) = get_mean_var(threadIdx.x);\n }\n }\n }\n\n float frag_mean[VEC_ELEMS / GCD_VEC_CPG];\n float frag_var[VEC_ELEMS / GCD_VEC_CPG];\n for (int k = 0; k < VEC_ELEMS; k += GCD_VEC_CPG) {\n frag_mean[k / GCD_VEC_CPG] = get_mean_var((thread_channel_start + k) / CPG - block_group_start).x;\n frag_var[k / GCD_VEC_CPG] = get_mean_var((thread_channel_start + k) / CPG - block_group_start).y;\n }\n\n for (int j = 0; j < ROWS_PER_BLOCK / ROWS_PER_IO; j++) {\n int64_t input_idx =\n n_loop * HW * C +\n (virtual_block_idx_y * ROWS_PER_BLOCK + j * ROWS_PER_IO + threadIdx.x / (C_PER_BLOCK / VEC_ELEMS)) * C +\n thread_channel_start;\n U val;\n if constexpr (LOAD_TWICE) {\n val = *reinterpret_cast(&x[input_idx]);\n } else {\n val = frag[j];\n }\n for (int k = 0; k < VEC_ELEMS; k++) {\n float f = ((float)val.data[k] - frag_mean[k / GCD_VEC_CPG]) * rsqrtf(frag_var[k / GCD_VEC_CPG] + eps) *\n (float)uw.data[k] +\n (float)ub.data[k];\n if constexpr (SILU) f = f / (1.f + expf(-f));\n val.data[k] = f;\n }\n *reinterpret_cast(&out[input_idx]) = val;\n }\n\n if constexpr (!STORE_MEAN_VAR_IN_SHARED_RED_BUFFER && USE_SHARED_RED_BUFFER && VIRTUAL_CLUSTER_SIZE > 1) {\n if constexpr (PERSISTENT) {\n if (nc_scheduler.at_end(n)) {\n cg::this_cluster().barrier_wait();\n }\n } else {\n cg::this_cluster().barrier_wait();\n }\n }\n\n if constexpr (!PERSISTENT) {\n break;\n }\n step ^= 1;\n }\n }\n}\n\nenum WgradSyncMethod {\n WGRAD_ARRIVE_AND_WAIT_GRID = 0, // grid arrive after the last virtual cluster sync\n WGRAD_ARRIVE_AND_WAIT_GROUP, // group arrive after the last virtual cluster sync (a group sync means synchronizing\n // all clusters cooperating on the same groups)\n WGRAD_REUSE_SUM_SYNC_GRID, // grid sync together with the last virtual cluster sync\n WGRAD_REUSE_SUM_SYNC_GROUP, // group sync together with the last virtual cluster sync\n WGRAD_SYNC_AT_LAST, // add a sync at the end of NC loops\n WGRAD_SYNC_UNSPECIFIED,\n};\n\ntemplate \n__global__ __launch_bounds__(BLOCK_DIM_X, BLOCKS_PER_SM) void gn_bwd_cuda_kernel(\n T* __restrict__ grad_input, T* __restrict__ grad_weight, T* __restrict__ grad_bias,\n T const* __restrict__ grad_output, T const* __restrict__ x, T const* __restrict__ w, T const* __restrict__ b,\n float const* __restrict__ mean_var, float eps, int64_t n, float* __restrict__ red_buffer,\n unsigned* __restrict__ barrier) {\n // Procedure Overview\n // 1. Thread sum: read from gmem, write partial sum to smem, store input in registers (if no LOAD_TWICE)\n // 2. Block sum: read from smem, write partial sum to gmem (or distributed shared memory if HARDWARE_CLUSTER is\n // used),\n // write wgrad to gmem at the last loop (at each loop if not CONSTANT_C_LOOP)\n // 3. Group sum: read from gmem, write mean&var to smem\n // 4. Scale: read mean&var from smem, read input from gmem (if LOAD_TWICE), write output to gmem\n // 5. Wgrad sum: read from gmem, write to gmem\n\n static_assert(BLOCK_DIM_X % 32 == 0, \"warp shuffle error\");\n\n constexpr int C = G * CPG;\n static_assert(C % C_PER_CLUSTER == 0, \"cannot divide channels into clusters\");\n static_assert(C_PER_CLUSTER % C_PER_BLOCK == 0, \"cannot divide a cluster into blocks\");\n static_assert(C_PER_CLUSTER % CPG == 0, \"no reduce between clusters, would produce incorrect results\");\n static_assert(!(C_PER_BLOCK % CPG == 0 && C_PER_CLUSTER != C_PER_BLOCK),\n \"inefficient configuration, please reduce C_PER_CLUSTER\");\n\n static_assert(ROWS_PER_BLOCK * C_PER_BLOCK % BLOCK_DIM_X == 0, \"cannot divide tile into threads\");\n struct alignas(VEC_ELEMS * sizeof(T)) U {\n T data[VEC_ELEMS];\n };\n\n // This function computes mean_dyw and mean_xdyw.\n // The function name is not changed because it has the same logic as the forward pass.\n auto compute_mean_var = [&](float2 sum) {\n float mean_dyw = sum.x / (HW * CPG);\n float mean_xdyw = sum.y / (HW * CPG);\n return float2{mean_dyw, mean_xdyw};\n };\n\n static_assert(HW % ROWS_PER_BLOCK == 0,\n \"HW must be divisible by ROWS_PER_BLOCK to determine the number of blocks on the HW axis\");\n constexpr int MAX_NUM_GROUPS_PER_BLOCK =\n C_PER_BLOCK % CPG == 0 ? C_PER_BLOCK / CPG : up_div(C_PER_BLOCK - gcd(C_PER_BLOCK, CPG), CPG) + 1;\n constexpr int VIRTUAL_CLUSTER_SIZE = (C_PER_CLUSTER / C_PER_BLOCK) * (HW / ROWS_PER_BLOCK);\n constexpr int virtual_cluster_dim_x = C_PER_CLUSTER / C_PER_BLOCK;\n constexpr int virtual_cluster_dim_y = HW / ROWS_PER_BLOCK;\n int virtual_block_idx_x = (blockIdx.x % VIRTUAL_CLUSTER_SIZE) % virtual_cluster_dim_x;\n int virtual_block_idx_y = (blockIdx.x % VIRTUAL_CLUSTER_SIZE) / virtual_cluster_dim_x;\n\n if constexpr (CompileCondition::matches()) {\n int step = 0;\n constexpr bool CONSTANT_C_LOOP = PERSISTENT && NUM_VIRTUAL_CLUSTERS % (C / C_PER_CLUSTER) == 0;\n if constexpr (!CONSTANT_C_LOOP) {\n static_assert(wgrad_sync_method != WGRAD_ARRIVE_AND_WAIT_GROUP && wgrad_sync_method != WGRAD_REUSE_SUM_SYNC_GROUP,\n \"grid sync is required when each block is responsible for multiple channel ranges\");\n }\n NCScheduler nc_scheduler(\n n); // TODO: I don't know why the template specialization with CONSTANT_C_LOOP=true is slower.\n\n [[maybe_unused]] int virtual_cluster_idx_c = blockIdx.y % (C / C_PER_CLUSTER);\n [[maybe_unused]] cg::grid_group::arrival_token wgrad_sync_token;\n [[maybe_unused]] float dw_thread[VEC_ELEMS];\n [[maybe_unused]] float db_thread[VEC_ELEMS];\n [[maybe_unused]] __shared__ union {\n float2 dwdb_block_buffer[BLOCK_DIM_X][VEC_ELEMS];\n struct {\n float wgrad_buffer[BLOCK_DIM_X / 32][32];\n float bgrad_buffer[BLOCK_DIM_X / 32][32];\n } transpose_buffer;\n } union_smem;\n if constexpr (REQUIRES_WGRAD && CONSTANT_C_LOOP) {\n for (int i = 0; i < VEC_ELEMS; i++) {\n dw_thread[i] = 0.f;\n db_thread[i] = 0.f;\n }\n }\n float* red_buffer_wgrad =\n &red_buffer[(2 * NUM_VIRTUAL_CLUSTERS * VIRTUAL_CLUSTER_SIZE * MAX_NUM_GROUPS_PER_BLOCK) * 2];\n unsigned* barrier_wgrad = barrier + NUM_VIRTUAL_CLUSTERS;\n if constexpr (REQUIRES_WGRAD && wgrad_sync_method != WGRAD_SYNC_AT_LAST) {\n if (nc_scheduler.at_end(n)) {\n static_assert(PERSISTENT, \"persistent is a must for reducing wgrad\");\n if constexpr (wgrad_sync_method == WGRAD_ARRIVE_AND_WAIT_GRID) {\n wgrad_sync_token = group_barrier_arrive(\n barrier_wgrad, blockIdx.x + blockIdx.y == 0);\n } else if constexpr (wgrad_sync_method == WGRAD_ARRIVE_AND_WAIT_GROUP) {\n wgrad_sync_token =\n group_barrier_arrive(\n barrier_wgrad + virtual_cluster_idx_c, blockIdx.x + blockIdx.y / (C / C_PER_CLUSTER) == 0);\n } else if constexpr (wgrad_sync_method == WGRAD_REUSE_SUM_SYNC_GRID) {\n wgrad_sync_token = group_barrier_arrive(\n barrier_wgrad, blockIdx.x + blockIdx.y == 0);\n group_barrier_wait(barrier_wgrad, wgrad_sync_token);\n } else if constexpr (wgrad_sync_method == WGRAD_REUSE_SUM_SYNC_GROUP) {\n wgrad_sync_token =\n group_barrier_arrive(\n barrier_wgrad + virtual_cluster_idx_c, blockIdx.x + blockIdx.y / (C / C_PER_CLUSTER) == 0);\n group_barrier_wait(barrier_wgrad + virtual_cluster_idx_c, wgrad_sync_token);\n }\n }\n }\n\n while (true) { // TODO: unroll the loop\n if constexpr (PERSISTENT) {\n if (nc_scheduler.at_end(n)) {\n break;\n }\n }\n auto [n_loop, c_loop] = nc_scheduler.get_nc();\n if constexpr (PERSISTENT) {\n nc_scheduler.next(n);\n }\n static_assert(C_PER_BLOCK % VEC_ELEMS == 0, \"cannot vectorize\");\n static_assert((BLOCK_DIM_X * VEC_ELEMS) % C_PER_BLOCK == 0,\n \"each block should load one or more C_PER_BLOCK at once\");\n constexpr int ROWS_PER_IO = BLOCK_DIM_X * VEC_ELEMS / C_PER_BLOCK;\n static_assert(ROWS_PER_BLOCK % ROWS_PER_IO == 0, \"cannot determine the IO times per batch\");\n int block_channel_start = virtual_block_idx_x * C_PER_BLOCK + c_loop * C_PER_CLUSTER;\n int block_group_start = block_channel_start / CPG;\n int thread_channel_start = block_channel_start + threadIdx.x % (C_PER_BLOCK / VEC_ELEMS) * VEC_ELEMS;\n U frag_x[ROWS_PER_BLOCK / ROWS_PER_IO];\n U frag_dy[ROWS_PER_BLOCK / ROWS_PER_IO];\n\n constexpr int GCD_VEC_CPG = gcd(VEC_ELEMS, CPG);\n\n constexpr bool SINGLE_GROUP_PER_BLOCK = CPG % C_PER_BLOCK == 0;\n [[maybe_unused]] __shared__ float2 sum_per_channel_multi_group[C_PER_BLOCK / GCD_VEC_CPG][relative_prime(\n 128 / (int)sizeof(float2), ROWS_PER_IO)];\n [[maybe_unused]] __shared__ float2 sum_per_channel_single_group[BLOCK_DIM_X / 32];\n\n float frag_mean[VEC_ELEMS / GCD_VEC_CPG];\n float frag_var[VEC_ELEMS / GCD_VEC_CPG];\n for (int k = 0; k < VEC_ELEMS; k += GCD_VEC_CPG) {\n float2 value = *reinterpret_cast(&mean_var[(n_loop * G + (thread_channel_start + k) / CPG) * 2]);\n frag_mean[k / GCD_VEC_CPG] = value.x;\n frag_var[k / GCD_VEC_CPG] = value.y;\n }\n\n U uw = *reinterpret_cast(&w[thread_channel_start]);\n U ub;\n if constexpr (SILU) {\n ub = *reinterpret_cast(&b[thread_channel_start]);\n }\n if constexpr (REQUIRES_WGRAD && !CONSTANT_C_LOOP) {\n for (int i = 0; i < VEC_ELEMS; i++) {\n dw_thread[i] = 0.f;\n db_thread[i] = 0.f;\n }\n }\n\n if constexpr (LOAD_TWICE) {\n float2 frag_sum_per_channel[VEC_ELEMS / GCD_VEC_CPG]{};\n for (int j = 0; j < ROWS_PER_BLOCK / ROWS_PER_IO; j++) {\n int64_t input_idx =\n n_loop * HW * C +\n (virtual_block_idx_y * ROWS_PER_BLOCK + j * ROWS_PER_IO + threadIdx.x / (C_PER_BLOCK / VEC_ELEMS)) * C +\n thread_channel_start;\n U ux = *reinterpret_cast(&x[input_idx]);\n U udy = *reinterpret_cast(&grad_output[input_idx]);\n for (int i = 0; i < VEC_ELEMS / GCD_VEC_CPG; i++) {\n float2 sum = frag_sum_per_channel[i];\n for (int k = 0; k < GCD_VEC_CPG; k++) {\n float rnorm = rsqrtf(frag_var[i] + eps);\n float x_norm =\n ((float)ux.data[i * GCD_VEC_CPG + k] - frag_mean[i]) * rnorm; // TODO: store rsqrtf in mean_var\n float grad_gn = udy.data[i * GCD_VEC_CPG + k];\n if constexpr (SILU) {\n float x_gn = x_norm * (float)uw.data[i * GCD_VEC_CPG + k] + (float)ub.data[i * GCD_VEC_CPG + k];\n float s = 1.f / (1.f + expf(-x_gn));\n grad_gn *= s * (1.f + x_gn * (1.f - s));\n }\n sum.x += grad_gn * (float)uw.data[i * GCD_VEC_CPG + k];\n sum.y += x_norm * (grad_gn * (float)uw.data[i * GCD_VEC_CPG + k]);\n if constexpr (REQUIRES_WGRAD) {\n dw_thread[i * GCD_VEC_CPG + k] += x_norm * grad_gn;\n db_thread[i * GCD_VEC_CPG + k] += grad_gn;\n }\n }\n frag_sum_per_channel[i] = sum;\n }\n }\n for (int i = 0; i < VEC_ELEMS / GCD_VEC_CPG; i++) {\n if constexpr (SINGLE_GROUP_PER_BLOCK) {\n for (int mask = 16; mask > 0; mask >>= 1) {\n frag_sum_per_channel[i].x += __shfl_xor_sync(FINAL_MASK, frag_sum_per_channel[i].x, mask, 32);\n frag_sum_per_channel[i].y += __shfl_xor_sync(FINAL_MASK, frag_sum_per_channel[i].y, mask, 32);\n }\n static_assert(VEC_ELEMS / GCD_VEC_CPG == 1, \"process only one element for each warp\");\n if (threadIdx.x % 32 == 0) {\n sum_per_channel_single_group[threadIdx.x / 32] = frag_sum_per_channel[i];\n }\n } else {\n sum_per_channel_multi_group[i * (C_PER_BLOCK / VEC_ELEMS) + threadIdx.x % (C_PER_BLOCK / VEC_ELEMS)]\n [threadIdx.x / (C_PER_BLOCK / VEC_ELEMS)] = frag_sum_per_channel[i];\n }\n }\n __syncthreads();\n } else {\n for (int j = 0; j < ROWS_PER_BLOCK / ROWS_PER_IO; j++) {\n int64_t input_idx =\n n_loop * HW * C +\n (virtual_block_idx_y * ROWS_PER_BLOCK + j * ROWS_PER_IO + threadIdx.x / (C_PER_BLOCK / VEC_ELEMS)) * C +\n thread_channel_start;\n frag_x[j] = *reinterpret_cast(&x[input_idx]);\n frag_dy[j] = *reinterpret_cast(&grad_output[input_idx]);\n }\n\n for (int i = 0; i < VEC_ELEMS / GCD_VEC_CPG; i++) {\n float2 sum = {0.f, 0.f};\n for (int j = 0; j < ROWS_PER_BLOCK / ROWS_PER_IO; j++) {\n for (int k = 0; k < GCD_VEC_CPG; k++) {\n float rnorm = rsqrtf(frag_var[i] + eps);\n float x_norm = ((float)frag_x[j].data[i * GCD_VEC_CPG + k] - frag_mean[i]) *\n rnorm; // TODO: store rsqrtf in mean_var\n float grad_gn = frag_dy[j].data[i * GCD_VEC_CPG + k];\n if constexpr (SILU) {\n float x_gn = x_norm * (float)uw.data[i * GCD_VEC_CPG + k] + (float)ub.data[i * GCD_VEC_CPG + k];\n float s = 1.f / (1.f + expf(-x_gn));\n grad_gn *= s * (1.f + x_gn * (1.f - s));\n }\n sum.x += grad_gn * (float)uw.data[i * GCD_VEC_CPG + k];\n sum.y += x_norm * (grad_gn * (float)uw.data[i * GCD_VEC_CPG + k]);\n if constexpr (REQUIRES_WGRAD) {\n dw_thread[i * GCD_VEC_CPG + k] += x_norm * grad_gn;\n db_thread[i * GCD_VEC_CPG + k] += grad_gn;\n }\n }\n }\n if constexpr (SINGLE_GROUP_PER_BLOCK) {\n for (int mask = 16; mask > 0; mask >>= 1) {\n sum.x += __shfl_xor_sync(FINAL_MASK, sum.x, mask, 32);\n sum.y += __shfl_xor_sync(FINAL_MASK, sum.y, mask, 32);\n }\n static_assert(VEC_ELEMS / GCD_VEC_CPG == 1, \"process only one element for each warp\");\n if (threadIdx.x % 32 == 0) {\n sum_per_channel_single_group[threadIdx.x / 32] = sum;\n }\n } else {\n sum_per_channel_multi_group[i * (C_PER_BLOCK / VEC_ELEMS) + threadIdx.x % (C_PER_BLOCK / VEC_ELEMS)]\n [threadIdx.x / (C_PER_BLOCK / VEC_ELEMS)] = sum;\n }\n }\n __syncthreads();\n }\n\n if ((CONSTANT_C_LOOP && nc_scheduler.at_end(n)) || !CONSTANT_C_LOOP) {\n constexpr int NT_C = max_divisor(C_PER_BLOCK, BLOCK_DIM_X); // Number of threads on the C axis\n constexpr int NT_R =\n 1; // std::min(32, (int)round_down_pow2(BLOCK_DIM_X / NT_C)); // Number of threads on the ROWS axis\n // TODO: swizzle for NT_R\n for (int i = 0; i < VEC_ELEMS; i++) {\n union_smem.dwdb_block_buffer[threadIdx.x][i ^ ((threadIdx.x / (16 / VEC_ELEMS)) & (VEC_ELEMS - 1))] =\n float2{dw_thread[i], db_thread[i]};\n }\n __syncthreads();\n static_assert(NT_C * NT_R <= BLOCK_DIM_X, \"not enough threads\");\n static_assert(C_PER_BLOCK % NT_C == 0, \"need to loop once more and check c < C_PER_BLOCK\");\n for (int i = 0; i < C_PER_BLOCK / NT_C; i++) {\n int c = i * NT_C + threadIdx.x / NT_R;\n float dw_block = 0.f;\n float db_block = 0.f;\n if (BLOCK_DIM_X == NT_C * NT_R || threadIdx.x < NT_C * NT_R) {\n for (int j = threadIdx.x % NT_R; j < ROWS_PER_IO; j += NT_R) {\n int src_thread = j * (C_PER_BLOCK / VEC_ELEMS) + c / VEC_ELEMS;\n float2 val = union_smem.dwdb_block_buffer[src_thread][(c % VEC_ELEMS) ^ ((src_thread / (16 / VEC_ELEMS)) &\n (VEC_ELEMS - 1))];\n dw_block += val.x;\n db_block += val.y;\n }\n }\n static_assert(32 % NT_R == 0, \"cannot shuffle\");\n for (int mask = NT_R / 2; mask > 0; mask >>= 1) {\n dw_block += __shfl_xor_sync(FINAL_MASK, dw_block, mask, 32);\n db_block += __shfl_xor_sync(FINAL_MASK, db_block, mask, 32);\n }\n if (BLOCK_DIM_X == NT_C * NT_R || threadIdx.x < NT_C * NT_R) {\n if (threadIdx.x % NT_R == 0) {\n if constexpr (CONSTANT_C_LOOP) {\n *reinterpret_cast(\n &red_buffer_wgrad\n [((blockIdx.y / (C / C_PER_CLUSTER) * virtual_cluster_dim_y + virtual_block_idx_y) * C +\n c_loop * C_PER_CLUSTER + virtual_block_idx_x * C_PER_BLOCK + c) *\n 2]) = float2{dw_block, db_block};\n } else {\n *reinterpret_cast(\n &red_buffer_wgrad[((n_loop * virtual_cluster_dim_y + virtual_block_idx_y) * C +\n c_loop * C_PER_CLUSTER + virtual_block_idx_x * C_PER_BLOCK + c) *\n 2]) = float2{dw_block, db_block};\n }\n }\n }\n }\n }\n\n constexpr bool USE_SHARED_RED_BUFFER = HARDWARE_CLUSTER || VIRTUAL_CLUSTER_SIZE == 1;\n constexpr bool STORE_MEAN_VAR_IN_SHARED_RED_BUFFER =\n VIRTUAL_CLUSTER_SIZE == 1 &&\n MAX_NUM_GROUPS_PER_BLOCK == 1; // MAX_NUM_GROUPS_PER_BLOCK > 1 is possible but not implemented\n [[maybe_unused]] __align__(16)\n __shared__ float2 shared_red_buffer[MAX_NUM_GROUPS_PER_BLOCK * (STORE_MEAN_VAR_IN_SHARED_RED_BUFFER ? 1 : 2)];\n\n // Block sum\n if constexpr (SINGLE_GROUP_PER_BLOCK) {\n // block reduce\n if (threadIdx.x < 32) {\n float2 sum_local_group =\n threadIdx.x < BLOCK_DIM_X / 32 ? sum_per_channel_single_group[threadIdx.x] : float2{0.f, 0.f};\n constexpr int warp_num_pow2 = round_up_pow2(BLOCK_DIM_X / 32);\n for (int mask = warp_num_pow2 / 2; mask > 0; mask >>= 1) {\n sum_local_group.x += __shfl_xor_sync(FINAL_MASK, sum_local_group.x, mask, 32);\n sum_local_group.y += __shfl_xor_sync(FINAL_MASK, sum_local_group.y, mask, 32);\n }\n if (threadIdx.x == 0) {\n if constexpr (USE_SHARED_RED_BUFFER) {\n if constexpr (STORE_MEAN_VAR_IN_SHARED_RED_BUFFER) {\n shared_red_buffer[0] = compute_mean_var(sum_local_group);\n } else {\n shared_red_buffer[step * MAX_NUM_GROUPS_PER_BLOCK + 0] = sum_local_group;\n }\n } else {\n *reinterpret_cast(\n &red_buffer[((step * gridDim.y + blockIdx.y) * VIRTUAL_CLUSTER_SIZE * MAX_NUM_GROUPS_PER_BLOCK +\n virtual_block_idx_x * virtual_cluster_dim_y * MAX_NUM_GROUPS_PER_BLOCK +\n // (threadIdx.x / THREADS_PER_GROUP) * virtual_cluster_dim_y +\n virtual_block_idx_y) *\n 2]) = sum_local_group;\n }\n }\n }\n } else {\n // The number of threads to calculate the sum of each group (should be a power of 2 for warp reduce)\n constexpr int THREADS_PER_GROUP = std::min(std::min(32U, round_up_pow2(ROWS_PER_IO)),\n round_up_pow2(BLOCK_DIM_X / MAX_NUM_GROUPS_PER_BLOCK / 2 + 1));\n static_assert(BLOCK_DIM_X >= MAX_NUM_GROUPS_PER_BLOCK * THREADS_PER_GROUP, \"not enough threads\");\n float2 sum_local_group = {0.f, 0.f};\n if (threadIdx.x / THREADS_PER_GROUP < MAX_NUM_GROUPS_PER_BLOCK) {\n int local_group_idx = block_group_start + threadIdx.x / THREADS_PER_GROUP;\n // TODO: map threads to both the CPG loop and the ROWS loop\n for (int local_c_loop = 0; local_c_loop < CPG; local_c_loop += GCD_VEC_CPG) {\n int c = local_group_idx * CPG + local_c_loop;\n if (C_PER_BLOCK % CPG == 0 || (c >= block_channel_start && c < block_channel_start + C_PER_BLOCK)) {\n for (int src_thread_tile_y = threadIdx.x % THREADS_PER_GROUP; src_thread_tile_y < ROWS_PER_IO;\n src_thread_tile_y += THREADS_PER_GROUP) {\n int channel_idx = (c - block_channel_start) / GCD_VEC_CPG;\n channel_idx = channel_idx % (VEC_ELEMS / GCD_VEC_CPG) * (C_PER_BLOCK / VEC_ELEMS) +\n channel_idx / (VEC_ELEMS / GCD_VEC_CPG);\n sum_local_group.x += sum_per_channel_multi_group[channel_idx][src_thread_tile_y].x;\n sum_local_group.y += sum_per_channel_multi_group[channel_idx][src_thread_tile_y].y;\n }\n }\n }\n }\n static_assert(32 % THREADS_PER_GROUP == 0, \"cannot shuffle\");\n for (int mask = THREADS_PER_GROUP / 2; mask > 0; mask >>= 1) {\n sum_local_group.x += __shfl_xor_sync(FINAL_MASK, sum_local_group.x, mask, 32);\n sum_local_group.y += __shfl_xor_sync(FINAL_MASK, sum_local_group.y, mask, 32);\n }\n if (threadIdx.x % THREADS_PER_GROUP == 0 && threadIdx.x / THREADS_PER_GROUP < MAX_NUM_GROUPS_PER_BLOCK) {\n if constexpr (USE_SHARED_RED_BUFFER) {\n static_assert(HARDWARE_CLUSTER || VIRTUAL_CLUSTER_SIZE == 1, \"no distributed shared memory\");\n if constexpr (STORE_MEAN_VAR_IN_SHARED_RED_BUFFER) {\n shared_red_buffer[threadIdx.x / THREADS_PER_GROUP] = compute_mean_var(sum_local_group);\n } else {\n shared_red_buffer[step * MAX_NUM_GROUPS_PER_BLOCK + threadIdx.x / THREADS_PER_GROUP] = sum_local_group;\n }\n } else {\n *reinterpret_cast(\n &red_buffer[((step * gridDim.y + blockIdx.y) * VIRTUAL_CLUSTER_SIZE * MAX_NUM_GROUPS_PER_BLOCK +\n virtual_block_idx_x * virtual_cluster_dim_y * MAX_NUM_GROUPS_PER_BLOCK +\n (threadIdx.x / THREADS_PER_GROUP) * virtual_cluster_dim_y + virtual_block_idx_y) *\n 2]) = sum_local_group;\n }\n }\n }\n\n if constexpr (REQUIRES_WGRAD && wgrad_sync_method != WGRAD_SYNC_AT_LAST) {\n if (nc_scheduler.at_end(n)) {\n static_assert(PERSISTENT, \"persistent is a must for reducing wgrad\");\n if constexpr (wgrad_sync_method == WGRAD_ARRIVE_AND_WAIT_GRID) {\n virtual_cluster_sync(barrier);\n wgrad_sync_token = group_barrier_arrive(\n barrier_wgrad, blockIdx.x + blockIdx.y == 0);\n } else if constexpr (wgrad_sync_method == WGRAD_ARRIVE_AND_WAIT_GROUP) {\n virtual_cluster_sync(barrier);\n wgrad_sync_token =\n group_barrier_arrive(\n barrier_wgrad + virtual_cluster_idx_c, blockIdx.x + blockIdx.y / (C / C_PER_CLUSTER) == 0);\n } else if constexpr (wgrad_sync_method == WGRAD_REUSE_SUM_SYNC_GRID) {\n static_assert(!HARDWARE_CLUSTER,\n \"Distributed smem sync cannot reuse gmem sync. Use WGRAD_ARRIVE_AND_WAIT_GRID instead.\");\n wgrad_sync_token = group_barrier_arrive(\n barrier_wgrad, blockIdx.x + blockIdx.y == 0);\n group_barrier_wait(barrier_wgrad, wgrad_sync_token);\n } else if constexpr (wgrad_sync_method == WGRAD_REUSE_SUM_SYNC_GROUP) {\n static_assert(!HARDWARE_CLUSTER,\n \"Distributed smem sync cannot reuse gmem sync. Use WGRAD_ARRIVE_AND_WAIT_GROUP instead.\");\n wgrad_sync_token =\n group_barrier_arrive(\n barrier_wgrad + virtual_cluster_idx_c, blockIdx.x + blockIdx.y / (C / C_PER_CLUSTER) == 0);\n group_barrier_wait(barrier_wgrad + virtual_cluster_idx_c, wgrad_sync_token);\n }\n } else {\n virtual_cluster_sync(barrier);\n }\n } else {\n virtual_cluster_sync(barrier);\n }\n\n // Group sum\n __shared__ float2 mean_var[MAX_NUM_GROUPS_PER_BLOCK];\n if constexpr (!STORE_MEAN_VAR_IN_SHARED_RED_BUFFER) {\n // The number of threads to calculate the sum of each group (should be a power of 2 for warp reduce)\n constexpr int THREADS_PER_GROUP = std::min(std::min(32U, round_up_pow2(virtual_cluster_dim_y)),\n round_up_pow2(BLOCK_DIM_X / MAX_NUM_GROUPS_PER_BLOCK / 2 + 1));\n static_assert(BLOCK_DIM_X >= MAX_NUM_GROUPS_PER_BLOCK * THREADS_PER_GROUP, \"not enough threads\");\n float2 sum_global_group = {0.f, 0.f};\n if (threadIdx.x / THREADS_PER_GROUP < MAX_NUM_GROUPS_PER_BLOCK) {\n if constexpr (C_PER_BLOCK % CPG == 0) {\n // Special case: no cross-virtual_cluster_dim_x reduction\n float2 buffer[up_div(virtual_cluster_dim_y, THREADS_PER_GROUP)];\n for (int i = threadIdx.x % THREADS_PER_GROUP; i < virtual_cluster_dim_y; i += THREADS_PER_GROUP) {\n float2 val;\n if constexpr (USE_SHARED_RED_BUFFER) {\n if constexpr (VIRTUAL_CLUSTER_SIZE == 1) {\n val = shared_red_buffer[step * MAX_NUM_GROUPS_PER_BLOCK + threadIdx.x / THREADS_PER_GROUP];\n } else {\n static_assert(HARDWARE_CLUSTER, \"no distributed shared memory\");\n float2 const* src_shared_red_buffer = cg::this_cluster().map_shared_rank(\n shared_red_buffer, i * virtual_cluster_dim_x + virtual_block_idx_x);\n val = src_shared_red_buffer[step * MAX_NUM_GROUPS_PER_BLOCK + threadIdx.x / THREADS_PER_GROUP];\n }\n } else {\n val = *reinterpret_cast(\n &red_buffer[((step * gridDim.y + blockIdx.y) * VIRTUAL_CLUSTER_SIZE * MAX_NUM_GROUPS_PER_BLOCK +\n virtual_block_idx_x * virtual_cluster_dim_y * MAX_NUM_GROUPS_PER_BLOCK +\n (threadIdx.x / THREADS_PER_GROUP) * virtual_cluster_dim_y + i) *\n 2]);\n }\n buffer[i / THREADS_PER_GROUP] = val;\n }\n for (int i = threadIdx.x % THREADS_PER_GROUP; i < virtual_cluster_dim_y; i += THREADS_PER_GROUP) {\n float2 val = buffer[i / THREADS_PER_GROUP];\n sum_global_group.x += val.x;\n sum_global_group.y += val.y;\n }\n } else {\n // Common case: cross-virtual_cluster_dim_x reduction\n int local_group_idx = block_group_start + threadIdx.x / THREADS_PER_GROUP;\n for (int i = threadIdx.x % THREADS_PER_GROUP; i < VIRTUAL_CLUSTER_SIZE; i += THREADS_PER_GROUP) {\n int src_virtual_block_idx_x = i % virtual_cluster_dim_x;\n int src_block_channel_start = src_virtual_block_idx_x * C_PER_BLOCK + c_loop * C_PER_CLUSTER;\n int src_block_group_start = src_block_channel_start / CPG;\n int relative_group_idx = local_group_idx - src_block_group_start;\n if (0 <= relative_group_idx && relative_group_idx < MAX_NUM_GROUPS_PER_BLOCK) {\n float2 val;\n if constexpr (USE_SHARED_RED_BUFFER) {\n static_assert(HARDWARE_CLUSTER, \"no distributed shared memory\");\n static_assert(VIRTUAL_CLUSTER_SIZE != 1,\n \"layout error: should not add (step * MAX_NUM_GROUPS_PER_BLOCK)\");\n float2 const* src_shared_red_buffer = cg::this_cluster().map_shared_rank(shared_red_buffer, i);\n val = src_shared_red_buffer[step * MAX_NUM_GROUPS_PER_BLOCK + relative_group_idx];\n } else {\n val = *reinterpret_cast(\n &red_buffer[((step * gridDim.y + blockIdx.y) * VIRTUAL_CLUSTER_SIZE * MAX_NUM_GROUPS_PER_BLOCK +\n src_virtual_block_idx_x * virtual_cluster_dim_y * MAX_NUM_GROUPS_PER_BLOCK +\n relative_group_idx * virtual_cluster_dim_y + i / virtual_cluster_dim_x) *\n 2]);\n }\n sum_global_group.x += val.x;\n sum_global_group.y += val.y;\n }\n }\n }\n }\n if constexpr (USE_SHARED_RED_BUFFER && VIRTUAL_CLUSTER_SIZE > 1) {\n // Need cluster sync after distributed shared memory access, otherwise behavior is undefined\n if constexpr (PERSISTENT) {\n if (nc_scheduler.at_end(n)) {\n cg::this_cluster().barrier_arrive();\n }\n } else {\n cg::this_cluster().barrier_arrive();\n }\n }\n static_assert(32 % THREADS_PER_GROUP == 0, \"cannot shuffle\");\n for (int mask = THREADS_PER_GROUP / 2; mask > 0; mask >>= 1) {\n sum_global_group.x += __shfl_xor_sync(FINAL_MASK, sum_global_group.x, mask, 32);\n sum_global_group.y += __shfl_xor_sync(FINAL_MASK, sum_global_group.y, mask, 32);\n }\n if (threadIdx.x % THREADS_PER_GROUP == 0 && threadIdx.x / THREADS_PER_GROUP < MAX_NUM_GROUPS_PER_BLOCK) {\n mean_var[threadIdx.x / THREADS_PER_GROUP] = compute_mean_var(sum_global_group);\n }\n __syncthreads();\n }\n\n auto get_mean_var = [&](int relative_group_idx) {\n return STORE_MEAN_VAR_IN_SHARED_RED_BUFFER ? shared_red_buffer[relative_group_idx]\n : mean_var[relative_group_idx];\n };\n\n float frag_dyw[VEC_ELEMS / GCD_VEC_CPG];\n float frag_xdyw[VEC_ELEMS / GCD_VEC_CPG];\n for (int k = 0; k < VEC_ELEMS; k += GCD_VEC_CPG) {\n frag_dyw[k / GCD_VEC_CPG] = get_mean_var((thread_channel_start + k) / CPG - block_group_start).x;\n frag_xdyw[k / GCD_VEC_CPG] = get_mean_var((thread_channel_start + k) / CPG - block_group_start).y;\n }\n\n for (int j = 0; j < ROWS_PER_BLOCK / ROWS_PER_IO; j++) {\n int64_t input_idx =\n n_loop * HW * C +\n (virtual_block_idx_y * ROWS_PER_BLOCK + j * ROWS_PER_IO + threadIdx.x / (C_PER_BLOCK / VEC_ELEMS)) * C +\n thread_channel_start;\n U ux;\n U udy;\n if constexpr (LOAD_TWICE) {\n ux = *reinterpret_cast(&x[input_idx]);\n udy = *reinterpret_cast(&grad_output[input_idx]);\n } else {\n ux = frag_x[j];\n udy = frag_dy[j];\n }\n U val;\n for (int k = 0; k < VEC_ELEMS; k++) {\n float rnorm = rsqrtf(frag_var[k / GCD_VEC_CPG] + eps);\n float x_norm = ((float)ux.data[k] - frag_mean[k / GCD_VEC_CPG]) * rnorm; // TODO: store rsqrtf in mean_var\n float grad_gn = udy.data[k];\n if constexpr (SILU) {\n float x_gn = x_norm * (float)uw.data[k] + (float)ub.data[k];\n float s = 1.f / (1.f + expf(-x_gn));\n grad_gn *= s * (1.f + x_gn * (1.f - s));\n }\n val.data[k] =\n (grad_gn * (float)uw.data[k] - frag_dyw[k / GCD_VEC_CPG] - frag_xdyw[k / GCD_VEC_CPG] * x_norm) * rnorm;\n }\n *reinterpret_cast(&grad_input[input_idx]) = val;\n }\n\n if constexpr (!STORE_MEAN_VAR_IN_SHARED_RED_BUFFER && USE_SHARED_RED_BUFFER && VIRTUAL_CLUSTER_SIZE > 1) {\n if constexpr (PERSISTENT) {\n if (nc_scheduler.at_end(n)) {\n cg::this_cluster().barrier_wait();\n }\n } else {\n cg::this_cluster().barrier_wait();\n }\n }\n\n if constexpr (!PERSISTENT) {\n break;\n }\n step ^= 1;\n }\n\n // Wgrad sum\n if constexpr (REQUIRES_WGRAD) {\n static_assert(PERSISTENT, \"cannot reduce wgrad\");\n static_assert(C % 32 == 0, \"cannot reduce wgrad\");\n if constexpr (wgrad_sync_method == WGRAD_ARRIVE_AND_WAIT_GRID) {\n group_barrier_wait(barrier_wgrad, wgrad_sync_token);\n } else if constexpr (wgrad_sync_method == WGRAD_ARRIVE_AND_WAIT_GROUP) {\n group_barrier_wait(barrier_wgrad + virtual_cluster_idx_c, wgrad_sync_token);\n } else if constexpr (wgrad_sync_method == WGRAD_SYNC_AT_LAST) {\n cg::this_grid().sync();\n }\n\n // If group sync, map blocks that are responsible for the same range of channels to these channels (named \"split\n // channels\"); otherwise, map all blocks to all channels.\n constexpr bool split_channels =\n wgrad_sync_method == WGRAD_ARRIVE_AND_WAIT_GROUP || wgrad_sync_method == WGRAD_REUSE_SUM_SYNC_GROUP;\n\n for (int c = split_channels ? virtual_cluster_idx_c * C_PER_CLUSTER +\n 32 * (blockIdx.y / (C / C_PER_CLUSTER) * VIRTUAL_CLUSTER_SIZE + blockIdx.x)\n : 32 * (blockIdx.y * VIRTUAL_CLUSTER_SIZE + blockIdx.x);\n split_channels ? c < (virtual_cluster_idx_c + 1) * C_PER_CLUSTER : c < C;\n c += split_channels ? 32 * (NUM_VIRTUAL_CLUSTERS / (C / C_PER_CLUSTER) * VIRTUAL_CLUSTER_SIZE)\n : 32 * (NUM_VIRTUAL_CLUSTERS * VIRTUAL_CLUSTER_SIZE)) {\n int64_t rows = (CONSTANT_C_LOOP ? std::min(n, (int64_t)NUM_VIRTUAL_CLUSTERS / (C / C_PER_CLUSTER)) : n) *\n virtual_cluster_dim_y;\n float sum_wgrad = 0.f;\n float sum_bgrad = 0.f;\n if ((split_channels &&\n (C_PER_CLUSTER % 32 == 0 || c + threadIdx.x % 32 < (virtual_cluster_idx_c + 1) * C_PER_CLUSTER)) ||\n (!split_channels && (C % 32 == 0 || c + threadIdx.x % 32 < C))) {\n for (int64_t i = threadIdx.x / 32; i < rows; i += BLOCK_DIM_X / 32) {\n float2 val = *reinterpret_cast(&red_buffer_wgrad[(i * C + c + threadIdx.x % 32) * 2]);\n sum_wgrad += val.x;\n sum_bgrad += val.y;\n }\n }\n constexpr int warp_num_pow2 = round_up_pow2(BLOCK_DIM_X / 32);\n union_smem.transpose_buffer\n .wgrad_buffer[threadIdx.x / 32][(threadIdx.x % 32) ^ ((threadIdx.x / 32) * (32 / warp_num_pow2))] =\n sum_wgrad;\n union_smem.transpose_buffer\n .bgrad_buffer[threadIdx.x / 32][(threadIdx.x % 32) ^ ((threadIdx.x / 32) * (32 / warp_num_pow2))] =\n sum_bgrad;\n __syncthreads();\n for (int i = threadIdx.x / warp_num_pow2;\n i < 32 &&\n ((split_channels && (C_PER_CLUSTER % 32 == 0 || c + i < (virtual_cluster_idx_c + 1) * C_PER_CLUSTER)) ||\n (!split_channels && (C % 32 == 0 || c + i < C)));\n i += BLOCK_DIM_X / warp_num_pow2) {\n int j = threadIdx.x % warp_num_pow2;\n float sum_wgrad =\n j < BLOCK_DIM_X / 32 ? union_smem.transpose_buffer.wgrad_buffer[j][i ^ (j * (32 / warp_num_pow2))] : 0.f;\n float sum_bgrad =\n j < BLOCK_DIM_X / 32 ? union_smem.transpose_buffer.bgrad_buffer[j][i ^ (j * (32 / warp_num_pow2))] : 0.f;\n for (int mask = warp_num_pow2 / 2; mask > 0; mask >>= 1) {\n sum_wgrad += __shfl_xor_sync((uint64_t(1) << warp_num_pow2) - 1, sum_wgrad, mask, warp_num_pow2);\n sum_bgrad += __shfl_xor_sync((uint64_t(1) << warp_num_pow2) - 1, sum_bgrad, mask, warp_num_pow2);\n }\n if (j == 0) {\n grad_weight[c + i] = sum_wgrad;\n grad_bias[c + i] = sum_bgrad;\n }\n }\n __syncthreads();\n }\n }\n }\n}\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_dispatch_hw_c.hpp", "content": "#pragma once\n\n#define DISPATCH_HW_C(hw, c, HW, C, ...) \\\n [&] { \\\n if (hw == 64 && c == 1280) { \\\n constexpr int HW = 64, C = 1280; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 64 && c == 2560) { \\\n constexpr int HW = 64, C = 2560; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 256 && c == 640) { \\\n constexpr int HW = 256, C = 640; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 256 && c == 1280) { \\\n constexpr int HW = 256, C = 1280; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 256 && c == 1920) { \\\n constexpr int HW = 256, C = 1920; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 256 && c == 2560) { \\\n constexpr int HW = 256, C = 2560; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 1024 && c == 320) { \\\n constexpr int HW = 1024, C = 320; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 1024 && c == 640) { \\\n constexpr int HW = 1024, C = 640; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 1024 && c == 960) { \\\n constexpr int HW = 1024, C = 960; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 1024 && c == 1280) { \\\n constexpr int HW = 1024, C = 1280; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 1024 && c == 1920) { \\\n constexpr int HW = 1024, C = 1920; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 4096 && c == 320) { \\\n constexpr int HW = 4096, C = 320; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 4096 && c == 640) { \\\n constexpr int HW = 4096, C = 640; \\\n return __VA_ARGS__(); \\\n } \\\n if (hw == 4096 && c == 960) { \\\n constexpr int HW = 4096, C = 960; \\\n return __VA_ARGS__(); \\\n } \\\n throw std::invalid_argument(\"DISPATCH_HW_C \" + std::to_string(hw) + \" \" + std::to_string(c)); \\\n }()\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_utils.cpp", "content": "#include \"gn_utils.hpp\"\n\n#include \n#include \n\nnamespace group_norm_v2 {\n\ncudaDeviceProp const& get_device_prop(int device_id) {\n static std::vector device_props;\n static std::once_flag flag;\n std::call_once(flag, [&] {\n int count;\n CUDA_CHECK(cudaGetDeviceCount(&count));\n device_props.resize(count);\n for (int i = 0; i < count; i++) {\n CUDA_CHECK(cudaGetDeviceProperties(&device_props[i], i));\n }\n });\n return device_props.at(device_id);\n}\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/group_norm_v2/gn_utils.hpp", "content": "#pragma once\n\n#include \n\n#include \n#include \n#include \n\n#include \"gn.hpp\"\n\n// Definition of CUDA_CHECK macro\n#define CUDA_CHECK(call) \\\n do { \\\n cudaError_t err_ = call; \\\n if (err_ != cudaSuccess) { \\\n fprintf(stderr, \"CUDA error at %s:%d code=%d(%s) \\\"%s\\\" \\n\", __FILE__, __LINE__, err_, cudaGetErrorString(err_), \\\n #call); \\\n exit(EXIT_FAILURE); \\\n } \\\n } while (0)\n\n#define GN_CUDA_HOST_PARAMS(T) \\\n T *out, T *x, T *w, T *b, float eps, bool silu, int64_t n, int64_t hw, int num_groups, int channels_per_group, \\\n float *mean_var_out, float *red_buffer, unsigned *barrier, int sm_margin, cudaStream_t stream, int device_id, \\\n Meta *meta_ptr, bool meta_only\n\n#define GN_BWD_CUDA_HOST_PARAMS(T) \\\n T *grad_input, T *grad_weight, T *grad_bias, T *grad_output, T *x, T *w, T *b, float *mean_var, float eps, \\\n bool silu, int64_t n, int64_t hw, int num_groups, int channels_per_group, float *red_buffer, unsigned *barrier, \\\n int sm_margin, cudaStream_t stream, int device_id, Meta *meta_ptr, bool meta_only\n\n#define GN_CUDA_HOST_ARGS \\\n out, x, w, b, eps, silu, n, hw, num_groups, channels_per_group, mean_var_out, red_buffer, barrier, sm_margin, \\\n stream, device_id, meta_ptr, meta_only\n\n#define GN_BWD_CUDA_HOST_ARGS \\\n grad_input, grad_weight, grad_bias, grad_output, x, w, b, mean_var, eps, silu, n, hw, num_groups, \\\n channels_per_group, red_buffer, barrier, sm_margin, stream, device_id, meta_ptr, meta_only\n\nnamespace group_norm_v2 {\n\ncudaDeviceProp const& get_device_prop(int device_id);\n\n#ifdef __CUDA_ARCH__\n\ntemplate \n__host__ __device__ inline int print_rank_0(char const* fmt, Ts&&... args) {\n if (threadIdx.x + threadIdx.y + threadIdx.z == 0 && blockIdx.x + blockIdx.y + blockIdx.z == 0) {\n return printf(fmt, std::forward(args)...);\n }\n return 0;\n}\n\n#endif\n\n} // namespace group_norm_v2\n" }, { "path": "apex/contrib/csrc/groupbn/batch_norm.cu", "content": "#include \n#include \n#include \n#include \n\n#include \"batch_norm.h\"\n\n#define cudaCheckErrors(msg) \\\n do { \\\n cudaError_t __err = cudaGetLastError(); \\\n if (__err != cudaSuccess) { \\\n fprintf(stderr, \"Fatal error: %s (%s at %s:%d)\\n\", msg, cudaGetErrorString(__err), __FILE__, __LINE__); \\\n fprintf(stderr, \"*** FAILED - ABORTING\\n\"); \\\n exit(1); \\\n } \\\n } while (0)\n\nstatic size_t round_up_to_multiple(size_t x, int multiple) { return ((x + multiple - 1) / multiple) * multiple; }\n\nstruct Workspace {\n Workspace(size_t size) : size(size), data(NULL) {\n auto& allocator = *::c10::cuda::CUDACachingAllocator::get();\n dataPtr = allocator.allocate(size);\n data = dataPtr.get();\n }\n Workspace(const Workspace&) = delete;\n Workspace(Workspace&&) = default;\n Workspace& operator=(Workspace&&) = default;\n ~Workspace() = default;\n\n size_t size;\n void* data;\n c10::DataPtr dataPtr;\n};\n\n// Return {y}\nat::Tensor nhwc_bn_fwd_train(const at::Tensor& x, const at::Tensor& scale, const at::Tensor& bias,\n const at::Tensor& running_mean, const at::Tensor& running_inv_var,\n const at::Tensor& minibatch_mean, const at::Tensor& minibatch_inv_var,\n const at::Tensor& ret_cta, const float momentum, const float epsilon, const bool fuse_relu,\n void* my_data, void* pair_data, void* pair_data2, void* pair_data3, const int bn_group,\n const at::Tensor& magic_tensor, const int occupancy, const int grid_dim_x,\n const bool coop) {\n const int N = x.size(0);\n const int H = x.size(1);\n const int W = x.size(2);\n const int C = x.size(3);\n\n // generating new magic number and use that for sync\n int* magic = magic_tensor.data_ptr();\n *magic = (*magic + 1) & 0xff;\n\n // Allocate output tensor\n at::Tensor y = at::empty({N, H, W, C}, x.options());\n\n // Create wrapper\n NhwcBatchNorm* bn = new NhwcBatchNorm();\n\n bn->setInputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W, bn_group);\n bn->setOutputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W);\n\n bn->setConstants(momentum, epsilon);\n\n // set pointers within the wrapper\n bn->setInputOutputPointers(x.data_ptr(), nullptr, y.data_ptr(), nullptr);\n\n bn->setWeightPointers({scale.data_ptr(), bias.data_ptr()}, {nullptr, nullptr});\n bn->setParameterPointers({running_mean.data_ptr(), running_inv_var.data_ptr()});\n\n // deal with workspace(s)\n auto workspace_bytes = bn->numWorkspaceBytes();\n // We'll create explicit tensors for the first 2 workspace ptrs, then allocate & offset\n // an allocated workspace for the others\n size_t total_workspace_bytes = 0;\n std::vector workspace_offsets;\n\n for (auto index = 3; index < workspace_bytes.size(); ++index) {\n total_workspace_bytes = round_up_to_multiple(total_workspace_bytes, 512);\n workspace_offsets.push_back(total_workspace_bytes);\n\n auto alloc_bytes = workspace_bytes[index];\n total_workspace_bytes += alloc_bytes;\n }\n\n // Allocate the workspace\n Workspace ws(total_workspace_bytes);\n\n std::vector workspace;\n workspace.push_back(minibatch_mean.data_ptr());\n workspace.push_back(minibatch_inv_var.data_ptr());\n\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n const int retired_cta_bytes = workspace_bytes[2];\n void* retired_ctas = ret_cta.data_ptr();\n assert(ret_cta.size(0) >= retired_cta_bytes);\n workspace.push_back(retired_ctas);\n\n for (auto index = 3; index < workspace_bytes.size(); ++index) {\n void* ptr = reinterpret_cast(ws.data) + workspace_offsets[index - 3];\n workspace.push_back(ptr);\n }\n\n bn->setWorkspacePointers(workspace, workspace_bytes);\n\n // Don't fuse in ReLU for now at least\n bn->fwd(stream, fuse_relu, my_data, pair_data, pair_data2, pair_data3, bn_group, *magic, occupancy, grid_dim_x, coop);\n\n return y;\n}\n\nat::Tensor nhwc_bn_fwd_eval(const at::Tensor& x, const at::Tensor& scale, const at::Tensor& bias,\n const at::Tensor& running_mean, const at::Tensor& running_inv_var,\n const at::Tensor& ret_cta, const int bn_group, const float momentum, const float epsilon,\n const bool fuse_relu) {\n const int N = x.size(0);\n const int H = x.size(1);\n const int W = x.size(2);\n const int C = x.size(3);\n\n // Allocate output tensor\n at::Tensor y = at::empty({N, H, W, C}, x.options());\n\n // Create wrapper\n NhwcBatchNorm* bn = new NhwcBatchNorm();\n\n bn->setInputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W, bn_group);\n bn->setOutputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W);\n\n bn->setConstants(momentum, epsilon);\n\n // set pointers within the wrapper\n bn->setInputOutputPointers(x.data_ptr(), nullptr, y.data_ptr(), nullptr);\n\n bn->setWeightPointers({scale.data_ptr(), bias.data_ptr()}, {nullptr, nullptr});\n bn->setParameterPointers({running_mean.data_ptr(), running_inv_var.data_ptr()});\n\n // deal with workspace(s)\n auto workspace_bytes = bn->numWorkspaceBytes();\n // We'll create explicit tensors for the first 2 workspace ptrs, then allocate & offset\n // an allocated workspace for the others\n size_t total_workspace_bytes = 0;\n std::vector workspace_offsets;\n\n for (auto index = 3; index < workspace_bytes.size(); ++index) {\n total_workspace_bytes = round_up_to_multiple(total_workspace_bytes, 512);\n workspace_offsets.push_back(total_workspace_bytes);\n\n auto alloc_bytes = workspace_bytes[index];\n total_workspace_bytes += alloc_bytes;\n }\n\n // Allocate the workspace\n Workspace ws(total_workspace_bytes);\n\n std::vector workspace;\n workspace.push_back(nullptr);\n workspace.push_back(nullptr);\n\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n const int retired_cta_bytes = workspace_bytes[2];\n void* retired_ctas = ret_cta.data_ptr();\n assert(ret_cta.size(0) >= retired_cta_bytes);\n workspace.push_back(retired_ctas);\n\n for (auto index = 3; index < workspace_bytes.size(); ++index) {\n void* ptr = reinterpret_cast(ws.data) + workspace_offsets[index - 3];\n workspace.push_back(ptr);\n }\n\n bn->setWorkspacePointers(workspace, workspace_bytes);\n\n // Don't fuse in ReLU for now at least\n bn->fwdInference(stream, fuse_relu);\n\n return y;\n}\n\nstd::vector nhwc_bn_bwd(const at::Tensor& x, const at::Tensor& dy, const at::Tensor& scale,\n const at::Tensor& bias, const at::Tensor& running_mean,\n const at::Tensor& running_inv_var, const at::Tensor& minibatch_mean,\n const at::Tensor& minibatch_inv_var, const at::Tensor& ret_cta,\n const float momentum, const float epsilon, const bool fuse_relu, void* my_data,\n void* pair_data, void* pair_data2, void* pair_data3, const int bn_group,\n const at::Tensor& magic_tensor, const int occupancy, const int grid_dim_x,\n const bool coop) {\n // shape\n const int N = x.size(0);\n const int H = x.size(1);\n const int W = x.size(2);\n const int C = x.size(3);\n\n // generating new magic number and use that for sync\n int* magic = magic_tensor.data_ptr();\n *magic = (*magic + 1) & 0xff;\n\n // outputs\n at::Tensor x_grad, scale_grad, bias_grad;\n\n // Allocate outputs\n x_grad = at::empty_like(x);\n scale_grad = at::empty_like(scale);\n bias_grad = at::empty_like(bias);\n\n // Create wrapper\n NhwcBatchNorm* bn = new NhwcBatchNorm();\n\n bn->setInputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W, bn_group);\n bn->setOutputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W);\n\n bn->setConstants(momentum, epsilon);\n\n // set pointers within the wrapper\n bn->setInputOutputPointers(x.data_ptr(), x_grad.data_ptr(), nullptr, dy.data_ptr());\n\n bn->setWeightPointers({scale.data_ptr(), bias.data_ptr()},\n {scale_grad.data_ptr(), bias_grad.data_ptr()});\n bn->setParameterPointers({running_mean.data_ptr(), running_inv_var.data_ptr()});\n\n // deal with workspace(s)\n auto workspace_bytes = bn->numWorkspaceBytes();\n // We'll create explicit tensors for the first 2 workspace ptrs, then allocate & offset\n // an allocated workspace for the others\n size_t total_workspace_bytes = 0;\n std::vector workspace_offsets;\n\n for (auto index = 3; index < workspace_bytes.size(); ++index) {\n total_workspace_bytes = round_up_to_multiple(total_workspace_bytes, 512);\n workspace_offsets.push_back(total_workspace_bytes);\n\n auto alloc_bytes = workspace_bytes[index];\n total_workspace_bytes += alloc_bytes;\n }\n\n // Allocate the workspace\n Workspace ws(total_workspace_bytes);\n\n std::vector workspace;\n workspace.push_back(minibatch_mean.data_ptr());\n workspace.push_back(minibatch_inv_var.data_ptr());\n\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n const int retired_cta_bytes = workspace_bytes[2];\n void* retired_ctas = ret_cta.data_ptr();\n assert(ret_cta.size(0) >= retired_cta_bytes);\n workspace.push_back(retired_ctas);\n\n for (auto index = 3; index < workspace_bytes.size(); ++index) {\n void* ptr = reinterpret_cast(ws.data) + workspace_offsets[index - 3];\n workspace.push_back(ptr);\n }\n\n bn->setWorkspacePointers(workspace, workspace_bytes);\n\n bn->dgrad(stream, fuse_relu, my_data, pair_data, pair_data2, pair_data3, bn_group, *magic, occupancy, grid_dim_x,\n coop);\n\n return std::vector{x_grad, scale_grad, bias_grad};\n}\n\nint nhwc_bn_fwd_occupancy() {\n int device_id = -1;\n cudaGetDevice(&device_id);\n\n // max occupancy supported by the code is 2\n return NhwcBatchNorm::smem_driven_fwd_occupancy(device_id, 2);\n}\n\nint nhwc_bn_bwd_occupancy() {\n int device_id = -1;\n cudaGetDevice(&device_id);\n\n // max occupancy supported by the code is 2\n return NhwcBatchNorm::smem_driven_bwd_occupancy(device_id, 2);\n}\n" }, { "path": "apex/contrib/csrc/groupbn/batch_norm.h", "content": "/*\n * Licensed to the Apache Software Foundation (ASF) under one\n * or more contributor license agreements. See the NOTICE file\n * distributed with this work for additional information\n * regarding copyright ownership. The ASF licenses this file\n * to you under the Apache License, Version 2.0 (the\n * \"License\"); you may not use this file except in compliance\n * with the License. You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing,\n * software distributed under the License is distributed on an\n * \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n * KIND, either express or implied. See the License for the\n * specific language governing permissions and limitations\n * under the License.\n */\n\n/*!\n * Copyright (c) 2018 by Contributors\n * \\file nhwc_batch_norm.h\n * \\brief CUDA NHWC Batch Normalization code\n * \\author Shankara Rao Thejaswi Nanditale, Dick Carter, Evgeni Krimer\n */\n#ifndef MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_H_\n#define MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_H_\n\n#include \n\n#include \n#include \n#include \n#include \n\n#include \"cuda_utils.h\"\n#include \"nhwc_batch_norm_kernel.h\"\n\n#define VERBOSE_DEFAULT false\n\nclass NhwcBatchNorm {\n public:\n NhwcBatchNorm() {\n name_ = \"nhwc_batchnorm\";\n createTensorDescriptor(&X_tensor_desc_);\n createTensorDescriptor(&Y_tensor_desc_);\n }\n\n ~NhwcBatchNorm() {\n destroyTensorDescriptor(X_tensor_desc_);\n destroyTensorDescriptor(Y_tensor_desc_);\n }\n\n void die() {\n std::cerr << \"batchnorm not initialized\" << std::endl;\n exit(-1);\n }\n\n void fwd(cudaStream_t stream, bool use_relu, void* my_data, void* pair_data, void* pair_data2, void* pair_data3,\n const int bn_group, const int magic, const int occupancy, const int grid_dim_x, const bool coop);\n void dgrad(cudaStream_t stream, bool use_relu, void* my_data, void* pair_data, void* pair_data2, void* pair_data3,\n const int bn_group, const int magic, const int occupancy, const int grid_dim_x, const bool coop);\n void fwdInference(cudaStream_t stream, bool use_relu);\n dim3 calc_fwd_grid(int* loop, const int grid_dim_x);\n dim3 calc_bwd_grid(int* loop, const int grid_dim_x);\n\n void setInputDescriptor(const cudnnTensorFormat_t format, const cudnnDataType_t data_type, int n, int c, int h, int w,\n int bn_group) {\n m_ = n * h * w;\n int m_bn_adjusted = m_ * bn_group;\n c_ = c;\n // factor to scale sum of squared errors to get saved variance. Must be 1/nhw.\n svar_inv_count_ = 1.f / m_bn_adjusted;\n // factor to scale sum of squared errors to get running variance. Should be 1/(nhw-1).\n int divisor = m_bn_adjusted - 1;\n // nhw == 1 is unlikely, but by setting the rvar_inv_count_ == 1.f, we avoid running var infs.\n rvar_inv_count_ = divisor == 0 ? 1.f : 1.f / divisor;\n setTensorDescriptor(X_tensor_desc_, format, data_type, n, c, h, w);\n }\n\n void setOutputDescriptor(const cudnnTensorFormat_t format, const cudnnDataType_t data_type, int n, int c, int h,\n int w) {\n setTensorDescriptor(Y_tensor_desc_, format, data_type, n, c, h, w);\n }\n\n const std::vector numWorkspaceBytes() const;\n\n void setWorkspacePointers(const std::vector& workspace, const std::vector& num_workspace_bytes);\n\n void setInputOutputPointers(void* X, void* dX, void* Y, void* dY) {\n X_ = X;\n dX_ = dX;\n Y_ = Y;\n dY_ = dY;\n }\n\n // Sets the pointers for the scale and weight (in that order) data and derivative buffers.\n void setWeightPointers(const std::vector& weight_pointers, const std::vector& deriv_pointers) {\n assert(weight_pointers.size() == 2);\n assert(deriv_pointers.size() == 2);\n scale_ = static_cast(weight_pointers[0]);\n bias_ = static_cast(weight_pointers[1]);\n dscale_ = static_cast(deriv_pointers[0]);\n dbias_ = static_cast(deriv_pointers[1]);\n }\n\n // Sets the pointers for the population mean and variance buffers, in that order.\n void setParameterPointers(const std::vector& param_pointers) {\n assert(param_pointers.size() == 2);\n population_mean_ = static_cast(param_pointers[0]);\n population_variance_ = static_cast(param_pointers[1]);\n }\n\n void setConstants(const double exp_avg_factor, const double eps) {\n exp_avg_factor_ = exp_avg_factor;\n eps_ = eps;\n }\n\n void processCudnnStatus(const cudnnStatus_t& status, const std::string& string = std::string(),\n bool verbose = VERBOSE_DEFAULT) {\n if (status != CUDNN_STATUS_SUCCESS)\n LOG(FATAL) << string << \" \" << cudnnGetErrorString(status);\n else if (verbose)\n LOG(INFO) << string << \" \" << cudnnGetErrorString(status);\n }\n\n void checkCudaStatus(const std::string& string = std::string(), bool verbose = VERBOSE_DEFAULT) {\n cudaError_t status = cudaGetLastError();\n if (status != cudaSuccess)\n LOG(FATAL) << string << \" \" << cudaGetErrorString(status);\n else if (verbose)\n LOG(INFO) << string << \" \" << cudaGetErrorString(status);\n }\n\n size_t size_retired_ctas(int grid_y) const {\n // Note that the value of max_grid_y to handle known GPUs is about 160.\n const int max_grid_y = 1024;\n if (grid_y > max_grid_y) LOG(INFO) << \"GPU capabilities exceeds assumptions.\";\n const int retired_cta_bytes = max_grid_y * 2 * sizeof(int);\n // Since the region will be initialized once and used for many kernels,\n // the idea is to return an ample size that will cover all uses.\n return retired_cta_bytes;\n }\n\n cudnnTensorDescriptor_t X_tensor_desc_ = nullptr;\n cudnnTensorDescriptor_t Y_tensor_desc_ = nullptr;\n\n void* X_ = nullptr;\n void* dX_ = nullptr;\n void* Y_ = nullptr;\n void* dY_ = nullptr;\n\n // Learned scale and bias weights.\n float* scale_ = nullptr;\n float* dscale_ = nullptr;\n float* bias_ = nullptr;\n float* dbias_ = nullptr;\n\n // Computed population mean and variance parameters.\n float* population_mean_ = nullptr;\n float* population_variance_ = nullptr;\n\n // Workspace buffers for minibatch mean and variance (computed in fwd, needed by bwd).\n float* minibatch_mean_ = nullptr;\n float* minibatch_variance_ = nullptr;\n\n int m_ = 0; // Number of values per channel that BN is normalizing.\n int c_ = 0; // Number of channels over which BN is normalizing.\n\n float svar_inv_count_ = 0.f; // factor to scale sum of squared errors to get saved variance\n float rvar_inv_count_ = 0.f; // factor to scale sum of squared errors to get running variance\n\n double exp_avg_factor_ = 0.;\n double eps_ = 0.;\n std::string name_;\n\n private:\n void setTensorDescriptor(cudnnTensorDescriptor_t descriptor, cudnnTensorFormat_t format, cudnnDataType_t data_type,\n int n, int c, int h, int w) {\n cudnnStatus_t status = CUDNN_STATUS_SUCCESS;\n status = cudnnSetTensor4dDescriptor(descriptor, format, data_type, n, c, h, w);\n processCudnnStatus(status, \"set tensor descriptor\");\n }\n\n void createTensorDescriptor(cudnnTensorDescriptor_t* descriptor) {\n cudnnStatus_t status = CUDNN_STATUS_SUCCESS;\n status = cudnnCreateTensorDescriptor(descriptor);\n processCudnnStatus(status, \"create tensor_descriptor\");\n }\n\n void destroyTensorDescriptor(cudnnTensorDescriptor_t descriptor) {\n cudnnStatus_t status = CUDNN_STATUS_SUCCESS;\n status = cudnnDestroyTensorDescriptor(descriptor);\n processCudnnStatus(status, \"destroy tensor_descriptor\");\n }\n\n protected:\n float* partial_sums_ = nullptr;\n int* partial_counts_ = nullptr;\n int* retired_ctas_ = nullptr;\n\n void _setFwdParams(NhwcBatchNormFwdParams* params) const;\n void _setFwdInferenceParams(NhwcBatchNormFwdInferenceParams* params) const;\n void _setBwdParams(NhwcBatchNormBwdParams* params) const;\n\n // @todo: ability to configure these?\n // Kernel params\n static const int USE_ONLINE_APPROACH = 1;\n static const int THREADS_PER_CTA = 512;\n static const int THREADS_PER_PIXEL = 16;\n static const int C_ELEMENTS_PER_CTA = 64;\n static const int ELEMENTS_PER_LDG = C_ELEMENTS_PER_CTA / THREADS_PER_PIXEL;\n static const int MAX_SMEM_WITHOUT_OPT_IN = 48 * 1024;\n\n typedef uint16_t StorageType;\n // typedef float StorageType;\n // increasing this to 6 causes spills in fwd kernel!\n static const int PIXELS_PER_THREAD_IN_REGISTERS_FWD = 5;\n static const int PIXELS_PER_THREAD_IN_REGISTERS_BWD = 3;\n static const int PIXELS_PER_THREAD_IN_SMEM_FWD = 10;\n static const int PIXELS_PER_THREAD_IN_SMEM_BWD = 5;\n\n static const int PIXELS_PER_THREAD_FWD = PIXELS_PER_THREAD_IN_REGISTERS_FWD + PIXELS_PER_THREAD_IN_SMEM_FWD;\n static const int PIXELS_PER_THREAD_BWD = PIXELS_PER_THREAD_IN_REGISTERS_BWD + PIXELS_PER_THREAD_IN_SMEM_BWD;\n static const int PIXELS_PER_THREAD_FWD_INFERENCE = 4;\n\n // Derived params\n static const size_t SMEM_SIZE_FWD =\n PIXELS_PER_THREAD_IN_SMEM_FWD * THREADS_PER_CTA * ELEMENTS_PER_LDG * sizeof(StorageType);\n static const size_t SMEM_SIZE_BWD =\n PIXELS_PER_THREAD_IN_SMEM_BWD * THREADS_PER_CTA * ELEMENTS_PER_LDG * 2 * sizeof(StorageType);\n static const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;\n static const int PIXELS_PER_CTA_FWD = THREADS_PER_CTA / THREADS_PER_PIXEL * PIXELS_PER_THREAD_FWD;\n static const int PIXELS_PER_CTA_BWD = THREADS_PER_CTA / THREADS_PER_PIXEL * PIXELS_PER_THREAD_BWD;\n static const int PIXELS_PER_CTA_FWD_INFERENCE = THREADS_PER_CTA / THREADS_PER_PIXEL * PIXELS_PER_THREAD_FWD_INFERENCE;\n\n // max grid.y in case of group bn is limited by exchange buffer size\n static const int MAX_GBN_BLOCK_Y = 256;\n\n // Helper function to launch the forward kernel.\n\n // We calculate (based on smem usage) the achievable occupancy and make sure we run a kernel\n // version that was compiled with that occupancy in its launch bounds. This way, we avoid\n // needless register spills.\n void _fwdKernelLauncher(cudaStream_t stream, NhwcBatchNormFwdParams params, dim3 grid_dim, int outer_loops,\n bool use_relu, const int occupancy, const bool coop) {\n#define LAUNCH_FWD_KERNEL(OUTER_LOOPS, USE_RELU, USE_ADD_RELU, COMPILED_FOR_OCCUPANCY, COOP) \\\n do { \\\n CHECK(SMEM_SIZE_FWD <= MAX_SMEM_WITHOUT_OPT_IN) << \"Nhwc batchnorm kernel smem too big.\"; \\\n auto fwd_func = \\\n nhwc_batch_norm_fwd; \\\n if (COMPILED_FOR_OCCUPANCY > 1) { \\\n cudaFuncSetAttribute(fwd_func, cudaFuncAttributePreferredSharedMemoryCarveout, 100); \\\n checkCudaStatus(name_ + \" fwd ser coop kernel (cudaFuncSetAttribute carveout)\"); \\\n } \\\n void* params_ptr = static_cast(¶ms); \\\n using FWD_FUNC = decltype(nhwc_batch_norm_fwd); \\\n if (COOP) { \\\n cudaLaunchCooperativeKernel(fwd_func, grid_dim, THREADS_PER_CTA, ¶ms_ptr, SMEM_SIZE_FWD, stream); \\\n } else { \\\n cudaLaunchKernel(fwd_func, grid_dim, THREADS_PER_CTA, ¶ms_ptr, SMEM_SIZE_FWD, stream); \\\n } \\\n checkCudaStatus(name_ + \" fwd ser coop kernel\"); \\\n } while (0)\n\n // Don't try for an occupancy > 2 as this will squeeze register use and create spills.\n if (outer_loops == 1 && use_relu) {\n if (occupancy >= 2)\n LAUNCH_FWD_KERNEL(1, true, false, 2, coop);\n else\n LAUNCH_FWD_KERNEL(1, true, false, 1, coop);\n } else if (outer_loops == 1 && !use_relu) {\n if (occupancy >= 2)\n LAUNCH_FWD_KERNEL(1, false, false, 2, coop);\n else\n LAUNCH_FWD_KERNEL(1, false, false, 1, coop);\n } else if (use_relu) {\n if (occupancy >= 2)\n LAUNCH_FWD_KERNEL(0, true, false, 2, coop);\n else\n LAUNCH_FWD_KERNEL(0, true, false, 1, coop);\n } else {\n if (occupancy >= 2)\n LAUNCH_FWD_KERNEL(0, false, false, 2, coop);\n else\n LAUNCH_FWD_KERNEL(0, false, false, 1, coop);\n }\n#undef LAUNCH_FWD_KERNEL\n }\n\n // Helper function to launch the backward kernel.\n\n void _bwdKernelLauncher(cudaStream_t stream, NhwcBatchNormBwdParams params, dim3 grid_dim, int outer_loops,\n bool use_relu, const int occupancy, const bool coop) {\n#define LAUNCH_BWD_KERNEL(OUTER_LOOPS, COMPILED_FOR_OCCUPANCY, COOP) \\\n do { \\\n CHECK(SMEM_SIZE_BWD <= MAX_SMEM_WITHOUT_OPT_IN) << \"Nhwc batchnorm kernel smem too big.\"; \\\n auto bwd_func = nhwc_batch_norm_bwd; \\\n if (COMPILED_FOR_OCCUPANCY > 1) { \\\n cudaFuncSetAttribute(bwd_func, cudaFuncAttributePreferredSharedMemoryCarveout, 100); \\\n checkCudaStatus(name_ + \" bwd coop serial kernel (cudaFuncSetAttribute carveout)\"); \\\n } \\\n void* params_ptr = static_cast(¶ms); \\\n using BWD_FUNC = \\\n decltype(nhwc_batch_norm_bwd); \\\n if (COOP) { \\\n cudaLaunchCooperativeKernel(bwd_func, grid_dim, THREADS_PER_CTA, ¶ms_ptr, SMEM_SIZE_BWD, stream); \\\n } else { \\\n cudaLaunchKernel(bwd_func, grid_dim, THREADS_PER_CTA, ¶ms_ptr, SMEM_SIZE_BWD, stream); \\\n } \\\n checkCudaStatus(name_ + \" bwd coop serial kernel\"); \\\n } while (0)\n\n#define LAUNCH_BWD_RELU_KERNEL(OUTER_LOOPS, COMPILED_FOR_OCCUPANCY, COOP) \\\n do { \\\n CHECK(SMEM_SIZE_BWD <= MAX_SMEM_WITHOUT_OPT_IN) << \"Nhwc batchnorm kernel smem too big.\"; \\\n auto bwd_relu_func = \\\n nhwc_batch_norm_bwd_relu; \\\n if (COMPILED_FOR_OCCUPANCY > 1) { \\\n cudaFuncSetAttribute(bwd_relu_func, cudaFuncAttributePreferredSharedMemoryCarveout, 100); \\\n checkCudaStatus(name_ + \" bwd-relu coop serial kernel (cudaFuncSetAttribute carveout)\"); \\\n } \\\n void* params_ptr = static_cast(¶ms); \\\n using BWD_RELU_FUNC = \\\n decltype(nhwc_batch_norm_bwd_relu); \\\n if (COOP) { \\\n cudaLaunchCooperativeKernel(bwd_relu_func, grid_dim, THREADS_PER_CTA, ¶ms_ptr, SMEM_SIZE_BWD, \\\n stream); \\\n } else { \\\n cudaLaunchKernel(bwd_relu_func, grid_dim, THREADS_PER_CTA, ¶ms_ptr, SMEM_SIZE_BWD, stream); \\\n } \\\n checkCudaStatus(name_ + \" bwd-relu coop serial kernel\"); \\\n } while (0)\n\n // Don't try for an occupancy > 2 as this will squeeze register use and create spills.\n if (outer_loops == 1 && use_relu) {\n if (occupancy >= 2)\n LAUNCH_BWD_RELU_KERNEL(1, 2, coop);\n else\n LAUNCH_BWD_RELU_KERNEL(1, 1, coop);\n } else if (outer_loops == 1 && !use_relu) {\n if (occupancy >= 2)\n LAUNCH_BWD_KERNEL(1, 2, coop);\n else\n LAUNCH_BWD_KERNEL(1, 1, coop);\n } else if (use_relu) {\n if (occupancy >= 2)\n LAUNCH_BWD_RELU_KERNEL(0, 2, coop);\n else\n LAUNCH_BWD_RELU_KERNEL(0, 1, coop);\n } else {\n if (occupancy >= 2)\n LAUNCH_BWD_KERNEL(0, 2, coop);\n else\n LAUNCH_BWD_KERNEL(0, 1, coop);\n }\n#undef LAUNCH_BWD_KERNEL\n }\n\n public:\n // Calculate the expected fwd kernel occupancy, as dictated by shared memory usage.\n static int smem_driven_fwd_occupancy(int device_id, const int max_cta_per_sm) {\n using namespace at::cuda::utils;\n int fwd_reduction_bytes = THREADS_PER_PIXEL * (THREADS_PER_CTA / 32) * ELEMENTS_PER_LDG * sizeof(float);\n int fwd_smem_bytes = SMEM_SIZE_FWD + fwd_reduction_bytes;\n int occupancy = MaxSharedMemoryPerMultiprocessor(device_id) / fwd_smem_bytes;\n return std::min(max_cta_per_sm, occupancy);\n }\n\n // Calculate the expected bwd kernel occupancy, as dictated by shared memory usage.\n static int smem_driven_bwd_occupancy(int device_id, const int max_cta_per_sm) {\n using namespace at::cuda::utils;\n int bwd_reduction_bytes = THREADS_PER_PIXEL * (THREADS_PER_CTA / 32) * ELEMENTS_PER_LDG * sizeof(float);\n int bwd_smem_bytes = SMEM_SIZE_BWD + bwd_reduction_bytes;\n int occupancy = MaxSharedMemoryPerMultiprocessor(device_id) / bwd_smem_bytes;\n return std::min(max_cta_per_sm, occupancy);\n }\n};\n\nconst std::vector NhwcBatchNorm::numWorkspaceBytes() const {\n assert(c_ > 0);\n\n // choose the max memory required between fwd/bwd passes\n int grid_x_fwd = div_up(m_, PIXELS_PER_CTA_FWD);\n int grid_x_bwd = div_up(m_, PIXELS_PER_CTA_BWD);\n int grid_x = max(grid_x_fwd, grid_x_bwd);\n int grid_y = div_up(c_, C_ELEMENTS_PER_CTA);\n\n const size_t num_mean_bytes = c_ * sizeof(float);\n const size_t num_variance_bytes = num_mean_bytes;\n const size_t size_sums = grid_y * grid_x * THREADS_PER_PIXEL * ELEMENTS_PER_LDG * 2 * sizeof(float);\n const size_t size_counts = grid_y * grid_x * sizeof(int);\n\n return {num_mean_bytes, num_variance_bytes, size_retired_ctas(grid_y), size_sums, size_counts};\n}\n\nvoid NhwcBatchNorm::setWorkspacePointers(const std::vector& workspace,\n const std::vector& num_workspace_bytes) {\n assert(workspace.size() == 5);\n assert(num_workspace_bytes.size() == 5);\n\n minibatch_mean_ = static_cast(workspace[0]);\n minibatch_variance_ = static_cast(workspace[1]);\n retired_ctas_ = static_cast(workspace[2]);\n partial_sums_ = static_cast(workspace[3]);\n partial_counts_ = static_cast(workspace[4]);\n}\n\nvoid NhwcBatchNorm::_setFwdParams(NhwcBatchNormFwdParams* params) const {\n params->gmem_src = static_cast(X_);\n params->gmem_dst = static_cast(Y_);\n params->gmem_src1 = nullptr;\n params->gmem_bias = bias_;\n params->gmem_scale = scale_;\n params->gmem_running_mean = population_mean_;\n params->gmem_running_var = population_variance_;\n params->gmem_saved_mean = minibatch_mean_;\n params->gmem_saved_var = minibatch_variance_;\n params->gmem_relu_bitmask = nullptr;\n params->nhw = m_;\n params->c = c_;\n params->svar_inv_count = svar_inv_count_;\n params->rvar_inv_count = rvar_inv_count_;\n params->gmem_sums = partial_sums_;\n params->gmem_counts = partial_counts_;\n params->gmem_retired_ctas = retired_ctas_;\n params->var_eps = eps_;\n params->outer_loops = 0;\n params->exp_avg_factor = static_cast(exp_avg_factor_);\n params->c_blks = div_up(c_, C_ELEMENTS_PER_CTA);\n}\n\nvoid NhwcBatchNorm::_setFwdInferenceParams(NhwcBatchNormFwdInferenceParams* params) const {\n params->gmem_src = static_cast(X_);\n params->gmem_dst = static_cast(Y_);\n params->gmem_src1 = nullptr;\n params->gmem_bias = bias_;\n params->gmem_scale = scale_;\n params->gmem_mean = population_mean_;\n params->gmem_var = population_variance_;\n params->nhw = m_;\n params->c = c_;\n params->var_eps = eps_;\n}\n\nvoid NhwcBatchNorm::_setBwdParams(NhwcBatchNormBwdParams* params) const {\n params->gmem_src = static_cast(X_);\n params->gmem_dy = static_cast(dY_);\n params->gmem_dst = static_cast(dX_);\n params->gmem_dst1 = nullptr;\n params->gmem_relu_bitmask = nullptr;\n params->gmem_dscale = dscale_;\n params->gmem_dbias = dbias_;\n params->gmem_scale = scale_;\n params->gmem_bias = bias_;\n params->gmem_saved_mean = minibatch_mean_;\n params->gmem_saved_var = minibatch_variance_;\n params->nhw = m_;\n params->c = c_;\n params->svar_inv_count = svar_inv_count_;\n params->gmem_sums = partial_sums_;\n params->gmem_retired_ctas = retired_ctas_;\n params->outer_loops = 0;\n params->c_blks = div_up(c_, C_ELEMENTS_PER_CTA);\n}\n\nvoid NhwcBatchNorm::fwdInference(cudaStream_t stream, bool use_relu) {\n bool ptrs_are_set = X_tensor_desc_ != nullptr && Y_tensor_desc_ != nullptr && scale_ != nullptr &&\n bias_ != nullptr\n // && minibatch_mean_ != nullptr\n // && minibatch_variance_ != nullptr\n && population_mean_ != nullptr && population_variance_ != nullptr &&\n X_ != nullptr\n // && dX_ != nullptr\n && Y_ != nullptr\n // && dY_ != nullptr\n // && dscale_ != nullptr\n // && dbias_ != nullptr\n && partial_sums_ != nullptr && partial_counts_ != nullptr;\n\n if (!ptrs_are_set) die();\n\n dim3 grid_dim;\n grid_dim.x = div_up(m_, PIXELS_PER_CTA_FWD_INFERENCE);\n grid_dim.y = div_up(c_, C_ELEMENTS_PER_CTA);\n\n // @todo: maybe just move this inside initialize routine?\n NhwcBatchNormFwdInferenceParams params;\n _setFwdInferenceParams(¶ms);\n\n if (use_relu) {\n nhwc_batch_norm_fwd_inference\n <<>>(params);\n checkCudaStatus(name_ + \" fwd_inference-relu kernel\");\n } else {\n nhwc_batch_norm_fwd_inference\n <<>>(params);\n checkCudaStatus(name_ + \" fwd_inference kernel\");\n }\n}\n\ndim3 NhwcBatchNorm::calc_fwd_grid(int* loop, const int grid_dim_x) {\n dim3 grid_dim;\n grid_dim.x = div_up(m_, PIXELS_PER_CTA_FWD);\n int c_blks = div_up(c_, C_ELEMENTS_PER_CTA);\n unsigned int max_grid_x = grid_dim_x;\n if (grid_dim.x <= max_grid_x) {\n *loop = 1;\n if (max_grid_x / grid_dim.x > 1) {\n grid_dim.y = std::min(c_blks, static_cast(max_grid_x / grid_dim.x));\n assert(grid_dim.y < MAX_GBN_BLOCK_Y); // FIXME: turn into a loop\n } else {\n grid_dim.y = 1;\n }\n } else {\n grid_dim.x = max_grid_x;\n grid_dim.y = 1;\n int nhw_in_regs = m_ - PIXELS_PER_THREAD_IN_SMEM_FWD * PIXELS_PER_LDG * grid_dim.x;\n int pixels_per_iteration = PIXELS_PER_THREAD_IN_REGISTERS_FWD * PIXELS_PER_LDG * grid_dim.x;\n *loop = div_up(nhw_in_regs, pixels_per_iteration);\n }\n return grid_dim;\n}\n\ndim3 NhwcBatchNorm::calc_bwd_grid(int* loop, const int grid_dim_x) {\n dim3 grid_dim;\n grid_dim.x = div_up(m_, PIXELS_PER_CTA_BWD);\n int c_blks = div_up(c_, C_ELEMENTS_PER_CTA);\n unsigned int max_grid_x = grid_dim_x;\n if (grid_dim.x <= max_grid_x) {\n *loop = 1;\n if (max_grid_x / grid_dim.x > 1) {\n grid_dim.y = std::min(c_blks, static_cast(max_grid_x / grid_dim.x));\n assert(grid_dim.y < MAX_GBN_BLOCK_Y); // FIXME: turn into a loop\n } else {\n grid_dim.y = 1;\n }\n } else {\n grid_dim.x = max_grid_x;\n grid_dim.y = 1;\n int nhw_in_regs = m_ - PIXELS_PER_THREAD_IN_SMEM_BWD * PIXELS_PER_LDG * grid_dim.x;\n int pixels_per_iteration = PIXELS_PER_THREAD_IN_REGISTERS_BWD * PIXELS_PER_LDG * grid_dim.x;\n *loop = div_up(nhw_in_regs, pixels_per_iteration);\n }\n return grid_dim;\n}\n\nvoid NhwcBatchNorm::fwd(cudaStream_t stream, bool use_relu, void* my_data, void* pair_data, void* pair_data2,\n void* pair_data3, const int bn_group, const int magic, const int occupancy,\n const int grid_dim_x, const bool coop) {\n bool ptrs_are_set = X_tensor_desc_ != nullptr && Y_tensor_desc_ != nullptr && scale_ != nullptr && bias_ != nullptr &&\n minibatch_mean_ != nullptr && minibatch_variance_ != nullptr && population_mean_ != nullptr &&\n population_variance_ != nullptr &&\n X_ != nullptr\n // && dX_ != nullptr\n && Y_ != nullptr\n // && dY_ != nullptr\n // && dscale_ != nullptr\n // && dbias_ != nullptr\n && partial_sums_ != nullptr && partial_counts_ != nullptr && retired_ctas_ != nullptr;\n\n if (!ptrs_are_set) die();\n\n // reset of retired_cta_count no longer needed\n\n NhwcBatchNormFwdParams params;\n _setFwdParams(¶ms);\n params.my_data = my_data;\n params.pair_datas[0] = pair_data;\n params.pair_datas[1] = pair_data2;\n params.pair_datas[2] = pair_data3;\n params.magic = magic;\n params.sync_iters = (bn_group == 8) ? 3 : (bn_group >> 1);\n\n dim3 grid_dim = calc_fwd_grid(¶ms.outer_loops, grid_dim_x);\n _fwdKernelLauncher(stream, params, grid_dim, params.outer_loops, use_relu, occupancy, coop);\n}\n\nvoid NhwcBatchNorm::dgrad(cudaStream_t stream, bool use_relu, void* my_data, void* pair_data, void* pair_data2,\n void* pair_data3, const int bn_group, const int magic, const int occupancy,\n const int grid_dim_x, const bool coop) {\n bool ptrs_are_set = X_tensor_desc_ != nullptr && Y_tensor_desc_ != nullptr && scale_ != nullptr &&\n (bias_ != nullptr || !use_relu) && minibatch_mean_ != nullptr &&\n minibatch_variance_ != nullptr\n // && population_mean_ != nullptr\n // && population_variance_ != nullptr\n && X_ != nullptr &&\n dX_ != nullptr\n // && Y_ != nullptr\n && dY_ != nullptr && dscale_ != nullptr && dbias_ != nullptr;\n\n if (!ptrs_are_set) die();\n\n // reset of retired_cta_count no longer needed\n\n NhwcBatchNormBwdParams params;\n _setBwdParams(¶ms);\n params.my_data = my_data;\n params.pair_datas[0] = pair_data;\n params.pair_datas[1] = pair_data2;\n params.pair_datas[2] = pair_data3;\n params.magic = magic;\n params.sync_iters = (bn_group == 8) ? 3 : (bn_group >> 1);\n params.wgrad_coeff = 1.0 / bn_group;\n\n dim3 grid_dim = calc_bwd_grid(¶ms.outer_loops, grid_dim_x);\n _bwdKernelLauncher(stream, params, grid_dim, params.outer_loops, use_relu, occupancy, coop);\n}\n\n#endif // MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_H_\n" }, { "path": "apex/contrib/csrc/groupbn/batch_norm_add_relu.cu", "content": "#include \n#include \n#include \n#include \n\n#include \"batch_norm_add_relu.h\"\n\n// FIXME move the common stuff to common h file\n#define cudaCheckErrors(msg) \\\n do { \\\n cudaError_t __err = cudaGetLastError(); \\\n if (__err != cudaSuccess) { \\\n fprintf(stderr, \"Fatal error: %s (%s at %s:%d)\\n\", msg, cudaGetErrorString(__err), __FILE__, __LINE__); \\\n fprintf(stderr, \"*** FAILED - ABORTING\\n\"); \\\n exit(1); \\\n } \\\n } while (0)\n\nstatic size_t round_up_to_multiple(size_t x, int multiple) { return ((x + multiple - 1) / multiple) * multiple; }\n\nstruct Workspace {\n Workspace(size_t size) : size(size), data(NULL) {\n auto& allocator = *::c10::cuda::CUDACachingAllocator::get();\n dataPtr = allocator.allocate(size);\n data = dataPtr.get();\n }\n Workspace(const Workspace&) = delete;\n Workspace(Workspace&&) = default;\n Workspace& operator=(Workspace&&) = default;\n ~Workspace() = default;\n\n size_t size;\n void* data;\n c10::DataPtr dataPtr;\n};\n\n// Return {y}\nat::Tensor nhwc_bn_addrelu_fwd_train(const at::Tensor& x, const at::Tensor& z, const at::Tensor& scale,\n const at::Tensor& bias, const at::Tensor& running_mean,\n const at::Tensor& running_inv_var, const at::Tensor& minibatch_mean,\n const at::Tensor& minibatch_inv_var, const at::Tensor& bitmask,\n const at::Tensor& ret_cta, const float momentum, const float epsilon,\n void* my_data, void* pair_data, void* pair_data2, void* pair_data3,\n const int bn_group, const at::Tensor& magic_tensor, const int occupancy,\n const int grid_dim_x, const bool coop) {\n const int N = x.size(0);\n const int H = x.size(1);\n const int W = x.size(2);\n const int C = x.size(3);\n\n // generating new magic number and use that for sync\n int* magic = magic_tensor.data_ptr();\n *magic = (*magic + 1) & 0xff;\n\n // Allocate output tensor\n at::Tensor y = at::empty({N, H, W, C}, x.options());\n\n // Create wrapper\n NhwcBatchNormAddRelu* bn = new NhwcBatchNormAddRelu();\n\n bn->setInputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W, bn_group);\n bn->setOutputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W);\n\n bn->setConstants(momentum, epsilon);\n\n // set pointers within the wrapper\n bn->setInputOutputPointers(x.data_ptr(), nullptr, y.data_ptr(), nullptr, z.data_ptr(),\n nullptr);\n\n bn->setWeightPointers({scale.data_ptr(), bias.data_ptr()}, {nullptr, nullptr});\n bn->setParameterPointers({running_mean.data_ptr(), running_inv_var.data_ptr()});\n\n // deal with workspace(s)\n auto workspace_bytes = bn->numWorkspaceBytes();\n // We'll create explicit tensors for the first 2 workspace ptrs, then allocate & offset\n // an allocated workspace for the others\n size_t total_workspace_bytes = 0;\n std::vector workspace_offsets;\n\n for (auto index = 4; index < workspace_bytes.size(); ++index) {\n total_workspace_bytes = round_up_to_multiple(total_workspace_bytes, 512);\n workspace_offsets.push_back(total_workspace_bytes);\n\n auto alloc_bytes = workspace_bytes[index];\n total_workspace_bytes += alloc_bytes;\n }\n\n // Allocate the workspace\n Workspace ws(total_workspace_bytes);\n\n std::vector workspace;\n workspace.push_back(minibatch_mean.data_ptr());\n workspace.push_back(minibatch_inv_var.data_ptr());\n workspace.push_back(bitmask.data_ptr());\n\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n const int retired_cta_bytes = workspace_bytes[3];\n void* retired_ctas = ret_cta.data_ptr();\n assert(ret_cta.size(0) >= retired_cta_bytes);\n\n workspace.push_back(retired_ctas);\n\n for (auto index = 4; index < workspace_bytes.size(); ++index) {\n void* ptr = reinterpret_cast(ws.data) + workspace_offsets[index - 4];\n workspace.push_back(ptr);\n }\n\n bn->setWorkspacePointers(workspace, workspace_bytes);\n\n // Don't fuse in ReLU for now at least\n bn->fwd(stream, my_data, pair_data, pair_data2, pair_data3, bn_group, *magic, occupancy, grid_dim_x, coop);\n\n return y;\n}\n\nat::Tensor nhwc_bn_addrelu_fwd_eval(const at::Tensor& x, const at::Tensor& z, const at::Tensor& scale,\n const at::Tensor& bias, const at::Tensor& running_mean,\n const at::Tensor& running_inv_var, const at::Tensor& ret_cta, const int bn_group,\n const float momentum, const float epsilon) {\n const int N = x.size(0);\n const int H = x.size(1);\n const int W = x.size(2);\n const int C = x.size(3);\n\n // Allocate output tensor\n at::Tensor y = at::empty({N, H, W, C}, x.options());\n\n // Create wrapper\n NhwcBatchNormAddRelu* bn = new NhwcBatchNormAddRelu();\n\n bn->setInputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W, bn_group);\n bn->setOutputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W);\n\n bn->setConstants(momentum, epsilon);\n\n // set pointers within the wrapper\n bn->setInputOutputPointers(x.data_ptr(), nullptr, y.data_ptr(), nullptr, z.data_ptr(),\n nullptr);\n\n bn->setWeightPointers({scale.data_ptr(), bias.data_ptr()}, {nullptr, nullptr});\n bn->setParameterPointers({running_mean.data_ptr(), running_inv_var.data_ptr()});\n\n // deal with workspace(s)\n auto workspace_bytes = bn->numWorkspaceBytes();\n // We'll create explicit tensors for the first 2 workspace ptrs, then allocate & offset\n // an allocated workspace for the others\n size_t total_workspace_bytes = 0;\n std::vector workspace_offsets;\n\n for (auto index = 4; index < workspace_bytes.size(); ++index) {\n total_workspace_bytes = round_up_to_multiple(total_workspace_bytes, 512);\n workspace_offsets.push_back(total_workspace_bytes);\n\n auto alloc_bytes = workspace_bytes[index];\n total_workspace_bytes += alloc_bytes;\n }\n\n // Allocate the workspace\n Workspace ws(total_workspace_bytes);\n\n std::vector workspace;\n workspace.push_back(nullptr);\n workspace.push_back(nullptr);\n workspace.push_back(nullptr);\n\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n const int retired_cta_bytes = workspace_bytes[3];\n void* retired_ctas = ret_cta.data_ptr();\n assert(ret_cta.size(0) >= retired_cta_bytes);\n workspace.push_back(retired_ctas);\n\n for (auto index = 4; index < workspace_bytes.size(); ++index) {\n void* ptr = reinterpret_cast(ws.data) + workspace_offsets[index - 4];\n workspace.push_back(ptr);\n }\n\n bn->setWorkspacePointers(workspace, workspace_bytes);\n\n // Don't fuse in ReLU for now at least\n bn->fwdInference(stream);\n\n return y;\n}\n\nstd::vector nhwc_bn_addrelu_bwd(const at::Tensor& x, const at::Tensor& dy, const at::Tensor& scale,\n const at::Tensor& bias, const at::Tensor& running_mean,\n const at::Tensor& running_inv_var, const at::Tensor& minibatch_mean,\n const at::Tensor& minibatch_inv_var, const at::Tensor& bitmask,\n const at::Tensor& ret_cta, const float momentum, const float epsilon,\n void* my_data, void* pair_data, void* pair_data2, void* pair_data3,\n const int bn_group, const at::Tensor& magic_tensor, const int occupancy,\n const int grid_dim_x, const bool coop) {\n // shape\n const int N = x.size(0);\n const int H = x.size(1);\n const int W = x.size(2);\n const int C = x.size(3);\n\n // generating new magic number and use that for sync\n int* magic = magic_tensor.data_ptr();\n *magic = (*magic + 1) & 0xff;\n\n // outputs\n at::Tensor x_grad, z_grad, scale_grad, bias_grad;\n\n // Allocate outputs\n x_grad = at::empty_like(x);\n z_grad = at::empty_like(x);\n scale_grad = at::empty_like(scale);\n bias_grad = at::empty_like(bias);\n\n // Create wrapper\n NhwcBatchNormAddRelu* bn = new NhwcBatchNormAddRelu();\n\n bn->setInputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W, bn_group);\n bn->setOutputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W);\n\n bn->setConstants(momentum, epsilon);\n\n // set pointers within the wrapper\n bn->setInputOutputPointers(x.data_ptr(), x_grad.data_ptr(), nullptr, dy.data_ptr(),\n nullptr, z_grad.data_ptr());\n\n bn->setWeightPointers({scale.data_ptr(), bias.data_ptr()},\n {scale_grad.data_ptr(), bias_grad.data_ptr()});\n bn->setParameterPointers({running_mean.data_ptr(), running_inv_var.data_ptr()});\n\n // deal with workspace(s)\n auto workspace_bytes = bn->numWorkspaceBytes();\n // We'll create explicit tensors for the first 2 workspace ptrs, then allocate & offset\n // an allocated workspace for the others\n size_t total_workspace_bytes = 0;\n std::vector workspace_offsets;\n\n for (auto index = 4; index < workspace_bytes.size(); ++index) {\n total_workspace_bytes = round_up_to_multiple(total_workspace_bytes, 512);\n workspace_offsets.push_back(total_workspace_bytes);\n\n auto alloc_bytes = workspace_bytes[index];\n total_workspace_bytes += alloc_bytes;\n }\n\n // Allocate the workspace\n Workspace ws(total_workspace_bytes);\n\n std::vector workspace;\n workspace.push_back(minibatch_mean.data_ptr());\n workspace.push_back(minibatch_inv_var.data_ptr());\n workspace.push_back(bitmask.data_ptr());\n\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n const int retired_cta_bytes = workspace_bytes[3];\n void* retired_ctas = ret_cta.data_ptr();\n assert(ret_cta.size(0) >= retired_cta_bytes);\n workspace.push_back(retired_ctas);\n\n for (auto index = 4; index < workspace_bytes.size(); ++index) {\n void* ptr = reinterpret_cast(ws.data) + workspace_offsets[index - 4];\n workspace.push_back(ptr);\n }\n\n bn->setWorkspacePointers(workspace, workspace_bytes);\n\n bn->dgrad(stream, my_data, pair_data, pair_data2, pair_data3, bn_group, *magic, occupancy, grid_dim_x, coop);\n\n return std::vector{x_grad, z_grad, scale_grad, bias_grad};\n}\n\nint nhwc_bn_addrelu_fwd_occupancy() {\n int device_id = -1;\n cudaGetDevice(&device_id);\n\n // max occupancy supported by the code is 2\n return NhwcBatchNormAddRelu::smem_driven_fwd_occupancy(device_id, 2);\n}\n\nint nhwc_bn_addrelu_bwd_occupancy() {\n int device_id = -1;\n cudaGetDevice(&device_id);\n\n // max occupancy supported by the code is 2\n return NhwcBatchNormAddRelu::smem_driven_bwd_occupancy(device_id, 2);\n}\n" }, { "path": "apex/contrib/csrc/groupbn/batch_norm_add_relu.h", "content": "/*\n * Licensed to the Apache Software Foundation (ASF) under one\n * or more contributor license agreements. See the NOTICE file\n * distributed with this work for additional information\n * regarding copyright ownership. The ASF licenses this file\n * to you under the Apache License, Version 2.0 (the\n * \"License\"); you may not use this file except in compliance\n * with the License. You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing,\n * software distributed under the License is distributed on an\n * \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n * KIND, either express or implied. See the License for the\n * specific language governing permissions and limitations\n * under the License.\n */\n\n/*!\n * Copyright (c) 2018 by Contributors\n * \\file nhwc_batch_norm_add_relu.h\n * \\brief CUDA NHWC Batch Normalization code with fused addition\n * \\author Shankara Rao Thejaswi Nanditale, Dick Carter, Maxim Milakov, Evgeni Krimer\n */\n#ifndef MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_ADD_RELU_H_\n#define MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_ADD_RELU_H_\n\n#include \n\n#include \n#include \n#include \n#include \n\n#include \"cuda_utils.h\"\n#include \"nhwc_batch_norm_kernel.h\"\n\n#define VERBOSE_DEFAULT false\n\nclass NhwcBatchNormAddRelu {\n public:\n NhwcBatchNormAddRelu() {\n name_ = \"nhwc_batchnormaddrelu\";\n createTensorDescriptor(&X_tensor_desc_);\n createTensorDescriptor(&Y_tensor_desc_);\n }\n\n ~NhwcBatchNormAddRelu() {\n destroyTensorDescriptor(X_tensor_desc_);\n destroyTensorDescriptor(Y_tensor_desc_);\n }\n\n void die() {\n std::cerr << \"batchnormaddrelu not initialized\" << std::endl;\n exit(-1);\n }\n\n void fwd(cudaStream_t stream, void* my_data, void* pair_data, void* pair_data2, void* pair_data3, const int bn_group,\n const int magic, const int occupancy, const int grid_dim_x, const bool coop);\n void dgrad(cudaStream_t stream, void* my_data, void* pair_data, void* pair_data2, void* pair_data3,\n const int bn_group, const int magic, const int occupancy, const int grid_dim_x, const bool coop);\n void fwdInference(cudaStream_t stream);\n dim3 calc_fwd_grid(int* loop, const int grid_dim_x);\n dim3 calc_bwd_grid(int* loop, const int grid_dim_x);\n\n void setInputDescriptor(const cudnnTensorFormat_t format, const cudnnDataType_t data_type, int n, int c, int h, int w,\n int bn_group) {\n m_ = n * h * w;\n int m_bn_adjusted = m_ * bn_group;\n c_ = c;\n // factor to scale sum of squared errors to get saved variance. Must be 1/nhw.\n svar_inv_count_ = 1.f / m_bn_adjusted;\n // factor to scale sum of squared errors to get running variance. Should be 1/(nhw-1).\n int divisor = m_bn_adjusted - 1;\n // nhw == 1 is unlikely, but by setting the rvar_inv_count_ == 1.f, we avoid running var infs.\n rvar_inv_count_ = divisor == 0 ? 1.f : 1.f / divisor;\n setTensorDescriptor(X_tensor_desc_, format, data_type, n, c, h, w);\n }\n\n void setOutputDescriptor(const cudnnTensorFormat_t format, const cudnnDataType_t data_type, int n, int c, int h,\n int w) {\n setTensorDescriptor(Y_tensor_desc_, format, data_type, n, c, h, w);\n }\n\n const std::vector numWorkspaceBytes() const;\n\n void setWorkspacePointers(const std::vector& workspace, const std::vector& num_workspace_bytes);\n\n void setInputOutputPointers(void* X, void* dX, void* Y, void* dY, void* addend, void* dAddend) {\n X_ = X;\n dX_ = dX;\n Y_ = Y;\n dY_ = dY;\n addend_ = addend;\n dAddend_ = dAddend;\n }\n\n // Sets the pointers for the scale and weight (in that order) data and derivative buffers.\n void setWeightPointers(const std::vector& weight_pointers, const std::vector& deriv_pointers) {\n assert(weight_pointers.size() == 2);\n assert(deriv_pointers.size() == 2);\n scale_ = static_cast(weight_pointers[0]);\n bias_ = static_cast(weight_pointers[1]);\n dscale_ = static_cast(deriv_pointers[0]);\n dbias_ = static_cast(deriv_pointers[1]);\n }\n\n // Sets the pointers for the population mean and variance buffers, in that order.\n void setParameterPointers(const std::vector& param_pointers) {\n assert(param_pointers.size() == 2);\n population_mean_ = static_cast(param_pointers[0]);\n population_variance_ = static_cast(param_pointers[1]);\n }\n\n void setConstants(const double exp_avg_factor, const double eps) {\n exp_avg_factor_ = exp_avg_factor;\n eps_ = eps;\n }\n\n void processCudnnStatus(const cudnnStatus_t& status, const std::string& string = std::string(),\n bool verbose = VERBOSE_DEFAULT) {\n if (status != CUDNN_STATUS_SUCCESS)\n LOG(FATAL) << string << \" \" << cudnnGetErrorString(status);\n else if (verbose)\n LOG(INFO) << string << \" \" << cudnnGetErrorString(status);\n }\n\n void checkCudaStatus(const std::string& string = std::string(), bool verbose = VERBOSE_DEFAULT) {\n cudaError_t status = cudaGetLastError();\n if (status != cudaSuccess)\n LOG(FATAL) << string << \" \" << cudaGetErrorString(status);\n else if (verbose)\n LOG(INFO) << string << \" \" << cudaGetErrorString(status);\n }\n\n size_t size_retired_ctas(int grid_y) const {\n // Note that the value of max_grid_y to handle known GPUs is about 160.\n const int max_grid_y = 1024;\n if (grid_y > max_grid_y) LOG(INFO) << \"GPU capabilities exceeds assumptions.\";\n const int retired_cta_bytes = max_grid_y * 2 * sizeof(int);\n // Since the region will be initialized once and used for many kernels,\n // the idea is to return an ample size that will cover all uses.\n return retired_cta_bytes;\n }\n\n cudnnTensorDescriptor_t X_tensor_desc_ = nullptr;\n cudnnTensorDescriptor_t Y_tensor_desc_ = nullptr;\n\n void* X_ = nullptr;\n void* dX_ = nullptr;\n void* Y_ = nullptr;\n void* dY_ = nullptr;\n void* addend_ = nullptr;\n void* dAddend_ = nullptr;\n\n // Learned scale and bias weights.\n float* scale_ = nullptr;\n float* dscale_ = nullptr;\n float* bias_ = nullptr;\n float* dbias_ = nullptr;\n\n // Computed population mean and variance parameters.\n float* population_mean_ = nullptr;\n float* population_variance_ = nullptr;\n\n // Workspace buffers for minibatch mean and variance (computed in fwd, needed by bwd).\n float* minibatch_mean_ = nullptr;\n float* minibatch_variance_ = nullptr;\n\n int m_ = 0; // Number of values per channel that BN is normalizing.\n int c_ = 0; // Number of channels over which BN is normalizing.\n\n float svar_inv_count_ = 0.f; // factor to scale sum of squared errors to get saved variance\n float rvar_inv_count_ = 0.f; // factor to scale sum of squared errors to get running variance\n\n double exp_avg_factor_ = 0.;\n double eps_ = 0.;\n std::string name_;\n\n private:\n void setTensorDescriptor(cudnnTensorDescriptor_t descriptor, cudnnTensorFormat_t format, cudnnDataType_t data_type,\n int n, int c, int h, int w) {\n cudnnStatus_t status = CUDNN_STATUS_SUCCESS;\n status = cudnnSetTensor4dDescriptor(descriptor, format, data_type, n, c, h, w);\n processCudnnStatus(status, \"set tensor descriptor\");\n }\n\n void createTensorDescriptor(cudnnTensorDescriptor_t* descriptor) {\n cudnnStatus_t status = CUDNN_STATUS_SUCCESS;\n status = cudnnCreateTensorDescriptor(descriptor);\n processCudnnStatus(status, \"create tensor_descriptor\");\n }\n\n void destroyTensorDescriptor(cudnnTensorDescriptor_t descriptor) {\n cudnnStatus_t status = CUDNN_STATUS_SUCCESS;\n status = cudnnDestroyTensorDescriptor(descriptor);\n processCudnnStatus(status, \"destroy tensor_descriptor\");\n }\n\n protected:\n float* partial_sums_ = nullptr;\n int* partial_counts_ = nullptr;\n int* retired_ctas_ = nullptr;\n unsigned int* relu_bitmask_ = nullptr;\n\n void _setFwdParams(NhwcBatchNormFwdParams* params) const;\n void _setFwdInferenceParams(NhwcBatchNormFwdInferenceParams* params) const;\n void _setBwdParams(NhwcBatchNormBwdParams* params) const;\n\n // @todo: ability to configure these?\n // Kernel params\n static const int USE_ONLINE_APPROACH = 1;\n static const int THREADS_PER_CTA = 512;\n static const int THREADS_PER_PIXEL = 16;\n static const int C_ELEMENTS_PER_CTA = 64;\n static const int ELEMENTS_PER_LDG = C_ELEMENTS_PER_CTA / THREADS_PER_PIXEL;\n static const int MAX_SMEM_WITHOUT_OPT_IN = 48 * 1024;\n\n typedef uint16_t StorageType;\n // increasing this to 6 causes spills in fwd kernel!\n static const int PIXELS_PER_THREAD_IN_REGISTERS_FWD = 5;\n static const int PIXELS_PER_THREAD_IN_REGISTERS_BWD = 3;\n static const int PIXELS_PER_THREAD_IN_SMEM_FWD = 10;\n static const int PIXELS_PER_THREAD_IN_SMEM_BWD = 5;\n\n static const int PIXELS_PER_THREAD_FWD = PIXELS_PER_THREAD_IN_REGISTERS_FWD + PIXELS_PER_THREAD_IN_SMEM_FWD;\n static const int PIXELS_PER_THREAD_BWD = PIXELS_PER_THREAD_IN_REGISTERS_BWD + PIXELS_PER_THREAD_IN_SMEM_BWD;\n static const int PIXELS_PER_THREAD_FWD_INFERENCE = 4;\n\n // Derived params\n static const size_t SMEM_SIZE_FWD =\n PIXELS_PER_THREAD_IN_SMEM_FWD * THREADS_PER_CTA * ELEMENTS_PER_LDG * sizeof(StorageType);\n static const size_t SMEM_SIZE_BWD =\n PIXELS_PER_THREAD_IN_SMEM_BWD * THREADS_PER_CTA * ELEMENTS_PER_LDG * 2 * sizeof(StorageType);\n static const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;\n static const int PIXELS_PER_CTA_FWD = THREADS_PER_CTA / THREADS_PER_PIXEL * PIXELS_PER_THREAD_FWD;\n static const int PIXELS_PER_CTA_BWD = THREADS_PER_CTA / THREADS_PER_PIXEL * PIXELS_PER_THREAD_BWD;\n static const int PIXELS_PER_CTA_FWD_INFERENCE = THREADS_PER_CTA / THREADS_PER_PIXEL * PIXELS_PER_THREAD_FWD_INFERENCE;\n\n // max grid.y in case of group bn is limited by exchange buffer size\n static const int MAX_GBN_BLOCK_Y = 256;\n\n // Helper function to launch the forward kernel.\n\n // We calculate (based on smem usage) the achievable occupancy and make sure we run a kernel\n // version that was compiled with that occupancy in its launch bounds. This way, we avoid\n // needless register spills.\n void _fwdKernelLauncher(cudaStream_t stream, NhwcBatchNormFwdParams params, dim3 grid_dim, int outer_loops,\n const int occupancy, const bool coop) {\n#define LAUNCH_FWD_KERNEL(OUTER_LOOPS, USE_RELU, USE_ADD_RELU, COMPILED_FOR_OCCUPANCY, COOP) \\\n do { \\\n CHECK(SMEM_SIZE_FWD <= MAX_SMEM_WITHOUT_OPT_IN) << \"Nhwc batchnormaddrelu kernel smem too big.\"; \\\n auto fwd_func = \\\n nhwc_batch_norm_fwd; \\\n if (COMPILED_FOR_OCCUPANCY > 1) { \\\n cudaFuncSetAttribute(fwd_func, cudaFuncAttributePreferredSharedMemoryCarveout, 100); \\\n checkCudaStatus(name_ + \" fwd ser coop kernel (cudaFuncSetAttribute carveout)\"); \\\n } \\\n void* params_ptr = static_cast(¶ms); \\\n using FWD_FUNC = decltype(nhwc_batch_norm_fwd); \\\n if (COOP) { \\\n cudaLaunchCooperativeKernel(fwd_func, grid_dim, THREADS_PER_CTA, ¶ms_ptr, SMEM_SIZE_FWD, stream); \\\n } else { \\\n cudaLaunchKernel(fwd_func, grid_dim, THREADS_PER_CTA, ¶ms_ptr, SMEM_SIZE_FWD, stream); \\\n } \\\n checkCudaStatus(name_ + \" fwd ser coop kernel\"); \\\n } while (0)\n\n // Don't try for an occupancy > 2 as this will squeeze register use and create spills.\n if (outer_loops == 1) {\n if (occupancy >= 2)\n LAUNCH_FWD_KERNEL(1, false, true, 2, coop);\n else\n LAUNCH_FWD_KERNEL(1, false, true, 1, coop);\n } else {\n if (occupancy >= 2)\n LAUNCH_FWD_KERNEL(0, false, true, 2, coop);\n else\n LAUNCH_FWD_KERNEL(0, false, true, 1, coop);\n }\n#undef LAUNCH_FWD_KERNEL\n }\n\n // Helper function to launch the backward kernel.\n\n void _bwdKernelLauncher(cudaStream_t stream, NhwcBatchNormBwdParams params, dim3 grid_dim, int outer_loops,\n const int occupancy, const bool coop) {\n#define LAUNCH_BWD_ADD_RELU_KERNEL(OUTER_LOOPS, COMPILED_FOR_OCCUPANCY, COOP) \\\n do { \\\n CHECK(SMEM_SIZE_BWD <= MAX_SMEM_WITHOUT_OPT_IN) << \"Nhwc batchnormaddrelu kernel smem too big.\"; \\\n auto bwd_add_relu_func = \\\n nhwc_batch_norm_bwd_add_relu; \\\n if (COMPILED_FOR_OCCUPANCY > 1) { \\\n cudaFuncSetAttribute(bwd_add_relu_func, cudaFuncAttributePreferredSharedMemoryCarveout, 100); \\\n checkCudaStatus(name_ + \" bwd-add-relu coop serial kernel (cudaFuncSetAttribute carveout)\"); \\\n } \\\n void* params_ptr = static_cast(¶ms); \\\n using BWD_ADD_RELU_FUNC = \\\n decltype(nhwc_batch_norm_bwd_add_relu); \\\n if (COOP) { \\\n cudaLaunchCooperativeKernel(bwd_add_relu_func, grid_dim, THREADS_PER_CTA, ¶ms_ptr, \\\n SMEM_SIZE_BWD, stream); \\\n } else { \\\n cudaLaunchKernel(bwd_add_relu_func, grid_dim, THREADS_PER_CTA, ¶ms_ptr, SMEM_SIZE_BWD, \\\n stream); \\\n } \\\n checkCudaStatus(name_ + \" bwd-add-relu coop serial kernel\"); \\\n } while (0)\n\n // Don't try for an occupancy > 2 as this will squeeze register use and create spills.\n if (outer_loops == 1) {\n if (occupancy >= 2)\n LAUNCH_BWD_ADD_RELU_KERNEL(1, 2, coop);\n else\n LAUNCH_BWD_ADD_RELU_KERNEL(1, 1, coop);\n } else {\n if (occupancy >= 2)\n LAUNCH_BWD_ADD_RELU_KERNEL(0, 2, coop);\n else\n LAUNCH_BWD_ADD_RELU_KERNEL(0, 1, coop);\n }\n#undef LAUNCH_BWD_KERNEL\n }\n\n public:\n // Calculate the expected fwd kernel occupancy, as dictated by shared memory usage.\n static int smem_driven_fwd_occupancy(int device_id, const int max_cta_per_sm) {\n using namespace at::cuda::utils;\n int fwd_reduction_bytes = THREADS_PER_PIXEL * (THREADS_PER_CTA / 32) * ELEMENTS_PER_LDG * sizeof(float);\n int fwd_smem_bytes = SMEM_SIZE_FWD + fwd_reduction_bytes;\n int occupancy = MaxSharedMemoryPerMultiprocessor(device_id) / fwd_smem_bytes;\n return std::min(max_cta_per_sm, occupancy);\n }\n\n // Calculate the expected bwd kernel occupancy, as dictated by shared memory usage.\n static int smem_driven_bwd_occupancy(int device_id, const int max_cta_per_sm) {\n using namespace at::cuda::utils;\n int bwd_reduction_bytes = THREADS_PER_PIXEL * (THREADS_PER_CTA / 32) * ELEMENTS_PER_LDG * sizeof(float);\n int bwd_smem_bytes = SMEM_SIZE_BWD + bwd_reduction_bytes;\n int occupancy = MaxSharedMemoryPerMultiprocessor(device_id) / bwd_smem_bytes;\n return std::min(max_cta_per_sm, occupancy);\n }\n};\n\nconst std::vector NhwcBatchNormAddRelu::numWorkspaceBytes() const {\n assert(c_ > 0);\n\n // choose the max memory required between fwd/bwd passes\n int grid_x_fwd = div_up(m_, PIXELS_PER_CTA_FWD);\n int grid_x_bwd = div_up(m_, PIXELS_PER_CTA_BWD);\n int grid_x = max(grid_x_fwd, grid_x_bwd);\n int grid_y = div_up(c_, C_ELEMENTS_PER_CTA);\n\n const size_t num_mean_bytes = c_ * sizeof(float);\n const size_t num_variance_bytes = num_mean_bytes;\n\n int elems_per_group = ((m_ + 31) & ~31) * 2;\n int group_count = div_up(c_, C_ELEMENTS_PER_CTA);\n const size_t bitmask_bytes = elems_per_group * group_count * sizeof(unsigned int);\n\n const size_t size_sums = grid_y * grid_x * THREADS_PER_PIXEL * ELEMENTS_PER_LDG * 2 * sizeof(float);\n const size_t size_counts = grid_y * grid_x * sizeof(int);\n\n return {num_mean_bytes, num_variance_bytes, bitmask_bytes, size_retired_ctas(grid_y), size_sums, size_counts};\n}\n\nvoid NhwcBatchNormAddRelu::setWorkspacePointers(const std::vector& workspace,\n const std::vector& num_workspace_bytes) {\n assert(workspace.size() == 6);\n assert(num_workspace_bytes.size() == 6);\n\n minibatch_mean_ = static_cast(workspace[0]);\n minibatch_variance_ = static_cast(workspace[1]);\n relu_bitmask_ = static_cast(workspace[2]);\n retired_ctas_ = static_cast(workspace[3]);\n partial_sums_ = static_cast(workspace[4]);\n partial_counts_ = static_cast(workspace[5]);\n}\n\nvoid NhwcBatchNormAddRelu::_setFwdParams(NhwcBatchNormFwdParams* params) const {\n params->gmem_src = static_cast(X_);\n params->gmem_dst = static_cast(Y_);\n params->gmem_src1 = static_cast(addend_);\n params->gmem_bias = bias_;\n params->gmem_scale = scale_;\n params->gmem_running_mean = population_mean_;\n params->gmem_running_var = population_variance_;\n params->gmem_saved_mean = minibatch_mean_;\n params->gmem_saved_var = minibatch_variance_;\n params->gmem_relu_bitmask = relu_bitmask_;\n params->nhw = m_;\n params->c = c_;\n params->svar_inv_count = svar_inv_count_;\n params->rvar_inv_count = rvar_inv_count_;\n params->gmem_sums = partial_sums_;\n params->gmem_counts = partial_counts_;\n params->gmem_retired_ctas = retired_ctas_;\n params->var_eps = eps_;\n params->outer_loops = 0;\n params->exp_avg_factor = static_cast(exp_avg_factor_);\n params->c_blks = div_up(c_, C_ELEMENTS_PER_CTA);\n}\n\nvoid NhwcBatchNormAddRelu::_setFwdInferenceParams(NhwcBatchNormFwdInferenceParams* params) const {\n params->gmem_src = static_cast(X_);\n params->gmem_dst = static_cast(Y_);\n params->gmem_src1 = static_cast(addend_);\n params->gmem_bias = bias_;\n params->gmem_scale = scale_;\n params->gmem_mean = population_mean_;\n params->gmem_var = population_variance_;\n params->nhw = m_;\n params->c = c_;\n params->var_eps = eps_;\n}\n\nvoid NhwcBatchNormAddRelu::_setBwdParams(NhwcBatchNormBwdParams* params) const {\n params->gmem_src = static_cast(X_);\n params->gmem_dy = static_cast(dY_);\n params->gmem_dst = static_cast(dX_);\n params->gmem_dst1 = static_cast(dAddend_);\n params->gmem_relu_bitmask = relu_bitmask_;\n params->gmem_dscale = dscale_;\n params->gmem_dbias = dbias_;\n params->gmem_scale = scale_;\n params->gmem_bias = bias_;\n params->gmem_saved_mean = minibatch_mean_;\n params->gmem_saved_var = minibatch_variance_;\n params->nhw = m_;\n params->c = c_;\n params->svar_inv_count = svar_inv_count_;\n params->gmem_sums = partial_sums_;\n params->gmem_retired_ctas = retired_ctas_;\n params->outer_loops = 0;\n params->c_blks = div_up(c_, C_ELEMENTS_PER_CTA);\n}\n\nvoid NhwcBatchNormAddRelu::fwdInference(cudaStream_t stream) {\n bool ptrs_are_set = X_tensor_desc_ != nullptr && Y_tensor_desc_ != nullptr && scale_ != nullptr &&\n bias_ != nullptr\n // && minibatch_mean_ != nullptr\n // && minibatch_variance_ != nullptr\n && population_mean_ != nullptr && population_variance_ != nullptr &&\n X_ != nullptr\n // && dX_ != nullptr\n && Y_ != nullptr &&\n addend_ != nullptr\n // && dY_ != nullptr\n // && dscale_ != nullptr\n // && dbias_ != nullptr\n && partial_sums_ != nullptr && partial_counts_ != nullptr;\n\n if (!ptrs_are_set) die();\n\n dim3 grid_dim;\n grid_dim.x = div_up(m_, PIXELS_PER_CTA_FWD_INFERENCE);\n grid_dim.y = div_up(c_, C_ELEMENTS_PER_CTA);\n\n // @todo: maybe just move this inside initialize routine?\n NhwcBatchNormFwdInferenceParams params;\n _setFwdInferenceParams(¶ms);\n\n nhwc_batch_norm_fwd_inference\n <<>>(params);\n checkCudaStatus(name_ + \" fwd_inference-relu kernel\");\n}\n\ndim3 NhwcBatchNormAddRelu::calc_fwd_grid(int* loop, const int grid_dim_x) {\n dim3 grid_dim;\n grid_dim.x = div_up(m_, PIXELS_PER_CTA_FWD);\n int c_blks = div_up(c_, C_ELEMENTS_PER_CTA);\n unsigned int max_grid_x = grid_dim_x;\n if (grid_dim.x <= max_grid_x) {\n *loop = 1;\n if (max_grid_x / grid_dim.x > 1) {\n grid_dim.y = std::min(c_blks, static_cast(max_grid_x / grid_dim.x));\n assert(grid_dim.y < MAX_GBN_BLOCK_Y); // FIXME: turn into a loop\n } else {\n grid_dim.y = 1;\n }\n } else {\n grid_dim.x = max_grid_x;\n grid_dim.y = 1;\n int nhw_in_regs = m_ - PIXELS_PER_THREAD_IN_SMEM_FWD * PIXELS_PER_LDG * grid_dim.x;\n int pixels_per_iteration = PIXELS_PER_THREAD_IN_REGISTERS_FWD * PIXELS_PER_LDG * grid_dim.x;\n *loop = div_up(nhw_in_regs, pixels_per_iteration);\n }\n return grid_dim;\n}\n\ndim3 NhwcBatchNormAddRelu::calc_bwd_grid(int* loop, const int grid_dim_x) {\n dim3 grid_dim;\n grid_dim.x = div_up(m_, PIXELS_PER_CTA_BWD);\n int c_blks = div_up(c_, C_ELEMENTS_PER_CTA);\n unsigned int max_grid_x = grid_dim_x;\n if (grid_dim.x <= max_grid_x) {\n *loop = 1;\n if (max_grid_x / grid_dim.x > 1) {\n grid_dim.y = std::min(c_blks, static_cast(max_grid_x / grid_dim.x));\n assert(grid_dim.y < MAX_GBN_BLOCK_Y); // FIXME: turn into a loop\n } else {\n grid_dim.y = 1;\n }\n } else {\n grid_dim.x = max_grid_x;\n grid_dim.y = 1;\n int nhw_in_regs = m_ - PIXELS_PER_THREAD_IN_SMEM_BWD * PIXELS_PER_LDG * grid_dim.x;\n int pixels_per_iteration = PIXELS_PER_THREAD_IN_REGISTERS_BWD * PIXELS_PER_LDG * grid_dim.x;\n *loop = div_up(nhw_in_regs, pixels_per_iteration);\n }\n return grid_dim;\n}\n\nvoid NhwcBatchNormAddRelu::fwd(cudaStream_t stream, void* my_data, void* pair_data, void* pair_data2, void* pair_data3,\n const int bn_group, const int magic, const int occupancy, const int grid_dim_x,\n const bool coop) {\n bool ptrs_are_set = X_tensor_desc_ != nullptr && Y_tensor_desc_ != nullptr && scale_ != nullptr && bias_ != nullptr &&\n minibatch_mean_ != nullptr && minibatch_variance_ != nullptr && relu_bitmask_ != nullptr &&\n population_mean_ != nullptr && population_variance_ != nullptr &&\n X_ != nullptr\n // && dX_ != nullptr\n && Y_ != nullptr &&\n addend_ != nullptr\n // && dY_ != nullptr\n // && dscale_ != nullptr\n // && dbias_ != nullptr\n && partial_sums_ != nullptr && partial_counts_ != nullptr && retired_ctas_ != nullptr;\n\n if (!ptrs_are_set) die();\n\n // reset of retired_cta_count no longer needed\n\n NhwcBatchNormFwdParams params;\n _setFwdParams(¶ms);\n\n params.my_data = my_data;\n params.pair_datas[0] = pair_data;\n params.pair_datas[1] = pair_data2;\n params.pair_datas[2] = pair_data3;\n params.magic = magic;\n params.sync_iters = (bn_group == 8) ? 3 : (bn_group >> 1);\n\n dim3 grid_dim = calc_fwd_grid(¶ms.outer_loops, grid_dim_x);\n _fwdKernelLauncher(stream, params, grid_dim, params.outer_loops, occupancy, coop);\n}\n\nvoid NhwcBatchNormAddRelu::dgrad(cudaStream_t stream, void* my_data, void* pair_data, void* pair_data2,\n void* pair_data3, const int bn_group, const int magic, const int occupancy,\n const int grid_dim_x, const bool coop) {\n bool ptrs_are_set = X_tensor_desc_ != nullptr && Y_tensor_desc_ != nullptr && scale_ != nullptr && bias_ != nullptr &&\n minibatch_mean_ != nullptr && minibatch_variance_ != nullptr &&\n relu_bitmask_ != nullptr\n // && population_mean_ != nullptr\n // && population_variance_ != nullptr\n && X_ != nullptr &&\n dX_ != nullptr\n // && Y_ != nullptr\n && dY_ != nullptr && dAddend_ != nullptr && dscale_ != nullptr && dbias_ != nullptr &&\n retired_ctas_ != nullptr;\n\n if (!ptrs_are_set) die();\n\n // reset of retired_cta_count no longer needed\n\n NhwcBatchNormBwdParams params;\n _setBwdParams(¶ms);\n\n params.my_data = my_data;\n params.pair_datas[0] = pair_data;\n params.pair_datas[1] = pair_data2;\n params.pair_datas[2] = pair_data3;\n params.magic = magic;\n params.sync_iters = (bn_group == 8) ? 3 : (bn_group >> 1);\n params.wgrad_coeff = 1.0 / bn_group;\n\n dim3 grid_dim = calc_bwd_grid(¶ms.outer_loops, grid_dim_x);\n _bwdKernelLauncher(stream, params, grid_dim, params.outer_loops, occupancy, coop);\n}\n\n#endif // MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_ADD_RELU_H_\n" }, { "path": "apex/contrib/csrc/groupbn/cuda_utils.h", "content": "#include \n#ifndef CUDA_UTILS_H\n#define CUDA_UTILS_H\n\nnamespace at {\nnamespace cuda {\n\nnamespace utils {\n\nstatic inline int MaxSharedMemoryPerMultiprocessor(int device_id) {\n return getDeviceProperties(device_id)->sharedMemPerMultiprocessor;\n}\n\n} // namespace utils\n} // namespace cuda\n} // namespace at\n\n#endif\n" }, { "path": "apex/contrib/csrc/groupbn/interface.cpp", "content": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"ATen/Generator.h\"\n#include \"ATen/Scalar.h\"\n#include \"ATen/Storage.h\"\n#include \"ATen/Tensor.h\"\n\nnamespace py = pybind11;\n\nint64_t get_buffer_size(const int bn_sync_steps);\n\nvoid* get_data_ptr(const at::Tensor& data);\n\nvoid* get_remote_data_ptr(const at::Tensor& handle, const int64_t offset);\n\nvoid close_remote_data(const at::Tensor& handle);\n\nat::Tensor nhwc_bn_fwd_train(const at::Tensor& x, const at::Tensor& scale, const at::Tensor& bias,\n const at::Tensor& running_mean, const at::Tensor& running_inv_var,\n const at::Tensor& minibatch_mean, const at::Tensor& minibatch_inv_var,\n const at::Tensor& ret_cta, const float momentum, const float epsilon, const bool fuse_relu,\n void* my_data, void* pair_data, void* pair_data2, void* pair_data3, const int bn_group,\n const at::Tensor& magic_tensor, const int occupancy, const int grid_dim_x,\n const bool coop);\n\nat::Tensor nhwc_bn_fwd_eval(const at::Tensor& x, const at::Tensor& scale, const at::Tensor& bias,\n const at::Tensor& running_mean, const at::Tensor& running_inv_var,\n const at::Tensor& ret_cta, const int bn_group, const float momentum, const float epsilon,\n const bool fuse_relu);\n\nstd::vector nhwc_bn_bwd(const at::Tensor& x, const at::Tensor& dy, const at::Tensor& scale,\n const at::Tensor& bias, const at::Tensor& running_mean,\n const at::Tensor& running_inv_var, const at::Tensor& minibatch_mean,\n const at::Tensor& minibatch_inv_var, const at::Tensor& ret_cta,\n const float momentum, const float epsilon, const bool fuse_relu, void* my_data,\n void* pair_data, void* pair_data2, void* pair_data3, const int bn_group,\n const at::Tensor& magic_tensor, const int occupancy, const int grid_dim_x,\n const bool coop);\n\nat::Tensor nhwc_bn_addrelu_fwd_train(const at::Tensor& x, const at::Tensor& z, const at::Tensor& scale,\n const at::Tensor& bias, const at::Tensor& running_mean,\n const at::Tensor& running_inv_var, const at::Tensor& minibatch_mean,\n const at::Tensor& minibatch_inv_var, const at::Tensor& bitmask,\n const at::Tensor& ret_cta, const float momentum, const float epsilon,\n void* my_data, void* pair_data, void* pair_data2, void* pair_data3,\n const int bn_group, const at::Tensor& magic_tensor, const int occupancy,\n const int grid_dim_x, const bool coop);\n\nat::Tensor nhwc_bn_addrelu_fwd_eval(const at::Tensor& x, const at::Tensor& z, const at::Tensor& scale,\n const at::Tensor& bias, const at::Tensor& running_mean,\n const at::Tensor& running_inv_var, const at::Tensor& ret_cta, const int bn_group,\n const float momentum, const float epsilon);\n\nstd::vector nhwc_bn_addrelu_bwd(const at::Tensor& x, const at::Tensor& dy, const at::Tensor& scale,\n const at::Tensor& bias, const at::Tensor& running_mean,\n const at::Tensor& running_inv_var, const at::Tensor& minibatch_mean,\n const at::Tensor& minibatch_inv_var, const at::Tensor& bitmask,\n const at::Tensor& ret_cta, const float momentum, const float epsilon,\n void* my_data, void* pair_data, void* pair_data2, void* pair_data3,\n const int bn_group, const at::Tensor& magic_tensor, const int occupancy,\n const int grid_dim_x, const bool coop);\n\nint nhwc_bn_fwd_occupancy();\nint nhwc_bn_bwd_occupancy();\n\nint nhwc_bn_addrelu_fwd_occupancy();\nint nhwc_bn_addrelu_bwd_occupancy();\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"get_buffer_size\", &get_buffer_size, \"get_buffer_size\", py::call_guard());\n m.def(\"get_data_ptr\", &get_data_ptr, \"get_data_ptr\", py::call_guard());\n m.def(\"get_remote_data_ptr\", &get_remote_data_ptr, \"get_remote_data_ptr\", py::call_guard());\n m.def(\"close_remote_data\", &close_remote_data, \"close_remote_data\", py::call_guard());\n\n m.def(\"bn_fwd_nhwc\", &nhwc_bn_fwd_train, \"bn_fwd_nhwc\", py::call_guard());\n m.def(\"bn_fwd_eval_nhwc\", &nhwc_bn_fwd_eval, \"bn_fwd_eval_nhwc\", py::call_guard());\n m.def(\"bn_bwd_nhwc\", &nhwc_bn_bwd, \"bn_bwd_nhwc\", py::call_guard());\n\n m.def(\"bn_fwd_nhwc_occupancy\", &nhwc_bn_fwd_occupancy, \"bn_fwd_nhwc_occupancy\",\n py::call_guard());\n m.def(\"bn_bwd_nhwc_occupancy\", &nhwc_bn_bwd_occupancy, \"bn_bwd_nhwc_occupancy\",\n py::call_guard());\n\n m.def(\"bn_addrelu_fwd_nhwc\", &nhwc_bn_addrelu_fwd_train, \"bn_addrelu_fwd_nhwc\",\n py::call_guard());\n m.def(\"bn_addrelu_fwd_eval_nhwc\", &nhwc_bn_addrelu_fwd_eval, \"bn_addrelu_fwd_eval_nhwc\",\n py::call_guard());\n m.def(\"bn_addrelu_bwd_nhwc\", &nhwc_bn_addrelu_bwd, \"bn_addrelu_bwd_nhwc\", py::call_guard());\n\n m.def(\"bn_addrelu_fwd_nhwc_occupancy\", &nhwc_bn_addrelu_fwd_occupancy, \"bn_addrelu_fwd_nhwc_occupancy\",\n py::call_guard());\n m.def(\"bn_addrelu_bwd_nhwc_occupancy\", &nhwc_bn_addrelu_bwd_occupancy, \"bn_addrelu_bwd_nhwc_occupancy\",\n py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/groupbn/ipc.cu", "content": "#include \n#include \n#include \n\n#define cudaCheckErrors(msg) \\\n do { \\\n cudaError_t __err = cudaGetLastError(); \\\n if (__err != cudaSuccess) { \\\n fprintf(stderr, \"Fatal error: %s (%s at %s:%d)\\n\", msg, cudaGetErrorString(__err), __FILE__, __LINE__); \\\n fprintf(stderr, \"*** FAILED - ABORTING\\n\"); \\\n exit(1); \\\n } \\\n } while (0)\n\ntemplate <>\nstruct std::hash {\n size_t operator()(const cudaIpcMemHandle_t& handle) const {\n size_t hash = 0;\n uint8_t* ptr = (uint8_t*)&handle;\n assert(sizeof(uint8_t) == 1);\n for (int i = 0; i < sizeof(cudaIpcMemHandle_t); i++) {\n hash += *ptr;\n ptr++;\n }\n return hash;\n }\n};\n\ntemplate <>\nstruct std::equal_to {\n bool operator()(const cudaIpcMemHandle_t& lhs, const cudaIpcMemHandle_t& rhs) const {\n return (std::memcmp((void*)&lhs, (void*)&rhs, sizeof(cudaIpcMemHandle_t)) == 0);\n }\n};\n\nnamespace {\n\nnamespace gpuipc {\n// from: src/operator/nn/cudnn/nhwc_batch_norm_kernel.h\n// The number of threads per pixel.\nconst int THREADS_PER_PIXEL = 16;\n// The number of elements per ldg.\nconst int ELEMENTS_PER_LDG = 4;\n// The number of reducing ops, each uses its own space : mean, var, dscale, dbias\nconst int REDUCE_OPS = 4;\n// Maximum block.y supported - limited due to buffer allocation\nconst int MAX_BLOCK_Y = 256;\nconst int MAX_OFFSET = REDUCE_OPS * MAX_BLOCK_Y;\nconst int BYTES_PER_ELEM = 4;\n// Buffer size per sync step\nconst int SINGLE_SYNC_BUFFER_BYTES = MAX_OFFSET * THREADS_PER_PIXEL * 2 * ELEMENTS_PER_LDG * BYTES_PER_ELEM;\n}; // namespace gpuipc\n\nclass IpcMemHandleRegistry {\n public:\n void* getPtr(const cudaIpcMemHandle_t& handle, int64_t offset) {\n if (registry_.count(handle) == 0) {\n registry_.insert(std::make_pair(handle, RegistryEntry()));\n registry_[handle].dev_ptr = ipcOpenMem(handle);\n }\n registry_[handle].ref_count++;\n return (((uint8_t*)registry_[handle].dev_ptr) + offset);\n }\n\n void releasePtr(const cudaIpcMemHandle_t& handle) {\n if (registry_.count(handle) == 0) {\n }\n if (--registry_[handle].ref_count == 0) {\n ipcCloseMem(registry_[handle].dev_ptr);\n registry_.erase(handle);\n }\n }\n\n struct RegistryEntry {\n void* dev_ptr;\n int ref_count;\n RegistryEntry() : dev_ptr(NULL), ref_count(0) {}\n };\n\n protected:\n std::unordered_map registry_;\n\n void* ipcOpenMem(const cudaIpcMemHandle_t& handle) {\n void* data;\n cudaIpcOpenMemHandle(&data, handle, cudaIpcMemLazyEnablePeerAccess);\n cudaCheckErrors(\"ipc init\");\n return data;\n }\n\n void ipcCloseMem(void* dev_ptr) {\n cudaIpcCloseMemHandle(dev_ptr);\n cudaCheckErrors(\"ipc close\");\n }\n};\n\n} // namespace\n\nstatic IpcMemHandleRegistry ipc_mem_registry;\n\nint64_t get_buffer_size(const int bn_sync_steps) { return bn_sync_steps * gpuipc::SINGLE_SYNC_BUFFER_BYTES; }\n\nvoid* get_remote_data_ptr(const at::Tensor& handle, const int64_t offset) {\n cudaIpcMemHandle_t my_handle;\n memcpy((unsigned char*)(&my_handle), handle.data_ptr(), sizeof(my_handle));\n return ipc_mem_registry.getPtr(my_handle, offset);\n}\n\nvoid close_remote_data(const at::Tensor& handle) {\n cudaIpcMemHandle_t my_handle;\n memcpy((unsigned char*)(&my_handle), handle.data_ptr(), sizeof(my_handle));\n ipc_mem_registry.releasePtr(my_handle);\n}\n\nvoid* get_data_ptr(const at::Tensor& data) { return data.data_ptr(); }\n" }, { "path": "apex/contrib/csrc/groupbn/nhwc_batch_norm_kernel.h", "content": "/*\n * Licensed to the Apache Software Foundation (ASF) under one\n * or more contributor license agreements. See the NOTICE file\n * distributed with this work for additional information\n * regarding copyright ownership. The ASF licenses this file\n * to you under the Apache License, Version 2.0 (the\n * \"License\"); you may not use this file except in compliance\n * with the License. You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing,\n * software distributed under the License is distributed on an\n * \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n * KIND, either express or implied. See the License for the\n * specific language governing permissions and limitations\n * under the License.\n */\n\n/*!\n * Copyright (c) 2018 by Contributors\n * \\file nhwc_batch_norm_kernel.h\n * \\brief CUDA NHWC Batch Normalization code\n * \\author Shankara Rao Thejaswi Nanditale, Dick Carter, Maxim Milakov, Evgeni Krimer\n */\n#ifndef MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_KERNEL_H_\n#define MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_KERNEL_H_\n\n#include \n\n#include \n\n#define DEVICE_FUNCTION static inline __device__\n\n// CTA margin used by cooperative launch. Can be overridden by env var NHWC_BATCHNORM_LAUNCH_MARGIN.\n#define NHWC_BATCHNORM_LAUNCH_MARGIN_MIN 3\n#define NHWC_BATCHNORM_LAUNCH_MARGIN_DEFAULT NHWC_BATCHNORM_LAUNCH_MARGIN_MIN\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct PackedStorage {\n enum { PACKED_ELEMENTS_PER_LDG = ELEMENTS_PER_LDG };\n typedef T Type;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct PackedStorage {\n enum { PACKED_ELEMENTS_PER_LDG = ELEMENTS_PER_LDG / 2 };\n typedef int Type;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void from_float(int (&dst)[N], const float (&src)[2 * N]) {\n#pragma unroll\n for (int i = 0; i < N; ++i) {\n uint16_t lo, hi;\n asm volatile(\"cvt.rn.f16.f32 %0, %1;\" : \"=h\"(lo) : \"f\"(src[2 * i + 0]));\n asm volatile(\"cvt.rn.f16.f32 %0, %1;\" : \"=h\"(hi) : \"f\"(src[2 * i + 1]));\n asm volatile(\"mov.b32 %0, {%1, %2};\" : \"=r\"(dst[i]) : \"h\"(lo), \"h\"(hi));\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void from_float(float (&dst)[N], const float (&src)[N]) {\n#pragma unroll\n for (int i = 0; i < N; ++i) {\n dst[i] = src[i];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void to_float(float (&dst)[2 * N], int (&src)[N]) {\n#pragma unroll\n for (int i = 0; i < N; ++i) {\n uint16_t lo, hi;\n asm volatile(\"mov.b32 {%0, %1}, %2;\" : \"=h\"(lo), \"=h\"(hi) : \"r\"(src[i]));\n asm volatile(\"cvt.f32.f16 %0, %1;\" : \"=f\"(dst[2 * i + 0]) : \"h\"(lo));\n asm volatile(\"cvt.f32.f16 %0, %1;\" : \"=f\"(dst[2 * i + 1]) : \"h\"(hi));\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void to_float(float (&dst)[N], float (&src)[N]) {\n#pragma unroll\n for (int i = 0; i < N; ++i) {\n dst[i] = src[i];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void ldg(int (&dst)[1], const uint16_t* gmem) { dst[0] = __ldg((const int*)gmem); }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void ldg_stream(int (&dst)[1], const uint16_t* gmem) {\n unsigned int tmp;\n asm volatile(\"ld.global.cs.nc.s32 %0, [%1];\" : \"=r\"(tmp) : \"l\"((const uint*)gmem));\n dst[0] = tmp;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void ldg(int (&dst)[2], const uint16_t* gmem) {\n int2 tmp = __ldg((const int2*)gmem);\n dst[0] = tmp.x;\n dst[1] = tmp.y;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void ldg_stream(int (&dst)[2], const uint16_t* gmem) {\n int2 tmp;\n asm volatile(\"ld.global.cs.nc.v2.s32 {%0,%1}, [%2];\" : \"=r\"(tmp.x), \"=r\"(tmp.y) : \"l\"((const int2*)gmem));\n dst[0] = tmp.x;\n dst[1] = tmp.y;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void ldg(float (&dst)[N], const uint16_t* gmem) {\n int tmp[N / 2];\n ldg(tmp, gmem);\n to_float(dst, tmp);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void ldg_stream(float (&dst)[N], const uint16_t* gmem) {\n int tmp[N / 2];\n ldg_stream(tmp, gmem);\n to_float(dst, tmp);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void stg(uint16_t* gmem, int (&src)[1]) { reinterpret_cast(gmem)[0] = src[0]; }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void stg_stream(uint16_t* gmem, int (&src)[1]) {\n unsigned int tmp = src[0];\n asm volatile(\"st.global.cs.s32 [%0], %1;\" ::\"l\"((uint*)gmem), \"r\"(tmp));\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void stg(uint16_t* gmem, int (&src)[2]) {\n reinterpret_cast(gmem)[0] = make_int2(src[0], src[1]);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void stg_stream(uint16_t* gmem, int (&src)[2]) {\n asm volatile(\"st.global.cs.v2.s32 [%0], {%1,%2};\" ::\"l\"((uint*)gmem), \"r\"(src[0]), \"r\"(src[1]));\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void stg(uint16_t* gmem, float (&src)[N]) {\n int tmp[N / 2];\n from_float(tmp, src);\n stg(gmem, tmp);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void stg_stream(uint16_t* gmem, float (&src)[N]) {\n int tmp[N / 2];\n from_float(tmp, src);\n stg_stream(gmem, tmp);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void read_from_gmem(float (&dst)[2], const float* gmem, int idx) {\n float2 tmp = __ldg(reinterpret_cast(&gmem[2 * idx]));\n dst[0] = tmp.x;\n dst[1] = tmp.y;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void read_from_gmem(float (&dst)[4], const float* gmem, int idx) {\n float4 tmp = __ldg(reinterpret_cast(&gmem[4 * idx]));\n dst[0] = tmp.x;\n dst[1] = tmp.y;\n dst[2] = tmp.z;\n dst[3] = tmp.w;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void read_from_smem(float (&x)[2], const float* smem, int idx) {\n float2 tmp = *(const float2*)&smem[2 * idx];\n x[0] = tmp.x;\n x[1] = tmp.y;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void read_from_smem(int (&x)[1], const int* smem, int idx) { x[0] = smem[idx]; }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void read_from_smem(float (&x)[4], const float* smem, int idx) {\n float4 tmp = *(const float4*)&smem[4 * idx];\n x[0] = tmp.x;\n x[1] = tmp.y;\n x[2] = tmp.z;\n x[3] = tmp.w;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void read_from_smem(int (&x)[2], const int* smem, int idx) {\n int2 tmp = *(const int2*)&smem[2 * idx];\n x[0] = tmp.x;\n x[1] = tmp.y;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void write_to_gmem(float* gmem, int idx, const float (&src)[2]) {\n reinterpret_cast(&gmem[2 * idx])[0] = make_float2(src[0], src[1]);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void write_to_gmem(float* gmem, int idx, const float (&src)[4]) {\n reinterpret_cast(&gmem[4 * idx])[0] = make_float4(src[0], src[1], src[2], src[3]);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void scaled_write_to_gmem(float* gmem, int idx, const float (&src)[4], const float coeff) {\n reinterpret_cast(&gmem[4 * idx])[0] =\n make_float4(src[0] * coeff, src[1] * coeff, src[2] * coeff, src[3] * coeff);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void write_to_smem(float* smem, int idx, const float (&x)[2]) {\n reinterpret_cast(&smem[2 * idx])[0] = make_float2(x[0], x[1]);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void write_to_smem(int* smem, int idx, const int (&x)[1]) { smem[idx] = x[0]; }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void write_to_smem(float* smem, int idx, const float (&x)[4]) {\n reinterpret_cast(&smem[4 * idx])[0] = make_float4(x[0], x[1], x[2], x[3]);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nDEVICE_FUNCTION void write_to_smem(int* smem, int idx, const int (&x)[2]) {\n reinterpret_cast(&smem[2 * idx])[0] = make_int2(x[0], x[1]);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void zero_array(int (&dst)[N]) {\n#pragma unroll\n for (int i = 0; i < N; ++i) {\n dst[i] = 0;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void zero_array(float (&dst)[N]) {\n#pragma unroll\n for (int i = 0; i < N; ++i) {\n dst[i] = 0.f;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void add(float (&x)[N], const float (&y)[N]) {\n#pragma unroll\n for (int i = 0; i < N; ++i) {\n x[i] += y[i];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void multiply(float (&x)[N], const float (&y)[N]) {\n#pragma unroll\n for (int i = 0; i < N; ++i) {\n x[i] *= y[i];\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void scale_(float (&x)[N], float scalar) {\n#pragma unroll\n for (int i = 0; i < N; ++i) {\n x[i] *= scalar;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void normalize(float (&x)[N], const float (&bias)[N], const float (&scale)[N], const float (&m1)[N]) {\n#pragma unroll\n for (int i = 0; i < N; ++i) {\n x[i] = bias[i] + scale[i] * (x[i] - m1[i]);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION Storage relu(Storage in) {\n Storage zero = (Storage)0.f;\n return (in < zero) ? zero : in;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void relu_activation(float (&x)[N]) {\n#pragma unroll\n for (int i = 0; i < N; ++i) {\n x[i] = relu(x[i]);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\ntemplate \nDEVICE_FUNCTION void parallel_sums_16x2(float* smem, float (&x)[4], int nhw, void* params_my_data,\n void** params_pair_datas, int off, const int magic, const int sync_iters) {\n // The size of a warp.\n const int THREADS_PER_WARP = 32;\n // The number of warps in a CTA.\n const int WARPS_PER_CTA = THREADS_PER_CTA / THREADS_PER_WARP;\n // The number of threads per pixel.\n const int THREADS_PER_PIXEL = 16;\n // The number of elements per ldg.\n const int ELEMENTS_PER_LDG = 4;\n // The number of reducing ops, each uses its own space : mean, var, dscale, dbias\n const int REDUCE_OPS = 4;\n // Maximum block.y supported - limited due to buffer allocation\n const int MAX_BLOCK_Y = 256;\n const int MAX_OFFSET = REDUCE_OPS * MAX_BLOCK_Y;\n // The warp decomposition.\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const int lane_id = threadIdx.x % THREADS_PER_WARP;\n // total size of data per sync iter\n const int data_total = MAX_OFFSET * THREADS_PER_PIXEL * ELEMENTS_PER_LDG * 2;\n\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL + lane_id);\n }\n\n // The warp leaders, write to SMEM.\n if (lane_id < THREADS_PER_PIXEL) {\n write_to_smem(smem, warp_id * THREADS_PER_PIXEL + lane_id, x);\n }\n\n // The data is in SMEM. Do the final reduction.\n __syncthreads();\n\n // The 1st warp does all the work.\n // We do the final reduction each half-warp sequentially reduces the final values.\n if (warp_id == 0) {\n read_from_smem(x, smem, threadIdx.x);\n\n#pragma unroll\n for (int offset = 1; offset < WARPS_PER_CTA / (THREADS_PER_WARP / THREADS_PER_PIXEL); ++offset) {\n float y[ELEMENTS_PER_LDG];\n // Read the mean and variance from the other pixel.\n read_from_smem(y, smem, threadIdx.x + offset * THREADS_PER_WARP);\n // Compute the updated sum.\n add(x, y);\n }\n\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL + lane_id);\n }\n\n // Make sure the data was read from SMEM.\n __syncwarp();\n\n // Store the final values.\n if (threadIdx.x < THREADS_PER_PIXEL) {\n // probably could do it earlier, before sync\n\n for (int sync_iter = 0; sync_iter < sync_iters; ++sync_iter) {\n // float* params_pair_data = (reinterpret_cast(params_pair_datas))[sync_iter];\n void* params_pair_data = params_pair_datas[sync_iter];\n\n // skip the space consumed by previous sync iterations\n const int xbuf_offset = sync_iter * data_total;\n // data starts after flags, but have to skip previous\n const int data_offset =\n xbuf_offset + off * ELEMENTS_PER_LDG * THREADS_PER_PIXEL * 2 + ELEMENTS_PER_LDG * threadIdx.x * 2;\n\n // after sums for this GPU were computed, let CTA0 broadcast the sum to over GPU\n if (blockIdx.x == 0) {\n volatile float* write_data = &((reinterpret_cast(params_pair_data))[data_offset]);\n\n // write the data to memory region to be reflected to other GPU\n asm volatile(\"st.global.wt.v4.b32 [%0], {%1,%2,%3,%4};\" ::\"l\"(write_data), \"f\"(x[0]), \"r\"(magic), \"f\"(x[2]),\n \"r\"(magic));\n\n asm volatile(\"st.global.wt.v4.b32 [%0], {%1,%2,%3,%4};\" ::\"l\"(write_data + 4), \"f\"(x[1]), \"r\"(magic),\n \"f\"(x[3]), \"r\"(magic));\n }\n\n // now each CTA (on each GPU) reads the data written by CTA 0 of the other GPU\n volatile float* read_data = &((reinterpret_cast(params_my_data))[data_offset]);\n\n float other[4];\n uint32_t other_flag_a, other_flag_b;\n do {\n asm volatile(\"ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];\"\n : \"=f\"(other[0]), \"=r\"(other_flag_a), \"=f\"(other[2]), \"=r\"(other_flag_b)\n : \"l\"(read_data));\n } while ((other_flag_a != magic) || (other_flag_b != magic));\n\n do {\n asm volatile(\"ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];\"\n : \"=f\"(other[1]), \"=r\"(other_flag_a), \"=f\"(other[3]), \"=r\"(other_flag_b)\n : \"l\"(read_data + 4));\n } while ((other_flag_a != magic) || (other_flag_b != magic));\n\n add(x, other);\n }\n // finally, after syncing up and accounting for partial sums from\n // other GPUs as required, write the result\n\n write_to_smem(smem, threadIdx.x, x);\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void parallel_sums_8x4(float* smem, float (&x)[4], int nhw) {\n // The size of a warp.\n const int THREADS_PER_WARP = 32;\n // The number of warps in a CTA.\n const int WARPS_PER_CTA = THREADS_PER_CTA / THREADS_PER_WARP;\n // The number of threads per pixel.\n const int THREADS_PER_PIXEL = 8;\n // The number of elements per ldg.\n const int ELEMENTS_PER_LDG = 4;\n // The warp decomposition.\n const int warp_id = threadIdx.x / THREADS_PER_WARP;\n const int lane_id = threadIdx.x % THREADS_PER_WARP;\n\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL + lane_id);\n x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL * 2 + lane_id);\n }\n\n // The warp leaders, write to SMEM.\n if (lane_id < THREADS_PER_PIXEL) {\n write_to_smem(smem, warp_id * THREADS_PER_PIXEL + lane_id, x);\n }\n\n // The data is in SMEM. Do the final reduction.\n __syncthreads();\n\n // The 1st warp does all the work.\n // We do the final reduction each half-warp sequentially reduces the final values.\n if (warp_id == 0) {\n read_from_smem(x, smem, threadIdx.x);\n\n#pragma unroll\n for (int offset = 1; offset < WARPS_PER_CTA / (THREADS_PER_WARP / THREADS_PER_PIXEL); ++offset) {\n float y[ELEMENTS_PER_LDG];\n // Read the mean and variance from the other pixel.\n read_from_smem(y, smem, threadIdx.x + offset * THREADS_PER_WARP);\n // Compute the updated sum.\n add(x, y);\n }\n\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL + lane_id);\n x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL * 2 + lane_id);\n }\n\n // Make sure the data was read from SMEM.\n __syncwarp();\n\n // Store the final values.\n if (threadIdx.x < THREADS_PER_PIXEL) {\n write_to_smem(smem, threadIdx.x, x);\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void parallel_sums(float* smem, float (&x)[ELEMENTS_PER_LDG], int nhw) {\n // The size of a warp.\n const int THREADS_PER_WARP = 32;\n // The number of warps in a CTA.\n const int WARPS_PER_CTA = THREADS_PER_CTA / THREADS_PER_WARP;\n // The number of pixels computed by a single warp.\n const int PIXELS_PER_WARP = THREADS_PER_WARP / THREADS_PER_PIXEL;\n\n // The position in the warp.\n const int nhw_in_warp = nhw % PIXELS_PER_WARP;\n // The C in the warp.\n const int c_in_warp = threadIdx.x % THREADS_PER_PIXEL;\n\n // Store the values to shared memory.\n write_to_smem(smem, threadIdx.x, x);\n\n // Compute the parallel sums.\n for (int offset = PIXELS_PER_WARP / 2; offset > 0; offset /= 2) {\n // NOP.\n __syncwarp();\n\n // Read the running sum from the other thread.\n float y[ELEMENTS_PER_LDG];\n if (nhw_in_warp < offset) {\n read_from_smem(y, smem, threadIdx.x + offset * THREADS_PER_PIXEL);\n }\n\n // Compute the updated sum.\n add(x, y);\n\n // NOP.\n __syncwarp();\n\n // Update the sum in SMEM.\n if (offset > 1 && nhw_in_warp < offset) {\n write_to_smem(smem, threadIdx.x, x);\n }\n }\n\n // The warps are done. Do the final reduction at the CTA level.\n __syncthreads();\n\n // The warp leaders, write to SMEM.\n const int idx = (threadIdx.x / THREADS_PER_WARP) * THREADS_PER_PIXEL + c_in_warp;\n if (nhw_in_warp == 0) {\n write_to_smem(smem, idx, x);\n }\n\n // The data is in SMEM. Do the final reduction.\n __syncthreads();\n\n // Read the 1st element to prepare the work.\n if (nhw < WARPS_PER_CTA / 2) {\n read_from_smem(x, smem, threadIdx.x);\n }\n\n // We have the running mean and running m2. Let's build the mean/var of the CTA.\n for (int offset = WARPS_PER_CTA / 2; offset > 0; offset /= 2) {\n // NOP.\n __syncwarp();\n\n // Read the mean and variance from the other pixel.\n float y[ELEMENTS_PER_LDG];\n if (nhw < offset) {\n read_from_smem(y, smem, threadIdx.x + offset * THREADS_PER_PIXEL);\n }\n\n // Compute the updated sum.\n add(x, y);\n\n // NOP.\n __syncwarp();\n\n // Store the mean/var for the different pixels.\n if (nhw < offset) {\n write_to_smem(smem, threadIdx.x, x);\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct ParallelSums {\n template \n DEVICE_FUNCTION void dispatch(float* smem, float (&x)[ELEMENTS_PER_LDG], int nhw) {\n parallel_sums(smem, x, nhw);\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate <>\nstruct ParallelSums<16, 4> {\n template \n DEVICE_FUNCTION void dispatch(float* smem, float (&x)[4], int nhw) {\n parallel_sums_16x2(smem, x, nhw, 0, 0, 0, 0, 0);\n }\n\n template \n DEVICE_FUNCTION void dispatchX(float* smem, float (&x)[4], int nhw, void* params_my_data, void** params_pair_datas,\n int off, const int magic, const unsigned int& sync_iters) {\n parallel_sums_16x2(smem, x, nhw, params_my_data, params_pair_datas, off, magic, sync_iters);\n }\n};\n\ntemplate <>\nstruct ParallelSums<8, 4> {\n template \n DEVICE_FUNCTION void dispatch(float* smem, float (&x)[4], int nhw) {\n parallel_sums_8x4(smem, x, nhw);\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstatic inline int div_up(int m, int n) { return (m + n - 1) / n; }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n// It is expected that all threads in the CTA enter this function!\nDEVICE_FUNCTION void inter_block_sync(int* gmem_retired_ctas, int expected_count, bool master) {\n // Register the CTA.\n if (threadIdx.x == 0) {\n // Issue the membar.\n __threadfence();\n // Notify that the CTA is done.\n int val_to_add = 1;\n if (master) {\n val_to_add = -(expected_count - 1);\n }\n atomicAdd(gmem_retired_ctas, val_to_add);\n }\n\n // Are all CTAs done?\n if (threadIdx.x == 0) {\n int retired_ctas = -1;\n do {\n __threadfence();\n asm volatile(\"ld.global.cg.b32 %0, [%1];\" : \"=r\"(retired_ctas) : \"l\"(gmem_retired_ctas));\n } while (retired_ctas != 0);\n }\n __syncthreads();\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct NhwcBatchNormFwdInferenceParams {\n // The input/output tensors.\n uint16_t *gmem_src, *gmem_dst, *gmem_src1;\n // the final mean and variance as calculated during the training process\n float *gmem_mean, *gmem_var;\n // The bias/scale.\n float *gmem_bias, *gmem_scale;\n // The dimensions.\n int nhw, c;\n // epsilon\n float var_eps;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n// No DESIRED_OCCUPANCY launch bounds needed, as this is not launched cooperatively\ntemplate \n__global__ __launch_bounds__(THREADS_PER_CTA) void nhwc_batch_norm_fwd_inference(\n NhwcBatchNormFwdInferenceParams params) {\n // The number of pixels loaded in a single LDG.\n const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;\n // The number of C elements per CTA.\n const int C_ELEMENTS_PER_CTA = THREADS_PER_PIXEL * ELEMENTS_PER_LDG;\n\n // The start position in the NHW dimension where the CTA starts.\n const int cta_nhw_stride = gridDim.x * PIXELS_PER_LDG;\n // Compute the NHW coordinate of the thread in the CTA.\n const int thread_in_cta_nhw = threadIdx.x / THREADS_PER_PIXEL;\n // thread's starting point in NHW\n const int thread_nhw = thread_in_cta_nhw + blockIdx.x * PIXELS_PER_LDG;\n\n // The position in the C dimension where the CTA starts.\n const int cta_c = blockIdx.y * C_ELEMENTS_PER_CTA;\n // Compute the C coordinate of the thread in the CTA.\n const int thread_in_cta_c = threadIdx.x % THREADS_PER_PIXEL;\n // Compute the C coordinate of the thread.\n const int thread_c = cta_c + thread_in_cta_c * ELEMENTS_PER_LDG;\n\n // Is the thread working on a valid C dimension?\n const int is_valid_c = thread_c < params.c;\n\n float mean[ELEMENTS_PER_LDG], var[ELEMENTS_PER_LDG];\n float scale[ELEMENTS_PER_LDG], bias[ELEMENTS_PER_LDG];\n zero_array(mean);\n zero_array(var);\n zero_array(scale);\n zero_array(bias);\n if (is_valid_c) {\n read_from_gmem(var, ¶ms.gmem_var[cta_c], thread_in_cta_c);\n read_from_gmem(scale, ¶ms.gmem_scale[cta_c], thread_in_cta_c);\n read_from_gmem(mean, ¶ms.gmem_mean[cta_c], thread_in_cta_c);\n read_from_gmem(bias, ¶ms.gmem_bias[cta_c], thread_in_cta_c);\n }\n\n// Update the scale with the stddev and eps.\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n scale[i] *= rsqrtf(var[i] + params.var_eps);\n }\n\n // The base pointers for reading/writing\n uint16_t* const gmem_src = ¶ms.gmem_src[thread_c];\n uint16_t* const gmem_dst = ¶ms.gmem_dst[thread_c];\n const uint16_t* gmem_src1 = nullptr;\n if (USE_ADD_RELU) {\n gmem_src1 = ¶ms.gmem_src1[thread_c];\n }\n\n // apply BN\n for (int nhw = thread_nhw; nhw < params.nhw; nhw += cta_nhw_stride) {\n float x_math[ELEMENTS_PER_LDG];\n zero_array(x_math);\n if (is_valid_c) {\n ldg(x_math, &gmem_src[nhw * params.c]);\n }\n\n // Normalize and apply activation function\n normalize(x_math, bias, scale, mean);\n if (USE_ADD_RELU) {\n float x1_math[ELEMENTS_PER_LDG];\n ldg(x1_math, &gmem_src1[nhw * params.c]);\n add(x_math, x1_math);\n relu_activation(x_math);\n } else if (USE_RELU) {\n relu_activation(x_math);\n }\n\n if (is_valid_c) {\n stg(&gmem_dst[nhw * params.c], x_math);\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct NhwcBatchNormFwdParams {\n // The input/output tensors.\n uint16_t *gmem_src, *gmem_dst, *gmem_src1;\n // The bias/scale.\n float *gmem_bias, *gmem_scale;\n // running mean/var (refer BN API from cudnn doc)\n float *gmem_running_mean, *gmem_running_var;\n // saved mean/var (refer BN API from cudnn doc)\n float *gmem_saved_mean, *gmem_saved_var;\n // ReLU bitmask\n unsigned int* gmem_relu_bitmask;\n // The dimensions.\n int nhw, c;\n // factor to scale sum of squared errors to get saved variance. Must be 1/nhw.\n float svar_inv_count;\n // factor to scale sum of squared errors to get running variance. Should be 1/nhw or 1/(nhw-1).\n float rvar_inv_count;\n // The buffer to do the reduction for mean, stddev and count.\n float* gmem_sums;\n // The buffer to count items in the different CTAs.\n int* gmem_counts;\n // The counters of retired CTAs.\n int* gmem_retired_ctas;\n // The epsilon to apply to the computation of the variance.\n float var_eps;\n // outer loop count\n int outer_loops;\n // exponential average factor\n float exp_avg_factor;\n // number of CTAs along .x dimension\n int c_blks;\n\n void* my_data;\n void* pair_datas[4];\n int magic;\n int sync_iters;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ __launch_bounds__(THREADS_PER_CTA,\n DESIRED_OCCUPANCY) void nhwc_batch_norm_fwd(NhwcBatchNormFwdParams params) {\n // The number of pixels loaded in a single LDG.\n const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;\n // The number of pixels computed per CTA stored in registers.\n const int PIXELS_PER_CTA_IN_REGISTERS = PIXELS_PER_THREAD_IN_REGISTERS * PIXELS_PER_LDG;\n // The number of pixels computed per CTA stored in SMEM.\n const int PIXELS_PER_CTA_IN_SMEM = PIXELS_PER_THREAD_IN_SMEM * PIXELS_PER_LDG;\n // The number of C elements per CTA.\n const int C_ELEMENTS_PER_CTA = THREADS_PER_PIXEL * ELEMENTS_PER_LDG;\n\n // Shared memory to do CTA-wide parallel sums.\n __shared__ float smem[THREADS_PER_PIXEL * (THREADS_PER_CTA / 32) * ELEMENTS_PER_LDG];\n\n // Compute the NHW coordinate of the thread in the CTA.\n const int thread_in_cta_nhw = threadIdx.x / THREADS_PER_PIXEL;\n\n // The adapter for the storage.\n typedef PackedStorage PackedStorage_;\n // The data type for packed storage in SMEM.\n typedef typename PackedStorage_::Type PackedStorageType;\n // The number of elements in the packed storage.\n const int PACKED_ELEMENTS_PER_LDG = PackedStorage_::PACKED_ELEMENTS_PER_LDG;\n // Registers to keep the data live for the persistent approach.\n PackedStorageType x_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];\n\n // Shared memory buffer to store the extra pixels.\n extern __shared__ PackedStorageType smem_storage_packed[];\n\n for (int c_blk_index = blockIdx.y; c_blk_index < params.c_blks; c_blk_index += gridDim.y) {\n // The position in the NHW dimension where the CTA starts.\n int cta_nhw_regs = blockIdx.x * PIXELS_PER_CTA_IN_REGISTERS;\n // The position in the NHW dimension where the CTA starts for the portion in SMEM.\n int cta_nhw_smem = blockIdx.x * PIXELS_PER_CTA_IN_SMEM;\n\n // The position in the C dimension where the CTA starts.\n const int cta_c = c_blk_index * C_ELEMENTS_PER_CTA;\n // Compute the C coordinate of the thread in the CTA.\n const int thread_in_cta_c = threadIdx.x % THREADS_PER_PIXEL;\n // Compute the C coordinate of the thread.\n int thread_c = cta_c + thread_in_cta_c * ELEMENTS_PER_LDG;\n\n // Is the thread working on a valid C dimension?\n const int is_valid_c = thread_c < params.c;\n\n // Clamp thread_c so that we load from valid locations even if we don't use the value\n if (!is_valid_c) thread_c = params.c - 4;\n\n // Single pass numerically stable algorithm, see:\n // https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Online_algorithm\n //\n // n = 0, mean = 0.0, M2 = 0.0\n //\n // for x in data:\n // n += 1\n // delta = x - mean\n // mean += delta/n\n // delta2 = x - mean\n // M2 += delta*delta2\n //\n // if n < 2:\n // return float('nan')\n // else:\n // return M2 / (n - 1)\n\n // Register to store the number of elements read so far.\n float count = 0.f, mean[ELEMENTS_PER_LDG], m2[ELEMENTS_PER_LDG];\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n mean[i] = 0.f;\n m2[i] = 0.f;\n }\n\n // The number of elements loaded by this CTA.\n int cta_count = 0;\n // The base pointer to load from.\n const uint16_t* gmem_src = ¶ms.gmem_src[thread_c];\n\n // outer loops\n int OUTER_LOOPS = OUTER_LOOPS_ == 1 ? 1 : params.outer_loops;\n // Load the batch of elements. Compute the mean/var across those elements.\n const int pixels_per_iteration = PIXELS_PER_CTA_IN_REGISTERS * gridDim.x;\n\n if (OUTER_LOOPS_ != 1) {\n // We cannot load everything to store persistently, so let's makes sure registers and\n // smem are fully utilized, offset is evenly divisible by 32\n int offset = (pixels_per_iteration * OUTER_LOOPS + PIXELS_PER_CTA_IN_SMEM * gridDim.x - params.nhw) & ~31;\n cta_nhw_regs -= offset;\n cta_nhw_smem -= offset;\n }\n\n#pragma unroll 1\n for (int loop_i = 0; loop_i < OUTER_LOOPS; ++loop_i) {\n // The nhw position.\n int nhw_regs = cta_nhw_regs + loop_i * pixels_per_iteration;\n // Update the number of elements loaded by this CTA. TODO: Skip if <= 0!!!\n cta_count += max(min(nhw_regs + PIXELS_PER_CTA_IN_REGISTERS, params.nhw) - max(nhw_regs, 0), 0);\n\n // Load the data and compute the local mean/sum and the variance.\n if (USE_ONLINE_APPROACH) {\n // Read the elements from memory.\n float is_valid[PIXELS_PER_THREAD_IN_REGISTERS];\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n const int idx = nhw_regs + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n zero_array(x_storage[i]);\n is_valid[i] = 0.f;\n if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {\n if (loop_i == OUTER_LOOPS - 1) {\n ldg_stream(x_storage[i], &gmem_src[idx * params.c]);\n } else {\n ldg(x_storage[i], &gmem_src[idx * params.c]);\n }\n is_valid[i] = 1.f;\n }\n }\n\n// Do the math.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n // Convert to float.\n float x_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage[i]);\n\n // Update the count.\n count += is_valid[i];\n // Invert the count.\n float inv_count = is_valid[i] ? 1.f / count : 0.f;\n\n// Update the mean and m2 using deltas.\n#pragma unroll\n for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {\n float delta0 = x_math[j] - mean[j];\n mean[j] += delta0 * inv_count;\n float delta1 = x_math[j] - mean[j];\n m2[j] += delta0 * delta1 * is_valid[i];\n }\n }\n } else {\n// Read the elements from memory.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n const int idx = nhw_regs + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n zero_array(x_storage[i]);\n if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {\n if (loop_i == OUTER_LOOPS - 1) {\n ldg_stream(x_storage[i], &gmem_src[idx * params.c]);\n } else {\n ldg(x_storage[i], &gmem_src[idx * params.c]);\n }\n count += 1.f;\n }\n }\n\n// Sum the elements in registers.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n // Convert to float.\n float x_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage[i]);\n\n// Update the mean and m2 using deltas.\n#pragma unroll\n for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {\n mean[j] += x_math[j];\n }\n }\n\n // Compute the mean.\n float inv_count = 1.f / count;\n#pragma unroll\n for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {\n mean[j] *= inv_count;\n }\n\n// Compute the variance.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n // Convert to float.\n float x_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage[i]);\n\n // Is it a valid pixel?\n float is_valid = i < static_cast(count) ? 1.f : 0.f;\n// Update the mean and m2 using deltas.\n#pragma unroll\n for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {\n m2[j] += (x_math[j] - mean[j]) * (x_math[j] - mean[j]) * is_valid;\n }\n }\n }\n }\n\n // The elements to load and store in SMEM.\n int smem_nhw = OUTER_LOOPS * pixels_per_iteration + cta_nhw_smem;\n // Load elements from SMEM, update the CTA count.\n int pixels_in_smem = min(smem_nhw + PIXELS_PER_CTA_IN_SMEM, params.nhw) - max(smem_nhw, 0);\n if (pixels_in_smem > 0) {\n cta_count += pixels_in_smem;\n for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {\n const int idx = smem_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n float is_pixel_valid = (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) ? 1.f : 0.f;\n\n PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG];\n ldg_stream(x_storage_local, &gmem_src[(is_pixel_valid ? idx : 0) * params.c]);\n\n // The offset to store in SMEM.\n const int offset = i * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n // Store in SMEM.\n write_to_smem(&smem_storage_packed[offset], threadIdx.x, x_storage_local);\n // Update the count.\n count += is_pixel_valid;\n // Invert the count.\n float inv_count = is_pixel_valid ? 1.f / count : 0.f;\n\n float x_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage_local);\n// Update the mean and m2 using deltas.\n#pragma unroll\n for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {\n float delta0 = x_math[j] - mean[j];\n mean[j] += delta0 * inv_count;\n float delta1 = x_math[j] - mean[j];\n m2[j] += delta0 * delta1 * is_pixel_valid;\n }\n }\n }\n\n // We scale the mean by the number of elements. It brings more stability.\n float m1[ELEMENTS_PER_LDG];\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n m1[i] = mean[i] * count;\n }\n\n // Run the parallel sum accross the CTA to get the local sum.\n ParallelSums::dispatch(smem, m1, thread_in_cta_nhw);\n __syncthreads();\n\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(m1, smem, thread_in_cta_c);\n __syncthreads();\n\n // Adjust the variance.\n float inv_cta_count = 1.f / static_cast(cta_count);\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n float mean_diff = m1[i] * inv_cta_count - mean[i];\n m2[i] = m2[i] + mean_diff * mean_diff * count;\n }\n\n // Run the parallel sum accross the CTA to get the local adjusted variance.\n ParallelSums::dispatch(smem, m2, thread_in_cta_nhw);\n\n // The workspace in global memory is distributed across the different CTA.\n int gmem_sums_offset = c_blk_index * gridDim.x * C_ELEMENTS_PER_CTA * 2;\n\n // Write the data for the CTA to global memory.\n float* gmem_sums = ¶ms.gmem_sums[gmem_sums_offset];\n if (threadIdx.x < THREADS_PER_PIXEL) {\n const int idx = blockIdx.x * THREADS_PER_PIXEL + threadIdx.x;\n write_to_gmem(&gmem_sums[0], idx, m1);\n write_to_gmem(&gmem_sums[C_ELEMENTS_PER_CTA * gridDim.x], idx, m2);\n }\n\n // The memory location to store the number of pixels per CTA.\n int* gmem_counts = ¶ms.gmem_counts[c_blk_index * gridDim.x];\n if (threadIdx.x == 0) {\n gmem_counts[blockIdx.x] = cta_count;\n }\n\n // Read the bias and scale.\n float bias[ELEMENTS_PER_LDG], scale[ELEMENTS_PER_LDG];\n if (is_valid_c) {\n read_from_gmem(bias, ¶ms.gmem_bias[cta_c], thread_in_cta_c);\n read_from_gmem(scale, ¶ms.gmem_scale[cta_c], thread_in_cta_c);\n }\n\n // The counters to count how many CTAs have retired at this point.\n // A given cta uses the same counter every other time through the outer loop.\n int* gmem_retired_ctas = ¶ms.gmem_retired_ctas[c_blk_index % (2 * gridDim.y)];\n inter_block_sync(gmem_retired_ctas, gridDim.x, blockIdx.x == 0);\n\n// Reset the mean to compute the global mean.\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n m1[i] = 0.f;\n }\n\n// Build the global mean.\n#pragma unroll 1\n for (int idx = threadIdx.x; idx < THREADS_PER_PIXEL * gridDim.x; idx += THREADS_PER_CTA) {\n float tmp[ELEMENTS_PER_LDG];\n read_from_gmem(tmp, gmem_sums, idx);\n add(m1, tmp);\n }\n\n if (params.sync_iters > 0) {\n ParallelSums::dispatchX(\n smem, m1, thread_in_cta_nhw, params.my_data, params.pair_datas, 4 * c_blk_index + 3, params.magic,\n params.sync_iters);\n } else {\n ParallelSums::dispatch(smem, m1, thread_in_cta_nhw);\n }\n __syncthreads();\n\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(m1, smem, thread_in_cta_c);\n __syncthreads();\n\n// Normalize the mean.\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n m1[i] = m1[i] * params.svar_inv_count;\n }\n\n// Reset the variance.\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n m2[i] = 0.f;\n }\n\n // for add+relu fusion\n const uint16_t* gmem_src1 = nullptr;\n if (USE_ADD_RELU) {\n gmem_src1 = ¶ms.gmem_src1[thread_c];\n }\n\n// Build the global variance.\n#pragma unroll 1\n for (int idx = threadIdx.x; idx < THREADS_PER_PIXEL * gridDim.x; idx += THREADS_PER_CTA) {\n // Read the means computed by different CTAs (again). Reuse tmp if we have 1 iteration.\n float tmp_mean[ELEMENTS_PER_LDG], tmp_var[ELEMENTS_PER_LDG];\n read_from_gmem(tmp_mean, &gmem_sums[0], idx);\n read_from_gmem(tmp_var, &gmem_sums[C_ELEMENTS_PER_CTA * gridDim.x], idx);\n\n // Read the number of pixels visited by a given CTA.\n cta_count = __ldg(&gmem_counts[idx / THREADS_PER_PIXEL]);\n\n // Compute the diff to update the variance.\n float mean_diff[ELEMENTS_PER_LDG], inv_cta_count = 1.f / static_cast(cta_count);\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n mean_diff[i] = m1[i] - tmp_mean[i] * inv_cta_count;\n }\n\n// Update the variance.\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n m2[i] += tmp_var[i] + mean_diff[i] * mean_diff[i] * static_cast(cta_count);\n }\n }\n\n if (params.sync_iters > 0) {\n ParallelSums::dispatchX(\n smem, m2, thread_in_cta_nhw, params.my_data, params.pair_datas, 4 * c_blk_index + 2, params.magic,\n params.sync_iters);\n } else {\n ParallelSums::dispatch(smem, m2, thread_in_cta_nhw);\n }\n __syncthreads();\n\n read_from_smem(m2, smem, thread_in_cta_c);\n\n // Finalize the stddev.\n // becasue saved var and running var may have different denominator, we don't do it here\n // scale_(m2, inv_count);\n\n // store the saved mean/var\n float svarinv[ELEMENTS_PER_LDG];\n bool is_valid_for_saving = is_valid_c && blockIdx.x == 0 && thread_in_cta_nhw == 0;\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n svarinv[i] = rsqrtf(m2[i] * params.svar_inv_count + params.var_eps);\n }\n if (is_valid_for_saving) {\n write_to_gmem(params.gmem_saved_mean, thread_c / ELEMENTS_PER_LDG, m1);\n write_to_gmem(params.gmem_saved_var, thread_c / ELEMENTS_PER_LDG, svarinv);\n }\n\n // store the running mean/var\n float rmean[ELEMENTS_PER_LDG], rvar[ELEMENTS_PER_LDG];\n zero_array(rmean);\n zero_array(rvar);\n if (params.exp_avg_factor != 1.f && is_valid_for_saving) {\n read_from_gmem(rmean, params.gmem_running_mean, thread_c / ELEMENTS_PER_LDG);\n read_from_gmem(rvar, params.gmem_running_var, thread_c / ELEMENTS_PER_LDG);\n }\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n rmean[i] = (1.f - params.exp_avg_factor) * rmean[i] + params.exp_avg_factor * m1[i];\n rvar[i] = (1.f - params.exp_avg_factor) * rvar[i] + params.exp_avg_factor * (m2[i] * params.rvar_inv_count);\n }\n if (is_valid_for_saving) {\n write_to_gmem(params.gmem_running_mean, thread_c / ELEMENTS_PER_LDG, rmean);\n write_to_gmem(params.gmem_running_var, thread_c / ELEMENTS_PER_LDG, rvar);\n }\n\n // Update the scale with the stddev and eps.\n multiply(scale, svarinv);\n\n // The base pointer to write to.\n uint16_t* const gmem_dst = ¶ms.gmem_dst[thread_c];\n\n unsigned int* const gmem_relu_bitmask = params.gmem_relu_bitmask + ((params.nhw + 31) & ~31) * 2 * c_blk_index;\n\n// Store the elements in registers.\n#pragma unroll 1\n for (int loop_i = OUTER_LOOPS - 1; loop_i >= 0; --loop_i) {\n // The value for nhw.\n int out_nhw = cta_nhw_regs + loop_i * pixels_per_iteration;\n\n// Normalize the elements and write to memory.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n const int idx = out_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n const bool is_valid_nhw = static_cast(idx) < static_cast(params.nhw);\n const bool is_valid = is_valid_nhw && is_valid_c;\n // Convert to float.\n float x_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage[i]);\n\n // Normalize and apply activation function\n normalize(x_math, bias, scale, m1);\n if (USE_ADD_RELU) {\n float x1_math[ELEMENTS_PER_LDG];\n ldg_stream(x1_math, &gmem_src1[(is_valid ? idx : 0) * params.c]);\n add(x_math, x1_math);\n unsigned int relu_mask;\n int lane_id = threadIdx.x & 31;\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n bool rectified = x_math[i] < 0.0F;\n unsigned int local_relu_mask = __ballot_sync(0xFFFFFFFFU, rectified);\n if (lane_id == i) {\n // Thread 0 remembers the relu_mask from the first time through this\n // loop, Thread 1 the next, Thread 2 the next, and Thread 3 the last.\n relu_mask = local_relu_mask;\n }\n if (rectified) {\n x_math[i] = 0.0F;\n }\n }\n if (is_valid_nhw && (lane_id < ELEMENTS_PER_LDG)) {\n gmem_relu_bitmask[idx * 2 + lane_id] = relu_mask;\n }\n } else if (USE_RELU) {\n relu_activation(x_math);\n }\n\n // Write back.\n if (is_valid) {\n stg_stream(&gmem_dst[idx * params.c], x_math);\n }\n }\n\n // The next value of nhw.\n out_nhw -= pixels_per_iteration;\n\n// Read the next elements from memory.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n const int idx = out_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {\n ldg_stream(x_storage[i], &gmem_src[idx * params.c]);\n }\n }\n }\n\n // Normalize the elements from SMEM and write them out.\n if (pixels_in_smem > 0) {\n#pragma unroll 2\n for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {\n const int idx = smem_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n const bool is_valid_nhw = static_cast(idx) < static_cast(params.nhw);\n const bool is_valid = is_valid_nhw && is_valid_c;\n\n // Read from SMEM.\n const int offset = i * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG];\n read_from_smem(x_storage_local, &smem_storage_packed[offset], threadIdx.x);\n float x_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage_local);\n\n // Normalize and apply activation function\n normalize(x_math, bias, scale, m1);\n if (USE_ADD_RELU) {\n float x1_math[ELEMENTS_PER_LDG];\n ldg_stream(x1_math, &gmem_src1[(is_valid ? idx : 0) * params.c]);\n add(x_math, x1_math);\n unsigned int relu_mask;\n int lane_id = threadIdx.x & 31;\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n bool rectified = x_math[i] < 0.0F;\n unsigned int local_relu_mask = __ballot_sync(0xFFFFFFFFU, rectified);\n if (lane_id == i) {\n relu_mask = local_relu_mask;\n }\n if (rectified) {\n x_math[i] = 0.0F;\n }\n }\n if (is_valid_nhw && (lane_id < ELEMENTS_PER_LDG)) {\n gmem_relu_bitmask[idx * 2 + lane_id] = relu_mask;\n }\n } else if (USE_RELU) {\n relu_activation(x_math);\n }\n\n // Write back.\n if (is_valid) {\n stg_stream(&gmem_dst[idx * params.c], x_math);\n }\n }\n }\n // We're about to start on the next c-blk. Needed?\n __syncthreads();\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct NhwcBatchNormBwdParams {\n // The input/output tensors.\n uint16_t *gmem_src, *gmem_dy, *gmem_dst, *gmem_dst1;\n // dscale/dbias\n float *gmem_dscale, *gmem_dbias;\n // The scale and bias.\n float *gmem_scale, *gmem_bias;\n // The mean/inv-var saved from fwd pass\n float *gmem_saved_mean, *gmem_saved_var;\n // ReLU bitmask\n unsigned int* gmem_relu_bitmask;\n // The dimensions.\n int nhw, c;\n // factor to scale sum of squared errors to get saved variance. Must be 1/nhw.\n float svar_inv_count;\n // The buffer to do the reduction for dscale and dbias\n float* gmem_sums;\n // The counters of retired CTAs.\n int* gmem_retired_ctas;\n // outer loop count\n int outer_loops;\n // number of CTAs along .x dimension\n int c_blks;\n\n void* my_data;\n void* pair_datas[4];\n int magic;\n int sync_iters;\n float wgrad_coeff;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void relu_bwd(float (&dy)[N], const float (&x)[N], const float (&mean_var_scale_bias)[N],\n const float (&var_scale)[N], bool valid_data) {\n#pragma unroll\n for (int j = 0; j < N; ++j) {\n float y = (x[j] * var_scale[j]) + mean_var_scale_bias[j];\n if ((y <= 0.f) && valid_data) {\n dy[j] = 0.f;\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void relu_bwd(float (&dy)[N], const float (&y)[N], bool valid_data) {\n#pragma unroll\n for (int j = 0; j < N; ++j) {\n if ((y[j] <= 0.f) && valid_data) {\n dy[j] = 0.f;\n }\n }\n}\n\ntemplate \nDEVICE_FUNCTION void relu_bwd(float (&dy)[N], const bool (&rectified)[N], bool valid_data) {\n#pragma unroll\n for (int j = 0; j < N; ++j) {\n if (rectified[j] && valid_data) {\n dy[j] = 0.f;\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void relu_bwd_for_dx(float (&dy)[N], const float (&x)[N], const float (&mean_var_scale_bias)[N],\n const float (&var_scale)[N]) {\n#pragma unroll\n for (int j = 0; j < N; ++j) {\n float y = (x[j] * var_scale[j]) + mean_var_scale_bias[j];\n if (y <= 0.f) {\n dy[j] = 0.f;\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void relu_bwd_for_dx(float (&dy)[N], const float (&y)[N]) {\n#pragma unroll\n for (int j = 0; j < N; ++j) {\n if (y[j] <= 0.f) {\n dy[j] = 0.f;\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void bwd_update(float (&dscale)[N], float (&dbias)[N], const float (&dy)[N], const float (&x)[N],\n const float (&mean)[N], float inv_count) {\n#pragma unroll\n for (int j = 0; j < N; ++j) {\n float delta0 = dy[j] - dbias[j];\n dbias[j] += delta0 * inv_count;\n delta0 = (dy[j] * (x[j] - mean[j])) - dscale[j];\n dscale[j] += delta0 * inv_count;\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nDEVICE_FUNCTION void bwd_dx(float (&dx)[N], const float (&dy)[N], const float (&var)[N], const float (&x)[N],\n const float (&mean)[N], const float (&dscale)[N], const float (&dbias)[N],\n float inv_count) {\n#pragma unroll\n for (int j = 0; j < N; ++j) {\n float tmp1 = dy[j] - (dbias[j] * inv_count);\n float tmp2 = dscale[j] * inv_count;\n float tmp3 = x[j] - mean[j];\n dx[j] = var[j] * (tmp1 - (tmp2 * tmp3));\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ __launch_bounds__(THREADS_PER_CTA,\n DESIRED_OCCUPANCY) void nhwc_batch_norm_bwd(NhwcBatchNormBwdParams params) {\n // The number of pixels loaded in a single LDG.\n const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;\n // The number of pixels computed per CTA stored in registers.\n const int PIXELS_PER_CTA_IN_REGISTERS = PIXELS_PER_THREAD_IN_REGISTERS * PIXELS_PER_LDG;\n // The number of pixels computed per CTA stored in SMEM.\n const int PIXELS_PER_CTA_IN_SMEM = PIXELS_PER_THREAD_IN_SMEM * PIXELS_PER_LDG;\n // The number of C elements per CTA.\n const int C_ELEMENTS_PER_CTA = THREADS_PER_PIXEL * ELEMENTS_PER_LDG;\n\n // Shared memory to do CTA-wide parallel sums.\n __shared__ float smem[THREADS_PER_PIXEL * (THREADS_PER_CTA / 32) * ELEMENTS_PER_LDG];\n\n // The adapter for the storage.\n typedef PackedStorage PackedStorage_;\n // The data type for packed storage in SMEM.\n typedef typename PackedStorage_::Type PackedStorageType;\n // The number of elements in the packed storage.\n const int PACKED_ELEMENTS_PER_LDG = PackedStorage_::PACKED_ELEMENTS_PER_LDG;\n // Registers to keep the data live for the persistent approach.\n PackedStorageType x_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];\n PackedStorageType dy_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];\n\n // Shared memory buffer to store the extra pixels.\n extern __shared__ PackedStorageType smem_storage_packed[];\n\n for (int c_blk_index = blockIdx.y; c_blk_index < params.c_blks; c_blk_index += gridDim.y) {\n // The position in the NHW dimension where the CTA starts.\n int cta_nhw_regs = blockIdx.x * PIXELS_PER_CTA_IN_REGISTERS;\n // The position in the NHW dimension where the CTA starts for the portion in SMEM.\n int cta_nhw_smem = blockIdx.x * PIXELS_PER_CTA_IN_SMEM;\n // Compute the NHW coordinate of the thread in the CTA.\n const int thread_in_cta_nhw = threadIdx.x / THREADS_PER_PIXEL;\n\n // The position in the C dimension where the CTA starts.\n const int cta_c = c_blk_index * C_ELEMENTS_PER_CTA;\n // Compute the C coordinate of the thread in the CTA.\n const int thread_in_cta_c = threadIdx.x % THREADS_PER_PIXEL;\n // Compute the C coordinate of the thread.\n const int thread_c = cta_c + thread_in_cta_c * ELEMENTS_PER_LDG;\n\n // Is the thread working on a valid C dimension?\n const int is_valid_c = thread_c < params.c;\n\n // Registers to store the mean used for entire duration\n float mean[ELEMENTS_PER_LDG];\n zero_array(mean);\n if (is_valid_c) {\n read_from_gmem(mean, params.gmem_saved_mean, thread_c / ELEMENTS_PER_LDG);\n }\n\n // accumulation related registers\n float count = 0.f, dscale[ELEMENTS_PER_LDG], dbias[ELEMENTS_PER_LDG];\n zero_array(dscale);\n zero_array(dbias);\n\n // The number of elements loaded by this CTA.\n int cta_count = 0;\n // The base pointers to load from.\n const uint16_t* gmem_src = ¶ms.gmem_src[thread_c];\n const uint16_t* gmem_dy = ¶ms.gmem_dy[thread_c];\n\n // outer loops\n int OUTER_LOOPS = OUTER_LOOPS_ == 1 ? 1 : params.outer_loops;\n // Load the batch of elements. Compute sum across them\n const int pixels_per_iteration = PIXELS_PER_CTA_IN_REGISTERS * gridDim.x;\n\n if (OUTER_LOOPS_ != 1) {\n // We cannot load everything to store persistently, so let's makes sure registers and\n // smem are fully utilized\n int offset = params.nhw - pixels_per_iteration * OUTER_LOOPS - PIXELS_PER_CTA_IN_SMEM * gridDim.x;\n cta_nhw_regs += offset;\n cta_nhw_smem += offset;\n }\n\n#pragma unroll 1\n for (int loop_i = 0; loop_i < OUTER_LOOPS; ++loop_i) {\n // The nhw position.\n int nhw_regs = cta_nhw_regs + loop_i * pixels_per_iteration;\n // Update the number of elements loaded by this CTA. TODO: Skip if <= 0!!!\n cta_count += max(0, min(PIXELS_PER_CTA_IN_REGISTERS, params.nhw - nhw_regs));\n\n // Read the elements from memory.\n float is_valid[PIXELS_PER_THREAD_IN_REGISTERS];\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n const int idx = nhw_regs + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n zero_array(x_storage[i]);\n zero_array(dy_storage[i]);\n is_valid[i] = 0.f;\n if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {\n if (loop_i == OUTER_LOOPS - 1) {\n ldg_stream(x_storage[i], &gmem_src[idx * params.c]);\n ldg_stream(dy_storage[i], &gmem_dy[idx * params.c]);\n } else {\n ldg(x_storage[i], &gmem_src[idx * params.c]);\n ldg(dy_storage[i], &gmem_dy[idx * params.c]);\n }\n is_valid[i] = 1.f;\n }\n }\n\n// Do the math.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n // Convert to float and update\n float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage[i]);\n to_float(dy_math, dy_storage[i]);\n\n // Update the count.\n count += is_valid[i];\n // Invert the count.\n float inv_count = is_valid[i] ? 1.f / count : 0.f;\n\n bwd_update(dscale, dbias, dy_math, x_math, mean, inv_count);\n }\n }\n\n // The elements to load and store in SMEM.\n int smem_nhw = OUTER_LOOPS * pixels_per_iteration + cta_nhw_smem;\n // Load elements from SMEM, update the CTA count.\n int pixels_in_smem = min(PIXELS_PER_CTA_IN_SMEM, params.nhw - smem_nhw);\n if (pixels_in_smem > 0) {\n cta_count += pixels_in_smem;\n for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {\n const int idx = smem_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n bool is_pixel_valid = (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c);\n PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG], dy_storage_local[PACKED_ELEMENTS_PER_LDG];\n zero_array(x_storage_local);\n zero_array(dy_storage_local);\n if (is_pixel_valid) {\n ldg_stream(x_storage_local, &gmem_src[idx * params.c]);\n ldg_stream(dy_storage_local, &gmem_dy[idx * params.c]);\n }\n\n // The offset to store in SMEM.\n int offset = i * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n // Store in SMEM.\n write_to_smem(&smem_storage_packed[offset], threadIdx.x, x_storage_local);\n offset += PIXELS_PER_THREAD_IN_SMEM * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n write_to_smem(&smem_storage_packed[offset], threadIdx.x, dy_storage_local);\n // Update the count.\n count += is_pixel_valid;\n // Invert the count.\n float inv_count = is_pixel_valid ? 1.f / count : 0.f;\n\n float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage_local);\n to_float(dy_math, dy_storage_local);\n\n bwd_update(dscale, dbias, dy_math, x_math, mean, inv_count);\n }\n }\n\n// We scale the mean by the number of elements. It brings more stability.\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n dbias[i] *= count;\n dscale[i] *= count;\n }\n\n // dscale parallel sum\n ParallelSums::dispatch(smem, dscale, thread_in_cta_nhw);\n __syncthreads();\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(dscale, smem, thread_in_cta_c);\n __syncthreads();\n\n // dbias parallel sum\n ParallelSums::dispatch(smem, dbias, thread_in_cta_nhw);\n __syncthreads();\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(dbias, smem, thread_in_cta_c);\n __syncthreads();\n\n // The workspace in global memory is distributed across the different CTA.\n int gmem_sums_offset = c_blk_index * gridDim.x * C_ELEMENTS_PER_CTA * 2;\n // Write the data for the CTA to global memory.\n float* gmem_sums = ¶ms.gmem_sums[gmem_sums_offset];\n if (threadIdx.x < THREADS_PER_PIXEL) {\n const int idx = blockIdx.x * THREADS_PER_PIXEL + threadIdx.x;\n write_to_gmem(&gmem_sums[0], idx, dscale);\n write_to_gmem(&gmem_sums[C_ELEMENTS_PER_CTA * gridDim.x], idx, dbias);\n }\n\n // The counters to count how many CTAs have retired at this point.\n // A given cta uses the same counter every other time through the outer loop.\n int* gmem_retired_ctas = ¶ms.gmem_retired_ctas[c_blk_index % (2 * gridDim.y)];\n inter_block_sync(gmem_retired_ctas, gridDim.x, blockIdx.x == 0);\n\n // Reset the accumulators for global summation\n zero_array(dscale);\n zero_array(dbias);\n\n// Build the global accumulation\n#pragma unroll 1\n for (int idx = threadIdx.x; idx < THREADS_PER_PIXEL * gridDim.x; idx += THREADS_PER_CTA) {\n float tmp1[ELEMENTS_PER_LDG], tmp2[ELEMENTS_PER_LDG];\n read_from_gmem(tmp1, gmem_sums, idx);\n read_from_gmem(tmp2, gmem_sums + C_ELEMENTS_PER_CTA * gridDim.x, idx);\n\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n dscale[i] += tmp1[i];\n dbias[i] += tmp2[i];\n }\n }\n\n // dscale parallel sum\n if (params.sync_iters > 0) {\n ParallelSums::dispatchX(\n smem, dscale, thread_in_cta_nhw, params.my_data, params.pair_datas, 4 * c_blk_index + 1, params.magic,\n params.sync_iters);\n } else {\n ParallelSums::dispatch(smem, dscale, thread_in_cta_nhw);\n }\n\n __syncthreads();\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(dscale, smem, thread_in_cta_c);\n __syncthreads();\n\n // dbias parallel sum\n if (params.sync_iters > 0) {\n ParallelSums::dispatchX(\n smem, dbias, thread_in_cta_nhw, params.my_data, params.pair_datas, 4 * c_blk_index + 0, params.magic,\n params.sync_iters);\n } else {\n ParallelSums::dispatch(smem, dbias, thread_in_cta_nhw);\n }\n\n __syncthreads();\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(dbias, smem, thread_in_cta_c);\n\n // inv-var\n float var[ELEMENTS_PER_LDG];\n zero_array(var);\n if (is_valid_c) {\n read_from_gmem(var, params.gmem_saved_var, thread_c / ELEMENTS_PER_LDG);\n }\n\n // Normalize the dscale.\n multiply(dscale, var);\n\n // store dscale/dbias\n bool is_valid_for_saving = is_valid_c && blockIdx.x == 0 && thread_in_cta_nhw == 0;\n if (is_valid_for_saving) {\n if (params.sync_iters > 0) {\n scaled_write_to_gmem(params.gmem_dscale, thread_c / ELEMENTS_PER_LDG, dscale, params.wgrad_coeff);\n scaled_write_to_gmem(params.gmem_dbias, thread_c / ELEMENTS_PER_LDG, dbias, params.wgrad_coeff);\n } else {\n write_to_gmem(params.gmem_dscale, thread_c / ELEMENTS_PER_LDG, dscale);\n write_to_gmem(params.gmem_dbias, thread_c / ELEMENTS_PER_LDG, dbias);\n }\n }\n\n // scale\n float scale[ELEMENTS_PER_LDG];\n zero_array(scale);\n if (is_valid_c) {\n read_from_gmem(scale, params.gmem_scale, thread_c / ELEMENTS_PER_LDG);\n }\n\n // Further normalize the dscale to be used in dx calculation\n multiply(dscale, var);\n // scale the inv-var as well, afterwards\n multiply(var, scale);\n\n // inverse count\n float inv_count = params.svar_inv_count;\n\n // The base pointer to write to.\n uint16_t* const gmem_dst = ¶ms.gmem_dst[thread_c];\n\n// Store the elements in registers.\n#pragma unroll 1\n for (int loop_i = OUTER_LOOPS - 1; loop_i >= 0; --loop_i) {\n // The value for nhw.\n int out_nhw = cta_nhw_regs + loop_i * pixels_per_iteration;\n\n// Normalize the elements and write to memory.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n // Convert to float.\n float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage[i]);\n to_float(dy_math, dy_storage[i]);\n\n float dx[ELEMENTS_PER_LDG];\n bwd_dx(dx, dy_math, var, x_math, mean, dscale, dbias, inv_count);\n\n // Write back.\n const int idx = out_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {\n stg_stream(&gmem_dst[idx * params.c], dx);\n }\n }\n\n // The next value of nhw.\n out_nhw -= pixels_per_iteration;\n\n// Read the next elements from memory.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n const int idx = out_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {\n ldg_stream(x_storage[i], &gmem_src[idx * params.c]);\n ldg_stream(dy_storage[i], &gmem_dy[idx * params.c]);\n }\n }\n }\n\n // Normalize the elements from SMEM and write them out.\n if (pixels_in_smem > 0) {\n for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {\n const int idx = smem_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n const bool is_valid = ((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c;\n if (is_valid) {\n // Read from SMEM.\n int offset = i * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG], dy_storage_local[PACKED_ELEMENTS_PER_LDG];\n read_from_smem(x_storage_local, &smem_storage_packed[offset], threadIdx.x);\n offset += PIXELS_PER_THREAD_IN_SMEM * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n read_from_smem(dy_storage_local, &smem_storage_packed[offset], threadIdx.x);\n float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage_local);\n to_float(dy_math, dy_storage_local);\n\n float dx[ELEMENTS_PER_LDG];\n bwd_dx(dx, dy_math, var, x_math, mean, dscale, dbias, inv_count);\n\n // Write back.\n stg_stream(&gmem_dst[idx * params.c], dx);\n }\n }\n }\n // We're about to start on the next c-blk. Needed?\n __syncthreads();\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ __launch_bounds__(THREADS_PER_CTA,\n DESIRED_OCCUPANCY) void nhwc_batch_norm_bwd_relu(NhwcBatchNormBwdParams params) {\n // The number of pixels loaded in a single LDG.\n const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;\n // The number of pixels computed per CTA stored in registers.\n const int PIXELS_PER_CTA_IN_REGISTERS = PIXELS_PER_THREAD_IN_REGISTERS * PIXELS_PER_LDG;\n // The number of pixels computed per CTA stored in SMEM.\n const int PIXELS_PER_CTA_IN_SMEM = PIXELS_PER_THREAD_IN_SMEM * PIXELS_PER_LDG;\n // The number of C elements per CTA.\n const int C_ELEMENTS_PER_CTA = THREADS_PER_PIXEL * ELEMENTS_PER_LDG;\n\n // Shared memory to do CTA-wide parallel sums.\n __shared__ float smem[THREADS_PER_PIXEL * (THREADS_PER_CTA / 32) * ELEMENTS_PER_LDG];\n\n // The adapter for the storage.\n typedef PackedStorage PackedStorage_;\n // The data type for packed storage in SMEM.\n typedef typename PackedStorage_::Type PackedStorageType;\n // The number of elements in the packed storage.\n const int PACKED_ELEMENTS_PER_LDG = PackedStorage_::PACKED_ELEMENTS_PER_LDG;\n // Registers to keep the data live for the persistent approach.\n PackedStorageType x_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];\n PackedStorageType dy_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];\n\n // Shared memory buffer to store the extra pixels.\n extern __shared__ PackedStorageType smem_storage_packed[];\n\n for (int c_blk_index = blockIdx.y; c_blk_index < params.c_blks; c_blk_index += gridDim.y) {\n // The position in the NHW dimension where the CTA starts.\n int cta_nhw_regs = blockIdx.x * PIXELS_PER_CTA_IN_REGISTERS;\n // The position in the NHW dimension where the CTA starts for the portion in SMEM.\n int cta_nhw_smem = blockIdx.x * PIXELS_PER_CTA_IN_SMEM;\n // Compute the NHW coordinate of the thread in the CTA.\n const int thread_in_cta_nhw = threadIdx.x / THREADS_PER_PIXEL;\n\n // The position in the C dimension where the CTA starts.\n const int cta_c = c_blk_index * C_ELEMENTS_PER_CTA;\n // Compute the C coordinate of the thread in the CTA.\n const int thread_in_cta_c = threadIdx.x % THREADS_PER_PIXEL;\n // Compute the C coordinate of the thread.\n const int thread_c = cta_c + thread_in_cta_c * ELEMENTS_PER_LDG;\n\n // Is the thread working on a valid C dimension?\n const int is_valid_c = thread_c < params.c;\n\n // Registers to store the mean/var/scale/bias used for the entire duration\n // Register usage optimizations:\n // 1. Can combine bias - (mean * var * scale) into a single register\n // 2. Can combine var * scale into a single register\n float varscale[ELEMENTS_PER_LDG];\n zero_array(varscale);\n if (is_valid_c) {\n read_from_gmem(varscale, params.gmem_saved_var, thread_c / ELEMENTS_PER_LDG);\n }\n float tmp[ELEMENTS_PER_LDG];\n zero_array(tmp);\n if (is_valid_c) {\n read_from_gmem(tmp, params.gmem_scale, thread_c / ELEMENTS_PER_LDG);\n }\n multiply(varscale, tmp);\n float mean[ELEMENTS_PER_LDG];\n zero_array(mean);\n if (is_valid_c) {\n read_from_gmem(mean, params.gmem_saved_mean, thread_c / ELEMENTS_PER_LDG);\n }\n zero_array(tmp);\n if (is_valid_c) {\n read_from_gmem(tmp, params.gmem_bias, thread_c / ELEMENTS_PER_LDG);\n }\n float mean_var_scale_bias[ELEMENTS_PER_LDG];\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n mean_var_scale_bias[i] = tmp[i] - (mean[i] * varscale[i]);\n }\n\n // accumulation related registers\n float count = 0.f, dscale[ELEMENTS_PER_LDG], dbias[ELEMENTS_PER_LDG];\n zero_array(dscale);\n zero_array(dbias);\n\n // The number of elements loaded by this CTA.\n int cta_count = 0;\n // The base pointers to load from.\n const uint16_t* gmem_src = ¶ms.gmem_src[thread_c];\n const uint16_t* gmem_dy = ¶ms.gmem_dy[thread_c];\n\n // outer loops\n int OUTER_LOOPS = OUTER_LOOPS_ == 1 ? 1 : params.outer_loops;\n // Load the batch of elements. Compute sum across them\n const int pixels_per_iteration = PIXELS_PER_CTA_IN_REGISTERS * gridDim.x;\n\n if (OUTER_LOOPS_ != 1) {\n // We cannot load everything to store persistently, so let's makes sure registers and\n // smem are fully utilized\n int offset = params.nhw - pixels_per_iteration * OUTER_LOOPS - PIXELS_PER_CTA_IN_SMEM * gridDim.x;\n cta_nhw_regs += offset;\n cta_nhw_smem += offset;\n }\n\n#pragma unroll 1\n for (int loop_i = 0; loop_i < OUTER_LOOPS; ++loop_i) {\n // The nhw position.\n int nhw_regs = cta_nhw_regs + loop_i * pixels_per_iteration;\n // Update the number of elements loaded by this CTA. TODO: Skip if <= 0!!!\n cta_count += max(0, min(PIXELS_PER_CTA_IN_REGISTERS, params.nhw - nhw_regs));\n\n // Read the elements from memory.\n float is_valid[PIXELS_PER_THREAD_IN_REGISTERS];\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n const int idx = nhw_regs + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n zero_array(x_storage[i]);\n zero_array(dy_storage[i]);\n is_valid[i] = 0.f;\n if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {\n if (loop_i == OUTER_LOOPS - 1) {\n ldg_stream(x_storage[i], &gmem_src[idx * params.c]);\n ldg_stream(dy_storage[i], &gmem_dy[idx * params.c]);\n } else {\n ldg(x_storage[i], &gmem_src[idx * params.c]);\n ldg(dy_storage[i], &gmem_dy[idx * params.c]);\n }\n is_valid[i] = 1.f;\n }\n }\n\n// Do the math.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n // Convert to float and update\n float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage[i]);\n to_float(dy_math, dy_storage[i]);\n\n // Update the count.\n count += is_valid[i];\n // Invert the count.\n float inv_count = is_valid[i] ? 1.f / count : 0.f;\n\n relu_bwd(dy_math, x_math, mean_var_scale_bias, varscale, is_valid[i]);\n bwd_update(dscale, dbias, dy_math, x_math, mean, inv_count);\n }\n }\n\n // The elements to load and store in SMEM.\n int smem_nhw = OUTER_LOOPS * pixels_per_iteration + cta_nhw_smem;\n // Load elements from SMEM, update the CTA count.\n int pixels_in_smem = min(PIXELS_PER_CTA_IN_SMEM, params.nhw - smem_nhw);\n if (pixels_in_smem > 0) {\n cta_count += pixels_in_smem;\n for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {\n const int idx = smem_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n bool is_pixel_valid = (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c);\n PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG], dy_storage_local[PACKED_ELEMENTS_PER_LDG];\n zero_array(x_storage_local);\n zero_array(dy_storage_local);\n if (is_pixel_valid) {\n ldg_stream(x_storage_local, &gmem_src[idx * params.c]);\n ldg_stream(dy_storage_local, &gmem_dy[idx * params.c]);\n }\n\n // The offset to store in SMEM.\n int offset = i * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n // Store in SMEM.\n write_to_smem(&smem_storage_packed[offset], threadIdx.x, x_storage_local);\n offset += PIXELS_PER_THREAD_IN_SMEM * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n write_to_smem(&smem_storage_packed[offset], threadIdx.x, dy_storage_local);\n // Update the count.\n count += is_pixel_valid;\n // Invert the count.\n float inv_count = is_pixel_valid ? 1.f / count : 0.f;\n\n float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage_local);\n to_float(dy_math, dy_storage_local);\n\n relu_bwd(dy_math, x_math, mean_var_scale_bias, varscale, is_pixel_valid);\n bwd_update(dscale, dbias, dy_math, x_math, mean, inv_count);\n }\n }\n\n// We scale the mean by the number of elements. It brings more stability.\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n dbias[i] *= count;\n dscale[i] *= count;\n }\n\n // dscale parallel sum\n ParallelSums::dispatch(smem, dscale, thread_in_cta_nhw);\n __syncthreads();\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(dscale, smem, thread_in_cta_c);\n __syncthreads();\n\n // dbias parallel sum\n ParallelSums::dispatch(smem, dbias, thread_in_cta_nhw);\n __syncthreads();\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(dbias, smem, thread_in_cta_c);\n __syncthreads();\n\n // The workspace in global memory is distributed across the different CTA.\n int gmem_sums_offset = c_blk_index * gridDim.x * C_ELEMENTS_PER_CTA * 2;\n // Write the data for the CTA to global memory.\n float* gmem_sums = ¶ms.gmem_sums[gmem_sums_offset];\n if (threadIdx.x < THREADS_PER_PIXEL) {\n const int idx = blockIdx.x * THREADS_PER_PIXEL + threadIdx.x;\n write_to_gmem(&gmem_sums[0], idx, dscale);\n write_to_gmem(&gmem_sums[C_ELEMENTS_PER_CTA * gridDim.x], idx, dbias);\n }\n\n // The counters to count how many CTAs have retired at this point.\n // A given cta uses the same counter every other time through the outer loop.\n int* gmem_retired_ctas = ¶ms.gmem_retired_ctas[c_blk_index % (2 * gridDim.y)];\n inter_block_sync(gmem_retired_ctas, gridDim.x, blockIdx.x == 0);\n\n // Reset the accumulators for global summation\n zero_array(dscale);\n zero_array(dbias);\n\n// Build the global accumulation\n#pragma unroll 1\n for (int idx = threadIdx.x; idx < THREADS_PER_PIXEL * gridDim.x; idx += THREADS_PER_CTA) {\n float tmp1[ELEMENTS_PER_LDG], tmp2[ELEMENTS_PER_LDG];\n read_from_gmem(tmp1, gmem_sums, idx);\n read_from_gmem(tmp2, gmem_sums + C_ELEMENTS_PER_CTA * gridDim.x, idx);\n\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n dscale[i] += tmp1[i];\n dbias[i] += tmp2[i];\n }\n }\n\n // dscale parallel sum\n if (params.sync_iters > 0) {\n ParallelSums::dispatchX(\n smem, dscale, thread_in_cta_nhw, params.my_data, params.pair_datas, 4 * c_blk_index + 1, params.magic,\n params.sync_iters);\n } else {\n ParallelSums::dispatch(smem, dscale, thread_in_cta_nhw);\n }\n\n __syncthreads();\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(dscale, smem, thread_in_cta_c);\n __syncthreads();\n\n // dbias parallel sum\n if (params.sync_iters > 0) {\n ParallelSums::dispatchX(\n smem, dbias, thread_in_cta_nhw, params.my_data, params.pair_datas, 4 * c_blk_index + 0, params.magic,\n params.sync_iters);\n } else {\n ParallelSums::dispatch(smem, dbias, thread_in_cta_nhw);\n }\n\n __syncthreads();\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(dbias, smem, thread_in_cta_c);\n\n // Normalize the dscale.\n float var[ELEMENTS_PER_LDG];\n zero_array(var);\n if (is_valid_c) {\n read_from_gmem(var, params.gmem_saved_var, thread_c / ELEMENTS_PER_LDG);\n }\n multiply(dscale, var);\n\n // store dscale/dbias\n bool is_valid_for_saving = is_valid_c && blockIdx.x == 0 && thread_in_cta_nhw == 0;\n if (is_valid_for_saving) {\n if (params.sync_iters > 0) {\n scaled_write_to_gmem(params.gmem_dscale, thread_c / ELEMENTS_PER_LDG, dscale, params.wgrad_coeff);\n scaled_write_to_gmem(params.gmem_dbias, thread_c / ELEMENTS_PER_LDG, dbias, params.wgrad_coeff);\n } else {\n write_to_gmem(params.gmem_dscale, thread_c / ELEMENTS_PER_LDG, dscale);\n write_to_gmem(params.gmem_dbias, thread_c / ELEMENTS_PER_LDG, dbias);\n }\n }\n\n // Further normalize the dscale to be used in dx calculation\n float scale[ELEMENTS_PER_LDG];\n zero_array(scale);\n if (is_valid_c) {\n read_from_gmem(scale, params.gmem_scale, thread_c / ELEMENTS_PER_LDG);\n }\n multiply(dscale, var);\n // scale the inv-var as well, afterwards\n multiply(var, scale);\n\n // inverse count\n float inv_count = params.svar_inv_count;\n\n // The base pointer to write to.\n uint16_t* const gmem_dst = ¶ms.gmem_dst[thread_c];\n\n// Store the elements in registers.\n#pragma unroll 1\n for (int loop_i = OUTER_LOOPS - 1; loop_i >= 0; --loop_i) {\n // The value for nhw.\n int out_nhw = cta_nhw_regs + loop_i * pixels_per_iteration;\n\n// Normalize the elements and write to memory.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n // Convert to float.\n float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage[i]);\n to_float(dy_math, dy_storage[i]);\n relu_bwd_for_dx(dy_math, x_math, mean_var_scale_bias, var);\n\n float dx[ELEMENTS_PER_LDG];\n bwd_dx(dx, dy_math, var, x_math, mean, dscale, dbias, inv_count);\n\n // Write back.\n const int idx = out_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {\n stg_stream(&gmem_dst[idx * params.c], dx);\n }\n }\n\n // The next value of nhw.\n out_nhw -= pixels_per_iteration;\n\n// Read the next elements from memory.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n const int idx = out_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {\n ldg_stream(x_storage[i], &gmem_src[idx * params.c]);\n ldg_stream(dy_storage[i], &gmem_dy[idx * params.c]);\n }\n }\n }\n\n // Normalize the elements from SMEM and write them out.\n if (pixels_in_smem > 0) {\n for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {\n const int idx = smem_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n const bool is_valid = ((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c;\n if (is_valid) {\n // Read from SMEM.\n int offset = i * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG], dy_storage_local[PACKED_ELEMENTS_PER_LDG];\n read_from_smem(x_storage_local, &smem_storage_packed[offset], threadIdx.x);\n offset += PIXELS_PER_THREAD_IN_SMEM * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n read_from_smem(dy_storage_local, &smem_storage_packed[offset], threadIdx.x);\n float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage_local);\n to_float(dy_math, dy_storage_local);\n relu_bwd_for_dx(dy_math, x_math, mean_var_scale_bias, var);\n\n float dx[ELEMENTS_PER_LDG];\n bwd_dx(dx, dy_math, var, x_math, mean, dscale, dbias, inv_count);\n\n // Write back.\n stg_stream(&gmem_dst[idx * params.c], dx);\n }\n }\n }\n // We're about to start on the next c-blk. Needed?\n __syncthreads();\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ __launch_bounds__(THREADS_PER_CTA,\n DESIRED_OCCUPANCY) void nhwc_batch_norm_bwd_add_relu(NhwcBatchNormBwdParams params) {\n // The number of pixels loaded in a single LDG.\n const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;\n // The number of pixels computed per CTA stored in registers.\n const int PIXELS_PER_CTA_IN_REGISTERS = PIXELS_PER_THREAD_IN_REGISTERS * PIXELS_PER_LDG;\n // The number of pixels computed per CTA stored in SMEM.\n const int PIXELS_PER_CTA_IN_SMEM = PIXELS_PER_THREAD_IN_SMEM * PIXELS_PER_LDG;\n // The number of C elements per CTA.\n const int C_ELEMENTS_PER_CTA = THREADS_PER_PIXEL * ELEMENTS_PER_LDG;\n\n // Shared memory to do CTA-wide parallel sums.\n __shared__ float smem[THREADS_PER_PIXEL * (THREADS_PER_CTA / 32) * ELEMENTS_PER_LDG];\n\n // The adapter for the storage.\n typedef PackedStorage PackedStorage_;\n // The data type for packed storage in SMEM.\n typedef typename PackedStorage_::Type PackedStorageType;\n // The number of elements in the packed storage.\n const int PACKED_ELEMENTS_PER_LDG = PackedStorage_::PACKED_ELEMENTS_PER_LDG;\n // Registers to keep the data live for the persistent approach.\n PackedStorageType x_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];\n PackedStorageType dy_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];\n\n // Shared memory buffer to store the extra pixels.\n extern __shared__ PackedStorageType smem_storage_packed[];\n\n for (int c_blk_index = blockIdx.y; c_blk_index < params.c_blks; c_blk_index += gridDim.y) {\n // The position in the NHW dimension where the CTA starts.\n int cta_nhw_regs = blockIdx.x * PIXELS_PER_CTA_IN_REGISTERS;\n // The position in the NHW dimension where the CTA starts for the portion in SMEM.\n int cta_nhw_smem = blockIdx.x * PIXELS_PER_CTA_IN_SMEM;\n // Compute the NHW coordinate of the thread in the CTA.\n const int thread_in_cta_nhw = threadIdx.x / THREADS_PER_PIXEL;\n\n // The position in the C dimension where the CTA starts.\n const int cta_c = c_blk_index * C_ELEMENTS_PER_CTA;\n // Compute the C coordinate of the thread in the CTA.\n const int thread_in_cta_c = threadIdx.x % THREADS_PER_PIXEL;\n // Compute the C coordinate of the thread.\n const int thread_c = cta_c + thread_in_cta_c * ELEMENTS_PER_LDG;\n\n // Is the thread working on a valid C dimension?\n const int is_valid_c = thread_c < params.c;\n\n float mean[ELEMENTS_PER_LDG];\n zero_array(mean);\n if (is_valid_c) {\n read_from_gmem(mean, params.gmem_saved_mean, thread_c / ELEMENTS_PER_LDG);\n }\n\n // accumulation related registers\n float count = 0.f, dscale[ELEMENTS_PER_LDG], dbias[ELEMENTS_PER_LDG];\n zero_array(dscale);\n zero_array(dbias);\n\n // The number of elements loaded by this CTA.\n int cta_count = 0;\n // The base pointers to load from.\n const uint16_t* gmem_src = ¶ms.gmem_src[thread_c];\n const uint16_t* gmem_dy = ¶ms.gmem_dy[thread_c];\n uint16_t* gmem_dst1 = ¶ms.gmem_dst1[thread_c];\n\n // outer loops\n int OUTER_LOOPS = OUTER_LOOPS_ == 1 ? 1 : params.outer_loops;\n // Load the batch of elements. Compute sum across them\n const int pixels_per_iteration = PIXELS_PER_CTA_IN_REGISTERS * gridDim.x;\n\n if (OUTER_LOOPS_ != 1) {\n // We cannot load everything to store persistently, so let's makes sure registers and\n // smem are fully utilized, offset is evenly divisible by 32\n int offset = (pixels_per_iteration * OUTER_LOOPS + PIXELS_PER_CTA_IN_SMEM * gridDim.x - params.nhw) & ~31;\n cta_nhw_regs -= offset;\n cta_nhw_smem -= offset;\n }\n\n const unsigned int* const gmem_relu_bitmask =\n params.gmem_relu_bitmask + ((params.nhw + 31) & ~31) * 2 * c_blk_index;\n\n#pragma unroll 1\n for (int loop_i = 0; loop_i < OUTER_LOOPS; ++loop_i) {\n // The nhw position.\n int nhw_regs = cta_nhw_regs + loop_i * pixels_per_iteration;\n // Update the number of elements loaded by this CTA. TODO: Skip if <= 0!!!\n cta_count += max(0, min(PIXELS_PER_CTA_IN_REGISTERS, params.nhw - nhw_regs));\n\n int lane_id = threadIdx.x & 31;\n\n // Read the elements from memory.\n float is_valid[PIXELS_PER_THREAD_IN_REGISTERS];\n unsigned int relu_mask[PIXELS_PER_THREAD_IN_REGISTERS];\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n const int idx = nhw_regs + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n zero_array(x_storage[i]);\n zero_array(dy_storage[i]);\n is_valid[i] = 0.f;\n const bool is_valid_nhw = static_cast(idx) < static_cast(params.nhw);\n if (is_valid_nhw) {\n if (is_valid_c) {\n if (loop_i == OUTER_LOOPS - 1) {\n ldg_stream(x_storage[i], &gmem_src[idx * params.c]);\n ldg_stream(dy_storage[i], &gmem_dy[idx * params.c]);\n } else {\n ldg(x_storage[i], &gmem_src[idx * params.c]);\n ldg(dy_storage[i], &gmem_dy[idx * params.c]);\n }\n is_valid[i] = 1.f;\n }\n\n if (lane_id < ELEMENTS_PER_LDG) {\n relu_mask[i] = gmem_relu_bitmask[idx * 2 + lane_id];\n }\n }\n }\n\n// Do the math.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n const int idx = nhw_regs + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n // Convert to float and update\n float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];\n bool rectified[ELEMENTS_PER_LDG];\n#pragma unroll\n for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {\n rectified[j] = ((__shfl_sync(0xFFFFFFFFU, relu_mask[i], j) & (1U << lane_id)) != 0);\n }\n to_float(x_math, x_storage[i]);\n to_float(dy_math, dy_storage[i]);\n\n // Update the count.\n count += is_valid[i];\n // Invert the count.\n float inv_count = is_valid[i] ? 1.f / count : 0.f;\n\n relu_bwd(dy_math, rectified, is_valid[i]);\n bwd_update(dscale, dbias, dy_math, x_math, mean, inv_count);\n\n // Lastly we need 'dy' only for BN, so store the 'relu-dgrad'ed version\n from_float(dy_storage[i], dy_math);\n\n // dZ for elementwise add\n if (is_valid[i]) {\n if (loop_i == OUTER_LOOPS - 1) {\n stg_stream(&gmem_dst1[idx * params.c], dy_storage[i]);\n } else {\n stg(&gmem_dst1[idx * params.c], dy_storage[i]);\n }\n }\n }\n }\n\n // The elements to load and store in SMEM.\n int smem_nhw = OUTER_LOOPS * pixels_per_iteration + cta_nhw_smem;\n // Load elements from SMEM, update the CTA count.\n int pixels_in_smem = min(PIXELS_PER_CTA_IN_SMEM, params.nhw - smem_nhw);\n if (pixels_in_smem > 0) {\n cta_count += pixels_in_smem;\n for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {\n const int idx = smem_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n const bool is_pixel_valid_nhw = static_cast(idx) < static_cast(params.nhw);\n const bool is_pixel_valid = is_pixel_valid_nhw && is_valid_c;\n PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG], dy_storage_local[PACKED_ELEMENTS_PER_LDG];\n unsigned int relu_mask;\n int lane_id = threadIdx.x & 31;\n zero_array(x_storage_local);\n zero_array(dy_storage_local);\n if (is_pixel_valid_nhw) {\n if (is_valid_c) {\n ldg_stream(x_storage_local, &gmem_src[idx * params.c]);\n ldg_stream(dy_storage_local, &gmem_dy[idx * params.c]);\n }\n if (lane_id < ELEMENTS_PER_LDG) {\n relu_mask = gmem_relu_bitmask[idx * 2 + lane_id];\n }\n }\n bool rectified[ELEMENTS_PER_LDG];\n#pragma unroll\n for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {\n rectified[j] = ((__shfl_sync(0xFFFFFFFFU, relu_mask, j) & (1U << lane_id)) != 0);\n }\n\n // The offset to store in SMEM.\n int offset = i * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n // Store in SMEM.\n write_to_smem(&smem_storage_packed[offset], threadIdx.x, x_storage_local);\n offset += PIXELS_PER_THREAD_IN_SMEM * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n // Update the count.\n count += is_pixel_valid;\n // Invert the count.\n float inv_count = is_pixel_valid ? 1.f / count : 0.f;\n\n float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage_local);\n to_float(dy_math, dy_storage_local);\n\n relu_bwd(dy_math, rectified, is_pixel_valid);\n bwd_update(dscale, dbias, dy_math, x_math, mean, inv_count);\n\n from_float(dy_storage_local, dy_math);\n // dZ for elementwise add\n if (is_pixel_valid) {\n stg_stream(&gmem_dst1[idx * params.c], dy_storage_local);\n }\n // only store the 'relu-dgrad'ed version!\n write_to_smem(&smem_storage_packed[offset], threadIdx.x, dy_storage_local);\n }\n }\n\n// We scale the mean by the number of elements. It brings more stability.\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n dbias[i] *= count;\n dscale[i] *= count;\n }\n\n // dscale parallel sum\n ParallelSums::dispatch(smem, dscale, thread_in_cta_nhw);\n __syncthreads();\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(dscale, smem, thread_in_cta_c);\n __syncthreads();\n\n // dbias parallel sum\n ParallelSums::dispatch(smem, dbias, thread_in_cta_nhw);\n __syncthreads();\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(dbias, smem, thread_in_cta_c);\n __syncthreads();\n\n // The workspace in global memory is distributed across the different CTA.\n int gmem_sums_offset = c_blk_index * gridDim.x * C_ELEMENTS_PER_CTA * 2;\n // Write the data for the CTA to global memory.\n float* gmem_sums = ¶ms.gmem_sums[gmem_sums_offset];\n if (threadIdx.x < THREADS_PER_PIXEL) {\n const int idx = blockIdx.x * THREADS_PER_PIXEL + threadIdx.x;\n write_to_gmem(&gmem_sums[0], idx, dscale);\n write_to_gmem(&gmem_sums[C_ELEMENTS_PER_CTA * gridDim.x], idx, dbias);\n }\n\n // The counters to count how many CTAs have retired at this point.\n // A given cta uses the same counter every other time through the outer loop.\n int* gmem_retired_ctas = ¶ms.gmem_retired_ctas[c_blk_index % (2 * gridDim.y)];\n inter_block_sync(gmem_retired_ctas, gridDim.x, blockIdx.x == 0);\n\n // Reset the accumulators for global summation\n zero_array(dscale);\n zero_array(dbias);\n\n// Build the global accumulation\n#pragma unroll 1\n for (int idx = threadIdx.x; idx < THREADS_PER_PIXEL * gridDim.x; idx += THREADS_PER_CTA) {\n float tmp1[ELEMENTS_PER_LDG], tmp2[ELEMENTS_PER_LDG];\n read_from_gmem(tmp1, gmem_sums, idx);\n read_from_gmem(tmp2, gmem_sums + C_ELEMENTS_PER_CTA * gridDim.x, idx);\n\n#pragma unroll\n for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {\n dscale[i] += tmp1[i];\n dbias[i] += tmp2[i];\n }\n }\n\n // dscale parallel sum\n if (params.sync_iters > 0) {\n ParallelSums::dispatchX(\n smem, dscale, thread_in_cta_nhw, params.my_data, params.pair_datas, 4 * c_blk_index + 1, params.magic,\n params.sync_iters);\n } else {\n ParallelSums::dispatch(smem, dscale, thread_in_cta_nhw);\n }\n\n __syncthreads();\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(dscale, smem, thread_in_cta_c);\n __syncthreads();\n\n // dbias parallel sum\n if (params.sync_iters > 0) {\n ParallelSums::dispatchX(\n smem, dbias, thread_in_cta_nhw, params.my_data, params.pair_datas, 4 * c_blk_index + 0, params.magic,\n params.sync_iters);\n } else {\n ParallelSums::dispatch(smem, dbias, thread_in_cta_nhw);\n }\n\n __syncthreads();\n // The values in shared memory correspond to the CTA-wide sums.\n read_from_smem(dbias, smem, thread_in_cta_c);\n\n // Normalize the dscale.\n float var[ELEMENTS_PER_LDG];\n zero_array(var);\n if (is_valid_c) {\n read_from_gmem(var, params.gmem_saved_var, thread_c / ELEMENTS_PER_LDG);\n }\n multiply(dscale, var);\n\n // store dscale/dbias\n bool is_valid_for_saving = is_valid_c && blockIdx.x == 0 && thread_in_cta_nhw == 0;\n if (is_valid_for_saving) {\n if (params.sync_iters > 0) {\n scaled_write_to_gmem(params.gmem_dscale, thread_c / ELEMENTS_PER_LDG, dscale, params.wgrad_coeff);\n scaled_write_to_gmem(params.gmem_dbias, thread_c / ELEMENTS_PER_LDG, dbias, params.wgrad_coeff);\n } else {\n write_to_gmem(params.gmem_dscale, thread_c / ELEMENTS_PER_LDG, dscale);\n write_to_gmem(params.gmem_dbias, thread_c / ELEMENTS_PER_LDG, dbias);\n }\n }\n\n // Further normalize the dscale to be used in dx calculation\n float scale[ELEMENTS_PER_LDG];\n zero_array(scale);\n if (is_valid_c) {\n read_from_gmem(scale, params.gmem_scale, thread_c / ELEMENTS_PER_LDG);\n }\n multiply(dscale, var);\n // scale the inv-var as well, afterwards\n multiply(var, scale);\n\n // inverse count\n float inv_count = params.svar_inv_count;\n\n // The base pointer to write to.\n uint16_t* const gmem_dst = ¶ms.gmem_dst[thread_c];\n\n// Store the elements in registers.\n#pragma unroll 1\n for (int loop_i = OUTER_LOOPS - 1; loop_i >= 0; --loop_i) {\n // The value for nhw.\n int out_nhw = cta_nhw_regs + loop_i * pixels_per_iteration;\n\n// Normalize the elements and write to memory.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n const int idx = out_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n const bool is_valid = ((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c;\n // Convert to float.\n float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage[i]);\n to_float(dy_math, dy_storage[i]);\n\n float dx[ELEMENTS_PER_LDG];\n bwd_dx(dx, dy_math, var, x_math, mean, dscale, dbias, inv_count);\n\n // Write back.\n if (is_valid) {\n stg_stream(&gmem_dst[idx * params.c], dx);\n }\n }\n\n // The next value of nhw.\n out_nhw -= pixels_per_iteration;\n\n// Read the next elements from memory.\n#pragma unroll\n for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {\n const int idx = out_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n float y[ELEMENTS_PER_LDG];\n zero_array(y);\n if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {\n ldg_stream(x_storage[i], &gmem_src[idx * params.c]);\n ldg_stream(dy_storage[i], &gmem_dst1[idx * params.c]);\n }\n }\n }\n\n // Normalize the elements from SMEM and write them out.\n if (pixels_in_smem > 0) {\n for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {\n const int idx = smem_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;\n const bool is_valid = ((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c;\n if (is_valid) {\n // Read from SMEM.\n int offset = i * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG], dy_storage_local[PACKED_ELEMENTS_PER_LDG];\n read_from_smem(x_storage_local, &smem_storage_packed[offset], threadIdx.x);\n offset += PIXELS_PER_THREAD_IN_SMEM * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;\n read_from_smem(dy_storage_local, &smem_storage_packed[offset], threadIdx.x);\n float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];\n to_float(x_math, x_storage_local);\n to_float(dy_math, dy_storage_local);\n\n float dx[ELEMENTS_PER_LDG];\n bwd_dx(dx, dy_math, var, x_math, mean, dscale, dbias, inv_count);\n\n // Write back.\n stg_stream(&gmem_dst[idx * params.c], dx);\n }\n }\n }\n // We're about to start on the next c-blk. Needed?\n __syncthreads();\n }\n}\n\n#endif // MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_KERNEL_H_\n" }, { "path": "apex/contrib/csrc/index_mul_2d/index_mul_2d_cuda.cpp", "content": "#include \n\n#include \n#include \n\nvoid index_mul_2d_float_foward_cuda(at::Tensor& out, const at::Tensor& in1, const at::Tensor& in2,\n const at::Tensor& idx1);\n\nvoid index_mul_2d_float_backward_cuda(at::Tensor& grad_in1, at::Tensor& grad_in2, const at::Tensor& grad_out,\n const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1);\n\nvoid index_mul_2d_float_backward_backward_cuda(at::Tensor& grad_grad_out, at::Tensor& grad_in1, at::Tensor& grad_in2,\n const at::Tensor& grad_out, const at::Tensor& grad_grad_in1,\n const at::Tensor& grad_grad_in2, const at::Tensor& in1,\n const at::Tensor& in2, const at::Tensor& idx1);\n\nvoid index_mul_2d_half_foward_cuda(at::Tensor& out, const at::Tensor& in1, const at::Tensor& in2,\n const at::Tensor& idx1);\n\nvoid index_mul_2d_half_backward_cuda(at::Tensor& grad_in1, at::Tensor& grad_in2, const at::Tensor& grad_out,\n const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1);\n\nvoid index_mul_2d_half_backward_backward_cuda(at::Tensor& grad_grad_out, at::Tensor& grad_in1, at::Tensor& grad_in2,\n const at::Tensor& grad_out, const at::Tensor& grad_grad_in1,\n const at::Tensor& grad_grad_in2, const at::Tensor& in1,\n const at::Tensor& in2, const at::Tensor& idx1);\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n#define CHECK_INPUT(x) \\\n CHECK_CUDA(x); \\\n CHECK_CONTIGUOUS(x)\n\nvoid index_mul_2d_float_forward(at::Tensor& out, const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1) {\n return index_mul_2d_float_foward_cuda(out, in1, in2, idx1);\n}\n\nvoid index_mul_2d_float_backward(at::Tensor& grad_in1, at::Tensor& grad_in2, const at::Tensor& grad_out,\n const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1) {\n return index_mul_2d_float_backward_cuda(grad_in1, grad_in2, grad_out, in1, in2, idx1);\n}\n\nvoid index_mul_2d_float_backwrad_backward(at::Tensor& grad_grad_out, at::Tensor& grad_in1, at::Tensor& grad_in2,\n const at::Tensor& grad_out, const at::Tensor& grad_grad_in1,\n const at::Tensor& grad_grad_in2, const at::Tensor& in1, const at::Tensor& in2,\n const at::Tensor& idx1) {\n return index_mul_2d_float_backward_backward_cuda(grad_grad_out, grad_in1, grad_in2, grad_out, grad_grad_in1,\n grad_grad_in2, in1, in2, idx1);\n}\n\nvoid index_mul_2d_half_forward(at::Tensor& out, const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1) {\n return index_mul_2d_half_foward_cuda(out, in1, in2, idx1);\n}\n\nvoid index_mul_2d_half_backward(at::Tensor& grad_in1, at::Tensor& grad_in2, const at::Tensor& grad_out,\n const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1) {\n return index_mul_2d_half_backward_cuda(grad_in1, grad_in2, grad_out, in1, in2, idx1);\n}\n\nvoid index_mul_2d_half_backwrad_backward(at::Tensor& grad_grad_out, at::Tensor& grad_in1, at::Tensor& grad_in2,\n const at::Tensor& grad_out, const at::Tensor& grad_grad_in1,\n const at::Tensor& grad_grad_in2, const at::Tensor& in1, const at::Tensor& in2,\n const at::Tensor& idx1) {\n return index_mul_2d_half_backward_backward_cuda(grad_grad_out, grad_in1, grad_in2, grad_out, grad_grad_in1,\n grad_grad_in2, in1, in2, idx1);\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"float_forward\", &index_mul_2d_float_forward, \"index mul float calculation forward (CUDA)\",\n py::call_guard());\n m.def(\"float_backward\", &index_mul_2d_float_backward, \"index mul float calculation backward (CUDA)\",\n py::call_guard());\n m.def(\"float_backward_backward\", &index_mul_2d_float_backwrad_backward,\n \"index mul float calculation backward backward (CUDA)\", py::call_guard());\n m.def(\"half_forward\", &index_mul_2d_half_forward, \"index mul half calculation forward (CUDA)\",\n py::call_guard());\n m.def(\"half_backward\", &index_mul_2d_half_backward, \"index mul half calculation backward (CUDA)\",\n py::call_guard());\n m.def(\"half_backward_backward\", &index_mul_2d_half_backwrad_backward,\n \"index mul half calculation backward backward (CUDA)\", py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/index_mul_2d/index_mul_2d_cuda_kernel.cu", "content": "#include \n#include \n#include \n\n#include \n\n__global__ void index_mul_2d_float_dim64(float* out, const float* in1, const float* in2, const int64_t* idx1,\n const int64_t size) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n constexpr int fea_dim = 64;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim) / 4 + tidx;\n int64_t vec_idx2 = (start_idx * fea_dim) / 4 + tidx;\n\n float4 res, src1, src2;\n src1 = reinterpret_cast(in1)[vec_idx1];\n src2 = reinterpret_cast(in2)[vec_idx2];\n res.x = src1.x * src2.x;\n res.y = src1.y * src2.y;\n res.z = src1.z * src2.z;\n res.w = src1.w * src2.w;\n reinterpret_cast(out)[vec_idx2] = res;\n }\n}\n\n__global__ void index_mul_2d_float(float* out, const float* in1, const float* in2, const int64_t* idx1,\n const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim);\n int64_t vec_idx2 = (start_idx * fea_dim);\n\n for (int i = tidx; i < fea_dim; i += stride) {\n out[vec_idx2 + i] = in1[vec_idx1 + i] * in2[vec_idx2 + i];\n }\n }\n}\n\n__global__ void index_mul_2d_half(at::Half* out, const at::Half* in1, const at::Half* in2, const int64_t* idx1,\n const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim);\n int64_t vec_idx2 = (start_idx * fea_dim);\n\n for (int i = tidx; i < fea_dim; i += stride) {\n out[vec_idx2 + i] = at::Half(static_cast(in1[vec_idx1 + i]) * static_cast(in2[vec_idx2 + i]));\n }\n }\n}\n\n__global__ void index_mul_2d_grad_float_dim64(float* grad_in1, float* grad_in2, const float* grad_out, const float* in1,\n const float* in2, const int64_t* idx1, const int64_t size) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n constexpr int fea_dim = 64;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim) / 4 + tidx;\n int64_t vec_idx2 = (start_idx * fea_dim) / 4 + tidx;\n\n float4 src_in1, src_in2, src_grad_out, dst_grad_in2;\n src_grad_out = reinterpret_cast(grad_out)[vec_idx2];\n src_in1 = reinterpret_cast(in1)[vec_idx1];\n src_in2 = reinterpret_cast(in2)[vec_idx2];\n int64_t grad_in1_base_idx = idx1[start_idx] * fea_dim + tidx * 4;\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 0, src_grad_out.x * src_in2.x);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 1, src_grad_out.y * src_in2.y);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 2, src_grad_out.z * src_in2.z);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 3, src_grad_out.w * src_in2.w);\n dst_grad_in2.x = src_grad_out.x * src_in1.x;\n dst_grad_in2.y = src_grad_out.y * src_in1.y;\n dst_grad_in2.z = src_grad_out.z * src_in1.z;\n dst_grad_in2.w = src_grad_out.w * src_in1.w;\n reinterpret_cast(grad_in2)[vec_idx2] = dst_grad_in2;\n }\n}\n\n__global__ void index_mul_2d_grad_float(float* grad_in1, float* grad_in2, const float* grad_out, const float* in1,\n const float* in2, const int64_t* idx1, const int64_t size,\n const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_in1 = in1[vec_idx1 + i];\n float src_in2 = in2[vec_idx2 + i];\n float src_grad_out = grad_out[vec_idx2 + i];\n grad_in2[vec_idx2 + i] = src_grad_out * src_in1;\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, src_grad_out * src_in2);\n }\n }\n}\n\n__global__ void index_mul_2d_grad_half(at::Half* grad_in1, at::Half* grad_in2, const at::Half* grad_out,\n const at::Half* in1, const at::Half* in2, const int64_t* idx1,\n const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_in1 = static_cast(in1[vec_idx1 + i]);\n float src_in2 = static_cast(in2[vec_idx2 + i]);\n float src_grad_out = static_cast(grad_out[vec_idx2 + i]);\n grad_in2[vec_idx2 + i] = at::Half(src_grad_out * src_in1);\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, at::Half(src_grad_out * src_in2));\n }\n }\n}\n\n__global__ void index_mul_2d_grad_grad_float_dim64(float* grad_grad_out, float* grad_in1, float* grad_in2,\n const float* grad_out, const float* grad_grad_in1,\n const float* grad_grad_in2, const float* in1, const float* in2,\n const int64_t* idx1, const int64_t size) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n constexpr int fea_dim = 64;\n\n if (start_idx < size) {\n int64_t vec_idx1 = (idx1[start_idx] * fea_dim) / 4 + tidx;\n int64_t vec_idx2 = (start_idx * fea_dim) / 4 + tidx;\n\n float4 src_grad_grad_in1, src_in1, src_grad_grad_in2, src_in2, src_grad_out;\n float4 dst_grad_grad_out, dst_grad_in2;\n src_grad_grad_in1 = reinterpret_cast(grad_grad_in1)[vec_idx1];\n src_in1 = reinterpret_cast(in1)[vec_idx1];\n src_grad_grad_in2 = reinterpret_cast(grad_grad_in2)[vec_idx2];\n src_in2 = reinterpret_cast(in2)[vec_idx2];\n dst_grad_grad_out.x = src_grad_grad_in1.x * src_in2.x + src_grad_grad_in2.x * src_in1.x;\n dst_grad_grad_out.y = src_grad_grad_in1.y * src_in2.y + src_grad_grad_in2.y * src_in1.y;\n dst_grad_grad_out.z = src_grad_grad_in1.z * src_in2.z + src_grad_grad_in2.z * src_in1.z;\n dst_grad_grad_out.w = src_grad_grad_in1.w * src_in2.w + src_grad_grad_in2.w * src_in1.w;\n reinterpret_cast(grad_grad_out)[vec_idx2] = dst_grad_grad_out;\n src_grad_out = reinterpret_cast(grad_out)[vec_idx2];\n int64_t grad_in1_base_idx = idx1[start_idx] * fea_dim + tidx * 4;\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 0, src_grad_grad_in2.x * src_grad_out.x);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 1, src_grad_grad_in2.y * src_grad_out.y);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 2, src_grad_grad_in2.z * src_grad_out.z);\n gpuAtomicAdd(grad_in1 + grad_in1_base_idx + 3, src_grad_grad_in2.w * src_grad_out.w);\n dst_grad_in2.x = src_grad_grad_in1.x * src_grad_out.x;\n dst_grad_in2.y = src_grad_grad_in1.y * src_grad_out.y;\n dst_grad_in2.z = src_grad_grad_in1.z * src_grad_out.z;\n dst_grad_in2.w = src_grad_grad_in1.w * src_grad_out.w;\n reinterpret_cast(grad_in2)[vec_idx2] = dst_grad_in2;\n }\n}\n\n__global__ void index_mul_2d_grad_grad_float(float* grad_grad_out, float* grad_in1, float* grad_in2,\n const float* grad_out, const float* grad_grad_in1,\n const float* grad_grad_in2, const float* in1, const float* in2,\n const int64_t* idx1, const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_grad_grad_in1 = grad_grad_in1[vec_idx1 + i];\n float src_grad_grad_in2 = grad_grad_in2[vec_idx2 + i];\n float src_in1 = in1[vec_idx1 + i];\n float src_in2 = in2[vec_idx2 + i];\n float src_grad_out = grad_out[vec_idx2 + i];\n grad_grad_out[vec_idx2 + i] = src_grad_grad_in1 * src_in2 + src_grad_grad_in2 * src_in1;\n grad_in2[vec_idx2 + i] = src_grad_grad_in1 * src_grad_out;\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, src_grad_grad_in2 * src_grad_out);\n }\n }\n}\n\n__global__ void index_mul_2d_grad_grad_half(at::Half* grad_grad_out, at::Half* grad_in1, at::Half* grad_in2,\n const at::Half* grad_out, const at::Half* grad_grad_in1,\n const at::Half* grad_grad_in2, const at::Half* in1, const at::Half* in2,\n const int64_t* idx1, const int64_t size, const int64_t fea_dim) {\n const int tidx = threadIdx.x;\n const int tidy = threadIdx.y;\n const int bidx = blockIdx.x;\n const int start_idx = bidx * blockDim.y + tidy;\n const int stride = blockDim.x;\n\n if (start_idx < size) {\n int64_t vec_idx1 = idx1[start_idx] * fea_dim;\n int64_t vec_idx2 = start_idx * fea_dim;\n\n for (int i = tidx; i < fea_dim; i += stride) {\n float src_grad_grad_in1 = static_cast(grad_grad_in1[vec_idx1 + i]);\n float src_grad_grad_in2 = static_cast(grad_grad_in2[vec_idx2 + i]);\n float src_in1 = static_cast(in1[vec_idx1 + i]);\n float src_in2 = static_cast(in2[vec_idx2 + i]);\n float src_grad_out = static_cast(grad_out[vec_idx2 + i]);\n grad_grad_out[vec_idx2 + i] = at::Half(src_grad_grad_in1 * src_in2 + src_grad_grad_in2 * src_in1);\n grad_in2[vec_idx2 + i] = at::Half(src_grad_grad_in1 * src_grad_out);\n gpuAtomicAdd(grad_in1 + vec_idx1 + i, at::Half(src_grad_grad_in2 * src_grad_out));\n }\n }\n}\n\nvoid index_mul_2d_float_foward_cuda(at::Tensor& out, const at::Tensor& in1, const at::Tensor& in2,\n const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fea_dim == 64) {\n const int BLOCK_THREADS_DIMX = 16;\n const int BLOCK_THREADS_DIMY = 16;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_float_dim64<<>>(\n out.data_ptr(), in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size);\n } else {\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_float<<>>(\n out.data_ptr(), in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n }\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid index_mul_2d_float_backward_cuda(at::Tensor& grad_in1, at::Tensor& grad_in2, const at::Tensor& grad_out,\n const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fea_dim == 64) {\n const int BLOCK_THREADS_DIMX = 16;\n const int BLOCK_THREADS_DIMY = 16;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_float_dim64<<>>(\n grad_in1.data_ptr(), grad_in2.data_ptr(), grad_out.data_ptr(), in1.data_ptr(),\n in2.data_ptr(), idx1.data_ptr(), size);\n\n AT_CUDA_CHECK(cudaGetLastError());\n } else {\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_float<<>>(\n grad_in1.data_ptr(), grad_in2.data_ptr(), grad_out.data_ptr(), in1.data_ptr(),\n in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n }\n}\n\nvoid index_mul_2d_float_backward_backward_cuda(at::Tensor& grad_grad_out, at::Tensor& grad_in1, at::Tensor& grad_in2,\n const at::Tensor& grad_out, const at::Tensor& grad_grad_in1,\n const at::Tensor& grad_grad_in2, const at::Tensor& in1,\n const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fea_dim == 64) {\n const int BLOCK_THREADS_DIMX = 16;\n const int BLOCK_THREADS_DIMY = 16;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_grad_float_dim64<<>>(\n grad_grad_out.data_ptr(), grad_in1.data_ptr(), grad_in2.data_ptr(),\n grad_out.data_ptr(), grad_grad_in1.data_ptr(), grad_grad_in2.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size);\n } else {\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_grad_float<<>>(\n grad_grad_out.data_ptr(), grad_in1.data_ptr(), grad_in2.data_ptr(),\n grad_out.data_ptr(), grad_grad_in1.data_ptr(), grad_grad_in2.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n }\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid index_mul_2d_half_foward_cuda(at::Tensor& out, const at::Tensor& in1, const at::Tensor& in2,\n const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_half<<>>(\n out.data_ptr(), in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size,\n fea_dim);\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid index_mul_2d_half_backward_cuda(at::Tensor& grad_in1, at::Tensor& grad_in2, const at::Tensor& grad_out,\n const at::Tensor& in1, const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_half<<>>(\n grad_in1.data_ptr(), grad_in2.data_ptr(), grad_out.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n}\n\nvoid index_mul_2d_half_backward_backward_cuda(at::Tensor& grad_grad_out, at::Tensor& grad_in1, at::Tensor& grad_in2,\n const at::Tensor& grad_out, const at::Tensor& grad_grad_in1,\n const at::Tensor& grad_grad_in2, const at::Tensor& in1,\n const at::Tensor& in2, const at::Tensor& idx1) {\n const int64_t size = in2.size(0);\n const int64_t fea_dim = in2.size(1);\n if (size < 0) {\n return;\n }\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const int BLOCK_THREADS_DIMX = 32;\n const int BLOCK_THREADS_DIMY = 8;\n const int BLOCK_NUMS = (size + BLOCK_THREADS_DIMY - 1) / BLOCK_THREADS_DIMY;\n\n index_mul_2d_grad_grad_half<<>>(\n grad_grad_out.data_ptr(), grad_in1.data_ptr(), grad_in2.data_ptr(),\n grad_out.data_ptr(), grad_grad_in1.data_ptr(), grad_grad_in2.data_ptr(),\n in1.data_ptr(), in2.data_ptr(), idx1.data_ptr(), size, fea_dim);\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n" }, { "path": "apex/contrib/csrc/layer_norm/ln.h", "content": "#pragma once\n\n#include \n#include \n#include \n#include \n\n#include \n\nnamespace layer_norm {\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct LaunchParams {\n size_t workspace_bytes;\n size_t barrier_size;\n\n cudaDeviceProp* props;\n\n cudaStream_t stream;\n\n Params params;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct FwdParams {\n FwdParams()\n : ctas_per_col(0),\n rows(0),\n cols(0),\n x(nullptr),\n z(nullptr),\n mu(nullptr),\n rs(nullptr),\n gamma(nullptr),\n beta(nullptr),\n workspace(nullptr),\n barrier(nullptr),\n epsilon(0.f) {}\n\n // For Multi-CTA, number of different CTA groups. Otherwise same as gridDim.x.\n int ctas_per_col;\n\n // Input is interpreted as matrix. We normalize across columns.\n int rows;\n int cols;\n\n // Common data pointers.\n void* x;\n void* z;\n void* mu;\n void* rs;\n void* gamma;\n void* beta;\n\n // Multi-CTA workspace in gmem.\n void* workspace;\n\n // Multi-CTA sync barriers in gmem.\n int* barrier;\n\n // Output of LN FWD.\n float epsilon;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct BwdParams : public FwdParams {\n BwdParams()\n : FwdParams(),\n dz(nullptr),\n dbeta_part(nullptr),\n dgamma_part(nullptr),\n dx(nullptr),\n dbeta(nullptr),\n dgamma(nullptr) {}\n // Input: gradient wrt. LN FWD output.\n void* dz;\n\n // Workspace for Wgrad pre-reduction.\n void* dbeta_part;\n void* dgamma_part;\n\n // Output: Dgrad.\n void* dx;\n // Output: Wgrad.\n void* dbeta;\n void* dgamma;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nusing FwdFunction = std::function&, const bool)>;\nusing BwdFunction = std::function&, const bool)>;\nusing FunctionKey = uint64_t;\nusing FwdRegistry = std::unordered_map;\nusing BwdRegistry = std::unordered_map;\n\nextern FwdRegistry FWD_FUNCS;\nextern BwdRegistry BWD_FUNCS;\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nusing fp32 = float;\nusing fp16 = half;\nusing bf16 = nv_bfloat16;\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct TypeId {};\n\ntemplate <>\nstruct TypeId {\n constexpr static uint32_t Value = 0;\n};\n\ntemplate <>\nstruct TypeId {\n constexpr static uint32_t Value = 1;\n};\n\ntemplate <>\nstruct TypeId {\n constexpr static uint32_t Value = 2;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Type2Key {\n constexpr static uint32_t Value = TypeId::Value << S;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct WeightType2Key : public Type2Key {};\n\ntemplate \nstruct InputType2Key : public Type2Key {};\n\ntemplate \nstruct OutputType2Key : public Type2Key {};\n\ntemplate \nstruct ComputeType2Key : public Type2Key {};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Types2Key {\n constexpr static uint32_t Value =\n WeightType2Key::Value | InputType2Key::Value | OutputType2Key::Value | ComputeType2Key::Value;\n constexpr static inline uint64_t get(const uint64_t hidden_size) {\n constexpr uint64_t type_key = Value;\n return (type_key << 32) | hidden_size;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct FwdRegistrar {\n FwdRegistrar(FwdFunction f) {\n uint64_t key = Types2Key::get(HIDDEN_SIZE);\n FWD_FUNCS.insert({key, f});\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct BwdRegistrar {\n BwdRegistrar(BwdFunction f) {\n uint64_t key = Types2Key::get(HIDDEN_SIZE);\n BWD_FUNCS.insert({key, f});\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace layer_norm\n" }, { "path": "apex/contrib/csrc/layer_norm/ln_api.cpp", "content": "#include \n\n#include \"ATen/cuda/CUDAContext.h\"\n#include \"ln.h\"\n\n/*\n\nSupported Type combinations:\n\ninput compute weights output\n=======================================\nfp32 fp32 fp32 fp32\nfp16 fp32 fp16 fp16\nbf16 fp32 bf16 bf16\nfp32 fp32 fp16 fp16\nfp32 fp32 bf16 bf16\n\nRemarks:\nOutput type = Weight type\nCompute always in FP32\n\n*/\n\nnamespace layer_norm {\n\n// Create registries and provide runtime versions of config hash functions.\n\nFwdRegistry FWD_FUNCS;\nBwdRegistry BWD_FUNCS;\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nuint32_t get_type_id(torch::Dtype dtype) {\n if (dtype == torch::kFloat16) {\n return TypeId::Value;\n } else if (dtype == torch::kBFloat16) {\n return TypeId::Value;\n } else if (dtype == torch::kFloat32) {\n return TypeId::Value;\n } else {\n TORCH_CHECK(false, \"Type not supported: \", dtype);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nuint64_t get_key(torch::Dtype wtype, torch::Dtype itype, torch::Dtype otype, torch::Dtype ctype, uint64_t hidden_size) {\n using namespace layer_norm;\n uint64_t type_key =\n get_type_id(wtype) | (get_type_id(itype) << 2) | (get_type_id(otype) << 4) | (get_type_id(ctype) << 6);\n uint64_t launcher_key = (type_key << 32) | hidden_size;\n return launcher_key;\n}\n\n} // namespace layer_norm\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nlayer_norm::FwdFunction& get_fwd_launcher(torch::Dtype wtype, torch::Dtype itype, torch::Dtype otype,\n torch::Dtype ctype, uint32_t hidden_size) {\n auto iter = layer_norm::FWD_FUNCS.find(layer_norm::get_key(wtype, itype, otype, ctype, hidden_size));\n if (iter != layer_norm::FWD_FUNCS.end()) {\n return iter->second;\n } else {\n TORCH_CHECK(false, \"FWD: Unsupported hidden_size or types: \", hidden_size, wtype, itype, otype, ctype);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nlayer_norm::BwdFunction& get_bwd_launcher(torch::Dtype wtype, torch::Dtype itype, torch::Dtype otype,\n torch::Dtype ctype, uint32_t hidden_size) {\n auto iter = layer_norm::BWD_FUNCS.find(layer_norm::get_key(wtype, itype, otype, ctype, hidden_size));\n if (iter != layer_norm::BWD_FUNCS.end()) {\n return iter->second;\n } else {\n TORCH_CHECK(false, \"BWD: Unsupported hidden_size or types: \", hidden_size, wtype, itype, otype, ctype);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstd::vector ln_fwd(const at::Tensor& x, // BxSxhidden_size\n const at::Tensor& gamma, // hidden_size\n const at::Tensor& beta, // hidden_size\n const float epsilon) {\n auto itype = x.scalar_type();\n auto wtype = gamma.scalar_type();\n auto otype = wtype;\n auto ctype = torch::kFloat32;\n\n TORCH_CHECK(beta.scalar_type() == wtype);\n\n TORCH_CHECK(x.is_cuda())\n TORCH_CHECK(gamma.is_cuda())\n TORCH_CHECK(beta.is_cuda())\n\n TORCH_CHECK(x.is_contiguous());\n auto sizes = x.sizes();\n TORCH_CHECK(sizes.size() == 2);\n\n const int rows = sizes[0];\n const int cols = sizes[1];\n auto hidden_size = gamma.numel();\n\n TORCH_CHECK(gamma.sizes() == beta.sizes());\n TORCH_CHECK(hidden_size == cols);\n\n TORCH_CHECK(epsilon >= 0.f);\n\n auto opts = x.options();\n\n auto z = torch::empty(sizes, opts.dtype(otype));\n\n auto mu = torch::empty({rows}, opts.dtype(ctype));\n auto rsigma = torch::empty({rows}, opts.dtype(ctype));\n\n layer_norm::LaunchParams launch_params;\n\n launch_params.props = at::cuda::getCurrentDeviceProperties();\n launch_params.stream = at::cuda::getCurrentCUDAStream().stream();\n\n // Request the kernel launcher.\n auto launcher = get_fwd_launcher(wtype, itype, otype, ctype, hidden_size);\n\n // Query the kernel-specific launch parameters.\n launcher(launch_params, true);\n\n at::Tensor workspace, barrier;\n\n // Set the kernel runtime parameters.\n layer_norm::FwdParams& params = launch_params.params;\n params.rows = rows;\n params.cols = cols;\n params.z = z.data_ptr();\n params.mu = mu.data_ptr();\n params.rs = rsigma.data_ptr();\n params.gamma = gamma.data_ptr();\n params.beta = beta.data_ptr();\n params.x = x.data_ptr();\n params.epsilon = epsilon;\n\n if (launch_params.barrier_size > 0) {\n auto options = x.options();\n barrier = torch::zeros(launch_params.barrier_size, options.dtype(torch::kInt32));\n workspace = torch::empty(launch_params.workspace_bytes, options.dtype(torch::kChar));\n params.workspace = workspace.data_ptr();\n params.barrier = barrier.data_ptr();\n }\n\n // Launch the kernel.\n launcher(launch_params, false);\n\n return {z, mu, rsigma};\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\nstd::vector ln_bwd(const at::Tensor& dz, // BxSxhidden_size\n const at::Tensor& x_or_z, // BxSxhidden_size\n c10::optional& mu_, // BxS, FP32!\n const at::Tensor& rsigma, // BxS, FP32!\n const at::Tensor& gamma, // hidden_size\n c10::optional& beta_, // hidden_size\n bool memory_efficient) {\n auto itype = x_or_z.scalar_type();\n auto wtype = gamma.scalar_type();\n auto otype = wtype;\n auto ctype = torch::kFloat32;\n\n TORCH_CHECK(dz.dtype() == otype);\n TORCH_CHECK(rsigma.dtype() == ctype);\n if (mu_.has_value()) {\n TORCH_CHECK(mu_.value().dtype() == ctype);\n }\n\n TORCH_CHECK(x_or_z.is_cuda());\n TORCH_CHECK(dz.is_cuda());\n TORCH_CHECK(rsigma.is_cuda());\n TORCH_CHECK(gamma.is_cuda());\n if (beta_.has_value()) {\n TORCH_CHECK(beta_.value().is_cuda());\n TORCH_CHECK(beta_.value().dtype() == wtype);\n }\n\n TORCH_CHECK(x_or_z.is_contiguous());\n TORCH_CHECK(dz.is_contiguous());\n\n auto sizes = x_or_z.sizes();\n TORCH_CHECK(sizes.size() == 2);\n TORCH_CHECK(dz.sizes() == sizes);\n auto rows = sizes[0];\n auto cols = sizes[1];\n\n auto hidden_size = gamma.numel();\n\n TORCH_CHECK(gamma.numel() == cols);\n if (beta_.has_value()) {\n TORCH_CHECK(beta_.value().numel() == cols);\n }\n\n auto options = x_or_z.options();\n\n auto dx = torch::empty_like(x_or_z);\n auto dgamma = torch::empty_like(gamma);\n auto dbeta = torch::empty_like(gamma);\n\n layer_norm::LaunchParams launch_params;\n launch_params.stream = at::cuda::getCurrentCUDAStream().stream();\n launch_params.props = at::cuda::getCurrentDeviceProperties();\n\n auto launcher = get_bwd_launcher(wtype, itype, otype, ctype, hidden_size);\n\n launcher(launch_params, true);\n\n auto dgamma_part = torch::empty({launch_params.params.ctas_per_col, hidden_size}, options.dtype(ctype));\n auto dbeta_part = torch::empty({launch_params.params.ctas_per_col, hidden_size}, options.dtype(ctype));\n at::Tensor workspace, barrier;\n\n layer_norm::BwdParams& params = launch_params.params;\n params.rows = rows;\n params.cols = cols;\n if (memory_efficient) {\n params.z = x_or_z.data_ptr();\n params.beta = beta_.value().data_ptr();\n } else {\n params.x = x_or_z.data_ptr();\n params.mu = mu_.value().data_ptr();\n }\n params.rs = rsigma.data_ptr();\n params.gamma = gamma.data_ptr();\n params.dz = dz.data_ptr();\n params.dx = dx.data_ptr();\n params.dbeta = dbeta.data_ptr();\n params.dgamma = dgamma.data_ptr();\n params.dbeta_part = dbeta_part.data_ptr();\n params.dgamma_part = dgamma_part.data_ptr();\n\n if (launch_params.barrier_size > 0) {\n // TODO Any way to avoid this?\n barrier = torch::zeros(launch_params.barrier_size, options.dtype(torch::kInt32));\n workspace = torch::empty(launch_params.workspace_bytes, options.dtype(torch::kChar));\n params.workspace = workspace.data_ptr();\n params.barrier = barrier.data_ptr();\n }\n\n launcher(launch_params, false);\n\n return {dx, dgamma, dbeta, dgamma_part, dbeta_part};\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.doc() = \"CUDA LayerNorm\";\n m.def(\"ln_fwd\", &ln_fwd, \"Run LayerNorm forward kernel\", py::call_guard());\n m.def(\"ln_bwd\", &ln_bwd, \"Run LayerNorm backward kernel\", py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/layer_norm/ln_bwd_kernels.cuh", "content": "#pragma once\n\n#include \"ln_utils.cuh\"\n\nnamespace layer_norm {\n\ntemplate \n__global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void ln_bwd_kernel(layer_norm::BwdParams params) {\n enum { ROWS_PER_CTA = Ktraits::ROWS_PER_CTA };\n enum { WARPS_M = Ktraits::WARPS_M };\n enum { WARPS_N = Ktraits::WARPS_N };\n enum { THREADS_PER_ROW = Ktraits::THREADS_PER_ROW };\n enum { COLS = Ktraits::COLS };\n enum { BYTES_PER_ROW = Ktraits::BYTES_PER_ROW };\n enum { LDGS = Ktraits::LDGS };\n enum { NUM_ELTS = Ktraits::ELTS_PER_LDG };\n enum { THREADS_PER_WARP = Ktraits::THREADS_PER_WARP };\n enum { CTAS_PER_ROW = Ktraits::CTAS_PER_ROW };\n\n using compute_t = typename Ktraits::compute_t;\n using index_t = typename Ktraits::index_t;\n using Ivec = typename Ktraits::Ivec;\n using Ovec = typename Ktraits::Ovec;\n using Wvec = typename Ktraits::Wvec;\n using Cvec = typename Ktraits::Cvec;\n using Reducer = typename Ktraits::Reducer;\n using reduce_t = typename Reducer::Type;\n\n extern __shared__ char smem_[];\n\n const index_t tidx = threadIdx.x;\n const index_t bidn = blockIdx.x % CTAS_PER_ROW;\n const index_t bidm = blockIdx.x / CTAS_PER_ROW;\n const index_t lane = tidx % THREADS_PER_WARP;\n const index_t warp = tidx / THREADS_PER_WARP;\n const index_t warp_m = warp / Ktraits::WARPS_N;\n const index_t warp_n = warp % Ktraits::WARPS_N;\n const index_t tid_r = warp_n * THREADS_PER_WARP + lane;\n\n const index_t r = bidm * Ktraits::ROWS_PER_CTA + warp_m;\n const index_t c = bidn * THREADS_PER_ROW + warp_n * THREADS_PER_WARP + lane;\n\n static_assert(COLS == THREADS_PER_ROW * LDGS * NUM_ELTS * CTAS_PER_ROW);\n\n Cvec dzy_sum[LDGS];\n Cvec dz_sum[LDGS];\n\n memset(dzy_sum, 0, sizeof(dzy_sum));\n memset(dz_sum, 0, sizeof(dz_sum));\n\n compute_t* smem_wgrad = reinterpret_cast(smem_);\n char* smem_dgrad = smem_ + Ktraits::SMEM_BYTES_WGRAD;\n\n Reducer reducer(params, bidm, bidn, warp_m, warp_n, lane, smem_dgrad);\n\n Sum sum;\n\n constexpr float rn = 1.f / float(COLS);\n Wvec gamma[LDGS];\n Wvec beta[LDGS];\n index_t idx = c;\n#pragma unroll\n for (int it = 0; it < LDGS; it++) {\n gamma[it].load_from(params.gamma, idx);\n if (params.z != nullptr) {\n beta[it].load_from(params.beta, idx);\n }\n idx += Ktraits::VEC_COLS_PER_LDG;\n }\n// TODO if ROWS_PER_CTA does not divide rows, we might get divergence in the\n// last blocks with syncthreads!\n// grid stride over rows\n#pragma unroll 1\n for (int row = r; row < params.rows; row += params.ctas_per_col * ROWS_PER_CTA) {\n const compute_t mu_r = params.z == nullptr ? static_cast(params.mu)[row] : 0.f;\n const compute_t rs_r = static_cast(params.rs)[row];\n Ivec x_or_z[LDGS];\n Ovec dz[LDGS];\n index_t idx = row * Ktraits::VEC_COLS + c;\n#pragma unroll\n for (int it = 0; it < LDGS; it++) {\n dz[it].load_from(params.dz, idx);\n if (params.z != nullptr) {\n x_or_z[it].load_from(params.z, idx);\n } else {\n x_or_z[it].load_from(params.x, idx);\n }\n idx += Ktraits::VEC_COLS_PER_LDG;\n }\n\n compute_t dy[LDGS * NUM_ELTS];\n compute_t y[LDGS * NUM_ELTS];\n\n compute_t mdy_local = 0.f;\n compute_t mdyy_local = 0.f;\n#pragma unroll\n for (int it = 0; it < LDGS; it++) {\n#pragma unroll\n for (int jt = 0; jt < NUM_ELTS; jt++) {\n compute_t gamma_tmp = compute_t(gamma[it].data.elt[jt]);\n compute_t beta_tmp = compute_t(beta[it].data.elt[jt]);\n compute_t x_or_z_tmp = compute_t(x_or_z[it].data.elt[jt]);\n compute_t y_tmp = params.z != nullptr ? (x_or_z_tmp - beta_tmp) / gamma_tmp : rs_r * (x_or_z_tmp - mu_r);\n compute_t dy_tmp = compute_t(dz[it].data.elt[jt]) * gamma_tmp;\n compute_t dz_tmp = dz[it].data.elt[jt];\n\n mdy_local += dy_tmp;\n mdyy_local += dy_tmp * y_tmp;\n\n dy[it * NUM_ELTS + jt] = dy_tmp;\n y[it * NUM_ELTS + jt] = y_tmp;\n\n dzy_sum[it].data.elt[jt] += dz_tmp * y_tmp;\n dz_sum[it].data.elt[jt] += dz_tmp;\n }\n }\n\n reduce_t result = reducer.allreduce({mdy_local, mdyy_local}, sum);\n mdy_local = layer_norm::Get<0>::of(result) * rn;\n mdyy_local = layer_norm::Get<1>::of(result) * rn;\n\n Ivec dx[LDGS];\n idx = row * Ktraits::VEC_COLS + c;\n#pragma unroll\n for (int it = 0; it < LDGS; it++) {\n#pragma unroll\n for (int jt = 0; jt < NUM_ELTS; jt++) {\n compute_t dy_tmp = dy[it * NUM_ELTS + jt];\n compute_t y_tmp = y[it * NUM_ELTS + jt];\n compute_t dx_tmp = rs_r * (dy_tmp - (mdyy_local * y_tmp + mdy_local));\n dx[it].data.elt[jt] = dx_tmp;\n }\n dx[it].store_to(params.dx, idx);\n idx += Ktraits::VEC_COLS_PER_LDG;\n }\n\n } // end: grid stride loop\n\n if (WARPS_M == 1) {\n idx = r * Ktraits::VEC_COLS + c;\n#pragma unroll\n for (int it = 0; it < LDGS; it++) {\n dz_sum[it].store_to(params.dbeta_part, idx);\n dzy_sum[it].store_to(params.dgamma_part, idx);\n idx += Ktraits::VEC_COLS_PER_LDG;\n }\n } else {\n static_assert(WARPS_M == 1 || Ktraits::CTAS_PER_ROW == 1, \"Multiple rows per CTA not supported for Multi-CTA.\");\n // Finalize reduction of part dgamma and dbeta for this CTA\n // by reducing over the rows held across the WARPS_M warps\n\n // Assumption: blockSize divides hidden size.\n enum { NUM_RES = COLS / Ktraits::THREADS_PER_CTA };\n static_assert(NUM_RES * Ktraits::THREADS_PER_CTA == COLS, \"\");\n\n idx = warp_m * Ktraits::VEC_COLS + tid_r;\n#pragma unroll\n for (int it = 0; it < LDGS; it++) {\n dz_sum[it].store_to(smem_wgrad, idx);\n idx += THREADS_PER_ROW;\n }\n __syncthreads();\n compute_t cta_dz_sum[NUM_RES];\n memset(cta_dz_sum, 0, sizeof(compute_t) * NUM_RES);\n for (int it = 0; it < ROWS_PER_CTA; it++) {\n for (int jt = 0; jt < NUM_RES; jt++) {\n cta_dz_sum[jt] += smem_wgrad[it * COLS + tidx + jt * Ktraits::THREADS_PER_CTA];\n }\n }\n __syncthreads();\n\n idx = warp_m * Ktraits::VEC_COLS + tid_r;\n#pragma unroll\n for (int it = 0; it < LDGS; it++) {\n dzy_sum[it].store_to(smem_wgrad, idx);\n idx += THREADS_PER_ROW;\n }\n __syncthreads();\n compute_t cta_dzy_sum[NUM_RES];\n memset(cta_dzy_sum, 0, sizeof(compute_t) * NUM_RES);\n for (int it = 0; it < ROWS_PER_CTA; it++) {\n for (int jt = 0; jt < NUM_RES; jt++) {\n cta_dzy_sum[jt] += smem_wgrad[it * COLS + tidx + jt * Ktraits::THREADS_PER_CTA];\n }\n }\n\n compute_t* dgamma_part = static_cast(params.dgamma_part) + bidm * COLS + tidx;\n for (int jt = 0; jt < NUM_RES; jt++) {\n *dgamma_part = cta_dzy_sum[jt];\n dgamma_part += Ktraits::THREADS_PER_CTA;\n }\n\n compute_t* dbeta_part = static_cast(params.dbeta_part) + bidm * COLS + tidx;\n for (int jt = 0; jt < NUM_RES; jt++) {\n *dbeta_part = cta_dz_sum[jt];\n dbeta_part += Ktraits::THREADS_PER_CTA;\n }\n }\n}\n\ntemplate \n__global__ __launch_bounds__(Kernel_traits::THREADS_PER_CTA) void ln_bwd_finalize_kernel(BwdParams params) {\n using compute_t = typename Kernel_traits::compute_t;\n using weight_t = typename Kernel_traits::weight_t;\n using index_t = typename Kernel_traits::index_t;\n using Reducer = typename Kernel_traits::Reducer;\n using reduce_t = typename Reducer::Type;\n\n Sum sum;\n enum { NUM_ELT = Kernel_traits::ELTS_PER_LDG };\n enum { THREADS_PER_WARP = Kernel_traits::THREADS_PER_WARP };\n\n __shared__ char smem_[Kernel_traits::SMEM_BYTES_PER_CTA];\n\n constexpr uint32_t bidm = 0;\n\n const uint32_t bidn = blockIdx.x;\n const uint32_t tidx = threadIdx.x;\n const uint32_t warp = tidx / THREADS_PER_WARP;\n const uint32_t lane = tidx % THREADS_PER_WARP;\n\n Reducer reducer(params, bidm, bidn, 0, 0, lane, smem_);\n\n const uint32_t c = bidn * THREADS_PER_WARP + lane;\n const uint32_t c_out = bidn * THREADS_PER_WARP / 2 + lane;\n constexpr uint32_t COL_STRIDE = Kernel_traits::CTAS * THREADS_PER_WARP;\n for (uint32_t col = c, col_out = c_out; col < Kernel_traits::COLS; col += COL_STRIDE, col_out += COL_STRIDE / 2) {\n // Each thread sums over NUM_ELT columns.\n Vec dbeta_local, dgamma_local;\n memset(&dgamma_local, 0, sizeof(dgamma_local));\n memset(&dbeta_local, 0, sizeof(dbeta_local));\n for (uint32_t row = warp; row < params.ctas_per_col; row += Kernel_traits::ROWS_PER_CTA) {\n index_t idx = row * Kernel_traits::COLS + col;\n\n Vec dbeta_part, dgamma_part;\n dbeta_part.load_from(params.dbeta_part, idx);\n dgamma_part.load_from(params.dgamma_part, idx);\n#pragma unroll\n for (int it = 0; it < NUM_ELT; it++) {\n dgamma_local.data.elt[it] += dgamma_part.data.elt[it];\n dbeta_local.data.elt[it] += dbeta_part.data.elt[it];\n }\n }\n\n void* smem_gamma = smem_;\n void* smem_beta = &smem_[Kernel_traits::SMEM_BYTES_TRANSPOSE];\n\n const int write_row = warp;\n const int write_col = lane ^ write_row;\n const int write_idx = write_row * THREADS_PER_WARP + write_col;\n\n dgamma_local.store_to(smem_gamma, write_idx);\n dbeta_local.store_to(smem_beta, write_idx);\n\n __syncthreads();\n\n // It would be probably safe to reuse the first row of smem_beta and smem_gamma\n void* smem_gamma_out = &smem_[2 * Kernel_traits::SMEM_BYTES_TRANSPOSE];\n void* smem_beta_out = &smem_[2 * Kernel_traits::SMEM_BYTES_TRANSPOSE + Kernel_traits::SMEM_BYTES_OUTPUT];\n\n // More than one iter iff ROWS_PER_CTA < 32.\n for (int w = warp; w < THREADS_PER_WARP; w += Kernel_traits::ROWS_PER_CTA) {\n const int read_row = lane;\n const int read_col = w ^ read_row;\n const int read_idx = read_row * THREADS_PER_WARP + read_col;\n\n memset(&dbeta_local, 0, sizeof(dbeta_local));\n memset(&dgamma_local, 0, sizeof(dgamma_local));\n\n // Load beta and gamma transposed\n if (read_row < Kernel_traits::ROWS_PER_CTA) {\n dbeta_local.load_from(smem_beta, read_idx);\n dgamma_local.load_from(smem_gamma, read_idx);\n }\n\n// Call reducer on the loaded value(s) and convert.\n#pragma unroll\n for (int it = 0; it < NUM_ELT; it++) {\n compute_t b_i = dbeta_local.data.elt[it];\n compute_t g_i = dgamma_local.data.elt[it];\n b_i = reducer.allreduce(b_i, sum);\n g_i = reducer.allreduce(g_i, sum);\n\n dgamma_local.data.elt[it] = g_i;\n dbeta_local.data.elt[it] = b_i;\n }\n\n // Leader stores the result at the current column.\n if (lane == 0) {\n dgamma_local.store_to(smem_gamma_out, w);\n dbeta_local.store_to(smem_beta_out, w);\n }\n }\n\n // All writes done.\n __syncthreads();\n\n // Pack and store: 2-wide stores with half the threads.\n if (warp == Kernel_traits::ROWS_PER_CTA - 1 && lane < THREADS_PER_WARP / 2) {\n using src_t = typename TypeToVec2::Type;\n using dst_t = typename TypeToVec2::Type;\n Vec dbeta_vec2, dgamma_vec2;\n Vec dbeta_out2, dgamma_out2;\n\n dgamma_vec2.load_from(smem_gamma_out, lane);\n dbeta_vec2.load_from(smem_beta_out, lane);\n#pragma unroll\n for (int it = 0; it < NUM_ELT; it++) {\n dgamma_out2.data.elt[it] = Converter::convert(dgamma_vec2.data.elt[it]);\n dbeta_out2.data.elt[it] = Converter::convert(dbeta_vec2.data.elt[it]);\n }\n dgamma_out2.store_to(params.dgamma, col_out);\n dbeta_out2.store_to(params.dbeta, col_out);\n }\n }\n}\n} // namespace layer_norm\n" }, { "path": "apex/contrib/csrc/layer_norm/ln_bwd_semi_cuda_kernel.cu", "content": "#include \"ln.h\"\n#include \"ln_bwd_kernels.cuh\"\n#include \"ln_kernel_traits.h\"\n#include \"ln_utils.cuh\"\n\nusing namespace layer_norm;\n\ntemplate \nvoid launch_(LaunchParams& launch_params, const bool configure_params) {\n using Kernel_traits = Kernel_traits;\n auto kernel = &ln_bwd_kernel;\n\n if (configure_params) {\n int ctas_per_sm;\n cudaError status_ = cudaOccupancyMaxActiveBlocksPerMultiprocessor(\n &ctas_per_sm, kernel, Kernel_traits::THREADS_PER_CTA, Kernel_traits::SMEM_BYTES);\n launch_params.params.ctas_per_col =\n launch_params.props->multiProcessorCount * ctas_per_sm / Kernel_traits::CTAS_PER_ROW;\n launch_params.barrier_size = 0;\n launch_params.workspace_bytes = 0;\n if (Kernel_traits::CTAS_PER_ROW > 1) {\n launch_params.barrier_size = 2 * launch_params.params.ctas_per_col;\n launch_params.workspace_bytes = launch_params.params.ctas_per_col * Kernel_traits::WARPS_M *\n Kernel_traits::CTAS_PER_ROW * sizeof(typename Kernel_traits::reduce_t) * 2;\n }\n return;\n }\n\n if (Kernel_traits::SMEM_BYTES >= 48 * 1024) {\n CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, Kernel_traits::SMEM_BYTES));\n }\n auto stream = launch_params.stream;\n auto ctas_per_col = launch_params.params.ctas_per_col;\n\n if (Kernel_traits::CTAS_PER_ROW == 1) {\n kernel<<>>(launch_params.params);\n } else {\n dim3 grid(Kernel_traits::CTAS_PER_ROW * ctas_per_col);\n dim3 block(Kernel_traits::THREADS_PER_CTA);\n void* params_ = (void*)&launch_params.params;\n cudaLaunchCooperativeKernel((void*)kernel, grid, block, (void**)¶ms_, Kernel_traits::SMEM_BYTES, stream);\n }\n\n using Kernel_traits_f =\n layer_norm::Kernel_traits_finalize;\n\n auto kernel_f = &layer_norm::ln_bwd_finalize_kernel;\n kernel_f<<>>(launch_params.params);\n}\n\n// Create backward launch function and register. Macro signature:\n// HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG, BYTES_PER_LDG_FINAL\n\nREGISTER_BWD_LAUNCHER(768, fp32, fp32, fp32, fp32, 1, 4, 1, 16, 4);\nREGISTER_BWD_LAUNCHER(768, fp16, fp16, fp16, fp32, 1, 4, 1, 16, 4);\nREGISTER_BWD_LAUNCHER(768, fp16, fp32, fp16, fp32, 1, 4, 1, 16, 4);\nREGISTER_BWD_LAUNCHER(768, bf16, bf16, bf16, fp32, 1, 4, 1, 16, 4);\nREGISTER_BWD_LAUNCHER(768, bf16, fp32, bf16, fp32, 1, 4, 1, 16, 4);\n\nREGISTER_BWD_LAUNCHER(1024, fp32, fp32, fp32, fp32, 1, 4, 1, 16, 4);\nREGISTER_BWD_LAUNCHER(1024, fp16, fp16, fp16, fp32, 1, 4, 1, 16, 4);\nREGISTER_BWD_LAUNCHER(1024, fp16, fp32, fp16, fp32, 1, 4, 1, 16, 4);\nREGISTER_BWD_LAUNCHER(1024, bf16, bf16, bf16, fp32, 1, 4, 1, 16, 4);\nREGISTER_BWD_LAUNCHER(1024, bf16, fp32, bf16, fp32, 1, 4, 1, 16, 4);\n\nREGISTER_BWD_LAUNCHER(1536, fp32, fp32, fp32, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(1536, fp16, fp16, fp16, fp32, 1, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(1536, fp16, fp32, fp16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(1536, bf16, bf16, bf16, fp32, 1, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(1536, bf16, fp32, bf16, fp32, 1, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(2048, fp32, fp32, fp32, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(2048, fp16, fp16, fp16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(2048, fp16, fp32, fp16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(2048, bf16, bf16, bf16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(2048, bf16, fp32, bf16, fp32, 1, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(2304, fp32, fp32, fp32, fp32, 1, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(2304, fp16, fp16, fp16, fp32, 1, 1, 4, 4, 4);\nREGISTER_BWD_LAUNCHER(2304, fp16, fp32, fp16, fp32, 1, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(2304, bf16, bf16, bf16, fp32, 1, 1, 4, 4, 4);\nREGISTER_BWD_LAUNCHER(2304, bf16, fp32, bf16, fp32, 1, 1, 4, 8, 4);\n\nREGISTER_BWD_LAUNCHER(3072, fp32, fp32, fp32, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(3072, fp16, fp16, fp16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(3072, fp16, fp32, fp16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(3072, bf16, bf16, bf16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(3072, bf16, fp32, bf16, fp32, 1, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(3840, fp32, fp32, fp32, fp32, 1, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(3840, fp16, fp16, fp16, fp32, 1, 1, 4, 4, 4);\nREGISTER_BWD_LAUNCHER(3840, fp16, fp32, fp16, fp32, 1, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(3840, bf16, bf16, bf16, fp32, 1, 1, 4, 4, 4);\nREGISTER_BWD_LAUNCHER(3840, bf16, fp32, bf16, fp32, 1, 1, 4, 8, 4);\n\nREGISTER_BWD_LAUNCHER(4096, fp32, fp32, fp32, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(4096, fp16, fp16, fp16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(4096, fp16, fp32, fp16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(4096, bf16, bf16, bf16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(4096, bf16, fp32, bf16, fp32, 1, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(5120, fp32, fp32, fp32, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(5120, fp16, fp16, fp16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(5120, fp16, fp32, fp16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(5120, bf16, bf16, bf16, fp32, 1, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(5120, bf16, fp32, bf16, fp32, 1, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(6144, fp32, fp32, fp32, fp32, 1, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(6144, fp16, fp16, fp16, fp32, 1, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(6144, fp16, fp32, fp16, fp32, 1, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(6144, bf16, bf16, bf16, fp32, 1, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(6144, bf16, fp32, bf16, fp32, 1, 1, 8, 16, 4);\n\nREGISTER_BWD_LAUNCHER(8192, fp32, fp32, fp32, fp32, 2, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(8192, fp16, fp16, fp16, fp32, 2, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(8192, fp16, fp32, fp16, fp32, 2, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(8192, bf16, bf16, bf16, fp32, 2, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(8192, bf16, fp32, bf16, fp32, 2, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(10240, fp32, fp32, fp32, fp32, 2, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(10240, fp16, fp16, fp16, fp32, 2, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(10240, fp16, fp32, fp16, fp32, 2, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(10240, bf16, bf16, bf16, fp32, 2, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(10240, bf16, fp32, bf16, fp32, 2, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(12288, fp32, fp32, fp32, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(12288, fp16, fp16, fp16, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(12288, fp16, fp32, fp16, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(12288, bf16, bf16, bf16, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(12288, bf16, fp32, bf16, fp32, 4, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(12800, fp32, fp32, fp32, fp32, 5, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(12800, fp16, fp16, fp16, fp32, 5, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(12800, fp16, fp32, fp16, fp32, 5, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(12800, bf16, bf16, bf16, fp32, 5, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(12800, bf16, fp32, bf16, fp32, 5, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(14336, fp32, fp32, fp32, fp32, 4, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(14336, fp16, fp16, fp16, fp32, 4, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(14336, fp16, fp32, fp16, fp32, 4, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(14336, bf16, bf16, bf16, fp32, 4, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(14336, bf16, fp32, bf16, fp32, 4, 1, 4, 8, 4);\n\nREGISTER_BWD_LAUNCHER(15360, fp32, fp32, fp32, fp32, 4, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(15360, fp16, fp16, fp16, fp32, 4, 1, 4, 4, 4);\nREGISTER_BWD_LAUNCHER(15360, fp16, fp32, fp16, fp32, 4, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(15360, bf16, bf16, bf16, fp32, 4, 1, 4, 4, 4);\nREGISTER_BWD_LAUNCHER(15360, bf16, fp32, bf16, fp32, 4, 1, 4, 8, 4);\n\nREGISTER_BWD_LAUNCHER(16384, fp32, fp32, fp32, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(16384, fp16, fp16, fp16, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(16384, fp16, fp32, fp16, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(16384, bf16, bf16, bf16, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(16384, bf16, fp32, bf16, fp32, 4, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(18432, fp32, fp32, fp32, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(18432, fp16, fp16, fp16, fp32, 4, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(18432, fp16, fp32, fp16, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(18432, bf16, bf16, bf16, fp32, 4, 1, 4, 8, 4);\nREGISTER_BWD_LAUNCHER(18432, bf16, fp32, bf16, fp32, 4, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(20480, fp32, fp32, fp32, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(20480, fp16, fp16, fp16, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(20480, fp16, fp32, fp16, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(20480, bf16, bf16, bf16, fp32, 4, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(20480, bf16, fp32, bf16, fp32, 4, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(24576, fp32, fp32, fp32, fp32, 4, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(24576, fp16, fp16, fp16, fp32, 4, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(24576, fp16, fp32, fp16, fp32, 4, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(24576, bf16, bf16, bf16, fp32, 4, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(24576, bf16, fp32, bf16, fp32, 4, 1, 8, 16, 4);\n\nREGISTER_BWD_LAUNCHER(25600, fp32, fp32, fp32, fp32, 5, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(25600, fp16, fp16, fp16, fp32, 5, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(25600, fp16, fp32, fp16, fp32, 5, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(25600, bf16, bf16, bf16, fp32, 5, 1, 4, 16, 4);\nREGISTER_BWD_LAUNCHER(25600, bf16, fp32, bf16, fp32, 5, 1, 4, 16, 4);\n\nREGISTER_BWD_LAUNCHER(30720, fp32, fp32, fp32, fp32, 4, 1, 8, 8, 4);\nREGISTER_BWD_LAUNCHER(30720, fp16, fp16, fp16, fp32, 4, 1, 8, 4, 4);\nREGISTER_BWD_LAUNCHER(30720, fp16, fp32, fp16, fp32, 4, 1, 8, 8, 4);\nREGISTER_BWD_LAUNCHER(30720, bf16, bf16, bf16, fp32, 4, 1, 8, 4, 4);\nREGISTER_BWD_LAUNCHER(30720, bf16, fp32, bf16, fp32, 4, 1, 8, 8, 4);\n\nREGISTER_BWD_LAUNCHER(32768, fp32, fp32, fp32, fp32, 4, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(32768, fp16, fp16, fp16, fp32, 4, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(32768, fp16, fp32, fp16, fp32, 4, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(32768, bf16, bf16, bf16, fp32, 4, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(32768, bf16, fp32, bf16, fp32, 4, 1, 8, 16, 4);\n\nREGISTER_BWD_LAUNCHER(40960, fp32, fp32, fp32, fp32, 4, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(40960, fp16, fp16, fp16, fp32, 4, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(40960, fp16, fp32, fp16, fp32, 4, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(40960, bf16, bf16, bf16, fp32, 4, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(40960, bf16, fp32, bf16, fp32, 4, 1, 8, 16, 4);\n\nREGISTER_BWD_LAUNCHER(49152, fp32, fp32, fp32, fp32, 8, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(49152, fp16, fp16, fp16, fp32, 8, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(49152, fp16, fp32, fp16, fp32, 8, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(49152, bf16, bf16, bf16, fp32, 8, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(49152, bf16, fp32, bf16, fp32, 8, 1, 8, 16, 4);\n\nREGISTER_BWD_LAUNCHER(65536, fp32, fp32, fp32, fp32, 8, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(65536, fp16, fp16, fp16, fp32, 8, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(65536, fp16, fp32, fp16, fp32, 8, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(65536, bf16, bf16, bf16, fp32, 8, 1, 8, 16, 4);\nREGISTER_BWD_LAUNCHER(65536, bf16, fp32, bf16, fp32, 8, 1, 8, 16, 4);\n" }, { "path": "apex/contrib/csrc/layer_norm/ln_fwd_cuda_kernel.cu", "content": "#include \"ln.h\"\n#include \"ln_fwd_kernels.cuh\"\n#include \"ln_kernel_traits.h\"\n#include \"ln_utils.cuh\"\n\nusing namespace layer_norm;\n\ntemplate \nvoid launch_(LaunchParams& launch_params, const bool configure_params) {\n using Kernel_traits = Kernel_traits;\n auto kernel = &ln_fwd_kernel;\n\n if (configure_params) {\n int ctas_per_sm;\n cudaError status_ = cudaOccupancyMaxActiveBlocksPerMultiprocessor(\n &ctas_per_sm, kernel, Kernel_traits::THREADS_PER_CTA, Kernel_traits::SMEM_BYTES_FWD);\n launch_params.params.ctas_per_col =\n launch_params.props->multiProcessorCount * ctas_per_sm / Kernel_traits::CTAS_PER_ROW;\n launch_params.barrier_size = 0;\n launch_params.workspace_bytes = 0;\n if (Kernel_traits::CTAS_PER_ROW > 1) {\n launch_params.barrier_size = 2 * launch_params.params.ctas_per_col;\n launch_params.workspace_bytes = launch_params.params.ctas_per_col * Kernel_traits::WARPS_M *\n Kernel_traits::CTAS_PER_ROW * sizeof(typename Kernel_traits::Stats::stats_t) * 2;\n }\n return;\n }\n\n if (Kernel_traits::SMEM_BYTES_FWD >= 48 * 1024) {\n CHECK_CUDA(\n cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, Kernel_traits::SMEM_BYTES_FWD));\n }\n auto stream = launch_params.stream;\n auto ctas_per_col = launch_params.params.ctas_per_col;\n\n if (Kernel_traits::CTAS_PER_ROW == 1) {\n kernel<<>>(\n launch_params.params);\n } else {\n dim3 grid(Kernel_traits::CTAS_PER_ROW * ctas_per_col);\n dim3 block(Kernel_traits::THREADS_PER_CTA);\n void* params_ = (void*)&launch_params.params;\n cudaLaunchCooperativeKernel((void*)kernel, grid, block, (void**)¶ms_, Kernel_traits::SMEM_BYTES_FWD, stream);\n }\n}\n\n// Create forward launch function and register. Macro signature:\n// HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG\n\nREGISTER_FWD_LAUNCHER(768, fp32, fp32, fp32, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(768, fp16, fp16, fp16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(768, fp16, fp32, fp16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(768, bf16, bf16, bf16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(768, bf16, fp32, bf16, fp32, 1, 4, 1, 16);\n\nREGISTER_FWD_LAUNCHER(1024, fp32, fp32, fp32, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(1024, fp16, fp16, fp16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(1024, fp16, fp32, fp16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(1024, bf16, bf16, bf16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(1024, bf16, fp32, bf16, fp32, 1, 4, 1, 16);\n\nREGISTER_FWD_LAUNCHER(1536, fp32, fp32, fp32, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(1536, fp16, fp16, fp16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(1536, fp16, fp32, fp16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(1536, bf16, bf16, bf16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(1536, bf16, fp32, bf16, fp32, 1, 4, 1, 16);\n\nREGISTER_FWD_LAUNCHER(2048, fp32, fp32, fp32, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(2048, fp16, fp16, fp16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(2048, fp16, fp32, fp16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(2048, bf16, bf16, bf16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(2048, bf16, fp32, bf16, fp32, 1, 4, 1, 16);\n\nREGISTER_FWD_LAUNCHER(2304, fp32, fp32, fp32, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(2304, fp16, fp16, fp16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(2304, fp16, fp32, fp16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(2304, bf16, bf16, bf16, fp32, 1, 4, 1, 16);\nREGISTER_FWD_LAUNCHER(2304, bf16, fp32, bf16, fp32, 1, 4, 1, 16);\n\nREGISTER_FWD_LAUNCHER(3072, fp32, fp32, fp32, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(3072, fp16, fp16, fp16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(3072, fp16, fp32, fp16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(3072, bf16, bf16, bf16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(3072, bf16, fp32, bf16, fp32, 1, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(3840, fp32, fp32, fp32, fp32, 1, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(3840, fp16, fp16, fp16, fp32, 1, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(3840, fp16, fp32, fp16, fp32, 1, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(3840, bf16, bf16, bf16, fp32, 1, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(3840, bf16, fp32, bf16, fp32, 1, 1, 4, 4);\n\nREGISTER_FWD_LAUNCHER(4096, fp32, fp32, fp32, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(4096, fp16, fp16, fp16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(4096, fp16, fp32, fp16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(4096, bf16, bf16, bf16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(4096, bf16, fp32, bf16, fp32, 1, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(5120, fp32, fp32, fp32, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(5120, fp16, fp16, fp16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(5120, fp16, fp32, fp16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(5120, bf16, bf16, bf16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(5120, bf16, fp32, bf16, fp32, 1, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(6144, fp32, fp32, fp32, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(6144, fp16, fp16, fp16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(6144, fp16, fp32, fp16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(6144, bf16, bf16, bf16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(6144, bf16, fp32, bf16, fp32, 1, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(8192, fp32, fp32, fp32, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(8192, fp16, fp16, fp16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(8192, fp16, fp32, fp16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(8192, bf16, bf16, bf16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(8192, bf16, fp32, bf16, fp32, 1, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(10240, fp32, fp32, fp32, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(10240, fp16, fp16, fp16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(10240, fp16, fp32, fp16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(10240, bf16, bf16, bf16, fp32, 1, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(10240, bf16, fp32, bf16, fp32, 1, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(12288, fp32, fp32, fp32, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(12288, fp16, fp16, fp16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(12288, fp16, fp32, fp16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(12288, bf16, bf16, bf16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(12288, bf16, fp32, bf16, fp32, 2, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(12800, fp32, fp32, fp32, fp32, 2, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(12800, fp16, fp16, fp16, fp32, 2, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(12800, fp16, fp32, fp16, fp32, 2, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(12800, bf16, bf16, bf16, fp32, 2, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(12800, bf16, fp32, bf16, fp32, 2, 1, 4, 4);\n\nREGISTER_FWD_LAUNCHER(14336, fp32, fp32, fp32, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(14336, fp16, fp16, fp16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(14336, fp16, fp32, fp16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(14336, bf16, bf16, bf16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(14336, bf16, fp32, bf16, fp32, 2, 1, 4, 8);\n\nREGISTER_FWD_LAUNCHER(15360, fp32, fp32, fp32, fp32, 2, 1, 4, 8);\nREGISTER_FWD_LAUNCHER(15360, fp16, fp16, fp16, fp32, 2, 1, 4, 8);\nREGISTER_FWD_LAUNCHER(15360, fp16, fp32, fp16, fp32, 2, 1, 4, 8);\nREGISTER_FWD_LAUNCHER(15360, bf16, bf16, bf16, fp32, 2, 1, 4, 8);\nREGISTER_FWD_LAUNCHER(15360, bf16, fp32, bf16, fp32, 2, 1, 4, 8);\n\nREGISTER_FWD_LAUNCHER(16384, fp32, fp32, fp32, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(16384, fp16, fp16, fp16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(16384, fp16, fp32, fp16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(16384, bf16, bf16, bf16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(16384, bf16, fp32, bf16, fp32, 2, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(18432, fp32, fp32, fp32, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(18432, fp16, fp16, fp16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(18432, fp16, fp32, fp16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(18432, bf16, bf16, bf16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(18432, bf16, fp32, bf16, fp32, 4, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(20480, fp32, fp32, fp32, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(20480, fp16, fp16, fp16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(20480, fp16, fp32, fp16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(20480, bf16, bf16, bf16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(20480, bf16, fp32, bf16, fp32, 2, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(24576, fp32, fp32, fp32, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(24576, fp16, fp16, fp16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(24576, fp16, fp32, fp16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(24576, bf16, bf16, bf16, fp32, 2, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(24576, bf16, fp32, bf16, fp32, 2, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(25600, fp32, fp32, fp32, fp32, 4, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(25600, fp16, fp16, fp16, fp32, 2, 1, 4, 8);\nREGISTER_FWD_LAUNCHER(25600, fp16, fp32, fp16, fp32, 4, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(25600, bf16, bf16, bf16, fp32, 2, 1, 4, 8);\nREGISTER_FWD_LAUNCHER(25600, bf16, fp32, bf16, fp32, 4, 1, 4, 4);\n\nREGISTER_FWD_LAUNCHER(30720, fp32, fp32, fp32, fp32, 4, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(30720, fp16, fp16, fp16, fp32, 4, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(30720, fp16, fp32, fp16, fp32, 4, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(30720, bf16, bf16, bf16, fp32, 4, 1, 4, 4);\nREGISTER_FWD_LAUNCHER(30720, bf16, fp32, bf16, fp32, 4, 1, 4, 4);\n\nREGISTER_FWD_LAUNCHER(32768, fp32, fp32, fp32, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(32768, fp16, fp16, fp16, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(32768, fp16, fp32, fp16, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(32768, bf16, bf16, bf16, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(32768, bf16, fp32, bf16, fp32, 4, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(40960, fp32, fp32, fp32, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(40960, fp16, fp16, fp16, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(40960, fp16, fp32, fp16, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(40960, bf16, bf16, bf16, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(40960, bf16, fp32, bf16, fp32, 4, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(49152, fp32, fp32, fp32, fp32, 8, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(49152, fp16, fp16, fp16, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(49152, fp16, fp32, fp16, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(49152, bf16, bf16, bf16, fp32, 4, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(49152, bf16, fp32, bf16, fp32, 4, 1, 4, 16);\n\nREGISTER_FWD_LAUNCHER(65536, fp32, fp32, fp32, fp32, 8, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(65536, fp16, fp16, fp16, fp32, 8, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(65536, fp16, fp32, fp16, fp32, 8, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(65536, bf16, bf16, bf16, fp32, 8, 1, 4, 16);\nREGISTER_FWD_LAUNCHER(65536, bf16, fp32, bf16, fp32, 8, 1, 4, 16);\n" }, { "path": "apex/contrib/csrc/layer_norm/ln_fwd_kernels.cuh", "content": "#pragma once\n\n#include \"ln.h\"\n#include \"ln_utils.cuh\"\n\nnamespace layer_norm {\n\ntemplate \n__global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void ln_fwd_kernel(FwdParams params) {\n enum { ROWS_PER_CTA = Ktraits::ROWS_PER_CTA };\n enum { WARPS_N = Ktraits::WARPS_N };\n enum { WARPS_M = Ktraits::WARPS_M };\n enum { THREADS_PER_ROW = Ktraits::THREADS_PER_ROW };\n enum { VEC_COLS_PER_LDG = Ktraits::VEC_COLS_PER_LDG };\n enum { BYTES_PER_ROW = Ktraits::BYTES_PER_ROW };\n enum { LDGS = Ktraits::LDGS };\n enum { NUM_ELTS = Ktraits::NUM_ELTS };\n enum { CTAS_PER_ROW = Ktraits::CTAS_PER_ROW };\n enum { THREADS_PER_WARP = Ktraits::THREADS_PER_WARP };\n\n using output_t = typename Ktraits::output_t;\n using index_t = typename Ktraits::index_t;\n using compute_t = typename Ktraits::compute_t;\n using Ivec = typename Ktraits::Ivec;\n using Ovec = typename Ktraits::Ovec;\n using Wvec = typename Ktraits::Wvec;\n using Cvec = typename Ktraits::Cvec;\n\n using Stats = typename Ktraits::Stats;\n using stats_t = typename Stats::stats_t;\n\n extern __shared__ char smem_[];\n\n const index_t tidx = threadIdx.x;\n const index_t bidn = blockIdx.x % CTAS_PER_ROW;\n const index_t bidm = blockIdx.x / CTAS_PER_ROW;\n const index_t lane = tidx % THREADS_PER_WARP;\n const index_t warp = tidx / THREADS_PER_WARP;\n const index_t warp_m = warp / WARPS_N;\n const index_t warp_n = warp % WARPS_N;\n\n const index_t r = bidm * ROWS_PER_CTA + warp_m;\n const index_t c = bidn * THREADS_PER_ROW + warp_n * THREADS_PER_WARP + lane;\n\n Stats stats(params, bidm, bidn, warp_m, warp_n, lane, smem_);\n\n compute_t* mu_ptr = static_cast(params.mu);\n compute_t* rs_ptr = static_cast(params.rs);\n\n Wvec gamma[LDGS];\n Wvec beta[LDGS];\n index_t idx = c;\n#pragma unroll\n for (int it = 0; it < LDGS; it++) {\n gamma[it].load_from(params.gamma, idx);\n beta[it].load_from(params.beta, idx);\n idx += VEC_COLS_PER_LDG;\n }\n\n constexpr compute_t rn = 1.f / compute_t(Ktraits::COLS);\n\n for (int row = r; row < params.rows; row += params.ctas_per_col * ROWS_PER_CTA) {\n Ivec x[LDGS];\n index_t idx = row * Ktraits::VEC_COLS + c;\n compute_t xf[LDGS * NUM_ELTS];\n#pragma unroll\n for (int it = 0; it < LDGS; it++) {\n x[it].load_from(params.x, idx);\n#pragma unroll\n for (int jt = 0; jt < NUM_ELTS; jt++) {\n compute_t x_ij = compute_t(x[it].data.elt[jt]);\n xf[it * NUM_ELTS + jt] = x_ij;\n }\n idx += VEC_COLS_PER_LDG;\n }\n\n stats_t s = stats.compute(xf, rn);\n\n compute_t mu = layer_norm::Get<0>::of(s);\n compute_t m2 = layer_norm::Get<1>::of(s);\n\n if (bidn == 0 && warp_n == 0 && lane == 0) {\n mu_ptr[row] = mu;\n }\n\n compute_t rs = rsqrtf(rn * m2 + params.epsilon);\n\n if (bidn == 0 && warp_n == 0 && lane == 0) {\n rs_ptr[row] = rs;\n }\n\n Ovec z[LDGS];\n idx = row * Ktraits::VEC_COLS + c;\n#pragma unroll\n for (int it = 0; it < LDGS; it++) {\n#pragma unroll\n for (int jt = 0; jt < NUM_ELTS; jt++) {\n output_t y_ij = output_t(rs * (xf[it * NUM_ELTS + jt] - mu));\n output_t g_ij = gamma[it].data.elt[jt];\n output_t b_ij = beta[it].data.elt[jt];\n z[it].data.elt[jt] = (g_ij * y_ij + b_ij);\n }\n z[it].store_to(params.z, idx);\n idx += VEC_COLS_PER_LDG;\n }\n }\n}\n\n} // namespace layer_norm\n" }, { "path": "apex/contrib/csrc/layer_norm/ln_kernel_traits.h", "content": "#pragma once\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nnamespace layer_norm {\ntemplate \nstruct Kernel_traits_base {\n using weight_t = weight_t_;\n using input_t = input_t_;\n using output_t = output_t_;\n using compute_t = compute_t_;\n using index_t = index_t_;\n\n enum { HIDDEN_SIZE = HIDDEN_SIZE_ };\n enum { THREADS_PER_CTA = THREADS_PER_CTA_ };\n enum { THREADS_PER_WARP = 32 };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate >\nstruct Kernel_traits_finalize : public Base {\n enum { ROWS_PER_CTA = Base::THREADS_PER_CTA / Base::THREADS_PER_WARP };\n static_assert((int)ROWS_PER_CTA <= (int)Base::THREADS_PER_WARP);\n // Bytes per global load from the input.\n enum { BYTES_PER_LDG = BYTES_PER_LDG_ };\n // Number of elements fetched by a global load.\n enum { ELTS_PER_LDG = BYTES_PER_LDG / sizeof(compute_t_) };\n // Bytes per global store of the weights.\n enum { BYTES_PER_STG = ELTS_PER_LDG * sizeof(weight_t_) };\n static_assert(sizeof(BYTES_PER_LDG) == 4, \"Conflict-free smem transpose only implemented for 4B compute type!\");\n static_assert(Base::THREADS_PER_CTA == ROWS_PER_CTA * Base::THREADS_PER_WARP, \"We assume one warp per row!\");\n // The total number of BYTES_PER_LDG-wide words in a hidden vector.\n enum { COLS = HIDDEN_SIZE_ * sizeof(compute_t_) / BYTES_PER_LDG };\n static_assert(COLS * BYTES_PER_LDG == HIDDEN_SIZE_ * sizeof(compute_t_));\n\n // Shared memory size to transpose the CTA result.\n enum { SMEM_BYTES_TRANSPOSE = Base::THREADS_PER_CTA * BYTES_PER_LDG };\n // Shared memory size to coalsece the CTA result.\n enum { SMEM_BYTES_OUTPUT = Base::THREADS_PER_WARP * BYTES_PER_LDG };\n // Shared memory requirement per CTA.\n enum { SMEM_BYTES_PER_CTA = 2 * SMEM_BYTES_TRANSPOSE + 2 * SMEM_BYTES_OUTPUT };\n\n // The type of the reducer.\n using Reducer = layer_norm::Reducer;\n\n // Condition for the whole CTA to participate in syncthreads.\n static_assert(COLS % Base::THREADS_PER_WARP == 0);\n enum { CTAS = COLS / Base::THREADS_PER_WARP };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate >\nstruct Kernel_traits : public Base {\n using input_t = typename Base::input_t;\n using weight_t = typename Base::weight_t;\n using compute_t = typename Base::compute_t;\n using output_t = typename Base::output_t;\n using index_t = typename Base::index_t;\n\n enum { CTAS_PER_ROW = CTAS_PER_ROW_ };\n enum { WARPS_M = WARPS_M_ };\n enum { WARPS_N = WARPS_N_ };\n enum { COLS = HIDDEN_SIZE_ };\n enum { HIDDEN_SIZE = HIDDEN_SIZE_ };\n enum { BYTES_PER_LDG = BYTES_PER_LDG_ };\n enum { NUM_ELTS = BYTES_PER_LDG / sizeof(input_t) };\n\n enum { THREADS_PER_ROW = WARPS_N * THREADS_PER_WARP };\n enum { THREADS_PER_CTA = WARPS_M * THREADS_PER_ROW };\n enum { ROWS_PER_CTA = WARPS_M };\n\n enum { BYTES_PER_ROW = COLS * sizeof(input_t) };\n enum { BYTES_PER_ROW_PER_CTA = THREADS_PER_ROW * BYTES_PER_LDG };\n // Multi-row per CTA not supported for multi-CTA => no smem for WGRAD needed\n enum { SMEM_BYTES_WGRAD = CTAS_PER_ROW > 1 ? 0 : ROWS_PER_CTA * COLS * sizeof(compute_t) };\n static_assert(WARPS_M == 1 || CTAS_PER_ROW == 1);\n\n using reduce_t = typename layer_norm::TypeToVec2::Type;\n using Reducer = layer_norm::Reducer;\n\n enum { SMEM_BYTES_DGRAD = Reducer::SMEM_BYTES };\n enum { SMEM_BYTES = SMEM_BYTES_DGRAD + SMEM_BYTES_WGRAD };\n\n using Ivec = layer_norm::Vec;\n using Ovec = layer_norm::Vec;\n using Wvec = layer_norm::Vec;\n using Cvec = layer_norm::Vec;\n enum { ELTS_PER_LDG = BYTES_PER_LDG / sizeof(input_t) };\n\n // Assume that each thread can handle the same number of elements in the output and weights as in the input.\n static_assert(sizeof(input_t) >= sizeof(output_t));\n static_assert(sizeof(input_t) >= sizeof(weight_t));\n // The number of columns fetched per load from input: one per thread.\n enum { VEC_COLS_PER_LDG = CTAS_PER_ROW * THREADS_PER_ROW };\n // The total number of vectorized loads/stores per hidden vector.\n enum { VEC_COLS = COLS / ELTS_PER_LDG };\n // The number of loads per thread for the input.\n enum { LDGS = VEC_COLS / VEC_COLS_PER_LDG };\n static_assert(LDGS * VEC_COLS_PER_LDG == VEC_COLS);\n // static_assert(LDGS * BYTES_PER_ROW_PER_CTA * CTAS_PER_ROW == BYTES_PER_ROW, \"\");\n\n using Stats = layer_norm::Stats;\n enum { SMEM_BYTES_FWD = Stats::SMEM_BYTES };\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace layer_norm\n" }, { "path": "apex/contrib/csrc/layer_norm/ln_utils.cuh", "content": "#pragma once\n\n#include \n#include \n\n#include \n\n#include \"ln.h\"\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nconstexpr uint32_t THREADS_PER_WARP = 32;\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline void check_cuda_(cudaError_t status, const char* file, int line) {\n if (status != cudaSuccess) {\n fprintf(stderr, \"CUDA Error: %s %s %d\\n\", cudaGetErrorString(status), file, line);\n exit(status);\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define CHECK_CUDA(ans) \\\n { \\\n check_cuda_((ans), __FILE__, __LINE__); \\\n }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define DIVUP(x, y) (((x) + ((y) - 1)) / (y))\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define REGISTER_FWD_LAUNCHER(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG) \\\n void ln_fwd_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE(LaunchParams& launch_params, \\\n const bool configure_params) { \\\n launch_( \\\n launch_params, configure_params); \\\n } \\\n static FwdRegistrar \\\n reg_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE( \\\n ln_fwd_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE)\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n#define REGISTER_BWD_LAUNCHER(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG, \\\n BYTES_PER_LDG_FINALIZE) \\\n void ln_bwd_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE(LaunchParams& launch_params, \\\n const bool configure_params) { \\\n launch_(launch_params, configure_params); \\\n } \\\n static BwdRegistrar \\\n reg_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE( \\\n ln_bwd_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE)\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ float2 operator+(const float2& a, const float2& b) { return {a.x + b.x, a.y + b.y}; }\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ninline __device__ void operator+=(float2& a, const float2& b) {\n a.x += b.x;\n a.y += b.y;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Sum {\n inline __device__ Sum() {}\n inline __device__ T operator()(const T& a, const T& b) { return a + b; }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ T warp_shuffle_xor(const T& x, uint32_t idx) {\n return __shfl_xor_sync(uint32_t(-1), x, idx);\n}\n\ntemplate <>\ninline __device__ float2 warp_shuffle_xor(const float2& x, uint32_t idx) {\n return {warp_shuffle_xor(x.x, idx), warp_shuffle_xor(x.y, idx)};\n}\n\ntemplate \ninline __device__ T warp_shuffle_down(const T& x, uint32_t idx) {\n return __shfl_down_sync(uint32_t(-1), x, idx);\n}\n\ntemplate <>\ninline __device__ float2 warp_shuffle_down(const float2& x, uint32_t idx) {\n return {warp_shuffle_down(x.x, idx), warp_shuffle_down(x.y, idx)};\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nnamespace layer_norm {\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct uint16 {\n uint4 u;\n uint4 v;\n uint4 s;\n uint4 t;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\nstruct uint8 {\n uint4 u;\n uint4 v;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct BytesToType {};\n\ntemplate <>\nstruct BytesToType<64> {\n using Type = uint16;\n static_assert(sizeof(Type) == 64);\n};\n\ntemplate <>\nstruct BytesToType<32> {\n using Type = uint8;\n static_assert(sizeof(Type) == 32);\n};\n\ntemplate <>\nstruct BytesToType<16> {\n using Type = uint4;\n static_assert(sizeof(Type) == 16);\n};\n\ntemplate <>\nstruct BytesToType<8> {\n using Type = uint64_t;\n static_assert(sizeof(Type) == 8);\n};\n\ntemplate <>\nstruct BytesToType<4> {\n using Type = uint32_t;\n static_assert(sizeof(Type) == 4);\n};\n\ntemplate <>\nstruct BytesToType<2> {\n using Type = uint16_t;\n static_assert(sizeof(Type) == 2);\n};\n\ntemplate <>\nstruct BytesToType<1> {\n using Type = uint8_t;\n static_assert(sizeof(Type) == 1);\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct TypeToVec2 {};\n\ntemplate <>\nstruct TypeToVec2 {\n using Type = float2;\n};\n\ntemplate <>\nstruct TypeToVec2 {\n using Type = half2;\n};\n\ntemplate <>\nstruct TypeToVec2 {\n using Type = nv_bfloat162;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Get {\n template \n static inline __device__ R of(const T& vec);\n};\n\ntemplate <>\ntemplate \ninline __device__ R Get<0>::of(const T& vec) {\n return vec.x;\n}\n\ntemplate <>\ntemplate \ninline __device__ R Get<1>::of(const T& vec) {\n return vec.y;\n}\n\ntemplate <>\ntemplate \ninline __device__ R Get<2>::of(const T& vec) {\n return vec.z;\n}\n\ntemplate <>\ntemplate \ninline __device__ R Get<3>::of(const T& vec) {\n return vec.w;\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Converter {\n static inline __device__ Dst convert(const Src& from) { return Dst(from); }\n};\n\ntemplate <>\nstruct Converter {\n static inline __device__ half2 convert(const float2& x) { return __float22half2_rn(x); }\n};\n\ntemplate <>\nstruct Converter {\n static inline __device__ nv_bfloat162 convert(const float2& x) {\n#if __CUDA_ARCH__ >= 800\n return __float22bfloat162_rn(x);\n#else\n union {\n nv_bfloat162 raw;\n nv_bfloat16 x;\n nv_bfloat16 y;\n } tmp;\n tmp.x = __float2bfloat16_rn(x.x);\n tmp.y = __float2bfloat16_rn(x.y);\n return tmp.raw;\n#endif\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Zeros {\n static inline __device__ T get() { return T(0.f); }\n};\n\ntemplate <>\nstruct Zeros {\n static inline __device__ float2 get() { return make_float2(0.f, 0.f); }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Vec {\n enum { BYTES = NUM_ELT * sizeof(Elt_type) };\n\n using Vec_type = typename BytesToType::Type;\n\n using Alias_type = union {\n Vec_type vec;\n Elt_type elt[NUM_ELT];\n };\n\n Alias_type data;\n\n template \n inline __device__ void to(Vec& other) {\n#pragma unroll\n for (int it = 0; it < NUM_ELT; it++) {\n other.data.elt[it] = S(this->data.elt[it]);\n }\n }\n\n template \n inline __device__ void assign(const Op& op) {\n#pragma unroll\n for (int it = 0; it < NUM_ELT; it++) {\n this->data.elt[it] = op(it);\n }\n }\n\n inline __device__ void load_from(const void* base_ptr, const size_t idx) {\n this->data.vec = static_cast(base_ptr)[idx];\n }\n\n inline __device__ void store_to(void* base_ptr, const size_t idx) {\n static_cast(base_ptr)[idx] = this->data.vec;\n }\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct InterCTASync {\n template \n inline __device__ InterCTASync(Params& params, uint32_t bidm, uint32_t bidn)\n : phase_counter_(0),\n b0_(params.barrier + bidm) // The barrier for this group of CTAs.\n ,\n b1_(params.barrier + bidm + params.ctas_per_col) // The barrier for this group of CTAs.\n {\n // BARRIERS ARE ASSUMED TO BE INITIALIZED TO 0!\n }\n\n inline __device__ void spin_wait_(int* barrier, int step, int expected) {\n asm volatile(\"red.release.gpu.global.add.s32 [%0], %1;\" ::\"l\"(barrier), \"r\"(step));\n for (int found = -1; found != expected;) {\n asm volatile(\"ld.global.acquire.gpu.b32 %0, [%1];\" : \"=r\"(found) : \"l\"(barrier));\n }\n }\n\n inline __device__ void sync() {\n // ALL THREADS MUST ENTER!\n\n // We switch barrier every iteration.\n int* barrier = phase_counter_ & 0x1 ? b1_ : b0_;\n // We decrement every other iteration.\n bool dec = phase_counter_ & 0x2;\n int step = dec ? -1 : 1;\n int expected = dec ? 0 : CTAS_PER_ROW;\n // There are only 4 phases: up/down for b0/b1.\n phase_counter_ = (phase_counter_ + 1) & 0x3;\n\n if (threadIdx.x == 0) {\n spin_wait_(barrier, step, expected);\n }\n // CTA waits for thread 0\n __syncthreads();\n }\n\n int phase_counter_;\n int* b0_;\n int* b1_;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Reducer : public Reducer {\n using InterCTASync = InterCTASync;\n using Base = Reducer;\n using Type = typename Base::Type;\n\n enum { SMEM_BYTES = Base::SMEM_BYTES };\n\n enum { WS_BARRIER_BYTES = 2 * sizeof(int) };\n enum { WS_DATA_BYTES = WARPS_M * CTAS_PER_ROW * sizeof(T) };\n\n // size of the barriers + temporary result per CTA (multiply with CTAS_PER_ROW to get total)\n enum { WORKSPACE_BYTES_PER_GROUP = Base::WORKSPACE_BYTES_PER_GROUP + WS_BARRIER_BYTES + WS_DATA_BYTES };\n\n template \n inline __device__ Reducer(Params& params, uint32_t bidm, uint32_t bidn, uint32_t warp_m, uint32_t warp_n,\n uint32_t lane, void* smem)\n : Base(params, bidm, bidn, warp_m, warp_n, lane, smem),\n inter_cta_(params, bidm, bidn),\n bidn_(bidn) // CTA id within the group.\n ,\n w0_(static_cast(params.workspace) + (bidm * WARPS_M + warp_m) * CTAS_PER_ROW),\n w1_(w0_ + params.ctas_per_col * WARPS_M * CTAS_PER_ROW) {}\n\n template \n inline __device__ T allreduce(T data, Op& op) {\n data = Base::reduce(data, op);\n // We switch workspace every iteration.\n T* workspace = inter_cta_.phase_counter_ & 0x1 ? w1_ : w0_;\n\n // Warp leaders 0 hold the CTA-local results.\n if (this->warp_n_ == 0 && this->lane_ == 0) {\n workspace[bidn_] = data;\n }\n inter_cta_.sync();\n static_assert(CTAS_PER_ROW <= 32);\n T total = Zeros::get();\n if (this->lane_ < CTAS_PER_ROW) {\n total = workspace[this->lane_];\n }\n total = Reducer::allreduce_(total, op);\n\n return total;\n }\n\n InterCTASync inter_cta_;\n\n T* w0_;\n T* w1_;\n int bidn_;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Reducer {\n using Type = T;\n enum { SMEM_BYTES = 0 };\n enum { WORKSPACE_BYTES_PER_GROUP = 0 };\n\n enum { THREADS_PER_WARP = 32 };\n\n template \n inline __device__ Reducer(Params& params, uint32_t bidm, uint32_t bidn, uint32_t warp_m, uint32_t warp_n,\n uint32_t lane, void* smem)\n : warp_n_(warp_n), lane_(lane) {}\n\n template \n static inline __device__ T allreduce_(T data, Op& op) {\n#pragma unroll\n for (int it = 1; it < THREADS_PER_WARP; it *= 2) {\n data = op(data, warp_shuffle_xor(data, it));\n }\n return data;\n }\n\n template \n inline __device__ T allreduce(T data, Op& op) {\n return allreduce_(data, op);\n }\n\n template \n inline __device__ T reduce(T data, Op& op) {\n// only lane 0 holds the result!\n#pragma unroll\n for (int it = THREADS_PER_WARP / 2; it > 0; it /= 2) {\n data = op(data, warp_shuffle_down(data, it));\n }\n return data;\n }\n int warp_n_;\n int lane_;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Reducer : public Reducer {\n using Base = Reducer;\n\n using Type = T;\n\n enum { SMEM_BYTES = Base::SMEM_BYTES + WARPS_M * WARPS_N * sizeof(T) * 2 };\n enum { WORKSPACE_BYTES_PER_GROUP = 0 };\n\n enum { THREADS_PER_WARP = 32 };\n\n template \n inline __device__ Reducer(Params& params, uint32_t bidm, uint32_t bidn, uint32_t warp_m, uint32_t warp_n,\n uint32_t lane, void* smem)\n : Base(params, bidm, bidn, warp_m, warp_n, lane, smem), use0_(true) {\n smem0_ = &static_cast(smem)[warp_m * WARPS_N];\n smem1_ = smem0_ + WARPS_M * WARPS_N;\n }\n\n template \n inline __device__ T allreduce(T data, Op& op) {\n T* smem = use0_ ? smem0_ : smem1_;\n use0_ = !use0_;\n data = Base::reduce(data, op);\n if (this->lane_ == 0) {\n smem[this->warp_n_] = data;\n }\n __syncthreads();\n T out = Zeros::get();\n#pragma unroll\n for (int it = 0; it < WARPS_N; it++) {\n out = op(out, smem[it]);\n }\n return out;\n }\n\n template \n inline __device__ T reduce(T data, Op& op) {\n T* smem = use0_ ? smem0_ : smem1_;\n use0_ = !use0_;\n // only intra-CTA group leader holds the result!\n data = Base::reduce(data, op);\n if (this->lane_ == 0) {\n smem[this->warp_n_] = data;\n }\n __syncthreads();\n T out = Zeros::get();\n if (this->warp_n_ == 0 && this->lane_ == 0) {\n#pragma unroll\n for (int it = 0; it < WARPS_N; it++) {\n out = op(out, smem[it]);\n }\n }\n return out;\n }\n\n T* smem0_;\n T* smem1_;\n bool use0_;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \ninline __device__ void warp_chan_upd_dynamic(T& m_a, T& m2_a, T& n_a, int num_active) {\n // Assume at least leftmost is valid and init: step = next_pow2(num_active) / 2 (might get NaN otherwise)\n int highest_bit_set = (8 * sizeof(num_active)) - __clz(num_active - 1);\n\n#pragma unroll\n for (int step = (1 << (highest_bit_set - 1)); step > 0; step /= 2) {\n // Exchange\n T n_b = warp_shuffle_down(n_a, step);\n T m_b = warp_shuffle_down(m_a, step);\n T m2_b = warp_shuffle_down(m2_a, step);\n\n // Update\n const T n_ab = n_a + n_b; // We can handle one of them being 0, not both.\n const T rn_ab = 1.f / n_ab; // Might have different n per thread, otherwise this would simplify :(\n const T delta = m_a - m_b;\n const float m2_ab = m2_a + m2_b + delta * delta * n_a * n_b * rn_ab;\n const float m_ab = (n_a * m_a + n_b * m_b) * rn_ab;\n\n n_a = n_ab;\n m_a = m_ab;\n m2_a = m2_ab;\n }\n // Intra-warp broadcast (only lane 0 has valid stats).\n m_a = __shfl_sync(uint32_t(-1), m_a, 0);\n m2_a = __shfl_sync(uint32_t(-1), m2_a, 0);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Stats {\n // This could be done generically with the Reducer. But then we would have to exchange 3 instead of 2 fields.\n\n using InterCTASync = InterCTASync;\n using BlockStats = Stats;\n using stats_t = typename BlockStats::stats_t;\n\n enum { SMEM_BYTES = BlockStats::SMEM_BYTES };\n\n template \n inline __device__ Stats(Params& params, uint32_t bidm, uint32_t bidn, uint32_t warp_m, uint32_t warp_n, uint32_t lane,\n void* smem)\n : inter_cta_(params, bidm, bidn),\n block_stats_(params, bidm, bidn, warp_m, warp_n, lane, smem),\n bidn_(bidn) // CTA id within the group.\n ,\n w0_(static_cast(params.workspace) + (bidm * WARPS_M + warp_m) * CTAS_PER_ROW),\n w1_(w0_ + params.ctas_per_col * WARPS_M * CTAS_PER_ROW),\n warp_n_(warp_n),\n lane_(lane) {}\n\n template \n inline __device__ stats_t compute(const T (&elts)[N], const T rn) {\n constexpr T ELTS_PER_ROW_PER_CTA = N * WARPS_N * THREADS_PER_WARP;\n // TODO rn is not really needed here..\n constexpr T block_rn = 1.f / T(ELTS_PER_ROW_PER_CTA);\n stats_t block_stats = block_stats_.compute(elts, block_rn);\n\n stats_t* workspace = inter_cta_.phase_counter_ & 0x1 ? w1_ : w0_;\n\n if (warp_n_ == 0 && lane_ == 0) {\n workspace[bidn_] = block_stats;\n }\n\n // Wait for all CTAS_PER_ROW CTAS in the group to have written their result.\n inter_cta_.sync();\n\n T n = Zeros::get();\n T m = Zeros::get();\n T m2 = Zeros::get();\n\n // Assume CTA group size in N less than 32, such that we can finalize with a single warp.\n static_assert(CTAS_PER_ROW <= 32);\n\n // Every warp does the final reduction locally.\n if (lane_ < CTAS_PER_ROW) {\n stats_t result = workspace[lane_];\n n = ELTS_PER_ROW_PER_CTA;\n m = layer_norm::Get<0>::of(result);\n m2 = layer_norm::Get<1>::of(result);\n }\n\n warp_chan_upd_dynamic(m, m2, n, CTAS_PER_ROW);\n\n return {m, m2};\n }\n\n InterCTASync inter_cta_;\n BlockStats block_stats_;\n\n stats_t* w0_;\n stats_t* w1_;\n int bidn_;\n int warp_n_;\n int lane_;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Stats {\n using WarpStats = Stats;\n using stats_t = typename WarpStats::stats_t;\n\n enum { SMEM_BYTES = WARPS_M * WARPS_N * sizeof(stats_t) * 2 };\n\n template \n inline __device__ Stats(Params& params, uint32_t bidm, uint32_t bidn, uint32_t warp_m, uint32_t warp_n, uint32_t lane,\n void* smem)\n : warp_stats_(params, bidm, bidn, warp_m, warp_n, lane, smem), use0_(true) {\n smem0_ = static_cast(smem) + warp_m * WARPS_N;\n smem1_ = smem0_ + WARPS_M * WARPS_N;\n }\n\n template \n inline __device__ stats_t compute(const T (&elts)[N], const T rn) {\n stats_t* smem = use0_ ? smem0_ : smem1_;\n use0_ = !use0_;\n // Compute warp local for all WARPS_N\n constexpr T warp_rn = 1.f / T(N * THREADS_PER_WARP);\n stats_t warp_stats = warp_stats_.compute(elts, warp_rn);\n\n // Each warp warp leader stores its stats\n const auto warp_n = warp_stats_.reducer_.warp_n_;\n const auto lane = warp_stats_.reducer_.lane_;\n if (lane == 0) {\n smem[warp_n] = warp_stats;\n }\n __syncthreads();\n\n T n = Zeros::get();\n T m = Zeros::get();\n T m2 = Zeros::get();\n\n // Assume that there are less than 32 warps, such that we can finalize with a single warp\n static_assert(WARPS_N <= 32);\n if (lane < WARPS_N) {\n stats_t result = smem[lane];\n n = N * THREADS_PER_WARP;\n m = layer_norm::Get<0>::of(result);\n m2 = layer_norm::Get<1>::of(result);\n }\n\n warp_chan_upd_dynamic(m, m2, n, WARPS_N);\n\n return {m, m2};\n }\n WarpStats warp_stats_;\n stats_t* smem0_;\n stats_t* smem1_;\n bool use0_;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct Stats {\n using stats_t = typename TypeToVec2::Type;\n // The simple Warp reducer.\n using Reducer = Reducer;\n\n enum { SMEM_BYTES = 0 };\n\n template \n inline __device__ Stats(Params& params, uint32_t bidm, uint32_t bidn, uint32_t warp_m, uint32_t warp_n, uint32_t lane,\n void* smem)\n : reducer_(params, bidm, bidn, warp_m, warp_n, lane, smem) {}\n\n template \n inline __device__ stats_t compute(const T (&elts)[N], const T rn) {\n auto sum = Sum();\n\n T m = Zeros::get();\n#pragma unroll\n for (int it = 0; it < N; it++) {\n m += elts[it];\n }\n m = reducer_.allreduce(m, sum) * rn;\n\n T m2 = Zeros::get();\n#pragma unroll\n for (int it = 0; it < N; it++) {\n T diff = (elts[it] - m);\n m2 += diff * diff;\n }\n m2 = reducer_.allreduce(m2, sum);\n\n return {m, m2};\n }\n\n Reducer reducer_;\n};\n\n////////////////////////////////////////////////////////////////////////////////////////////////////\n\n} // namespace layer_norm\n" }, { "path": "apex/contrib/csrc/multihead_attn/additive_masked_softmax_dropout_cuda.cu", "content": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"dropout.cuh\"\n#include \"softmax.cuh\"\n\n// symbol to be automatically resolved by PyTorch libs\n\nnamespace multihead_attn {\nnamespace fused_softmax {\nnamespace additive_mask_softmax_dropout {\n\nstd::vector fwd_cuda(bool is_training, int heads, torch::Tensor const& input, const half* pad_mask,\n float dropout_prob) {\n const int attn_batches = input.size(0);\n const int sequences = attn_batches / heads;\n const int q_seq_len = input.size(1);\n const int k_seq_len = q_seq_len;\n // const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n\n // There is no reason to use more than one stream as every kernel is\n // sequentially dependent\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // 3 Intermediate Results + Output (Note: dropout intermediates are generated\n // by ATen library code)\n auto act_options = input.options().requires_grad(false);\n auto mask_options = act_options.dtype(torch::kUInt8);\n\n torch::Tensor softmax_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_mask = torch::empty({attn_batches, q_seq_len, k_seq_len}, mask_options);\n\n // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)\n void* input_ptr = static_cast(input.data_ptr());\n void* softmax_results_ptr = static_cast(softmax_results.data_ptr());\n\n // Padded Softmax\n [[maybe_unused]] bool softmax_success = false;\n if (pad_mask == nullptr) {\n softmax_success = dispatch_softmax(reinterpret_cast(softmax_results_ptr),\n reinterpret_cast(input_ptr), k_seq_len,\n k_seq_len, attn_batches * q_seq_len);\n } else {\n softmax_success = dispatch_additive_masked_softmax(\n reinterpret_cast(softmax_results_ptr), reinterpret_cast(input_ptr), pad_mask, k_seq_len,\n k_seq_len, attn_batches * q_seq_len, attn_batches * q_seq_len / sequences);\n }\n\n if (is_training) {\n // use at:: function so that C++ version generates the same random mask as\n // python version\n auto dropout_tuple = at::_fused_dropout(softmax_results, 1.0f - dropout_prob);\n dropout_results = std::get<0>(dropout_tuple);\n dropout_mask = std::get<1>(dropout_tuple);\n }\n\n // Matmul2\n\n return {dropout_results, dropout_mask, softmax_results};\n}\n\ntorch::Tensor bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& softmax_results,\n torch::Tensor const& dropout_mask, float dropout_prob) {\n const int attn_batches = output_grads.size(0);\n const int q_seq_len = output_grads.size(1);\n const int k_seq_len = q_seq_len;\n // const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n // TODO: Streams can be used in Backprop but I haven't added more than one\n // in my first attempt to create the code\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // Output Tensor Allocations\n // torch::Tensor input_grads = torch::empty_like(output_grads);\n\n // Apply Dropout Mask and Scale by Dropout Probability\n // Softmax Grad\n dispatch_masked_scale_softmax_backward_stream(\n static_cast(output_grads.data_ptr()), static_cast(output_grads.data_ptr()),\n reinterpret_cast(softmax_results.data_ptr()), static_cast(dropout_mask.data_ptr()),\n 1.0 / (1.0 - dropout_prob), k_seq_len, k_seq_len, attn_batches * q_seq_len, stream);\n // backward pass is completely in-place\n return output_grads;\n}\n} // namespace additive_mask_softmax_dropout\n} // namespace fused_softmax\n} // namespace multihead_attn\n" }, { "path": "apex/contrib/csrc/multihead_attn/dropout.cuh", "content": "#pragma once\n#include \n\n#ifdef OLD_GENERATOR_PATH\n#include \n#else\n#include \n#endif\n\n#include \n#include \n\nnamespace {\nconstexpr int UNROLL = 4;\n} // namespace\n\ntemplate \n__global__ void apex_fused_dropout_kernel(scalar_t const* inputs, scalar_t* outputs, uint8_t* mask,\n IndexType totalElements, accscalar_t p, std::pair seeds) {\n accscalar_t pinv = accscalar_t(1) / p;\n IndexType idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n curandStatePhilox4_32_10_t state;\n curand_init(seeds.first, idx, seeds.second, &state);\n\n IndexType rounded_size =\n ((totalElements - 1) / (blockDim.x * gridDim.x * UNROLL) + 1) * blockDim.x * gridDim.x * UNROLL;\n for (IndexType linearIndex = idx; linearIndex < rounded_size; linearIndex += gridDim.x * blockDim.x * UNROLL) {\n float4 rand = curand_uniform4(&state);\n scalar_t src[UNROLL];\n rand.x = rand.x <= p;\n rand.y = rand.y <= p;\n rand.z = rand.z <= p;\n rand.w = rand.w <= p;\n\n for (int ii = 0; ii < UNROLL; ii++) {\n IndexType li = linearIndex + blockDim.x * gridDim.x * ii;\n if (li < totalElements) {\n src[ii] = inputs[li];\n }\n }\n for (int ii = 0; ii < UNROLL; ii++) {\n IndexType li = linearIndex + blockDim.x * gridDim.x * ii;\n if (li < totalElements) {\n outputs[li] = src[ii] * (&rand.x)[ii] * pinv;\n mask[li] = (uint8_t)(&rand.x)[ii];\n }\n }\n __syncthreads();\n }\n}\n\ntemplate \n__global__ void apex_dropout_add_kernel(scalar_t const* inputs, scalar_t const* add_inputs, scalar_t* outputs,\n uint8_t* mask, IndexType totalElements, accscalar_t p,\n std::pair seeds) {\n accscalar_t pinv = accscalar_t(1) / p;\n IndexType idx = blockIdx.x * blockDim.x + threadIdx.x;\n\n curandStatePhilox4_32_10_t state;\n curand_init(seeds.first, idx, seeds.second, &state);\n\n IndexType rounded_size =\n ((totalElements - 1) / (blockDim.x * gridDim.x * UNROLL) + 1) * blockDim.x * gridDim.x * UNROLL;\n for (IndexType linearIndex = idx; linearIndex < rounded_size; linearIndex += gridDim.x * blockDim.x * UNROLL) {\n float4 rand = curand_uniform4(&state);\n scalar_t src[UNROLL];\n scalar_t add_src[UNROLL];\n rand.x = rand.x <= p;\n rand.y = rand.y <= p;\n rand.z = rand.z <= p;\n rand.w = rand.w <= p;\n for (int ii = 0; ii < UNROLL; ii++) {\n IndexType li = linearIndex + blockDim.x * gridDim.x * ii;\n if (li < totalElements) {\n src[ii] = inputs[li];\n add_src[ii] = add_inputs[li];\n }\n }\n for (int ii = 0; ii < UNROLL; ii++) {\n IndexType li = linearIndex + blockDim.x * gridDim.x * ii;\n if (li < totalElements) {\n accscalar_t int1 = src[ii] * (&rand.x)[ii] * pinv;\n outputs[li] = static_cast(static_cast(add_src[ii]) + int1);\n mask[li] = (uint8_t)(&rand.x)[ii];\n }\n }\n __syncthreads();\n }\n}\n\ntemplate \n__global__ void apex_add_kernel(scalar_t const* inputs, scalar_t const* add_inputs, scalar_t* outputs,\n IndexType totalElements) {\n IndexType idx = blockIdx.x * blockDim.x + threadIdx.x;\n IndexType rounded_size =\n ((totalElements - 1) / (blockDim.x * gridDim.x * UNROLL) + 1) * blockDim.x * gridDim.x * UNROLL;\n for (IndexType linearIndex = idx; linearIndex < rounded_size; linearIndex += gridDim.x * blockDim.x * UNROLL) {\n scalar_t src[UNROLL];\n scalar_t add_src[UNROLL];\n for (int ii = 0; ii < UNROLL; ii++) {\n IndexType li = linearIndex + blockDim.x * gridDim.x * ii;\n if (li < totalElements) {\n src[ii] = inputs[li];\n add_src[ii] = add_inputs[li];\n }\n }\n for (int ii = 0; ii < UNROLL; ii++) {\n IndexType li = linearIndex + blockDim.x * gridDim.x * ii;\n if (li < totalElements) {\n outputs[li] = src[ii] + add_src[ii];\n }\n }\n __syncthreads();\n }\n}\n\ntemplate \n__global__ void apex_masked_scale_kernel(scalar_t const* inputs, scalar_t* outputs, uint8_t const* mask,\n IndexType totalElements, accscalar_t scale) {\n IndexType idx = blockIdx.x * blockDim.x + threadIdx.x;\n IndexType rounded_size =\n ((totalElements - 1) / (blockDim.x * gridDim.x * UNROLL) + 1) * blockDim.x * gridDim.x * UNROLL;\n for (IndexType linearIndex = idx; linearIndex < rounded_size; linearIndex += gridDim.x * blockDim.x * UNROLL) {\n scalar_t src[UNROLL];\n scalar_t msk[UNROLL];\n for (int ii = 0; ii < UNROLL; ii++) {\n IndexType li = linearIndex + blockDim.x * gridDim.x * ii;\n if (li < totalElements) {\n src[ii] = static_cast(inputs[li]);\n msk[ii] = static_cast(mask[li]);\n }\n }\n for (int ii = 0; ii < UNROLL; ii++) {\n IndexType li = linearIndex + blockDim.x * gridDim.x * ii;\n if (li < totalElements) {\n outputs[li] = static_cast(src[ii]) * scale * static_cast(msk[ii]);\n }\n }\n }\n}\n\ntemplate \nvoid apex_fused_dropout_cuda(scalar_t const* inputs, scalar_t* outputs, uint8_t* mask, IndexType totalElements,\n accscalar_t p) {\n auto gen = at::cuda::detail::getDefaultCUDAGenerator();\n\n int block_size = 256;\n dim3 dim_block(block_size);\n dim3 grid((totalElements + block_size - 1) / block_size);\n unsigned int blocks_per_sm = at::cuda::getCurrentDeviceProperties()->maxThreadsPerMultiProcessor / block_size;\n grid.x = std::min((unsigned int)at::cuda::getCurrentDeviceProperties()->multiProcessorCount * blocks_per_sm, grid.x);\n\n // number of times random will be generated per thread, to offset philox\n // counter in the random state\n int64_t counter_offset = ((totalElements - 1) / (block_size * grid.x * UNROLL) + 1) * UNROLL;\n std::pair rng_engine_inputs;\n {\n // See Note [Acquire lock when using random generators]\n std::lock_guard lock(gen.mutex());\n rng_engine_inputs = at::check_generator(gen)->philox_engine_inputs(counter_offset);\n }\n\n apex_fused_dropout_kernel<<>>(\n inputs, outputs, mask, totalElements, p, rng_engine_inputs);\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\ntemplate \nvoid apex_dropout_add_cuda(scalar_t const* inputs, scalar_t const* add_inputs, scalar_t* outputs, uint8_t* mask,\n IndexType totalElements, accscalar_t p) {\n auto gen = at::cuda::detail::getDefaultCUDAGenerator();\n\n int block_size = 256;\n dim3 dim_block(block_size);\n dim3 grid((totalElements + block_size - 1) / block_size);\n unsigned int blocks_per_sm = at::cuda::getCurrentDeviceProperties()->maxThreadsPerMultiProcessor / block_size;\n grid.x = std::min((unsigned int)at::cuda::getCurrentDeviceProperties()->multiProcessorCount * blocks_per_sm, grid.x);\n\n // number of times random will be generated per thread, to offset philox\n // counter in the random state\n int64_t counter_offset = ((totalElements - 1) / (block_size * grid.x * UNROLL) + 1) * UNROLL;\n std::pair rng_engine_inputs;\n {\n // See Note [Acquire lock when using random generators]\n std::lock_guard lock(gen.mutex());\n rng_engine_inputs = at::check_generator(gen)->philox_engine_inputs(counter_offset);\n }\n\n apex_dropout_add_kernel<<>>(\n inputs, add_inputs, outputs, mask, totalElements, p, rng_engine_inputs);\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\ntemplate \nvoid apex_add_cuda(scalar_t const* inputs, scalar_t const* add_inputs, scalar_t* outputs, IndexType totalElements) {\n int block_size = 256;\n dim3 dim_block(block_size);\n dim3 grid((totalElements + block_size - 1) / block_size);\n unsigned int blocks_per_sm = at::cuda::getCurrentDeviceProperties()->maxThreadsPerMultiProcessor / block_size;\n grid.x = std::min((unsigned int)at::cuda::getCurrentDeviceProperties()->multiProcessorCount * blocks_per_sm, grid.x);\n\n apex_add_kernel\n <<>>(inputs, add_inputs, outputs, totalElements);\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\ntemplate \nvoid apex_masked_scale_cuda(scalar_t const* inputs, scalar_t* outputs, uint8_t const* mask, IndexType totalElements,\n accscalar_t scale) {\n int block_size = 256;\n dim3 dim_block(block_size);\n dim3 grid((totalElements + block_size - 1) / block_size);\n unsigned int blocks_per_sm = at::cuda::getCurrentDeviceProperties()->maxThreadsPerMultiProcessor / block_size;\n grid.x = std::min((unsigned int)at::cuda::getCurrentDeviceProperties()->multiProcessorCount * blocks_per_sm, grid.x);\n\n apex_masked_scale_kernel\n <<>>(inputs, outputs, mask, totalElements, scale);\n C10_CUDA_CHECK(cudaGetLastError());\n}\n" }, { "path": "apex/contrib/csrc/multihead_attn/encdec_multihead_attn_cuda.cu", "content": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"dropout.cuh\"\n#include \"softmax.cuh\"\n#include \"strided_batched_gemm.cuh\"\n\nnamespace multihead_attn {\nnamespace encdec {\nnamespace cublas_gemmex {\n\nstd::vector fwd_cuda(bool use_time_mask, bool is_training, int heads, torch::Tensor const& inputs_q,\n torch::Tensor const& inputs_kv, torch::Tensor const& input_weights_q,\n torch::Tensor const& input_weights_kv, torch::Tensor const& output_weights,\n const uint8_t* pad_mask, float dropout_prob) {\n const int embed_dim = inputs_q.size(2);\n const int sequences = inputs_q.size(1);\n const int q_seq_len = inputs_q.size(0);\n const int k_seq_len = inputs_kv.size(0);\n const int batches_q = sequences * q_seq_len;\n const int batches_kv = sequences * k_seq_len;\n const int head_dim = embed_dim / heads;\n const int output_lin_q_dim = embed_dim;\n const int output_lin_kv_dim = 2 * embed_dim;\n const int attn_batches = heads * sequences;\n const int lead_dim_q = attn_batches * head_dim;\n const int lead_dim_kv = attn_batches * 2 * head_dim;\n const int batch_stride_q = head_dim;\n const int batch_stride_kv = 2 * head_dim;\n const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n const float alpha = 1.0;\n const float beta = 0.0;\n const float scale = 1.0 / sqrt(static_cast(head_dim));\n\n // There is no reason to use more than one stream as every kernel is\n // sequentially dependent\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // 3 Intermediate Results + Output (Note: dropout intermediates are generated\n // by ATen library code)\n auto act_options = inputs_q.options().requires_grad(false);\n auto mask_options = act_options.dtype(torch::kUInt8);\n\n torch::Tensor input_lin_q_results = torch::empty({q_seq_len, sequences, output_lin_q_dim}, act_options);\n torch::Tensor input_lin_kv_results = torch::empty({k_seq_len, sequences, output_lin_kv_dim}, act_options);\n torch::Tensor softmax_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_mask = torch::empty({attn_batches, q_seq_len, k_seq_len}, mask_options);\n torch::Tensor matmul2_results = torch::empty({q_seq_len, attn_batches, head_dim}, act_options);\n torch::Tensor outputs = torch::empty_like(inputs_q, act_options);\n\n // Input Linear Results Pointers to Q, K, and V of interviewed activations\n void* q_lin_results_ptr = static_cast(input_lin_q_results.data_ptr());\n void* k_lin_results_ptr = static_cast(input_lin_kv_results.data_ptr());\n void* v_lin_results_ptr = static_cast(static_cast(input_lin_kv_results.data_ptr()) + head_dim);\n\n // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)\n void* softmax_results_ptr = static_cast(softmax_results.data_ptr());\n\n char a_layout_t{'t'};\n char a_layout_n{'n'};\n char b_layout_n{'n'};\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));\n // Input Linear Q Fwd\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, output_lin_q_dim, batches_q, embed_dim, static_cast(&alpha),\n static_cast(input_weights_q.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(inputs_q.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta),\n q_lin_results_ptr, CUDA_R_16F, output_lin_q_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Input Linear KV Fwd\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, output_lin_kv_dim, batches_kv, embed_dim, static_cast(&alpha),\n static_cast(input_weights_kv.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(inputs_kv.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta),\n k_lin_results_ptr, CUDA_R_16F, output_lin_kv_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // MatMul1 of Dot-Product Attention Plus scaling by 1/Sqrt(head size)\n gemm_switch_fp32accum(a_layout_t, b_layout_n, k_seq_len, q_seq_len, head_dim, scale,\n static_cast(k_lin_results_ptr), lead_dim_kv, batch_stride_kv,\n static_cast(q_lin_results_ptr), lead_dim_q, batch_stride_q, beta,\n static_cast(softmax_results_ptr), k_seq_len, k_seq_len * q_seq_len, attn_batches);\n\n // Padded Softmax\n bool softmax_success = false;\n if (pad_mask == nullptr) {\n softmax_success = dispatch_softmax(reinterpret_cast(softmax_results_ptr),\n reinterpret_cast(softmax_results_ptr), k_seq_len,\n k_seq_len, attn_batches * q_seq_len);\n } else {\n if (use_time_mask) {\n softmax_success = dispatch_time_masked_softmax(\n reinterpret_cast(softmax_results_ptr), reinterpret_cast(softmax_results_ptr), pad_mask,\n k_seq_len, k_seq_len, attn_batches * q_seq_len, q_seq_len);\n } else {\n softmax_success = dispatch_masked_softmax(\n reinterpret_cast(softmax_results_ptr), reinterpret_cast(softmax_results_ptr), pad_mask,\n k_seq_len, k_seq_len, attn_batches * q_seq_len, attn_batches * q_seq_len / sequences);\n }\n }\n assert(softmax_success);\n\n if (is_training) {\n apex_fused_dropout_cuda(\n static_cast(softmax_results.data_ptr()), static_cast(dropout_results.data_ptr()),\n static_cast(dropout_mask.data_ptr()), dropout_elems, (1.0f - dropout_prob));\n }\n\n // Matmul2\n gemm_switch_fp32accum(a_layout_n, b_layout_n, head_dim, q_seq_len, k_seq_len, alpha,\n static_cast(v_lin_results_ptr), lead_dim_kv, batch_stride_kv,\n (is_training) ? static_cast(dropout_results.data_ptr())\n : static_cast(softmax_results.data_ptr()),\n k_seq_len, k_seq_len * q_seq_len, beta, static_cast(matmul2_results.data_ptr()),\n head_dim * attn_batches, head_dim, attn_batches);\n\n // Output Linear\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, embed_dim, batches_q, embed_dim, static_cast(&alpha),\n static_cast(output_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(matmul2_results.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta),\n static_cast(outputs.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F,\n // CUBLAS_GEMM_ALGO1_TENSOR_OP));\n CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));\n\n return {input_lin_q_results, input_lin_kv_results, softmax_results, dropout_results,\n dropout_mask, matmul2_results, outputs};\n}\n\nstd::vector bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_q_results, torch::Tensor const& input_lin_kv_results,\n torch::Tensor const& inputs_q, torch::Tensor const& inputs_kv,\n torch::Tensor const& input_weights_q, torch::Tensor const& input_weights_kv,\n torch::Tensor const& output_weights, torch::Tensor const& dropout_mask,\n float dropout_prob) {\n const int embed_dim = inputs_q.size(2);\n const int sequences = inputs_q.size(1);\n const int q_seq_len = inputs_q.size(0);\n const int k_seq_len = inputs_kv.size(0);\n const int batches_q = sequences * q_seq_len;\n const int batches_kv = sequences * k_seq_len;\n const int head_dim = embed_dim / heads;\n const int output_lin_q_dim = embed_dim;\n const int output_lin_kv_dim = 2 * embed_dim;\n const int attn_batches = heads * sequences;\n const int lead_dim_q = attn_batches * head_dim;\n const int lead_dim_kv = attn_batches * 2 * head_dim;\n const int batch_stride_q = head_dim;\n const int batch_stride_kv = 2 * head_dim;\n const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n const float alpha = 1.0;\n const float beta = 0.0;\n const float scale = 1.0 / sqrt(static_cast(head_dim));\n\n // TODO: Streams can be used in Backprop but I haven't added more than one\n // in my first attempt to create the code\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // Output Tensor Allocations\n torch::Tensor input_q_grads = torch::empty_like(inputs_q);\n torch::Tensor input_kv_grads = torch::empty_like(inputs_kv);\n torch::Tensor input_weight_q_grads = torch::empty_like(input_weights_q);\n torch::Tensor input_weight_kv_grads = torch::empty_like(input_weights_kv);\n torch::Tensor output_weight_grads = torch::empty_like(output_weights);\n // Intermediate Tensor Allocations\n at::Tensor output_lin_grads = torch::empty_like(matmul2_results);\n at::Tensor matmul2_grads = torch::empty_like(dropout_results);\n at::Tensor input_lin_q_output_grads = torch::empty_like(input_lin_q_results);\n at::Tensor input_lin_kv_output_grads = torch::empty_like(input_lin_kv_results);\n\n auto q_lin_results_ptr = static_cast(input_lin_q_results.data_ptr());\n auto k_lin_results_ptr = static_cast(input_lin_kv_results.data_ptr());\n auto v_lin_results_ptr = static_cast(input_lin_kv_results.data_ptr()) + head_dim;\n\n auto q_lin_grads_ptr = static_cast(input_lin_q_output_grads.data_ptr());\n auto k_lin_grads_ptr = static_cast(input_lin_kv_output_grads.data_ptr());\n auto v_lin_grads_ptr = static_cast(input_lin_kv_output_grads.data_ptr()) + head_dim;\n\n char a_layout_n{'n'};\n char a_layout_t{'t'};\n char b_layout_n{'n'};\n char b_layout_t{'t'};\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));\n\n // Output Linear Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches_q, embed_dim,\n static_cast(&alpha),\n static_cast(output_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(output_grads.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_lin_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Output Linear Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, embed_dim, batches_q,\n static_cast(&alpha),\n static_cast(matmul2_results.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(output_grads.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_weight_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // MatMul2 Dgrad1\n gemm_switch_fp32accum(a_layout_t, b_layout_n, k_seq_len, q_seq_len, head_dim, alpha,\n static_cast(v_lin_results_ptr), lead_dim_kv, batch_stride_kv,\n static_cast(output_lin_grads.data_ptr()), head_dim * attn_batches, head_dim, beta,\n static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len, attn_batches);\n\n // Matmul2 Dgrad2\n gemm_switch_fp32accum(a_layout_n, b_layout_t, head_dim, k_seq_len, q_seq_len, alpha,\n static_cast(output_lin_grads.data_ptr()), head_dim * attn_batches, head_dim,\n static_cast(dropout_results.data_ptr()), k_seq_len, k_seq_len * q_seq_len, beta,\n v_lin_grads_ptr, lead_dim_kv, batch_stride_kv, attn_batches);\n\n // Apply Dropout Mask and Scale by Dropout Probability\n apex_masked_scale_cuda(\n static_cast(matmul2_grads.data_ptr()), static_cast(matmul2_grads.data_ptr()),\n static_cast(dropout_mask.data_ptr()), dropout_elems, (1.0 / (1.0 - dropout_prob)));\n\n // Softmax Grad\n bool softmax_success = false;\n softmax_success = dispatch_softmax_backward(\n static_cast(matmul2_grads.data_ptr()), static_cast(matmul2_grads.data_ptr()),\n reinterpret_cast(softmax_results.data_ptr()), k_seq_len, k_seq_len, attn_batches * q_seq_len);\n assert(softmax_success);\n\n // Matmul1 Dgrad1\n gemm_switch_fp32accum(a_layout_n, b_layout_n, head_dim, q_seq_len, k_seq_len, scale, k_lin_results_ptr, lead_dim_kv,\n batch_stride_kv, static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta, q_lin_grads_ptr, lead_dim_q, batch_stride_q, attn_batches);\n\n // Matmul1 Dgrad2\n gemm_switch_fp32accum(a_layout_n, b_layout_t, head_dim, k_seq_len, q_seq_len, scale, q_lin_results_ptr, lead_dim_q,\n batch_stride_q, static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta, k_lin_grads_ptr, lead_dim_kv, batch_stride_kv, attn_batches);\n\n // Input Linear Q Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches_q, output_lin_q_dim, static_cast(&alpha),\n static_cast(input_weights_q.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(q_lin_grads_ptr), CUDA_R_16F, output_lin_q_dim, static_cast(&beta),\n static_cast(input_q_grads.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F,\n // CUBLAS_GEMM_ALGO10_TENSOR_OP));\n CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Input Linear Q Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, output_lin_q_dim, batches_q,\n static_cast(&alpha), static_cast(inputs_q.data_ptr()),\n CUDA_R_16F, embed_dim, static_cast(q_lin_grads_ptr), CUDA_R_16F,\n output_lin_q_dim, static_cast(&beta),\n static_cast(input_weight_q_grads.data_ptr()), CUDA_R_16F, embed_dim,\n CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Input Linear KV Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches_kv, output_lin_kv_dim, static_cast(&alpha),\n static_cast(input_weights_kv.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(k_lin_grads_ptr), CUDA_R_16F, output_lin_kv_dim, static_cast(&beta),\n static_cast(input_kv_grads.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F,\n // CUBLAS_GEMM_ALGO10_TENSOR_OP));\n CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Input Linear KV Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, output_lin_kv_dim, batches_kv,\n static_cast(&alpha), static_cast(inputs_kv.data_ptr()),\n CUDA_R_16F, embed_dim, static_cast(k_lin_grads_ptr), CUDA_R_16F,\n output_lin_kv_dim, static_cast(&beta),\n static_cast(input_weight_kv_grads.data_ptr()), CUDA_R_16F, embed_dim,\n CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));\n\n return {input_q_grads, input_kv_grads, input_weight_q_grads, input_weight_kv_grads, output_weight_grads};\n}\n\n} // end namespace cublas_gemmex\n} // end namespace encdec\n} // end namespace multihead_attn\n" }, { "path": "apex/contrib/csrc/multihead_attn/encdec_multihead_attn_norm_add_cuda.cu", "content": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"dropout.cuh\"\n#include \"layer_norm.cuh\"\n#include \"softmax.cuh\"\n#include \"strided_batched_gemm.cuh\"\n\nnamespace multihead_attn {\nnamespace encdec_norm_add {\nnamespace cublas_gemmex {\n\nstd::vector fwd_cuda(bool use_time_mask, bool is_training, int heads, torch::Tensor const& inputs_q,\n torch::Tensor const& inputs_kv, torch::Tensor const& lyr_nrm_gamma_weights,\n torch::Tensor const& lyr_nrm_beta_weights, torch::Tensor const& input_weights_q,\n torch::Tensor const& input_weights_kv, torch::Tensor const& output_weights,\n const uint8_t* pad_mask, float dropout_prob) {\n const int embed_dim = inputs_q.size(2);\n const int sequences = inputs_q.size(1);\n const int q_seq_len = inputs_q.size(0);\n const int k_seq_len = inputs_kv.size(0);\n const int batches_q = sequences * q_seq_len;\n const int batches_kv = sequences * k_seq_len;\n const int total_tokens_q = batches_q * embed_dim;\n const int head_dim = embed_dim / heads;\n const int output_lin_q_dim = embed_dim;\n const int output_lin_kv_dim = 2 * embed_dim;\n const int attn_batches = heads * sequences;\n const int lead_dim_q = attn_batches * head_dim;\n const int lead_dim_kv = attn_batches * 2 * head_dim;\n const int batch_stride_q = head_dim;\n const int batch_stride_kv = 2 * head_dim;\n const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n const float alpha = 1.0;\n const float beta = 0.0;\n const float scale = 1.0 / sqrt(static_cast(head_dim));\n\n // There is no reason to use more than one stream as every kernel is\n // sequentially dependent\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // 3 Intermediate Results + Output (Note: dropout intermediates are generated\n // by ATen library code)\n auto act_options = inputs_q.options().requires_grad(false);\n auto lyr_nrm_options = act_options.dtype(torch::kFloat32);\n auto mask_options = act_options.dtype(torch::kUInt8);\n\n torch::Tensor lyr_nrm_mean = torch::empty({batches_q}, lyr_nrm_options);\n torch::Tensor lyr_nrm_invvar = torch::empty({batches_q}, lyr_nrm_options);\n torch::Tensor lyr_nrm_results = torch::empty_like(inputs_q, act_options);\n\n torch::Tensor input_lin_q_results = torch::empty({q_seq_len, sequences, output_lin_q_dim}, act_options);\n torch::Tensor input_lin_kv_results = torch::empty({k_seq_len, sequences, output_lin_kv_dim}, act_options);\n torch::Tensor softmax_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_mask = torch::empty({attn_batches, q_seq_len, k_seq_len}, mask_options);\n torch::Tensor matmul2_results = torch::empty({q_seq_len, attn_batches, head_dim}, act_options);\n torch::Tensor output_lin_results = torch::empty_like(inputs_q, act_options);\n torch::Tensor dropout_add_mask = torch::empty_like(inputs_q, mask_options);\n torch::Tensor outputs = torch::empty_like(inputs_q, act_options);\n\n // Input Linear Results Pointers to Q, K, and V of interviewed activations\n void* q_lin_results_ptr = static_cast(input_lin_q_results.data_ptr());\n void* k_lin_results_ptr = static_cast(input_lin_kv_results.data_ptr());\n void* v_lin_results_ptr = static_cast(static_cast(input_lin_kv_results.data_ptr()) + head_dim);\n\n // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)\n void* softmax_results_ptr = static_cast(softmax_results.data_ptr());\n\n char a_layout_t{'t'};\n char a_layout_n{'n'};\n char b_layout_n{'n'};\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));\n // Layer Norm\n HostApplyLayerNorm(\n static_cast(lyr_nrm_results.data_ptr()), static_cast(lyr_nrm_mean.data_ptr()),\n static_cast(lyr_nrm_invvar.data_ptr()), static_cast(inputs_q.data_ptr()),\n static_cast(batches_q), // n1\n static_cast(embed_dim), // n2\n 1.0e-5, static_cast(lyr_nrm_gamma_weights.data_ptr()),\n static_cast(lyr_nrm_beta_weights.data_ptr()));\n\n // Input Linear Q Fwd\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, output_lin_q_dim, batches_q, embed_dim, static_cast(&alpha),\n static_cast(input_weights_q.data_ptr()), CUDA_R_16F, embed_dim,\n // static_cast(inputs_q.data_ptr()),\n static_cast(lyr_nrm_results.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta),\n q_lin_results_ptr, CUDA_R_16F, output_lin_q_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Input Linear KV Fwd\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, output_lin_kv_dim, batches_kv, embed_dim, static_cast(&alpha),\n static_cast(input_weights_kv.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(inputs_kv.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta),\n k_lin_results_ptr, CUDA_R_16F, output_lin_kv_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // MatMul1 of Dot-Product Attention Plus scaling by 1/Sqrt(head size)\n gemm_switch_fp32accum(a_layout_t, b_layout_n, k_seq_len, q_seq_len, head_dim, scale,\n static_cast(k_lin_results_ptr), lead_dim_kv, batch_stride_kv,\n static_cast(q_lin_results_ptr), lead_dim_q, batch_stride_q, beta,\n static_cast(softmax_results_ptr), k_seq_len, k_seq_len * q_seq_len, attn_batches);\n\n // Padded Softmax\n bool softmax_success = false;\n if (pad_mask == nullptr) {\n softmax_success = dispatch_softmax(reinterpret_cast(softmax_results_ptr),\n reinterpret_cast(softmax_results_ptr), k_seq_len,\n k_seq_len, attn_batches * q_seq_len);\n } else {\n if (use_time_mask) {\n softmax_success = dispatch_time_masked_softmax(\n reinterpret_cast(softmax_results_ptr), reinterpret_cast(softmax_results_ptr), pad_mask,\n k_seq_len, k_seq_len, attn_batches * q_seq_len, q_seq_len);\n } else {\n softmax_success = dispatch_masked_softmax(\n reinterpret_cast(softmax_results_ptr), reinterpret_cast(softmax_results_ptr), pad_mask,\n k_seq_len, k_seq_len, attn_batches * q_seq_len, attn_batches * q_seq_len / sequences);\n }\n }\n assert(softmax_success);\n\n if (is_training) {\n apex_fused_dropout_cuda(\n static_cast(softmax_results.data_ptr()), static_cast(dropout_results.data_ptr()),\n static_cast(dropout_mask.data_ptr()), dropout_elems, (1.0f - dropout_prob));\n }\n\n // Matmul2\n gemm_switch_fp32accum(a_layout_n, b_layout_n, head_dim, q_seq_len, k_seq_len, alpha,\n static_cast(v_lin_results_ptr), lead_dim_kv, batch_stride_kv,\n (is_training) ? static_cast(dropout_results.data_ptr())\n : static_cast(softmax_results.data_ptr()),\n // static_cast(dropout_results.data_ptr()),\n k_seq_len, k_seq_len * q_seq_len, beta, static_cast(matmul2_results.data_ptr()),\n head_dim * attn_batches, head_dim, attn_batches);\n\n // Output Linear\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, embed_dim, batches_q, embed_dim, static_cast(&alpha),\n static_cast(output_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(matmul2_results.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta),\n static_cast(output_lin_results.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F,\n // CUBLAS_GEMM_ALGO1_TENSOR_OP));\n CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // End-of-block Dropout-Add\n if (is_training) {\n apex_dropout_add_cuda(\n static_cast(output_lin_results.data_ptr()), static_cast(inputs_q.data_ptr()),\n static_cast(outputs.data_ptr()), static_cast(dropout_add_mask.data_ptr()), total_tokens_q,\n (1.0f - dropout_prob));\n } else {\n apex_add_cuda(static_cast(output_lin_results.data_ptr()),\n static_cast(inputs_q.data_ptr()),\n static_cast(outputs.data_ptr()), total_tokens_q);\n }\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));\n\n return {lyr_nrm_results, lyr_nrm_mean, lyr_nrm_invvar, input_lin_q_results,\n input_lin_kv_results, softmax_results, dropout_results, dropout_mask,\n matmul2_results, dropout_add_mask, outputs};\n}\n\nstd::vector bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_q_results, torch::Tensor const& input_lin_kv_results,\n torch::Tensor const& lyr_nrm_results, torch::Tensor const& lyr_nrm_mean,\n torch::Tensor const& lyr_nrm_invvar, torch::Tensor const& inputs_q,\n torch::Tensor const& inputs_kv, torch::Tensor const& lyr_nrm_gamma_weights,\n torch::Tensor const& lyr_nrm_beta_weights, torch::Tensor const& input_weights_q,\n torch::Tensor const& input_weights_kv, torch::Tensor const& output_weights,\n torch::Tensor const& dropout_mask, torch::Tensor const& dropout_add_mask,\n float dropout_prob) {\n const int embed_dim = inputs_q.size(2);\n const int sequences = inputs_q.size(1);\n const int q_seq_len = inputs_q.size(0);\n const int k_seq_len = inputs_kv.size(0);\n const int batches_q = sequences * q_seq_len;\n const int batches_kv = sequences * k_seq_len;\n const int total_tokens_q = batches_q * embed_dim;\n const int head_dim = embed_dim / heads;\n const int output_lin_q_dim = embed_dim;\n const int output_lin_kv_dim = 2 * embed_dim;\n const int attn_batches = heads * sequences;\n const int lead_dim_q = attn_batches * head_dim;\n const int lead_dim_kv = attn_batches * 2 * head_dim;\n const int batch_stride_q = head_dim;\n const int batch_stride_kv = 2 * head_dim;\n const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n const float alpha = 1.0;\n const float beta = 0.0;\n const float scale = 1.0 / sqrt(static_cast(head_dim));\n\n // TODO: Streams can be used in Backprop but I haven't added more than one\n // in my first attempt to create the code\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // Output Tensor Allocations\n torch::Tensor input_q_grads = torch::empty_like(inputs_q);\n torch::Tensor input_kv_grads = torch::empty_like(inputs_kv);\n torch::Tensor lyr_nrm_gamma_grads = torch::empty_like(lyr_nrm_gamma_weights);\n torch::Tensor lyr_nrm_beta_grads = torch::empty_like(lyr_nrm_beta_weights);\n torch::Tensor input_weight_q_grads = torch::empty_like(input_weights_q);\n torch::Tensor input_weight_kv_grads = torch::empty_like(input_weights_kv);\n torch::Tensor output_weight_grads = torch::empty_like(output_weights);\n // Intermediate Tensor Allocations\n at::Tensor dropout_add_grads = torch::empty_like(output_grads);\n at::Tensor output_lin_grads = torch::empty_like(matmul2_results);\n at::Tensor matmul2_grads = torch::empty_like(dropout_results);\n at::Tensor input_lin_q_output_grads = torch::empty_like(input_lin_q_results);\n at::Tensor input_lin_kv_output_grads = torch::empty_like(input_lin_kv_results);\n at::Tensor input_lin_q_grads = torch::empty_like(inputs_q);\n\n auto q_lin_results_ptr = static_cast(input_lin_q_results.data_ptr());\n auto k_lin_results_ptr = static_cast(input_lin_kv_results.data_ptr());\n auto v_lin_results_ptr = static_cast(input_lin_kv_results.data_ptr()) + head_dim;\n\n auto q_lin_grads_ptr = static_cast(input_lin_q_output_grads.data_ptr());\n auto k_lin_grads_ptr = static_cast(input_lin_kv_output_grads.data_ptr());\n auto v_lin_grads_ptr = static_cast(input_lin_kv_output_grads.data_ptr()) + head_dim;\n\n char a_layout_n{'n'};\n char a_layout_t{'t'};\n char b_layout_n{'n'};\n char b_layout_t{'t'};\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));\n\n // Dropout Add Backward\n apex_masked_scale_cuda(\n static_cast(output_grads.data_ptr()), static_cast(dropout_add_grads.data_ptr()),\n static_cast(dropout_add_mask.data_ptr()), total_tokens_q, (1.0 / (1.0 - dropout_prob)));\n\n // Output Linear Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches_q, embed_dim,\n static_cast(&alpha),\n static_cast(output_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(dropout_add_grads.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_lin_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Output Linear Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, embed_dim, batches_q,\n static_cast(&alpha),\n static_cast(matmul2_results.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(dropout_add_grads.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_weight_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // MatMul2 Dgrad1\n gemm_switch_fp32accum(a_layout_t, b_layout_n, k_seq_len, q_seq_len, head_dim, alpha,\n static_cast(v_lin_results_ptr), lead_dim_kv, batch_stride_kv,\n static_cast(output_lin_grads.data_ptr()), head_dim * attn_batches, head_dim, beta,\n static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len, attn_batches);\n\n // Matmul2 Dgrad2\n gemm_switch_fp32accum(a_layout_n, b_layout_t, head_dim, k_seq_len, q_seq_len, alpha,\n static_cast(output_lin_grads.data_ptr()), head_dim * attn_batches, head_dim,\n static_cast(dropout_results.data_ptr()), k_seq_len, k_seq_len * q_seq_len, beta,\n v_lin_grads_ptr, lead_dim_kv, batch_stride_kv, attn_batches);\n\n // Apply Dropout Mask and Scale by Dropout Probability\n apex_masked_scale_cuda(\n static_cast(matmul2_grads.data_ptr()), static_cast(matmul2_grads.data_ptr()),\n static_cast(dropout_mask.data_ptr()), dropout_elems, (1.0 / (1.0 - dropout_prob)));\n\n // Softmax Grad\n bool softmax_success = false;\n softmax_success = dispatch_softmax_backward(\n static_cast(matmul2_grads.data_ptr()), static_cast(matmul2_grads.data_ptr()),\n reinterpret_cast(softmax_results.data_ptr()), k_seq_len, k_seq_len, attn_batches * q_seq_len);\n assert(softmax_success);\n\n // Matmul1 Dgrad1\n gemm_switch_fp32accum(a_layout_n, b_layout_n, head_dim, q_seq_len, k_seq_len, scale, k_lin_results_ptr, lead_dim_kv,\n batch_stride_kv, static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta, q_lin_grads_ptr, lead_dim_q, batch_stride_q, attn_batches);\n\n // Matmul1 Dgrad2\n gemm_switch_fp32accum(a_layout_n, b_layout_t, head_dim, k_seq_len, q_seq_len, scale, q_lin_results_ptr, lead_dim_q,\n batch_stride_q, static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta, k_lin_grads_ptr, lead_dim_kv, batch_stride_kv, attn_batches);\n\n // Input Linear Q Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches_q, output_lin_q_dim, static_cast(&alpha),\n static_cast(input_weights_q.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(q_lin_grads_ptr), CUDA_R_16F, output_lin_q_dim, static_cast(&beta),\n // static_cast(input_q_grads.data_ptr()),\n static_cast(input_lin_q_grads.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F,\n // CUBLAS_GEMM_ALGO10_TENSOR_OP));\n CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Input Linear Q Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, output_lin_q_dim, batches_q,\n static_cast(&alpha), static_cast(inputs_q.data_ptr()),\n CUDA_R_16F, embed_dim, static_cast(q_lin_grads_ptr), CUDA_R_16F,\n output_lin_q_dim, static_cast(&beta),\n static_cast(input_weight_q_grads.data_ptr()), CUDA_R_16F, embed_dim,\n CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Input Linear KV Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches_kv, output_lin_kv_dim, static_cast(&alpha),\n static_cast(input_weights_kv.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(k_lin_grads_ptr), CUDA_R_16F, output_lin_kv_dim, static_cast(&beta),\n static_cast(input_kv_grads.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F,\n // CUBLAS_GEMM_ALGO10_TENSOR_OP));\n CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Input Linear KV Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, output_lin_kv_dim, batches_kv,\n static_cast(&alpha), static_cast(inputs_kv.data_ptr()),\n CUDA_R_16F, embed_dim, static_cast(k_lin_grads_ptr), CUDA_R_16F,\n output_lin_kv_dim, static_cast(&beta),\n static_cast(input_weight_kv_grads.data_ptr()), CUDA_R_16F, embed_dim,\n CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Fused Layer Norm Bwd with Residual Add\n HostLayerNormGradient(\n static_cast(input_lin_q_grads.data_ptr()), static_cast(output_grads.data_ptr()),\n static_cast(lyr_nrm_mean.data_ptr()), static_cast(lyr_nrm_invvar.data_ptr()),\n inputs_q,\n static_cast(batches_q), // n1\n static_cast(embed_dim), // n2\n static_cast(lyr_nrm_gamma_weights.data_ptr()),\n static_cast(lyr_nrm_beta_weights.data_ptr()), 1.0e-5, static_cast(input_q_grads.data_ptr()),\n static_cast(lyr_nrm_gamma_grads.data_ptr()), static_cast(lyr_nrm_beta_grads.data_ptr()));\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));\n\n return {input_q_grads, input_kv_grads, lyr_nrm_gamma_grads, lyr_nrm_beta_grads,\n input_weight_q_grads, input_weight_kv_grads, output_weight_grads};\n}\n\n} // end namespace cublas_gemmex\n} // end namespace encdec_norm_add\n} // end namespace multihead_attn\n" }, { "path": "apex/contrib/csrc/multihead_attn/layer_norm.cuh", "content": "#pragma once\n#include \n#include \n#include \n\n#include \n\nnamespace {\ntemplate \n__device__ void cuWelfordOnlineSum(const U curr, U& mu, U& sigma2, U& count) {\n count = count + U(1);\n U delta = curr - mu;\n U lmean = mu + delta / count;\n mu = lmean;\n U delta2 = curr - lmean;\n sigma2 = sigma2 + delta * delta2;\n}\n\ntemplate \n__device__ void cuChanOnlineSum(const U muB, const U sigma2B, const U countB, U& mu, U& sigma2, U& count) {\n U delta = muB - mu;\n U nA = count;\n U nB = countB;\n count = count + countB;\n U nX = count;\n if (nX > U(0)) {\n nA = nA / nX;\n nB = nB / nX;\n mu = nA * mu + nB * muB;\n sigma2 = sigma2 + sigma2B + delta * delta * nA * nB * nX;\n } else {\n mu = U(0);\n sigma2 = U(0);\n }\n}\n\ntemplate \n__device__ void cuWelfordMuSigma2(const T* __restrict__ vals, const int n1, const int n2, const int i1, U& mu,\n U& sigma2, U* buf) {\n // Assumptions:\n // 1) blockDim.x == warpSize\n // 2) Tensor is contiguous\n // 3) 2*blockDim.y*sizeof(U)+blockDim.y*sizeof(int) shared memory available.\n //\n // compute variance and mean over n2\n U count = U(0);\n mu = U(0);\n sigma2 = U(0);\n if (i1 < n1) {\n // one warp normalizes one n1 index,\n // synchronization is implicit\n // initialize with standard Welford algorithm\n const int numx = blockDim.x * blockDim.y;\n const int thrx = threadIdx.x + threadIdx.y * blockDim.x;\n const T* lvals = vals + i1 * n2;\n int l = 4 * thrx;\n for (; l + 3 < n2; l += 4 * numx) {\n for (int k = 0; k < 4; ++k) {\n U curr = static_cast(lvals[l + k]);\n cuWelfordOnlineSum(curr, mu, sigma2, count);\n }\n }\n for (; l < n2; ++l) {\n U curr = static_cast(lvals[l]);\n cuWelfordOnlineSum(curr, mu, sigma2, count);\n }\n // intra-warp reductions\n for (int l = 0; l <= 4; ++l) {\n int srcLaneB = (threadIdx.x + (1 << l)) & 31;\n U muB = WARP_SHFL(mu, srcLaneB);\n U countB = WARP_SHFL(count, srcLaneB);\n U sigma2B = WARP_SHFL(sigma2, srcLaneB);\n cuChanOnlineSum(muB, sigma2B, countB, mu, sigma2, count);\n }\n // threadIdx.x == 0 has correct values for each warp\n // inter-warp reductions\n if (blockDim.y > 1) {\n U* ubuf = (U*)buf;\n U* ibuf = (U*)(ubuf + blockDim.y);\n for (int offset = blockDim.y / 2; offset > 0; offset /= 2) {\n // upper half of warps write to shared\n if (threadIdx.x == 0 && threadIdx.y >= offset && threadIdx.y < 2 * offset) {\n const int wrt_y = threadIdx.y - offset;\n ubuf[2 * wrt_y] = mu;\n ubuf[2 * wrt_y + 1] = sigma2;\n ibuf[wrt_y] = count;\n }\n __syncthreads();\n // lower half merges\n if (threadIdx.x == 0 && threadIdx.y < offset) {\n U muB = ubuf[2 * threadIdx.y];\n U sigma2B = ubuf[2 * threadIdx.y + 1];\n U countB = ibuf[threadIdx.y];\n cuChanOnlineSum(muB, sigma2B, countB, mu, sigma2, count);\n }\n __syncthreads();\n }\n // threadIdx.x = 0 && threadIdx.y == 0 only thread that has correct values\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n ubuf[0] = mu;\n ubuf[1] = sigma2;\n }\n __syncthreads();\n mu = ubuf[0];\n sigma2 = ubuf[1] / U(n2);\n // don't care about final value of count, we know count == n2\n } else {\n mu = WARP_SHFL(mu, 0);\n sigma2 = WARP_SHFL(sigma2 / U(n2), 0);\n }\n }\n}\n\ntemplate <>\n__device__ void cuWelfordMuSigma2(const at::Half* __restrict__ vals, const int n1, const int n2, const int i1,\n float& mu, float& sigma2, float* buf) {\n // Assumptions:\n // 1) blockDim.x == warpSize\n // 2) Tensor is contiguous\n // 3) 2*blockDim.y*sizeof(U)+blockDim.y*sizeof(int) shared memory available.\n //\n // compute variance and mean over n2\n float count = 0.0f;\n mu = float(0);\n sigma2 = float(0);\n\n if (i1 < n1) {\n // one warp normalizes one n1 index,\n // synchronization is implicit\n // initialize with standard Welford algorithm\n const int numx = blockDim.x * blockDim.y;\n const int thrx = threadIdx.x + threadIdx.y * blockDim.x;\n const at::Half* lvals = vals + i1 * n2;\n int l = 8 * thrx;\n if ((((size_t)lvals) & 3) != 0) {\n // 16 bit alignment\n // first thread consumes first point\n if (thrx == 0) {\n float curr = static_cast(lvals[0]);\n cuWelfordOnlineSum(curr, mu, sigma2, count);\n }\n ++l;\n }\n // at this point, lvals[l] are 32 bit aligned for all threads.\n for (; l + 7 < n2; l += 8 * numx) {\n for (int k = 0; k < 8; k += 2) {\n float2 curr = __half22float2(*((__half2*)(lvals + l + k)));\n cuWelfordOnlineSum(curr.x, mu, sigma2, count);\n cuWelfordOnlineSum(curr.y, mu, sigma2, count);\n }\n }\n for (; l < n2; ++l) {\n float curr = static_cast(lvals[l]);\n cuWelfordOnlineSum(curr, mu, sigma2, count);\n }\n // intra-warp reductions\n for (int l = 0; l <= 4; ++l) {\n int srcLaneB = (threadIdx.x + (1 << l)) & 31;\n float muB = WARP_SHFL(mu, srcLaneB);\n float countB = WARP_SHFL(count, srcLaneB);\n float sigma2B = WARP_SHFL(sigma2, srcLaneB);\n cuChanOnlineSum(muB, sigma2B, countB, mu, sigma2, count);\n }\n // threadIdx.x == 0 has correct values for each warp\n // inter-warp reductions\n if (blockDim.y > 1) {\n float* ubuf = (float*)buf;\n float* ibuf = (float*)(ubuf + blockDim.y);\n for (int offset = blockDim.y / 2; offset > 0; offset /= 2) {\n // upper half of warps write to shared\n if (threadIdx.x == 0 && threadIdx.y >= offset && threadIdx.y < 2 * offset) {\n const int wrt_y = threadIdx.y - offset;\n ubuf[2 * wrt_y] = mu;\n ubuf[2 * wrt_y + 1] = sigma2;\n ibuf[wrt_y] = count;\n }\n __syncthreads();\n // lower half merges\n if (threadIdx.x == 0 && threadIdx.y < offset) {\n float muB = ubuf[2 * threadIdx.y];\n float sigma2B = ubuf[2 * threadIdx.y + 1];\n float countB = ibuf[threadIdx.y];\n cuChanOnlineSum(muB, sigma2B, countB, mu, sigma2, count);\n }\n __syncthreads();\n }\n // threadIdx.x = 0 && threadIdx.y == 0 only thread that has correct values\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n ubuf[0] = mu;\n ubuf[1] = sigma2;\n }\n __syncthreads();\n mu = ubuf[0];\n sigma2 = ubuf[1] / float(n2);\n // don't care about final value of count, we know count == n2\n } else {\n mu = WARP_SHFL(mu, 0);\n sigma2 = WARP_SHFL(sigma2 / float(n2), 0);\n }\n }\n}\n\ntemplate \n__device__ U rsqrt(U v) {\n return U(1) / sqrt(v);\n}\ntemplate <>\n__device__ float rsqrt(float v) {\n return rsqrtf(v);\n}\ntemplate <>\n__device__ double rsqrt(double v) {\n return rsqrt(v);\n}\n\n// This is the un-specialized struct. Note that we prevent instantiation of\n// this struct by putting an undefined symbol in the function body so it won't\n// compile.\n// template \n// struct SharedMemory\n// {\n// // Ensure that we won't compile any un-specialized types\n// __device__ T *getPointer()\n// {\n// extern __device__ void error(void);\n// error();\n// return NULL;\n// }\n// };\n// https://github.com/NVIDIA/apex/issues/246\ntemplate \nstruct SharedMemory;\ntemplate <>\nstruct SharedMemory {\n __device__ float* getPointer() {\n extern __shared__ float s_float[];\n return s_float;\n }\n};\n\ntemplate <>\nstruct SharedMemory {\n __device__ double* getPointer() {\n extern __shared__ double s_double[];\n return s_double;\n }\n};\n\ntemplate \n__global__ void cuApplyLayerNorm(T* __restrict__ output_vals, U* __restrict__ mean, U* __restrict__ invvar,\n const T* __restrict__ vals, const int n1, const int n2, const U epsilon,\n const T* __restrict__ gamma, const T* __restrict__ beta) {\n // Assumptions:\n // 1) blockDim.x == warpSize\n // 2) Tensors are contiguous\n //\n for (auto i1 = blockIdx.y; i1 < n1; i1 += gridDim.y) {\n SharedMemory shared;\n U* buf = shared.getPointer();\n U mu, sigma2;\n cuWelfordMuSigma2(vals, n1, n2, i1, mu, sigma2, buf);\n const T* lvals = vals + i1 * n2;\n T* ovals = output_vals + i1 * n2;\n U c_invvar = rsqrt(sigma2 + epsilon);\n const int numx = blockDim.x * blockDim.y;\n const int thrx = threadIdx.x + threadIdx.y * blockDim.x;\n if (gamma != NULL && beta != NULL) {\n for (int i = thrx; i < n2; i += numx) {\n U curr = static_cast(lvals[i]);\n ovals[i] = gamma[i] * static_cast(c_invvar * (curr - mu)) + beta[i];\n }\n } else {\n for (int i = thrx; i < n2; i += numx) {\n U curr = static_cast(lvals[i]);\n ovals[i] = static_cast(c_invvar * (curr - mu));\n }\n }\n if (threadIdx.x == 0 && threadIdx.y == 0) {\n mean[i1] = mu;\n invvar[i1] = c_invvar;\n }\n }\n}\n\ntemplate \n__device__ void cuLoadWriteStridedInputs(const int i1_block, const int thr_load_row_off, const int thr_load_col_off,\n const int i2_off, const int row_stride, U* warp_buf1, U* warp_buf2,\n const T* input, const T* dout, const int i1_end, const int n2,\n const U* __restrict__ mean, const U* __restrict__ invvar) {\n int i1 = i1_block + thr_load_row_off;\n if (i1 < i1_end) {\n U curr_mean = mean[i1];\n U curr_invvar = invvar[i1];\n for (int k = 0; k < blockDim.y; ++k) {\n int i2 = i2_off + k;\n int load_idx = i1 * n2 + i2;\n int write_idx = thr_load_row_off * row_stride + thr_load_col_off + k;\n if (i2 < n2) {\n U curr_input = static_cast(input[load_idx]);\n U curr_dout = static_cast(dout[load_idx]);\n warp_buf1[write_idx] = curr_dout;\n warp_buf2[write_idx] = curr_dout * (curr_input - curr_mean) * curr_invvar;\n } else {\n warp_buf1[write_idx] = U(0);\n warp_buf2[write_idx] = U(0);\n }\n }\n } else {\n for (int k = 0; k < blockDim.y; ++k) {\n int write_idx = thr_load_row_off * row_stride + thr_load_col_off + k;\n warp_buf1[write_idx] = U(0);\n warp_buf2[write_idx] = U(0);\n }\n }\n}\n\ntemplate \n__device__ void cuLoadAddStridedInputs(const int i1_block, const int thr_load_row_off, const int thr_load_col_off,\n const int i2_off, const int row_stride, U* warp_buf1, U* warp_buf2,\n const T* input, const T* dout, const int i1_end, const int n2,\n const U* __restrict__ mean, const U* __restrict__ invvar) {\n int i1 = i1_block + thr_load_row_off;\n if (i1 < i1_end) {\n U curr_mean = mean[i1];\n U curr_invvar = invvar[i1];\n for (int k = 0; k < blockDim.y; ++k) {\n int i2 = i2_off + k;\n int load_idx = i1 * n2 + i2;\n int write_idx = thr_load_row_off * row_stride + thr_load_col_off + k;\n if (i2 < n2) {\n U curr_input = static_cast(input[load_idx]);\n U curr_dout = static_cast(dout[load_idx]);\n warp_buf1[write_idx] += curr_dout;\n warp_buf2[write_idx] += curr_dout * (curr_input - curr_mean) * curr_invvar;\n }\n }\n }\n}\n\ntemplate \n__global__ void cuComputePartGradGammaBeta(const T* __restrict__ dout, const T* __restrict__ input, const int n1,\n const int n2, const U* __restrict__ mean, const U* __restrict__ invvar,\n U epsilon, U* part_grad_gamma, U* part_grad_beta) {\n const int numsegs_n1 = (n1 + blockDim.y * blockDim.y - 1) / (blockDim.y * blockDim.y);\n const int segs_per_block = (numsegs_n1 + gridDim.y - 1) / gridDim.y;\n const int i1_beg = blockIdx.y * segs_per_block * blockDim.y * blockDim.y;\n const int i1_beg_plus_one = (blockIdx.y + 1) * segs_per_block * blockDim.y * blockDim.y;\n const int i1_end = i1_beg_plus_one < n1 ? i1_beg_plus_one : n1;\n const int row_stride = blockDim.x + 1;\n const int thr_load_col_off = (threadIdx.x * blockDim.y) & (blockDim.x - 1);\n const int thr_load_row_off = (threadIdx.x * blockDim.y) / blockDim.x + threadIdx.y * blockDim.y;\n const int i2_off = blockIdx.x * blockDim.x + thr_load_col_off;\n SharedMemory shared;\n U* buf = shared.getPointer(); // buf has at least blockDim.x * blockDim.y *\n // blockDim.y + (blockDim.y -\n // 1)*(blockDim.x/blockDim.y) elements\n U* warp_buf1 = (U*)buf;\n U* warp_buf2 = warp_buf1 + blockDim.y * blockDim.y * row_stride;\n // compute partial sums from strided inputs\n // do this to increase number of loads in flight\n cuLoadWriteStridedInputs(i1_beg, thr_load_row_off, thr_load_col_off, i2_off, row_stride, warp_buf1, warp_buf2, input,\n dout, i1_end, n2, mean, invvar);\n for (int i1_block = i1_beg + blockDim.y * blockDim.y; i1_block < i1_end; i1_block += blockDim.y * blockDim.y) {\n cuLoadAddStridedInputs(i1_block, thr_load_row_off, thr_load_col_off, i2_off, row_stride, warp_buf1, warp_buf2,\n input, dout, i1_end, n2, mean, invvar);\n }\n __syncthreads();\n // inter-warp reductions\n // sum within each warp\n U acc1 = U(0);\n U acc2 = U(0);\n for (int k = 0; k < blockDim.y; ++k) {\n int row1 = threadIdx.y + k * blockDim.y;\n int idx1 = row1 * row_stride + threadIdx.x;\n acc1 += warp_buf1[idx1];\n acc2 += warp_buf2[idx1];\n }\n warp_buf1[threadIdx.y * row_stride + threadIdx.x] = acc1;\n warp_buf2[threadIdx.y * row_stride + threadIdx.x] = acc2;\n __syncthreads();\n // sum all warps\n for (int offset = blockDim.y / 2; offset > 1; offset /= 2) {\n if (threadIdx.y < offset) {\n int row1 = threadIdx.y;\n int row2 = threadIdx.y + offset;\n int idx1 = row1 * row_stride + threadIdx.x;\n int idx2 = row2 * row_stride + threadIdx.x;\n warp_buf1[idx1] += warp_buf1[idx2];\n warp_buf2[idx1] += warp_buf2[idx2];\n }\n __syncthreads();\n }\n int i2 = blockIdx.x * blockDim.x + threadIdx.x;\n if (threadIdx.y == 0 && i2 < n2) {\n int row1 = threadIdx.y;\n int row2 = threadIdx.y + 1;\n int idx1 = row1 * row_stride + threadIdx.x;\n int idx2 = row2 * row_stride + threadIdx.x;\n part_grad_beta[blockIdx.y * n2 + i2] = warp_buf1[idx1] + warp_buf1[idx2];\n part_grad_gamma[blockIdx.y * n2 + i2] = warp_buf2[idx1] + warp_buf2[idx2];\n }\n}\n\ntemplate \n__global__ void cuComputeGradGammaBeta(const U* part_grad_gamma, const U* part_grad_beta, const int part_size,\n const int n1, const int n2, T* grad_gamma, T* grad_beta) {\n // sum partial gradients for gamma and beta\n SharedMemory shared;\n U* buf = shared.getPointer();\n int i2 = blockIdx.x * blockDim.x + threadIdx.x;\n if (i2 < n2) {\n // each warp does sequential reductions until reduced part_size is num_warps\n int num_warp_reductions = part_size / blockDim.y;\n U sum_gamma = U(0);\n U sum_beta = U(0);\n const U* part_grad_gamma_ptr = part_grad_gamma + threadIdx.y * num_warp_reductions * n2 + i2;\n const U* part_grad_beta_ptr = part_grad_beta + threadIdx.y * num_warp_reductions * n2 + i2;\n for (int warp_offset = 0; warp_offset < num_warp_reductions; ++warp_offset) {\n sum_gamma += part_grad_gamma_ptr[warp_offset * n2];\n sum_beta += part_grad_beta_ptr[warp_offset * n2];\n }\n // inter-warp reductions\n const int nbsize3 = blockDim.x * blockDim.y / 2;\n for (int offset = blockDim.y / 2; offset >= 1; offset /= 2) {\n // top half write to shared memory\n if (threadIdx.y >= offset && threadIdx.y < 2 * offset) {\n const int write_idx = (threadIdx.y - offset) * blockDim.x + threadIdx.x;\n buf[write_idx] = sum_gamma;\n buf[write_idx + nbsize3] = sum_beta;\n }\n __syncthreads();\n // bottom half sums\n if (threadIdx.y < offset) {\n const int read_idx = threadIdx.y * blockDim.x + threadIdx.x;\n sum_gamma += buf[read_idx];\n sum_beta += buf[read_idx + nbsize3];\n }\n __syncthreads();\n }\n // write out fully summed gradients\n if (threadIdx.y == 0) {\n grad_gamma[i2] = sum_gamma;\n grad_beta[i2] = sum_beta;\n }\n }\n}\n\ntemplate \n__global__ void cuComputeGradInput(const T* __restrict__ dout, const T* __restrict__ dout_resid,\n const T* __restrict__ input, const int n1, const int n2, const U* __restrict__ mean,\n const U* __restrict__ invvar, U epsilon, const T* gamma, T* grad_input) {\n for (auto i1 = blockIdx.y; i1 < n1; i1 += gridDim.y) {\n U sum_loss1 = U(0);\n U sum_loss2 = U(0);\n const U c_mean = mean[i1];\n const U c_invvar = invvar[i1];\n const T* k_input = input + i1 * n2;\n const T* k_dout = dout + i1 * n2;\n const T* k_dout_resid = dout_resid + i1 * n2;\n const int numx = blockDim.x * blockDim.y;\n const int thrx = threadIdx.x + threadIdx.y * blockDim.x;\n if (gamma != NULL) {\n int l = 4 * thrx;\n for (; l + 3 < n2; l += 4 * numx) {\n for (int k = 0; k < 4; ++k) {\n const U c_h = static_cast(k_input[l + k]);\n const U c_loss = static_cast(k_dout[l + k]);\n sum_loss1 += c_loss * static_cast(gamma[l + k]);\n sum_loss2 += c_loss * static_cast(gamma[l + k]) * (c_h - c_mean) * c_invvar;\n }\n }\n for (; l < n2; ++l) {\n const U c_h = static_cast(k_input[l]);\n const U c_loss = static_cast(k_dout[l]);\n sum_loss1 += c_loss * static_cast(gamma[l]);\n sum_loss2 += c_loss * static_cast(gamma[l]) * (c_h - c_mean) * c_invvar;\n }\n } else {\n int l = 4 * thrx;\n for (; l + 3 < n2; l += 4 * numx) {\n for (int k = 0; k < 4; ++k) {\n const U c_h = static_cast(k_input[l + k]);\n const U c_loss = static_cast(k_dout[l + k]);\n sum_loss1 += c_loss;\n sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;\n }\n }\n for (; l < n2; ++l) {\n const U c_h = static_cast(k_input[l]);\n const U c_loss = static_cast(k_dout[l]);\n sum_loss1 += c_loss;\n sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;\n }\n }\n // intra-warp reductions\n for (int mask = blockDim.x / 2; mask > 0; mask /= 2) {\n sum_loss1 += WARP_SHFL_XOR(sum_loss1, mask);\n sum_loss2 += WARP_SHFL_XOR(sum_loss2, mask);\n }\n // inter-warp reductions\n if (blockDim.y > 1) {\n SharedMemory shared;\n U* buf = shared.getPointer();\n for (int offset = blockDim.y / 2; offset > 0; offset /= 2) {\n // upper half of warps write to shared\n if (threadIdx.y >= offset && threadIdx.y < 2 * offset) {\n const int wrt_i = (threadIdx.y - offset) * blockDim.x + threadIdx.x;\n buf[2 * wrt_i] = sum_loss1;\n buf[2 * wrt_i + 1] = sum_loss2;\n }\n __syncthreads();\n // lower half merges\n if (threadIdx.y < offset) {\n const int read_i = threadIdx.y * blockDim.x + threadIdx.x;\n sum_loss1 += buf[2 * read_i];\n sum_loss2 += buf[2 * read_i + 1];\n }\n __syncthreads();\n }\n if (threadIdx.y == 0) {\n buf[2 * threadIdx.x] = sum_loss1;\n buf[2 * threadIdx.x + 1] = sum_loss2;\n }\n __syncthreads();\n if (threadIdx.y != 0) {\n sum_loss1 = buf[2 * threadIdx.x];\n sum_loss2 = buf[2 * threadIdx.x + 1];\n }\n }\n // all threads now have the two sums over l\n U fH = (U)n2;\n U term1 = (U(1) / fH) * c_invvar;\n T* k_grad_input = grad_input + i1 * n2;\n if (gamma != NULL) {\n for (int l = thrx; l < n2; l += numx) {\n const U c_h = static_cast(k_input[l]);\n const U c_loss = static_cast(k_dout[l]);\n const T c_resid = static_cast(k_dout_resid[l]);\n U f_grad_input = fH * c_loss * static_cast(gamma[l]);\n f_grad_input -= sum_loss1;\n f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;\n f_grad_input *= term1;\n k_grad_input[l] = static_cast(f_grad_input) + c_resid;\n }\n } else {\n for (int l = thrx; l < n2; l += numx) {\n const U c_h = static_cast(k_input[l]);\n const U c_loss = static_cast(k_dout[l]);\n const T c_resid = static_cast(k_dout_resid[l]);\n U f_grad_input = fH * c_loss;\n f_grad_input -= sum_loss1;\n f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;\n f_grad_input *= term1;\n k_grad_input[l] = static_cast(f_grad_input) + c_resid;\n }\n }\n }\n}\n\ntemplate \nvoid HostApplyLayerNorm(T* output, U* mean, U* invvar, const T* input, int n1, int n2, double epsilon, const T* gamma,\n const T* beta) {\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n const dim3 threads(32, 4, 1);\n const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];\n const dim3 blocks(1, std::min((uint64_t)n1, maxGridY), 1);\n int nshared = threads.y > 1 ? threads.y * sizeof(U) + (threads.y / 2) * sizeof(U) : 0;\n cuApplyLayerNorm<<>>(output, mean, invvar, input, n1, n2, U(epsilon), gamma, beta);\n}\n\ntemplate \nvoid HostLayerNormGradient(const T* dout, const T* dout_resid, const U* mean, const U* invvar, const at::Tensor& input,\n int n1, int n2, const T* gamma, const T* beta, double epsilon, T* grad_input, T* grad_gamma,\n T* grad_beta) {\n auto stream = at::cuda::getCurrentCUDAStream().stream();\n\n if (gamma != NULL && beta != NULL) {\n // compute grad_gamma(j) and grad_beta(j)\n const int part_size = 16;\n const dim3 threads2(32, 4, 1);\n const dim3 blocks2((n2 + threads2.x - 1) / threads2.x, part_size, 1);\n const int nshared2_a = 2 * sizeof(U) * threads2.y * threads2.y * (threads2.x + 1);\n const int nshared2_b = threads2.x * threads2.y * sizeof(U);\n const int nshared2 = nshared2_a > nshared2_b ? nshared2_a : nshared2_b;\n at::Tensor part_grad_gamma = at::empty(\n {part_size, n2}, input.options().dtype(input.scalar_type() == at::ScalarType::Half ? at::ScalarType::Float\n : input.scalar_type()));\n at::Tensor part_grad_beta = at::empty_like(part_grad_gamma);\n cuComputePartGradGammaBeta<<>>(\n dout, static_cast(input.data_ptr()), n1, n2, mean, invvar, U(epsilon),\n static_cast(part_grad_gamma.data_ptr()), static_cast(part_grad_beta.data_ptr()));\n\n const dim3 threads3(32, 8, 1);\n const dim3 blocks3((n2 + threads2.x - 1) / threads2.x, 1, 1);\n const int nshared3 = threads3.x * threads3.y * sizeof(U);\n cuComputeGradGammaBeta<<>>(static_cast(part_grad_gamma.data_ptr()),\n static_cast(part_grad_beta.data_ptr()),\n part_size, n1, n2, grad_gamma, grad_beta);\n }\n\n // compute grad_input\n const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];\n const dim3 blocks1(1, std::min((uint64_t)n1, maxGridY), 1);\n const dim3 threads1(32, 4, 1);\n int nshared = threads1.y > 1 ? threads1.y * threads1.x * sizeof(U) : 0;\n cuComputeGradInput<<>>(dout, dout_resid, static_cast(input.data_ptr()), n1,\n n2, mean, invvar, U(epsilon), gamma, grad_input);\n}\n} // namespace\n" }, { "path": "apex/contrib/csrc/multihead_attn/masked_softmax_dropout_cuda.cu", "content": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"dropout.cuh\"\n#include \"softmax.cuh\"\n\nnamespace multihead_attn {\nnamespace fused_softmax {\nnamespace mask_softmax_dropout {\n\nstd::vector fwd_cuda(bool is_training, int heads, torch::Tensor const& input, const uint8_t* pad_mask,\n float dropout_prob) {\n const int attn_batches = input.size(0);\n const int sequences = attn_batches / heads;\n const int q_seq_len = input.size(1);\n const int k_seq_len = q_seq_len;\n const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n\n // There is no reason to use more than one stream as every kernel is\n // sequentially dependent\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // 3 Intermediate Results + Output (Note: dropout intermediates are generated\n // by ATen library code)\n auto act_options = input.options().requires_grad(false);\n auto mask_options = act_options.dtype(torch::kUInt8);\n\n torch::Tensor softmax_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_mask = torch::empty({attn_batches, q_seq_len, k_seq_len}, mask_options);\n\n // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)\n void* input_ptr = static_cast(input.data_ptr());\n void* softmax_results_ptr = static_cast(softmax_results.data_ptr());\n\n // Padded Softmax\n bool softmax_success = false;\n if (pad_mask == nullptr) {\n softmax_success = dispatch_softmax(reinterpret_cast(softmax_results_ptr),\n reinterpret_cast(input_ptr), k_seq_len,\n k_seq_len, attn_batches * q_seq_len);\n } else {\n softmax_success = dispatch_masked_softmax(\n reinterpret_cast(softmax_results_ptr), reinterpret_cast(input_ptr), pad_mask, k_seq_len,\n k_seq_len, attn_batches * q_seq_len, attn_batches * q_seq_len / sequences);\n }\n\n if (is_training) {\n // use at:: function so that C++ version generates the same random mask as\n // python version\n auto dropout_tuple = at::_fused_dropout(softmax_results, 1.0f - dropout_prob);\n dropout_results = std::get<0>(dropout_tuple);\n dropout_mask = std::get<1>(dropout_tuple);\n }\n\n // Matmul2\n\n return {dropout_results, dropout_mask, softmax_results};\n}\n\ntorch::Tensor bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& softmax_results,\n torch::Tensor const& dropout_mask, const uint8_t* padding_mask, float dropout_prob) {\n const int attn_batches = output_grads.size(0);\n const int q_seq_len = output_grads.size(1);\n const int k_seq_len = q_seq_len;\n const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n // TODO: Streams can be used in Backprop but I haven't added more than one\n // in my first attempt to create the code\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // Output Tensor Allocations\n // torch::Tensor input_grads = torch::empty_like(output_grads);\n\n // Apply Dropout Mask and Scale by Dropout Probability\n // Softmax Grad\n if (padding_mask == nullptr) {\n dispatch_masked_scale_softmax_backward_stream(\n static_cast(output_grads.data_ptr()), static_cast(output_grads.data_ptr()),\n reinterpret_cast(softmax_results.data_ptr()), static_cast(dropout_mask.data_ptr()),\n 1.0 / (1.0 - dropout_prob), k_seq_len, k_seq_len, attn_batches * q_seq_len, stream);\n } else {\n dispatch_masked_scale_softmax_backward_masked_out_stream(\n static_cast(output_grads.data_ptr()), static_cast(output_grads.data_ptr()),\n reinterpret_cast(softmax_results.data_ptr()), static_cast(dropout_mask.data_ptr()),\n static_cast(padding_mask), 1.0 / (1.0 - dropout_prob), k_seq_len, k_seq_len,\n attn_batches * q_seq_len, heads, stream);\n }\n // backward pass is completely in-place\n return output_grads;\n}\n} // namespace mask_softmax_dropout\n} // namespace fused_softmax\n} // namespace multihead_attn\n" }, { "path": "apex/contrib/csrc/multihead_attn/multihead_attn_frontend.cpp", "content": "#include \n#include \n\n#include \n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n#define CHECK_INPUT(x) \\\n CHECK_CUDA(x); \\\n CHECK_CONTIGUOUS(x)\n\nnamespace multihead_attn {\nnamespace fused_softmax {\nnamespace additive_mask_softmax_dropout {\n\nstd::vector fwd_cuda(bool is_training, int heads, torch::Tensor const& input, const half* pad_mask,\n float dropout_prob);\n\ntorch::Tensor bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& softmax_results,\n torch::Tensor const& dropout_mask, float dropout_prob);\n\nstd::vector fwd(bool use_mask, bool is_training, int heads, torch::Tensor const& input,\n torch::Tensor const& pad_mask, float dropout_prob) {\n TORCH_CHECK(input.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n if (use_mask) {\n TORCH_CHECK(pad_mask.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(pad_mask.scalar_type() == at::ScalarType::Half, \"Only BYTE is supported\");\n }\n\n return fwd_cuda(is_training, heads, input, use_mask ? static_cast(pad_mask.data_ptr()) : nullptr,\n dropout_prob);\n}\n\ntorch::Tensor bwd(bool use_mask, int heads, torch::Tensor const& output_grads, torch::Tensor const& softmax_results,\n torch::Tensor const& dropout_mask, float dropout_prob) {\n TORCH_CHECK(output_grads.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(softmax_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(dropout_mask.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(output_grads.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(softmax_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n // TORCH_CHECK(dropout_mask.scalar_type() == at::ScalarType::Byte,\n // \"Only BYTE is supported\");\n\n return bwd_cuda(heads, output_grads, softmax_results, dropout_mask, dropout_prob);\n}\n\n} // namespace additive_mask_softmax_dropout\nnamespace mask_softmax_dropout {\n\nstd::vector fwd_cuda(bool is_training, int heads, torch::Tensor const& input, const uint8_t* pad_mask,\n float dropout_prob);\n\ntorch::Tensor bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& softmax_results,\n torch::Tensor const& dropout_mask, const uint8_t* padding_mask, float dropout_prob);\n\nstd::vector fwd(bool use_mask, bool is_training, int heads, torch::Tensor const& input,\n torch::Tensor const& pad_mask, float dropout_prob) {\n TORCH_CHECK(input.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n\n if (use_mask) {\n TORCH_CHECK(pad_mask.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(pad_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n }\n\n return fwd_cuda(is_training, heads, input, use_mask ? static_cast(pad_mask.data_ptr()) : nullptr,\n dropout_prob);\n}\n\ntorch::Tensor bwd(bool use_mask, int heads, torch::Tensor const& output_grads, torch::Tensor const& softmax_results,\n torch::Tensor const& dropout_mask, torch::Tensor const& padding_mask, float dropout_prob) {\n TORCH_CHECK(output_grads.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(softmax_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(dropout_mask.dim() == 3, \"expected 3D tensor\");\n\n TORCH_CHECK(output_grads.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(softmax_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n // TORCH_CHECK(dropout_mask.scalar_type() == at::ScalarType::Byte,\n // \"Only BYTE is supported\");\n\n return bwd_cuda(heads, output_grads, softmax_results, dropout_mask,\n use_mask ? static_cast(padding_mask.data_ptr()) : nullptr, dropout_prob);\n}\n\n} // end namespace mask_softmax_dropout\n} // end namespace fused_softmax\n\nnamespace encdec {\nnamespace cublas_gemmex {\n\nstd::vector fwd_cuda(bool use_time_mask, bool is_training, int heads, torch::Tensor const& inputs_q,\n torch::Tensor const& inputs_kv, torch::Tensor const& input_weights_q,\n torch::Tensor const& input_weights_kv, torch::Tensor const& output_weights,\n const uint8_t* pad_mask, float dropout_prob);\nstd::vector bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_q_results, torch::Tensor const& input_lin_kv_results,\n torch::Tensor const& inputs_q, torch::Tensor const& inputs_kv,\n torch::Tensor const& input_weights_q, torch::Tensor const& input_weights_kv,\n torch::Tensor const& output_weights, torch::Tensor const& dropout_mask,\n float dropout_prob);\n\nstd::vector fwd(bool use_mask, bool use_time_mask, bool is_training, int heads,\n torch::Tensor const& inputs_q, torch::Tensor const& inputs_kv,\n torch::Tensor const& input_weights_q, torch::Tensor const& input_weights_kv,\n torch::Tensor const& output_weights, torch::Tensor const& pad_mask, float dropout_prob) {\n TORCH_CHECK(inputs_q.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(inputs_kv.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_weights_q.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(input_weights_kv.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(output_weights.dim() == 2, \"expected 2D tensor\");\n\n TORCH_CHECK(inputs_q.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(inputs_kv.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights_q.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights_kv.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(output_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n\n if (use_mask) {\n TORCH_CHECK(pad_mask.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(pad_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n }\n\n return fwd_cuda(use_time_mask, is_training, heads, inputs_q, inputs_kv, input_weights_q, input_weights_kv,\n output_weights, use_mask ? static_cast(pad_mask.data_ptr()) : nullptr, dropout_prob);\n}\n\nstd::vector bwd(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_q_results, torch::Tensor const& input_lin_kv_results,\n torch::Tensor const& inputs_q, torch::Tensor const& inputs_kv,\n torch::Tensor const& input_weights_q, torch::Tensor const& input_weights_kv,\n torch::Tensor const& output_weights, torch::Tensor const& dropout_mask,\n float dropout_prob) {\n TORCH_CHECK(output_grads.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(matmul2_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(dropout_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(softmax_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_lin_q_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_lin_kv_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(inputs_q.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(inputs_kv.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_weights_q.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(input_weights_kv.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(output_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(dropout_mask.dim() == 3, \"expected 3D tensor\");\n\n TORCH_CHECK(output_grads.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(matmul2_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(dropout_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(softmax_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_lin_q_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_lin_kv_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(inputs_q.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(inputs_kv.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights_q.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights_kv.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(output_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(dropout_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n\n return bwd_cuda(heads, output_grads, matmul2_results, dropout_results, softmax_results, input_lin_q_results,\n input_lin_kv_results, inputs_q, inputs_kv, input_weights_q, input_weights_kv, output_weights,\n dropout_mask, dropout_prob);\n}\n\n} // end namespace cublas_gemmex\n} // end namespace encdec\n\nnamespace encdec_norm_add {\nnamespace cublas_gemmex {\n\nstd::vector fwd_cuda(bool use_time_mask, bool is_training, int heads, torch::Tensor const& inputs_q,\n torch::Tensor const& inputs_kv, torch::Tensor const& lyr_nrm_gamma_weights,\n torch::Tensor const& lyr_nrm_beta_weights, torch::Tensor const& input_weights_q,\n torch::Tensor const& input_weights_kv, torch::Tensor const& output_weights,\n const uint8_t* pad_mask, float dropout_prob);\n\nstd::vector bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_q_results, torch::Tensor const& input_lin_kv_results,\n torch::Tensor const& lyr_nrm_results, torch::Tensor const& lyr_nrm_mean,\n torch::Tensor const& lyr_nrm_invvar, torch::Tensor const& inputs_q,\n torch::Tensor const& inputs_kv, torch::Tensor const& lyr_nrm_gamma_weights,\n torch::Tensor const& lyr_nrm_beta_weights, torch::Tensor const& input_weights_q,\n torch::Tensor const& input_weights_kv, torch::Tensor const& output_weights,\n torch::Tensor const& dropout_mask, torch::Tensor const& dropout_add_mask,\n float dropout_prob);\n\nstd::vector fwd(bool use_mask, bool use_time_mask, bool is_training, int heads,\n torch::Tensor const& inputs_q, torch::Tensor const& inputs_kv,\n torch::Tensor const& lyr_nrm_gamma_weights, torch::Tensor const& lyr_nrm_beta_weights,\n torch::Tensor const& input_weights_q, torch::Tensor const& input_weights_kv,\n torch::Tensor const& output_weights, torch::Tensor const& pad_mask, float dropout_prob) {\n TORCH_CHECK(inputs_q.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(inputs_kv.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(lyr_nrm_gamma_weights.dim() == 1, \"expected 1D tensor\");\n TORCH_CHECK(lyr_nrm_beta_weights.dim() == 1, \"expected 1D tensor\");\n TORCH_CHECK(input_weights_q.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(input_weights_kv.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(output_weights.dim() == 2, \"expected 2D tensor\");\n\n TORCH_CHECK(inputs_q.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(inputs_kv.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(lyr_nrm_gamma_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(lyr_nrm_beta_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights_q.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights_kv.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(output_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n\n if (use_mask) {\n TORCH_CHECK(pad_mask.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(pad_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n }\n\n return fwd_cuda(use_time_mask, is_training, heads, inputs_q, inputs_kv, lyr_nrm_gamma_weights, lyr_nrm_beta_weights,\n input_weights_q, input_weights_kv, output_weights,\n use_mask ? static_cast(pad_mask.data_ptr()) : nullptr, dropout_prob);\n}\n\nstd::vector bwd(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_q_results, torch::Tensor const& input_lin_kv_results,\n torch::Tensor const& lyr_nrm_results, torch::Tensor const& lyr_nrm_mean,\n torch::Tensor const& lyr_nrm_invvar, torch::Tensor const& inputs_q,\n torch::Tensor const& inputs_kv, torch::Tensor const& lyr_nrm_gamma_weights,\n torch::Tensor const& lyr_nrm_beta_weights, torch::Tensor const& input_weights_q,\n torch::Tensor const& input_weights_kv, torch::Tensor const& output_weights,\n torch::Tensor const& dropout_mask, torch::Tensor const& dropout_add_mask,\n float dropout_prob) {\n TORCH_CHECK(output_grads.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(matmul2_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(dropout_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(softmax_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_lin_q_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_lin_kv_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(lyr_nrm_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(lyr_nrm_mean.dim() == 1, \"expected 1D tensor\");\n TORCH_CHECK(lyr_nrm_invvar.dim() == 1, \"expected 1D tensor\");\n TORCH_CHECK(inputs_q.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(inputs_kv.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(lyr_nrm_gamma_weights.dim() == 1, \"expected 1D tensor\");\n TORCH_CHECK(lyr_nrm_beta_weights.dim() == 1, \"expected 1D tensor\");\n TORCH_CHECK(input_weights_q.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(input_weights_kv.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(output_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(dropout_mask.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(dropout_add_mask.dim() == 3, \"expected 3D tensor\");\n\n TORCH_CHECK(output_grads.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(matmul2_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(dropout_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(softmax_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_lin_q_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_lin_kv_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(lyr_nrm_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(lyr_nrm_mean.scalar_type() == at::ScalarType::Float, \"Only FLOAT is supported\");\n TORCH_CHECK(lyr_nrm_invvar.scalar_type() == at::ScalarType::Float, \"Only FLOAT is supported\");\n TORCH_CHECK(inputs_q.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(inputs_kv.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(lyr_nrm_gamma_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(lyr_nrm_beta_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights_q.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights_kv.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(output_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(dropout_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n TORCH_CHECK(dropout_add_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n\n return bwd_cuda(heads, output_grads, matmul2_results, dropout_results, softmax_results, input_lin_q_results,\n input_lin_kv_results, lyr_nrm_results, lyr_nrm_mean, lyr_nrm_invvar, inputs_q, inputs_kv,\n lyr_nrm_gamma_weights, lyr_nrm_beta_weights, input_weights_q, input_weights_kv, output_weights,\n dropout_mask, dropout_add_mask, dropout_prob);\n}\n\n} // end namespace cublas_gemmex\n} // end namespace encdec_norm_add\n\nnamespace self {\nnamespace cublas_gemmex {\n\nstd::vector fwd_cuda(bool use_time_mask, bool is_training, int heads, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n const uint8_t* pad_mask, float dropout_prob);\n\nstd::vector bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_results, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n torch::Tensor const& dropout_mask, float dropout_prob);\n\nstd::vector fwd(bool use_mask, bool use_time_mask, bool is_training, int heads,\n torch::Tensor const& inputs, torch::Tensor const& input_weights,\n torch::Tensor const& output_weights, torch::Tensor const& pad_mask, float dropout_prob) {\n TORCH_CHECK(inputs.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(output_weights.dim() == 2, \"expected 2D tensor\");\n\n TORCH_CHECK(inputs.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(output_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n\n if (use_mask) {\n TORCH_CHECK(pad_mask.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(pad_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n }\n\n return fwd_cuda(use_time_mask, is_training, heads, inputs, input_weights, output_weights,\n use_mask ? static_cast(pad_mask.data_ptr()) : nullptr, dropout_prob);\n}\n\nstd::vector bwd(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_results, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n torch::Tensor const& dropout_mask, float dropout_prob) {\n TORCH_CHECK(output_grads.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(matmul2_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(dropout_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(softmax_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_lin_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(inputs.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(output_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(dropout_mask.dim() == 3, \"expected 3D tensor\");\n\n TORCH_CHECK(output_grads.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(matmul2_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(dropout_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(softmax_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_lin_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(inputs.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(output_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(dropout_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n\n return bwd_cuda(heads, output_grads, matmul2_results, dropout_results, softmax_results, input_lin_results, inputs,\n input_weights, output_weights, dropout_mask, dropout_prob);\n}\n\n} // end namespace cublas_gemmex\n} // end namespace self\nnamespace self_bias {\nnamespace cublas_gemmex {\n\nstd::vector fwd_cuda(bool use_time_mask, bool is_training, int heads, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n torch::Tensor const& input_biases, torch::Tensor const& output_biases,\n const uint8_t* pad_mask, float dropout_prob);\n\nstd::vector bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_results, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n // torch::Tensor const& input_biases,\n // torch::Tensor const& output_biases,\n torch::Tensor const& dropout_mask, float dropout_prob);\n\nstd::vector fwd(bool use_mask, bool use_time_mask, bool is_training, int heads,\n torch::Tensor const& inputs, torch::Tensor const& input_weights,\n torch::Tensor const& output_weights, torch::Tensor const& input_biases,\n torch::Tensor const& output_biases, torch::Tensor const& pad_mask, float dropout_prob) {\n TORCH_CHECK(inputs.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(output_weights.dim() == 2, \"expected 2D tensor\");\n\n TORCH_CHECK(inputs.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(output_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n\n if (use_mask) {\n TORCH_CHECK(pad_mask.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(pad_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n }\n\n return fwd_cuda(use_time_mask, is_training, heads, inputs, input_weights, output_weights, input_biases, output_biases,\n use_mask ? static_cast(pad_mask.data_ptr()) : nullptr, dropout_prob);\n}\n\nstd::vector bwd(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_results, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n torch::Tensor const& dropout_mask, float dropout_prob) {\n TORCH_CHECK(output_grads.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(matmul2_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(dropout_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(softmax_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_lin_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(inputs.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(output_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(dropout_mask.dim() == 3, \"expected 3D tensor\");\n\n TORCH_CHECK(output_grads.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(matmul2_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(dropout_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(softmax_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_lin_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(inputs.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(output_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(dropout_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n\n return bwd_cuda(heads, output_grads, matmul2_results, dropout_results, softmax_results, input_lin_results, inputs,\n input_weights, output_weights, dropout_mask, dropout_prob);\n}\n\n} // end namespace cublas_gemmex\n} // namespace self_bias\nnamespace self_bias_additive_mask {\nnamespace cublas_gemmex {\n\nstd::vector fwd_cuda(bool use_time_mask, bool is_training, int heads, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n torch::Tensor const& input_biases, torch::Tensor const& output_biases,\n const half* pad_mask, float dropout_prob);\n\nstd::vector bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results,\n // torch::Tensor const& softmax_results,\n torch::Tensor const& bmm1_results, torch::Tensor const& pad_mask,\n torch::Tensor const& input_lin_results, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n // torch::Tensor const& input_biases,\n // torch::Tensor const& output_biases,\n torch::Tensor const& dropout_mask, float dropout_prob);\n\nstd::vector fwd(bool use_mask, bool use_time_mask, bool is_training, int heads,\n torch::Tensor const& inputs, torch::Tensor const& input_weights,\n torch::Tensor const& output_weights, torch::Tensor const& input_biases,\n torch::Tensor const& output_biases, torch::Tensor const& pad_mask, float dropout_prob) {\n TORCH_CHECK(inputs.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(output_weights.dim() == 2, \"expected 2D tensor\");\n\n TORCH_CHECK(inputs.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(output_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(use_mask, \"no mask is not supported\");\n\n if (use_mask) {\n TORCH_CHECK(pad_mask.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(pad_mask.scalar_type() == at::ScalarType::Half, \"Only Half is supported\");\n }\n\n return fwd_cuda(use_time_mask, is_training, heads, inputs, input_weights, output_weights, input_biases, output_biases,\n use_mask ? static_cast(pad_mask.data_ptr()) : nullptr, dropout_prob);\n}\n\nstd::vector bwd(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& bmm1_results,\n torch::Tensor const& pad_mask, torch::Tensor const& input_lin_results,\n torch::Tensor const& inputs, torch::Tensor const& input_weights,\n torch::Tensor const& output_weights, torch::Tensor const& dropout_mask,\n float dropout_prob) {\n TORCH_CHECK(output_grads.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(matmul2_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(dropout_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_lin_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(inputs.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(output_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(dropout_mask.dim() == 3, \"expected 3D tensor\");\n\n TORCH_CHECK(output_grads.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(matmul2_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(dropout_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_lin_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(inputs.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(output_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(dropout_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n\n return bwd_cuda(heads, output_grads, matmul2_results, dropout_results, bmm1_results, pad_mask, input_lin_results,\n inputs, input_weights, output_weights, dropout_mask, dropout_prob);\n}\n\n} // end namespace cublas_gemmex\n} // namespace self_bias_additive_mask\n\nnamespace self_norm_add {\nnamespace cublas_gemmex {\n\nstd::vector fwd_cuda(bool use_time_mask, bool is_training, int heads, torch::Tensor const& inputs,\n torch::Tensor const& lyr_nrm_gamma_weights,\n torch::Tensor const& lyr_nrm_beta_weights, torch::Tensor const& input_weights,\n torch::Tensor const& output_weights, const uint8_t* pad_mask, float dropout_prob);\n\nstd::vector bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_results, torch::Tensor const& lyr_nrm_results,\n torch::Tensor const& lyr_nrm_mean, torch::Tensor const& lyr_nrm_invvar,\n torch::Tensor const& inputs, torch::Tensor const& lyr_nrm_gamma_weights,\n torch::Tensor const& lyr_nrm_beta_weights, torch::Tensor const& input_weights,\n torch::Tensor const& output_weights, torch::Tensor const& dropout_mask,\n torch::Tensor const& dropout_add_mask, float dropout_prob);\n\nstd::vector fwd(bool use_mask, bool use_time_mask, bool is_training, int heads,\n torch::Tensor const& inputs, torch::Tensor const& lyr_nrm_gamma_weights,\n torch::Tensor const& lyr_nrm_beta_weights, torch::Tensor const& input_weights,\n torch::Tensor const& output_weights, torch::Tensor const& pad_mask, float dropout_prob) {\n TORCH_CHECK(inputs.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(lyr_nrm_gamma_weights.dim() == 1, \"expected 1D tensor\");\n TORCH_CHECK(lyr_nrm_beta_weights.dim() == 1, \"expected 1D tensor\");\n TORCH_CHECK(input_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(output_weights.dim() == 2, \"expected 2D tensor\");\n\n TORCH_CHECK(inputs.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(lyr_nrm_gamma_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(lyr_nrm_beta_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(output_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n\n if (use_mask) {\n TORCH_CHECK(pad_mask.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(pad_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n }\n\n return fwd_cuda(use_time_mask, is_training, heads, inputs, lyr_nrm_gamma_weights, lyr_nrm_beta_weights, input_weights,\n output_weights, use_mask ? static_cast(pad_mask.data_ptr()) : nullptr, dropout_prob);\n}\n\nstd::vector bwd(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_results, torch::Tensor const& lyr_nrm_results,\n torch::Tensor const& lyr_nrm_mean, torch::Tensor const& lyr_nrm_invvar,\n torch::Tensor const& inputs, torch::Tensor const& lyr_nrm_gamma_weights,\n torch::Tensor const& lyr_nrm_beta_weights, torch::Tensor const& input_weights,\n torch::Tensor const& output_weights, torch::Tensor const& dropout_mask,\n torch::Tensor const& dropout_add_mask, float dropout_prob) {\n TORCH_CHECK(output_grads.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(matmul2_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(dropout_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(softmax_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(input_lin_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(lyr_nrm_results.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(lyr_nrm_mean.dim() == 1, \"expected 1D tensor\");\n TORCH_CHECK(lyr_nrm_invvar.dim() == 1, \"expected 1D tensor\");\n TORCH_CHECK(inputs.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(lyr_nrm_gamma_weights.dim() == 1, \"expected 1D tensor\");\n TORCH_CHECK(lyr_nrm_beta_weights.dim() == 1, \"expected 1D tensor\");\n TORCH_CHECK(input_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(output_weights.dim() == 2, \"expected 2D tensor\");\n TORCH_CHECK(dropout_mask.dim() == 3, \"expected 3D tensor\");\n TORCH_CHECK(dropout_add_mask.dim() == 3, \"expected 3D tensor\");\n\n TORCH_CHECK(output_grads.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(matmul2_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(dropout_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(softmax_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_lin_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(lyr_nrm_results.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(lyr_nrm_mean.scalar_type() == at::ScalarType::Float, \"Only FLOAT is supported\");\n TORCH_CHECK(lyr_nrm_invvar.scalar_type() == at::ScalarType::Float, \"Only FLOAT is supported\");\n TORCH_CHECK(inputs.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(lyr_nrm_gamma_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(lyr_nrm_beta_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(input_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(output_weights.scalar_type() == at::ScalarType::Half, \"Only HALF is supported\");\n TORCH_CHECK(dropout_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n TORCH_CHECK(dropout_add_mask.scalar_type() == at::ScalarType::Byte, \"Only BYTE is supported\");\n\n return bwd_cuda(heads, output_grads, matmul2_results, dropout_results, softmax_results, input_lin_results,\n lyr_nrm_results, lyr_nrm_mean, lyr_nrm_invvar, inputs, lyr_nrm_gamma_weights, lyr_nrm_beta_weights,\n input_weights, output_weights, dropout_mask, dropout_add_mask, dropout_prob);\n}\n\n} // end namespace cublas_gemmex\n} // end namespace self_norm_add\n} // end namespace multihead_attn\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"additive_mask_softmax_dropout_forward\", &multihead_attn::fused_softmax::additive_mask_softmax_dropout::fwd,\n \"Self Multihead Attention masked softmax dropout -- Forward.\", py::call_guard());\n m.def(\"additive_mask_softmax_dropout_backward\", &multihead_attn::fused_softmax::additive_mask_softmax_dropout::bwd,\n \"Self Multihead Attention masked softmax dropout -- Backward.\", py::call_guard());\n m.def(\"mask_softmax_dropout_forward\", &multihead_attn::fused_softmax::mask_softmax_dropout::fwd,\n \"Self Multihead Attention masked softmax dropout -- Forward.\", py::call_guard());\n m.def(\"mask_softmax_dropout_backward\", &multihead_attn::fused_softmax::mask_softmax_dropout::bwd,\n \"Self Multihead Attention masked softmax dropout -- Backward.\", py::call_guard());\n m.def(\"encdec_multihead_attn_forward\", &multihead_attn::encdec::cublas_gemmex::fwd,\n \"Encdec Multihead Attention Forward.\", py::call_guard());\n m.def(\"encdec_multihead_attn_backward\", &multihead_attn::encdec::cublas_gemmex::bwd,\n \"Encdec Multihead Attention Backward.\", py::call_guard());\n m.def(\"encdec_multihead_attn_norm_add_forward\", &multihead_attn::encdec_norm_add::cublas_gemmex::fwd,\n \"Encdec Multihead Attention Plus Layer Norm and Residual Add Forward.\",\n py::call_guard());\n m.def(\"encdec_multihead_attn_norm_add_backward\", &multihead_attn::encdec_norm_add::cublas_gemmex::bwd,\n \"Encdec Multihead Attention Plus Layer Norm and Residual Add Backward.\",\n py::call_guard());\n m.def(\"self_attn_forward\", &multihead_attn::self::cublas_gemmex::fwd, \"Self Multihead Attention Forward.\",\n py::call_guard());\n m.def(\"self_attn_backward\", &multihead_attn::self::cublas_gemmex::bwd, \"Self Multihead Attention Backward.\",\n py::call_guard());\n m.def(\"self_attn_bias_forward\", &multihead_attn::self_bias::cublas_gemmex::fwd,\n \"Self Multihead Attention with Bias -- Forward.\", py::call_guard());\n m.def(\"self_attn_bias_backward\", &multihead_attn::self_bias::cublas_gemmex::bwd,\n \"Self Multihead Attention with Bias -- Backward.\", py::call_guard());\n m.def(\"self_attn_bias_additive_mask_forward\", &multihead_attn::self_bias_additive_mask::cublas_gemmex::fwd,\n \"Self Multihead Attention with Bias -- Forward.\", py::call_guard());\n m.def(\"self_attn_bias_additive_mask_backward\", &multihead_attn::self_bias_additive_mask::cublas_gemmex::bwd,\n \"Self Multihead Attention with Bias -- Backward.\", py::call_guard());\n m.def(\"self_attn_norm_add_forward\", &multihead_attn::self_norm_add::cublas_gemmex::fwd,\n \"Self Multihead Attention Plus Layer Norm and Residual Add Forward.\", py::call_guard());\n m.def(\"self_attn_norm_add_backward\", &multihead_attn::self_norm_add::cublas_gemmex::bwd,\n \"Self Multihead Attention Plus Layer Norm and Residual Add Backward.\",\n py::call_guard());\n}\n\n#undef CHECK_CUDA\n#undef CHECK_CONTIGUOUS\n#undef CHECK_INPUT\n" }, { "path": "apex/contrib/csrc/multihead_attn/philox.cuh", "content": "#pragma once\n// Philox CUDA.\n\nnamespace {\n\nclass Philox {\n public:\n __device__ inline Philox(unsigned long long seed, unsigned long long subsequence, unsigned long long offset)\n : STATE(0) {\n // key.x = (unsigned int)seed;\n // key.y = (unsigned int)(seed >> 32);\n // counter = make_uint4(0, 0, 0, 0);\n // counter.z = (unsigned int)(subsequence);\n // counter.w = (unsigned int)(subsequence >> 32);\n // STATE = 0;\n // incr_n(offset / 4);\n\n key = reinterpret_cast(seed);\n ull2* tmp = reinterpret_cast(&counter);\n tmp->x = offset / 4;\n tmp->y = subsequence;\n }\n __device__ inline uint4 operator()() {\n if (STATE == 0) {\n uint4 counter_ = counter;\n uint2 key_ = key;\n // 7-round philox\n for (int i = 0; i < 6; i++) {\n counter_ = single_round(counter_, key_);\n key_.x += (kPhilox10A);\n key_.y += (kPhilox10B);\n }\n output = single_round(counter_, key_);\n incr();\n }\n // return a float4 directly\n // unsigned long ret;\n // switch(STATE) {\n // case 0: ret = output.x; break;\n // case 1: ret = output.y; break;\n // case 2: ret = output.z; break;\n // case 3: ret = output.w; break;\n //}\n // STATE = (STATE + 1) % 4;\n return output;\n }\n\n private:\n struct ull2 {\n uint64_t x;\n uint64_t y;\n };\n uint4 counter;\n uint4 output;\n uint2 key;\n unsigned int STATE;\n __device__ inline void incr_n(unsigned long long n) {\n unsigned int nlo = (unsigned int)(n);\n unsigned int nhi = (unsigned int)(n >> 32);\n counter.x += nlo;\n if (counter.x < nlo) nhi++;\n counter.y += nhi;\n if (nhi <= counter.y) return;\n if (++counter.z) return;\n ++counter.w;\n }\n\n __device__ uint4 incr128(uint4 ctr) {\n uint4 res;\n asm(\"add.cc.u32 %0, %4, %8;\\n\\t\"\n \"addc.cc.u32 %1, %5, %9;\\n\\t\"\n \"addc.cc.u32 %2, %6, %10;\\n\\t\"\n \"addc.u32 %3, %7, %11;\\n\\t\"\n : \"=r\"(res.x), \"=r\"(res.y), \"=r\"(res.z), \"=r\"(res.w)\n : \"r\"(ctr.x), \"r\"(ctr.y), \"r\"(ctr.z), \"r\"(ctr.w), \"n\"(1), \"n\"(0), \"n\"(0), \"n\"(0));\n return res;\n }\n\n __device__ inline void incr() { counter = incr128(counter); }\n __device__ unsigned int mulhilo32(unsigned int a, unsigned int b, unsigned int* result_high) {\n *result_high = __umulhi(a, b);\n return a * b;\n }\n __device__ uint2 mulhilo32_v2(unsigned int a, unsigned int b) {\n uint2* res;\n unsigned long long tmp;\n asm(\"mul.wide.u32 %0, %1, %2;\\n\\t\" : \"=l\"(tmp) : \"r\"(a), \"r\"(b));\n res = (uint2*)(&tmp);\n return *res;\n }\n __device__ inline uint4 single_round(uint4 ctr, uint2 key) {\n // unsigned int hi0;\n // unsigned int hi1;\n // unsigned int lo0 = mulhilo32(kPhiloxSA, ctr.x, &hi0);\n // unsigned int lo1 = mulhilo32(kPhiloxSB, ctr.z, &hi1);\n // uint4 ret = {hi1 ^ ctr.y ^ key.x, lo1, hi0 ^ ctr.w ^ key.y, lo0};\n uint2 res0 = mulhilo32_v2(kPhiloxSA, ctr.x);\n uint2 res1 = mulhilo32_v2(kPhiloxSB, ctr.z);\n uint4 ret = {res1.y ^ ctr.y ^ key.x, res1.x, res0.y ^ ctr.w ^ key.y, res0.x};\n return ret;\n }\n static const unsigned long kPhilox10A = 0x9E3779B9;\n static const unsigned long kPhilox10B = 0xBB67AE85;\n static const unsigned long kPhiloxSA = 0xD2511F53;\n static const unsigned long kPhiloxSB = 0xCD9E8D57;\n};\n// Inverse of 2^32.\nconstexpr float M_RAN_INVM32 = 2.3283064e-10f;\n__device__ __inline__ float4 uniform4(uint4 x) {\n return make_float4(x.x * M_RAN_INVM32, x.y * M_RAN_INVM32, x.z * M_RAN_INVM32, x.w * M_RAN_INVM32);\n}\n\n} // namespace\n" }, { "path": "apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_additive_mask_cuda.cu", "content": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"dropout.cuh\"\n#include \"softmax.cuh\"\n#include \"strided_batched_gemm.cuh\"\n\nnamespace multihead_attn {\nnamespace self_bias_additive_mask {\nnamespace cublas_gemmex {\n\nstd::vector fwd_cuda(bool use_time_mask, bool is_training, int heads, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n torch::Tensor const& input_biases, torch::Tensor const& output_biases,\n const half* pad_mask, float dropout_prob) {\n const int embed_dim = inputs.size(2);\n const int sequences = inputs.size(1);\n const int q_seq_len = inputs.size(0);\n const int k_seq_len = q_seq_len;\n const int batches = sequences * q_seq_len;\n const int head_dim = embed_dim / heads;\n const int output_lin_dim = 3 * embed_dim;\n const int attn_batches = heads * sequences;\n const int lead_dim = attn_batches * 3 * head_dim;\n const int batch_stride = 3 * head_dim;\n [[maybe_unused]] const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n const float alpha = 1.0;\n const float beta_zero = 0.0;\n const float beta_one = 1.0;\n const float scale = 1.0 / sqrt(static_cast(head_dim));\n\n // There is no reason to use more than one stream as every kernel is\n // sequentially dependent\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // 3 Intermediate Results + Output (Note: dropout intermediates are generated\n // by ATen library code)\n auto act_options = inputs.options().requires_grad(false);\n auto mask_options = act_options.dtype(torch::kUInt8);\n\n torch::Tensor input_lin_results = torch::empty({q_seq_len, sequences, output_lin_dim}, act_options);\n torch::Tensor bmm1_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_mask = torch::empty({attn_batches, q_seq_len, k_seq_len}, mask_options);\n torch::Tensor matmul2_results = torch::empty({q_seq_len, attn_batches, head_dim}, act_options);\n torch::Tensor outputs = torch::empty_like(inputs, act_options);\n\n // Input Linear Results Pointers to Q, K, and V of interviewed activations\n void* q_lin_results_ptr = static_cast(input_lin_results.data_ptr());\n void* k_lin_results_ptr = static_cast(static_cast(input_lin_results.data_ptr()) + head_dim);\n void* v_lin_results_ptr = static_cast(static_cast(input_lin_results.data_ptr()) + 2 * head_dim);\n\n // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)\n void* bmm1_results_ptr = static_cast(bmm1_results.data_ptr());\n void* dropout_results_ptr = static_cast(dropout_results.data_ptr());\n\n char a_layout_t{'t'};\n char a_layout_n{'n'};\n char b_layout_n{'n'};\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));\n // Input Linear Fwd\n input_lin_results.copy_(input_biases);\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, output_lin_dim, batches, embed_dim, static_cast(&alpha),\n static_cast(input_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(inputs.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta_one),\n q_lin_results_ptr, CUDA_R_16F, output_lin_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // MatMul1 of Dot-Product Attention Plus scaling by 1/Sqrt(head size)\n gemm_switch_fp32accum(a_layout_t, b_layout_n, k_seq_len, q_seq_len, head_dim, scale,\n static_cast(k_lin_results_ptr), lead_dim, batch_stride,\n static_cast(q_lin_results_ptr), lead_dim, batch_stride, beta_zero,\n static_cast(bmm1_results_ptr), k_seq_len, k_seq_len * q_seq_len, attn_batches);\n // Padded Softmax\n [[maybe_unused]] bool softmax_success = false;\n if (is_training) {\n softmax_success = dispatch_additive_masked_softmax_dropout(\n reinterpret_cast(dropout_results_ptr),\n (is_training) ? reinterpret_cast(dropout_mask.data_ptr()) : nullptr,\n reinterpret_cast(bmm1_results_ptr), pad_mask, attn_batches * q_seq_len * q_seq_len, k_seq_len,\n k_seq_len, attn_batches * q_seq_len, attn_batches * q_seq_len / sequences, 1.0f - dropout_prob, stream);\n } else {\n softmax_success = dispatch_additive_masked_softmax(\n reinterpret_cast(dropout_results_ptr), // this is actually softmax results, but\n // making it consistent for the next function\n reinterpret_cast(bmm1_results_ptr), pad_mask, k_seq_len, k_seq_len, attn_batches * q_seq_len,\n attn_batches * q_seq_len / sequences);\n }\n\n // Matmul2\n gemm_switch_fp32accum(\n a_layout_n, b_layout_n, head_dim, q_seq_len, k_seq_len, alpha, static_cast(v_lin_results_ptr),\n lead_dim, batch_stride, static_cast(dropout_results.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta_zero, static_cast(matmul2_results.data_ptr()), head_dim * attn_batches, head_dim, attn_batches);\n\n outputs.copy_(output_biases);\n\n // Output Linear\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, embed_dim, batches, embed_dim, static_cast(&alpha),\n static_cast(output_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(matmul2_results.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta_one),\n static_cast(outputs.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F,\n // CUBLAS_GEMM_ALGO1_TENSOR_OP));\n CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));\n\n return {input_lin_results, bmm1_results, dropout_results, dropout_mask, matmul2_results, outputs};\n}\n\nstd::vector bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& bmm1_results,\n torch::Tensor const& pad_mask, torch::Tensor const& input_lin_results,\n torch::Tensor const& inputs, torch::Tensor const& input_weights,\n torch::Tensor const& output_weights, torch::Tensor const& dropout_mask,\n float dropout_prob) {\n const int embed_dim = inputs.size(2);\n const int sequences = inputs.size(1);\n const int q_seq_len = inputs.size(0);\n const int k_seq_len = q_seq_len;\n const int batches = sequences * q_seq_len;\n const int head_dim = embed_dim / heads;\n const int output_lin_dim = 3 * embed_dim;\n const int attn_batches = heads * sequences;\n const int lead_dim = attn_batches * 3 * head_dim;\n const int batch_stride = 3 * head_dim;\n [[maybe_unused]] const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n const float alpha = 1.0;\n const float beta = 0.0;\n const float scale = 1.0 / sqrt(static_cast(head_dim));\n\n // TODO: Streams can be used in Backprop but I haven't added more than one\n // in my first attempt to create the code\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // Output Tensor Allocations\n torch::Tensor input_grads = torch::empty_like(inputs);\n torch::Tensor input_weight_grads = torch::empty_like(input_weights);\n torch::Tensor output_weight_grads = torch::empty_like(output_weights);\n // Intermediate Tensor Allocations\n at::Tensor output_lin_grads = torch::empty_like(matmul2_results);\n at::Tensor matmul2_grads = torch::empty_like(dropout_results);\n at::Tensor input_lin_output_grads = torch::empty_like(input_lin_results);\n\n auto q_lin_results_ptr = static_cast(input_lin_results.data_ptr());\n auto k_lin_results_ptr = static_cast(input_lin_results.data_ptr()) + head_dim;\n auto v_lin_results_ptr = static_cast(input_lin_results.data_ptr()) + 2 * head_dim;\n\n auto q_lin_grads_ptr = static_cast(input_lin_output_grads.data_ptr());\n auto k_lin_grads_ptr = static_cast(input_lin_output_grads.data_ptr()) + head_dim;\n auto v_lin_grads_ptr = static_cast(input_lin_output_grads.data_ptr()) + 2 * head_dim;\n\n char a_layout_n{'n'};\n char a_layout_t{'t'};\n char b_layout_n{'n'};\n char b_layout_t{'t'};\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));\n\n // Output Linear Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches, embed_dim,\n static_cast(&alpha),\n static_cast(output_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(output_grads.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_lin_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n // Output Linear Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, embed_dim, batches,\n static_cast(&alpha),\n static_cast(matmul2_results.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(output_grads.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_weight_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n auto output_bias_grads = output_grads.view({-1, embed_dim}).sum(0, false);\n // MatMul2 Dgrad1\n gemm_switch_fp32accum(a_layout_t, b_layout_n, k_seq_len, q_seq_len, head_dim, alpha,\n static_cast(v_lin_results_ptr), lead_dim, batch_stride,\n static_cast(output_lin_grads.data_ptr()), head_dim * attn_batches, head_dim, beta,\n static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len, attn_batches);\n\n // Matmul2 Dgrad2\n gemm_switch_fp32accum(a_layout_n, b_layout_t, head_dim, k_seq_len, q_seq_len, alpha,\n static_cast(output_lin_grads.data_ptr()), head_dim * attn_batches, head_dim,\n static_cast(dropout_results.data_ptr()), k_seq_len, k_seq_len * q_seq_len, beta,\n v_lin_grads_ptr, lead_dim, batch_stride, attn_batches);\n\n // Apply Dropout Mask and Scale by Dropout Probability\n // Softmax Grad\n dispatch_masked_scale_softmax_backward_recompute(\n static_cast(matmul2_grads.data_ptr()), static_cast(matmul2_grads.data_ptr()),\n reinterpret_cast(bmm1_results.data_ptr()), reinterpret_cast(pad_mask.data_ptr()),\n static_cast(dropout_mask.data_ptr()), 1.0 / (1.0 - dropout_prob), k_seq_len, k_seq_len,\n attn_batches * q_seq_len / sequences, attn_batches * q_seq_len, stream);\n\n // Matmul1 Dgrad1\n gemm_switch_fp32accum(a_layout_n, b_layout_n, head_dim, q_seq_len, k_seq_len, scale, k_lin_results_ptr, lead_dim,\n batch_stride, static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta, q_lin_grads_ptr, lead_dim, batch_stride, attn_batches);\n\n // Matmul1 Dgrad2\n gemm_switch_fp32accum(a_layout_n, b_layout_t, head_dim, k_seq_len, q_seq_len, scale, q_lin_results_ptr, lead_dim,\n batch_stride, static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta, k_lin_grads_ptr, lead_dim, batch_stride, attn_batches);\n // Input Linear Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches, output_lin_dim,\n static_cast(&alpha),\n static_cast(input_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(input_lin_output_grads.data_ptr()),\n // static_cast(q_lin_grads_ptr),\n CUDA_R_16F, output_lin_dim, static_cast(&beta),\n static_cast(input_grads.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F,\n // CUBLAS_GEMM_ALGO10_TENSOR_OP));\n CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Input Linear Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, output_lin_dim, batches, static_cast(&alpha),\n static_cast(inputs.data_ptr()), CUDA_R_16F, embed_dim, static_cast(q_lin_grads_ptr),\n CUDA_R_16F, output_lin_dim, static_cast(&beta), static_cast(input_weight_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n auto input_bias_grads = input_lin_output_grads.view({-1, output_lin_dim}).sum(0, false);\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));\n\n return {input_grads, input_weight_grads, output_weight_grads, input_bias_grads, output_bias_grads};\n}\n\n} // end namespace cublas_gemmex\n} // namespace self_bias_additive_mask\n} // end namespace multihead_attn\n" }, { "path": "apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_cuda.cu", "content": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"dropout.cuh\"\n#include \"softmax.cuh\"\n#include \"strided_batched_gemm.cuh\"\n\nnamespace multihead_attn {\nnamespace self_bias {\nnamespace cublas_gemmex {\n\nstd::vector fwd_cuda(bool use_time_mask, bool is_training, int heads, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n torch::Tensor const& input_biases, torch::Tensor const& output_biases,\n const uint8_t* pad_mask, float dropout_prob) {\n const int embed_dim = inputs.size(2);\n const int sequences = inputs.size(1);\n const int q_seq_len = inputs.size(0);\n const int k_seq_len = q_seq_len;\n const int batches = sequences * q_seq_len;\n const int head_dim = embed_dim / heads;\n const int output_lin_dim = 3 * embed_dim;\n const int attn_batches = heads * sequences;\n const int lead_dim = attn_batches * 3 * head_dim;\n const int batch_stride = 3 * head_dim;\n [[maybe_unused]] const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n const float alpha = 1.0;\n const float beta_zero = 0.0;\n const float beta_one = 1.0;\n const float scale = 1.0 / sqrt(static_cast(head_dim));\n\n // There is no reason to use more than one stream as every kernel is\n // sequentially dependent\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // 3 Intermediate Results + Output (Note: dropout intermediates are generated\n // by ATen library code)\n auto act_options = inputs.options().requires_grad(false);\n auto mask_options = act_options.dtype(torch::kUInt8);\n\n torch::Tensor input_lin_results = torch::empty({q_seq_len, sequences, output_lin_dim}, act_options);\n torch::Tensor softmax_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_mask = torch::empty({attn_batches, q_seq_len, k_seq_len}, mask_options);\n torch::Tensor matmul2_results = torch::empty({q_seq_len, attn_batches, head_dim}, act_options);\n torch::Tensor outputs = torch::empty_like(inputs, act_options);\n\n // Input Linear Results Pointers to Q, K, and V of interviewed activations\n void* q_lin_results_ptr = static_cast(input_lin_results.data_ptr());\n void* k_lin_results_ptr = static_cast(static_cast(input_lin_results.data_ptr()) + head_dim);\n void* v_lin_results_ptr = static_cast(static_cast(input_lin_results.data_ptr()) + 2 * head_dim);\n\n // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)\n void* softmax_results_ptr = static_cast(softmax_results.data_ptr());\n\n char a_layout_t{'t'};\n char a_layout_n{'n'};\n char b_layout_n{'n'};\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));\n // Input Linear Fwd\n input_lin_results.copy_(input_biases);\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, output_lin_dim, batches, embed_dim, static_cast(&alpha),\n static_cast(input_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(inputs.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta_one),\n q_lin_results_ptr, CUDA_R_16F, output_lin_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // MatMul1 of Dot-Product Attention Plus scaling by 1/Sqrt(head size)\n gemm_switch_fp32accum(a_layout_t, b_layout_n, k_seq_len, q_seq_len, head_dim, scale,\n static_cast(k_lin_results_ptr), lead_dim, batch_stride,\n static_cast(q_lin_results_ptr), lead_dim, batch_stride, beta_zero,\n static_cast(softmax_results_ptr), k_seq_len, k_seq_len * q_seq_len, attn_batches);\n // Padded Softmax\n [[maybe_unused]] bool softmax_success = false;\n if (pad_mask == nullptr) {\n softmax_success = dispatch_softmax(reinterpret_cast(softmax_results_ptr),\n reinterpret_cast(softmax_results_ptr), k_seq_len,\n k_seq_len, attn_batches * q_seq_len);\n } else {\n if (use_time_mask) {\n softmax_success = dispatch_time_masked_softmax(\n reinterpret_cast(softmax_results_ptr), reinterpret_cast(softmax_results_ptr), pad_mask,\n k_seq_len, k_seq_len, attn_batches * q_seq_len, q_seq_len);\n } else {\n softmax_success = dispatch_masked_softmax(\n reinterpret_cast(softmax_results_ptr), reinterpret_cast(softmax_results_ptr), pad_mask,\n k_seq_len, k_seq_len, attn_batches * q_seq_len, attn_batches * q_seq_len / sequences);\n }\n }\n\n if (is_training) {\n // use at:: function so that C++ version generates the same random mask as\n // python version\n auto dropout_tuple = at::_fused_dropout(softmax_results, 1.0f - dropout_prob);\n dropout_results = std::get<0>(dropout_tuple);\n dropout_mask = std::get<1>(dropout_tuple);\n }\n\n // Matmul2\n gemm_switch_fp32accum(a_layout_n, b_layout_n, head_dim, q_seq_len, k_seq_len, alpha,\n static_cast(v_lin_results_ptr), lead_dim, batch_stride,\n (is_training) ? static_cast(dropout_results.data_ptr())\n : static_cast(softmax_results.data_ptr()),\n k_seq_len, k_seq_len * q_seq_len, beta_zero, static_cast(matmul2_results.data_ptr()),\n head_dim * attn_batches, head_dim, attn_batches);\n\n outputs.copy_(output_biases);\n\n // Output Linear\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, embed_dim, batches, embed_dim, static_cast(&alpha),\n static_cast(output_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(matmul2_results.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta_one),\n static_cast(outputs.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F,\n // CUBLAS_GEMM_ALGO1_TENSOR_OP));\n CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));\n\n return {input_lin_results, softmax_results, dropout_results, dropout_mask, matmul2_results, outputs};\n}\n\nstd::vector bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_results, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n torch::Tensor const& dropout_mask, float dropout_prob) {\n const int embed_dim = inputs.size(2);\n const int sequences = inputs.size(1);\n const int q_seq_len = inputs.size(0);\n const int k_seq_len = q_seq_len;\n const int batches = sequences * q_seq_len;\n const int head_dim = embed_dim / heads;\n const int output_lin_dim = 3 * embed_dim;\n const int attn_batches = heads * sequences;\n const int lead_dim = attn_batches * 3 * head_dim;\n const int batch_stride = 3 * head_dim;\n [[maybe_unused]] const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n const float alpha = 1.0;\n const float beta = 0.0;\n const float scale = 1.0 / sqrt(static_cast(head_dim));\n\n // TODO: Streams can be used in Backprop but I haven't added more than one\n // in my first attempt to create the code\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // Output Tensor Allocations\n torch::Tensor input_grads = torch::empty_like(inputs);\n torch::Tensor input_weight_grads = torch::empty_like(input_weights);\n torch::Tensor output_weight_grads = torch::empty_like(output_weights);\n // Intermediate Tensor Allocations\n at::Tensor output_lin_grads = torch::empty_like(matmul2_results);\n at::Tensor matmul2_grads = torch::empty_like(dropout_results);\n at::Tensor input_lin_output_grads = torch::empty_like(input_lin_results);\n\n auto q_lin_results_ptr = static_cast(input_lin_results.data_ptr());\n auto k_lin_results_ptr = static_cast(input_lin_results.data_ptr()) + head_dim;\n auto v_lin_results_ptr = static_cast(input_lin_results.data_ptr()) + 2 * head_dim;\n\n auto q_lin_grads_ptr = static_cast(input_lin_output_grads.data_ptr());\n auto k_lin_grads_ptr = static_cast(input_lin_output_grads.data_ptr()) + head_dim;\n auto v_lin_grads_ptr = static_cast(input_lin_output_grads.data_ptr()) + 2 * head_dim;\n\n char a_layout_n{'n'};\n char a_layout_t{'t'};\n char b_layout_n{'n'};\n char b_layout_t{'t'};\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));\n\n // Output Linear Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches, embed_dim,\n static_cast(&alpha),\n static_cast(output_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(output_grads.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_lin_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n // Output Linear Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, embed_dim, batches,\n static_cast(&alpha),\n static_cast(matmul2_results.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(output_grads.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_weight_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n auto output_bias_grads = output_grads.view({-1, embed_dim}).sum(0, false);\n // MatMul2 Dgrad1\n gemm_switch_fp32accum(a_layout_t, b_layout_n, k_seq_len, q_seq_len, head_dim, alpha,\n static_cast(v_lin_results_ptr), lead_dim, batch_stride,\n static_cast(output_lin_grads.data_ptr()), head_dim * attn_batches, head_dim, beta,\n static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len, attn_batches);\n\n // Matmul2 Dgrad2\n gemm_switch_fp32accum(a_layout_n, b_layout_t, head_dim, k_seq_len, q_seq_len, alpha,\n static_cast(output_lin_grads.data_ptr()), head_dim * attn_batches, head_dim,\n static_cast(dropout_results.data_ptr()), k_seq_len, k_seq_len * q_seq_len, beta,\n v_lin_grads_ptr, lead_dim, batch_stride, attn_batches);\n\n // Apply Dropout Mask and Scale by Dropout Probability\n // Softmax Grad\n dispatch_masked_scale_softmax_backward_stream(\n static_cast(matmul2_grads.data_ptr()), static_cast(matmul2_grads.data_ptr()),\n reinterpret_cast(softmax_results.data_ptr()), static_cast(dropout_mask.data_ptr()),\n 1.0 / (1.0 - dropout_prob), k_seq_len, k_seq_len, attn_batches * q_seq_len, stream);\n\n // Matmul1 Dgrad1\n gemm_switch_fp32accum(a_layout_n, b_layout_n, head_dim, q_seq_len, k_seq_len, scale, k_lin_results_ptr, lead_dim,\n batch_stride, static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta, q_lin_grads_ptr, lead_dim, batch_stride, attn_batches);\n\n // Matmul1 Dgrad2\n gemm_switch_fp32accum(a_layout_n, b_layout_t, head_dim, k_seq_len, q_seq_len, scale, q_lin_results_ptr, lead_dim,\n batch_stride, static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta, k_lin_grads_ptr, lead_dim, batch_stride, attn_batches);\n // Input Linear Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches, output_lin_dim,\n static_cast(&alpha),\n static_cast(input_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(input_lin_output_grads.data_ptr()),\n // static_cast(q_lin_grads_ptr),\n CUDA_R_16F, output_lin_dim, static_cast(&beta),\n static_cast(input_grads.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F,\n // CUBLAS_GEMM_ALGO10_TENSOR_OP));\n CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Input Linear Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, output_lin_dim, batches, static_cast(&alpha),\n static_cast(inputs.data_ptr()), CUDA_R_16F, embed_dim, static_cast(q_lin_grads_ptr),\n CUDA_R_16F, output_lin_dim, static_cast(&beta), static_cast(input_weight_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n auto input_bias_grads = input_lin_output_grads.view({-1, output_lin_dim}).sum(0, false);\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));\n\n return {input_grads, input_weight_grads, output_weight_grads, input_bias_grads, output_bias_grads};\n}\n\n} // end namespace cublas_gemmex\n} // namespace self_bias\n} // end namespace multihead_attn\n" }, { "path": "apex/contrib/csrc/multihead_attn/self_multihead_attn_cuda.cu", "content": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"dropout.cuh\"\n#include \"softmax.cuh\"\n#include \"strided_batched_gemm.cuh\"\n\nnamespace multihead_attn {\nnamespace self {\nnamespace cublas_gemmex {\n\nstd::vector fwd_cuda(bool use_time_mask, bool is_training, int heads, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n const uint8_t* pad_mask, float dropout_prob) {\n const int embed_dim = inputs.size(2);\n const int sequences = inputs.size(1);\n const int q_seq_len = inputs.size(0);\n const int k_seq_len = q_seq_len;\n const int batches = sequences * q_seq_len;\n const int head_dim = embed_dim / heads;\n const int output_lin_dim = 3 * embed_dim;\n const int attn_batches = heads * sequences;\n const int lead_dim = attn_batches * 3 * head_dim;\n const int batch_stride = 3 * head_dim;\n const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n const float alpha = 1.0;\n const float beta = 0.0;\n const float scale = 1.0 / sqrt(static_cast(head_dim));\n\n // There is no reason to use more than one stream as every kernel is\n // sequentially dependent\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // 3 Intermediate Results + Output (Note: dropout intermediates are generated\n // by ATen library code)\n auto act_options = inputs.options().requires_grad(false);\n auto mask_options = act_options.dtype(torch::kUInt8);\n\n torch::Tensor input_lin_results = torch::empty({q_seq_len, sequences, output_lin_dim}, act_options);\n torch::Tensor softmax_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_mask = torch::empty({attn_batches, q_seq_len, k_seq_len}, mask_options);\n torch::Tensor matmul2_results = torch::empty({q_seq_len, attn_batches, head_dim}, act_options);\n torch::Tensor outputs = torch::empty_like(inputs, act_options);\n\n // Input Linear Results Pointers to Q, K, and V of interviewed activations\n void* q_lin_results_ptr = static_cast(input_lin_results.data_ptr());\n void* k_lin_results_ptr = static_cast(static_cast(input_lin_results.data_ptr()) + head_dim);\n void* v_lin_results_ptr = static_cast(static_cast(input_lin_results.data_ptr()) + 2 * head_dim);\n\n // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)\n void* softmax_results_ptr = static_cast(softmax_results.data_ptr());\n\n char a_layout_t{'t'};\n char a_layout_n{'n'};\n char b_layout_n{'n'};\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));\n // Input Linear Fwd\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, output_lin_dim, batches, embed_dim, static_cast(&alpha),\n static_cast(input_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(inputs.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta),\n q_lin_results_ptr, CUDA_R_16F, output_lin_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // MatMul1 of Dot-Product Attention Plus scaling by 1/Sqrt(head size)\n gemm_switch_fp32accum(a_layout_t, b_layout_n, k_seq_len, q_seq_len, head_dim, scale,\n static_cast(k_lin_results_ptr), lead_dim, batch_stride,\n static_cast(q_lin_results_ptr), lead_dim, batch_stride, beta,\n static_cast(softmax_results_ptr), k_seq_len, k_seq_len * q_seq_len, attn_batches);\n\n // Padded Softmax\n bool softmax_success = false;\n if (pad_mask == nullptr) {\n softmax_success = dispatch_softmax(reinterpret_cast(softmax_results_ptr),\n reinterpret_cast(softmax_results_ptr), k_seq_len,\n k_seq_len, attn_batches * q_seq_len);\n } else {\n if (use_time_mask) {\n softmax_success = dispatch_time_masked_softmax(\n reinterpret_cast(softmax_results_ptr), reinterpret_cast(softmax_results_ptr), pad_mask,\n k_seq_len, k_seq_len, attn_batches * q_seq_len, q_seq_len);\n } else {\n softmax_success = dispatch_masked_softmax(\n reinterpret_cast(softmax_results_ptr), reinterpret_cast(softmax_results_ptr), pad_mask,\n k_seq_len, k_seq_len, attn_batches * q_seq_len, attn_batches * q_seq_len / sequences);\n }\n }\n assert(softmax_success);\n\n if (is_training) {\n apex_fused_dropout_cuda(\n static_cast(softmax_results.data_ptr()), static_cast(dropout_results.data_ptr()),\n static_cast(dropout_mask.data_ptr()), dropout_elems, (1.0f - dropout_prob));\n }\n\n // Matmul2\n gemm_switch_fp32accum(a_layout_n, b_layout_n, head_dim, q_seq_len, k_seq_len, alpha,\n static_cast(v_lin_results_ptr), lead_dim, batch_stride,\n (is_training) ? static_cast(dropout_results.data_ptr())\n : static_cast(softmax_results.data_ptr()),\n k_seq_len, k_seq_len * q_seq_len, beta, static_cast(matmul2_results.data_ptr()),\n head_dim * attn_batches, head_dim, attn_batches);\n\n // Output Linear\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, embed_dim, batches, embed_dim, static_cast(&alpha),\n static_cast(output_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(matmul2_results.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta),\n static_cast(outputs.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));\n\n return {input_lin_results, softmax_results, dropout_results, dropout_mask, matmul2_results, outputs};\n}\n\nstd::vector bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_results, torch::Tensor const& inputs,\n torch::Tensor const& input_weights, torch::Tensor const& output_weights,\n torch::Tensor const& dropout_mask, float dropout_prob) {\n const int embed_dim = inputs.size(2);\n const int sequences = inputs.size(1);\n const int q_seq_len = inputs.size(0);\n const int k_seq_len = q_seq_len;\n const int batches = sequences * q_seq_len;\n const int head_dim = embed_dim / heads;\n const int output_lin_dim = 3 * embed_dim;\n const int attn_batches = heads * sequences;\n const int lead_dim = attn_batches * 3 * head_dim;\n const int batch_stride = 3 * head_dim;\n const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n const float alpha = 1.0;\n const float beta = 0.0;\n const float scale = 1.0 / sqrt(static_cast(head_dim));\n\n // TODO: Streams can be used in Backprop but I haven't added more than one\n // in my first attempt to create the code\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // Output Tensor Allocations\n torch::Tensor input_grads = torch::empty_like(inputs);\n torch::Tensor input_weight_grads = torch::empty_like(input_weights);\n torch::Tensor output_weight_grads = torch::empty_like(output_weights);\n // Intermediate Tensor Allocations\n at::Tensor output_lin_grads = torch::empty_like(matmul2_results);\n at::Tensor matmul2_grads = torch::empty_like(dropout_results);\n at::Tensor input_lin_output_grads = torch::empty_like(input_lin_results);\n\n auto q_lin_results_ptr = static_cast(input_lin_results.data_ptr());\n auto k_lin_results_ptr = static_cast(input_lin_results.data_ptr()) + head_dim;\n auto v_lin_results_ptr = static_cast(input_lin_results.data_ptr()) + 2 * head_dim;\n\n auto q_lin_grads_ptr = static_cast(input_lin_output_grads.data_ptr());\n auto k_lin_grads_ptr = static_cast(input_lin_output_grads.data_ptr()) + head_dim;\n auto v_lin_grads_ptr = static_cast(input_lin_output_grads.data_ptr()) + 2 * head_dim;\n\n char a_layout_n{'n'};\n char a_layout_t{'t'};\n char b_layout_n{'n'};\n char b_layout_t{'t'};\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));\n\n // Output Linear Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches, embed_dim,\n static_cast(&alpha),\n static_cast(output_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(output_grads.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_lin_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Output Linear Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, embed_dim, batches,\n static_cast(&alpha),\n static_cast(matmul2_results.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(output_grads.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_weight_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // MatMul2 Dgrad1\n gemm_switch_fp32accum(a_layout_t, b_layout_n, k_seq_len, q_seq_len, head_dim, alpha,\n static_cast(v_lin_results_ptr), lead_dim, batch_stride,\n static_cast(output_lin_grads.data_ptr()), head_dim * attn_batches, head_dim, beta,\n static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len, attn_batches);\n\n // Matmul2 Dgrad2\n gemm_switch_fp32accum(a_layout_n, b_layout_t, head_dim, k_seq_len, q_seq_len, alpha,\n static_cast(output_lin_grads.data_ptr()), head_dim * attn_batches, head_dim,\n static_cast(dropout_results.data_ptr()), k_seq_len, k_seq_len * q_seq_len, beta,\n v_lin_grads_ptr, lead_dim, batch_stride, attn_batches);\n\n // Apply Dropout Mask and Scale by Dropout Probability\n apex_masked_scale_cuda(\n static_cast(matmul2_grads.data_ptr()), static_cast(matmul2_grads.data_ptr()),\n static_cast(dropout_mask.data_ptr()), dropout_elems, (1.0 / (1.0 - dropout_prob)));\n\n // Softmax Grad\n bool softmax_success = false;\n softmax_success = dispatch_softmax_backward(\n static_cast(matmul2_grads.data_ptr()), static_cast(matmul2_grads.data_ptr()),\n reinterpret_cast(softmax_results.data_ptr()), k_seq_len, k_seq_len, attn_batches * q_seq_len);\n assert(softmax_success);\n\n // Matmul1 Dgrad1\n gemm_switch_fp32accum(a_layout_n, b_layout_n, head_dim, q_seq_len, k_seq_len, scale, k_lin_results_ptr, lead_dim,\n batch_stride, static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta, q_lin_grads_ptr, lead_dim, batch_stride, attn_batches);\n\n // Matmul1 Dgrad2\n gemm_switch_fp32accum(a_layout_n, b_layout_t, head_dim, k_seq_len, q_seq_len, scale, q_lin_results_ptr, lead_dim,\n batch_stride, static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta, k_lin_grads_ptr, lead_dim, batch_stride, attn_batches);\n\n // Input Linear Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches, output_lin_dim, static_cast(&alpha),\n static_cast(input_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(q_lin_grads_ptr), CUDA_R_16F, output_lin_dim, static_cast(&beta),\n static_cast(input_grads.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Input Linear Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, output_lin_dim, batches, static_cast(&alpha),\n static_cast(inputs.data_ptr()), CUDA_R_16F, embed_dim, static_cast(q_lin_grads_ptr),\n CUDA_R_16F, output_lin_dim, static_cast(&beta), static_cast(input_weight_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));\n\n return {input_grads, input_weight_grads, output_weight_grads};\n}\n\n} // end namespace cublas_gemmex\n} // end namespace self\n} // end namespace multihead_attn\n" }, { "path": "apex/contrib/csrc/multihead_attn/self_multihead_attn_norm_add_cuda.cu", "content": "#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \"dropout.cuh\"\n#include \"layer_norm.cuh\"\n#include \"softmax.cuh\"\n#include \"strided_batched_gemm.cuh\"\n\nnamespace multihead_attn {\nnamespace self_norm_add {\nnamespace cublas_gemmex {\n\nstd::vector fwd_cuda(bool use_time_mask, bool is_training, int heads, torch::Tensor const& inputs,\n torch::Tensor const& lyr_nrm_gamma_weights,\n torch::Tensor const& lyr_nrm_beta_weights, torch::Tensor const& input_weights,\n torch::Tensor const& output_weights, const uint8_t* pad_mask, float dropout_prob) {\n const int embed_dim = inputs.size(2);\n const int sequences = inputs.size(1);\n const int q_seq_len = inputs.size(0);\n const int k_seq_len = q_seq_len;\n const int batches = sequences * q_seq_len;\n const int total_tokens = batches * embed_dim;\n const int head_dim = embed_dim / heads;\n const int output_lin_dim = 3 * embed_dim;\n const int attn_batches = heads * sequences;\n const int lead_dim = attn_batches * 3 * head_dim;\n const int batch_stride = 3 * head_dim;\n const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n const float alpha = 1.0;\n const float beta = 0.0;\n const float scale = 1.0 / sqrt(static_cast(head_dim));\n\n // There is no reason to use more than one stream as every kernel is\n // sequentially dependent\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // 3 Intermediate Results + Output (Note: dropout intermediates are generated\n // by ATen library code)\n auto act_options = inputs.options().requires_grad(false);\n auto lyr_nrm_options = act_options.dtype(torch::kFloat32);\n auto mask_options = act_options.dtype(torch::kUInt8);\n\n torch::Tensor lyr_nrm_mean = torch::empty({batches}, lyr_nrm_options);\n torch::Tensor lyr_nrm_invvar = torch::empty({batches}, lyr_nrm_options);\n torch::Tensor lyr_nrm_results = torch::empty_like(inputs, act_options);\n\n torch::Tensor input_lin_results = torch::empty({q_seq_len, sequences, output_lin_dim}, act_options);\n torch::Tensor softmax_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_results = torch::empty({attn_batches, q_seq_len, k_seq_len}, act_options);\n torch::Tensor dropout_mask = torch::empty({attn_batches, q_seq_len, k_seq_len}, mask_options);\n torch::Tensor matmul2_results = torch::empty({q_seq_len, attn_batches, head_dim}, act_options);\n torch::Tensor output_lin_results = torch::empty_like(inputs, act_options);\n torch::Tensor dropout_add_mask = torch::empty_like(inputs, mask_options);\n torch::Tensor outputs = torch::empty_like(inputs, act_options);\n\n // Input Linear Results Pointers to Q, K, and V of interviewed activations\n void* q_lin_results_ptr = static_cast(input_lin_results.data_ptr());\n void* k_lin_results_ptr = static_cast(static_cast(input_lin_results.data_ptr()) + head_dim);\n void* v_lin_results_ptr = static_cast(static_cast(input_lin_results.data_ptr()) + 2 * head_dim);\n\n // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)\n void* softmax_results_ptr = static_cast(softmax_results.data_ptr());\n\n char a_layout_t{'t'};\n char a_layout_n{'n'};\n char b_layout_n{'n'};\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));\n // Layer Norm\n HostApplyLayerNorm(\n static_cast(lyr_nrm_results.data_ptr()), static_cast(lyr_nrm_mean.data_ptr()),\n static_cast(lyr_nrm_invvar.data_ptr()), static_cast(inputs.data_ptr()),\n static_cast(batches), // n1\n static_cast(embed_dim), // n2\n 1.0e-5, static_cast(lyr_nrm_gamma_weights.data_ptr()),\n static_cast(lyr_nrm_beta_weights.data_ptr()));\n\n // Input Linear Fwd\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_T, CUBLAS_OP_N, output_lin_dim, batches, embed_dim, static_cast(&alpha),\n static_cast(input_weights.data_ptr()), CUDA_R_16F, embed_dim,\n // static_cast(inputs.data_ptr()),\n static_cast(lyr_nrm_results.data_ptr()), CUDA_R_16F, embed_dim, static_cast(&beta),\n q_lin_results_ptr, CUDA_R_16F, output_lin_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // MatMul1 of Dot-Product Attention Plus scaling by 1/Sqrt(head size)\n gemm_switch_fp32accum(a_layout_t, b_layout_n, k_seq_len, q_seq_len, head_dim, scale,\n static_cast(k_lin_results_ptr), lead_dim, batch_stride,\n static_cast(q_lin_results_ptr), lead_dim, batch_stride, beta,\n static_cast(softmax_results_ptr), k_seq_len, k_seq_len * q_seq_len, attn_batches);\n\n // Padded Softmax\n bool softmax_success = false;\n if (pad_mask == nullptr) {\n softmax_success = dispatch_softmax(reinterpret_cast(softmax_results_ptr),\n reinterpret_cast(softmax_results_ptr), k_seq_len,\n k_seq_len, attn_batches * q_seq_len);\n } else {\n if (use_time_mask) {\n softmax_success = dispatch_time_masked_softmax(\n reinterpret_cast(softmax_results_ptr), reinterpret_cast(softmax_results_ptr), pad_mask,\n k_seq_len, k_seq_len, attn_batches * q_seq_len, q_seq_len);\n } else {\n softmax_success = dispatch_masked_softmax(\n reinterpret_cast(softmax_results_ptr), reinterpret_cast(softmax_results_ptr), pad_mask,\n k_seq_len, k_seq_len, attn_batches * q_seq_len, attn_batches * q_seq_len / sequences);\n }\n }\n assert(softmax_success);\n\n if (is_training) {\n apex_fused_dropout_cuda(\n static_cast(softmax_results.data_ptr()), static_cast(dropout_results.data_ptr()),\n static_cast(dropout_mask.data_ptr()), dropout_elems, (1.0f - dropout_prob));\n }\n\n // Matmul2\n gemm_switch_fp32accum(a_layout_n, b_layout_n, head_dim, q_seq_len, k_seq_len, alpha,\n static_cast(v_lin_results_ptr), lead_dim, batch_stride,\n (is_training) ? static_cast(dropout_results.data_ptr())\n : static_cast(softmax_results.data_ptr()),\n // static_cast(dropout_results.data_ptr()),\n k_seq_len, k_seq_len * q_seq_len, beta, static_cast(matmul2_results.data_ptr()),\n head_dim * attn_batches, head_dim, attn_batches);\n\n // Output Linear\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, embed_dim, batches, embed_dim,\n static_cast(&alpha),\n static_cast(output_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(matmul2_results.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_lin_results.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // End-of-block Dropout-Add\n if (is_training) {\n apex_dropout_add_cuda(\n static_cast(output_lin_results.data_ptr()), static_cast(inputs.data_ptr()),\n static_cast(outputs.data_ptr()), static_cast(dropout_add_mask.data_ptr()), total_tokens,\n (1.0f - dropout_prob));\n } else {\n apex_add_cuda(static_cast(output_lin_results.data_ptr()),\n static_cast(inputs.data_ptr()),\n static_cast(outputs.data_ptr()), total_tokens);\n }\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));\n\n return {lyr_nrm_results, lyr_nrm_mean, lyr_nrm_invvar, input_lin_results, softmax_results,\n dropout_results, dropout_mask, matmul2_results, dropout_add_mask, outputs};\n}\n\nstd::vector bwd_cuda(int heads, torch::Tensor const& output_grads, torch::Tensor const& matmul2_results,\n torch::Tensor const& dropout_results, torch::Tensor const& softmax_results,\n torch::Tensor const& input_lin_results, torch::Tensor const& lyr_nrm_results,\n torch::Tensor const& lyr_nrm_mean, torch::Tensor const& lyr_nrm_invvar,\n torch::Tensor const& inputs, torch::Tensor const& lyr_nrm_gamma_weights,\n torch::Tensor const& lyr_nrm_beta_weights, torch::Tensor const& input_weights,\n torch::Tensor const& output_weights, torch::Tensor const& dropout_mask,\n torch::Tensor const& dropout_add_mask, float dropout_prob) {\n const int embed_dim = inputs.size(2);\n const int sequences = inputs.size(1);\n const int q_seq_len = inputs.size(0);\n const int k_seq_len = q_seq_len;\n const int batches = sequences * q_seq_len;\n const int total_tokens = batches * embed_dim;\n const int head_dim = embed_dim / heads;\n const int output_lin_dim = 3 * embed_dim;\n const int attn_batches = heads * sequences;\n const int lead_dim = attn_batches * 3 * head_dim;\n const int batch_stride = 3 * head_dim;\n const int dropout_elems = attn_batches * q_seq_len * k_seq_len;\n const float alpha = 1.0;\n const float beta = 0.0;\n const float scale = 1.0 / sqrt(static_cast(head_dim));\n\n // TODO: Streams can be used in Backprop but I haven't added more than one\n // in my first attempt to create the code\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n\n // Output Tensor Allocations\n torch::Tensor input_grads = torch::empty_like(inputs);\n torch::Tensor lyr_nrm_gamma_grads = torch::empty_like(lyr_nrm_gamma_weights);\n torch::Tensor lyr_nrm_beta_grads = torch::empty_like(lyr_nrm_beta_weights);\n torch::Tensor input_weight_grads = torch::empty_like(input_weights);\n torch::Tensor output_weight_grads = torch::empty_like(output_weights);\n // Intermediate Tensor Allocations\n torch::Tensor dropout_add_grads = torch::empty_like(output_grads);\n torch::Tensor output_lin_grads = torch::empty_like(matmul2_results);\n torch::Tensor matmul2_grads = torch::empty_like(dropout_results);\n torch::Tensor input_lin_output_grads = torch::empty_like(input_lin_results);\n torch::Tensor input_lin_grads = torch::empty_like(inputs);\n\n auto q_lin_results_ptr = static_cast(input_lin_results.data_ptr());\n auto k_lin_results_ptr = static_cast(input_lin_results.data_ptr()) + head_dim;\n auto v_lin_results_ptr = static_cast(input_lin_results.data_ptr()) + 2 * head_dim;\n\n auto q_lin_grads_ptr = static_cast(input_lin_output_grads.data_ptr());\n auto k_lin_grads_ptr = static_cast(input_lin_output_grads.data_ptr()) + head_dim;\n auto v_lin_grads_ptr = static_cast(input_lin_output_grads.data_ptr()) + 2 * head_dim;\n\n char a_layout_n{'n'};\n char a_layout_t{'t'};\n char b_layout_n{'n'};\n char b_layout_t{'t'};\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));\n\n // Dropout Add Backward\n apex_masked_scale_cuda(\n static_cast(output_grads.data_ptr()), static_cast(dropout_add_grads.data_ptr()),\n static_cast(dropout_add_mask.data_ptr()), total_tokens, (1.0 / (1.0 - dropout_prob)));\n\n // Output Linear Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches, embed_dim,\n static_cast(&alpha),\n static_cast(output_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(dropout_add_grads.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_lin_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Output Linear Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, embed_dim, batches,\n static_cast(&alpha),\n static_cast(matmul2_results.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(dropout_add_grads.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(&beta), static_cast(output_weight_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // MatMul2 Dgrad1\n gemm_switch_fp32accum(a_layout_t, b_layout_n, k_seq_len, q_seq_len, head_dim, alpha,\n static_cast(v_lin_results_ptr), lead_dim, batch_stride,\n static_cast(output_lin_grads.data_ptr()), head_dim * attn_batches, head_dim, beta,\n static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len, attn_batches);\n\n // Matmul2 Dgrad2\n gemm_switch_fp32accum(a_layout_n, b_layout_t, head_dim, k_seq_len, q_seq_len, alpha,\n static_cast(output_lin_grads.data_ptr()), head_dim * attn_batches, head_dim,\n static_cast(dropout_results.data_ptr()), k_seq_len, k_seq_len * q_seq_len, beta,\n v_lin_grads_ptr, lead_dim, batch_stride, attn_batches);\n\n // Apply Dropout Mask and Scale by Dropout Probability\n apex_masked_scale_cuda(\n static_cast(matmul2_grads.data_ptr()), static_cast(matmul2_grads.data_ptr()),\n static_cast(dropout_mask.data_ptr()), dropout_elems, (1.0 / (1.0 - dropout_prob)));\n\n // Softmax Grad\n bool softmax_success = false;\n softmax_success = dispatch_softmax_backward(\n static_cast(matmul2_grads.data_ptr()), static_cast(matmul2_grads.data_ptr()),\n reinterpret_cast(softmax_results.data_ptr()), k_seq_len, k_seq_len, attn_batches * q_seq_len);\n assert(softmax_success);\n\n // Matmul1 Dgrad1\n gemm_switch_fp32accum(a_layout_n, b_layout_n, head_dim, q_seq_len, k_seq_len, scale, k_lin_results_ptr, lead_dim,\n batch_stride, static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta, q_lin_grads_ptr, lead_dim, batch_stride, attn_batches);\n\n // Matmul1 Dgrad2\n gemm_switch_fp32accum(a_layout_n, b_layout_t, head_dim, k_seq_len, q_seq_len, scale, q_lin_results_ptr, lead_dim,\n batch_stride, static_cast(matmul2_grads.data_ptr()), k_seq_len, k_seq_len * q_seq_len,\n beta, k_lin_grads_ptr, lead_dim, batch_stride, attn_batches);\n\n // Input Linear Dgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(\n handle, CUBLAS_OP_N, CUBLAS_OP_N, embed_dim, batches, output_lin_dim, static_cast(&alpha),\n static_cast(input_weights.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(q_lin_grads_ptr), CUDA_R_16F, output_lin_dim, static_cast(&beta),\n // static_cast(input_grads.data_ptr()),\n static_cast(input_lin_grads.data_ptr()), CUDA_R_16F, embed_dim, CUDA_R_32F,\n // CUBLAS_GEMM_ALGO10_TENSOR_OP));\n CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Input Linear Wgrad\n TORCH_CUDABLAS_CHECK(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, embed_dim, output_lin_dim, batches,\n static_cast(&alpha),\n // static_cast(inputs.data_ptr()),\n static_cast(lyr_nrm_results.data_ptr()), CUDA_R_16F, embed_dim,\n static_cast(q_lin_grads_ptr), CUDA_R_16F, output_lin_dim,\n static_cast(&beta), static_cast(input_weight_grads.data_ptr()),\n CUDA_R_16F, embed_dim, CUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP));\n\n // Fused Layer Norm Bwd with Residual Add\n HostLayerNormGradient(\n static_cast(input_lin_grads.data_ptr()), static_cast(output_grads.data_ptr()),\n static_cast(lyr_nrm_mean.data_ptr()), static_cast(lyr_nrm_invvar.data_ptr()), inputs,\n static_cast(batches), // n1\n static_cast(embed_dim), // n2\n static_cast(lyr_nrm_gamma_weights.data_ptr()),\n static_cast(lyr_nrm_beta_weights.data_ptr()), 1.0e-5, static_cast(input_grads.data_ptr()),\n static_cast(lyr_nrm_gamma_grads.data_ptr()), static_cast(lyr_nrm_beta_grads.data_ptr()));\n\n TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));\n\n return {input_grads, lyr_nrm_gamma_grads, lyr_nrm_beta_grads, input_weight_grads, output_weight_grads};\n}\n\n} // end namespace cublas_gemmex\n} // end namespace self_norm_add\n} // end namespace multihead_attn\n" }, { "path": "apex/contrib/csrc/multihead_attn/softmax.cuh", "content": "#pragma once\n#include \n\n#include \n\n#include \"philox.cuh\"\n\n#ifdef OLD_GENERATOR_PATH\n#include \n#else\n#include \n#endif\n\n#include \n#include \n#include \n\n#include \n#include \n#include \n\nnamespace {\ntemplate \n__device__ __inline__ void copy_vector(Datatype* dst, const Datatype* src);\n\ntemplate \n__device__ __inline__ void apply_mask(Datatype* dst, Datatype value, const uint8_t* src);\n\ntemplate \n__device__ __inline__ void apply_additive_mask(Datatype* dst, const Datatype* additive_mask);\n\ntemplate <>\n__device__ __inline__ void copy_vector<__half, 1>(__half* dst, const __half* src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(float* dst, const float* src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector<__half, 4>(__half* dst, const __half* src) {\n *((float2*)dst) = *((float2*)src);\n}\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t* dst, const uint8_t* src) {\n *dst = *src;\n}\n\ntemplate <>\n__device__ __inline__ void copy_vector(uint8_t* dst, const uint8_t* src) {\n *((half2*)dst) = *((half2*)src);\n}\n\ntemplate <>\n__device__ __inline__ void apply_mask<__half, 1>(__half* dst, __half value, const uint8_t* src) {\n if (*src == 1) {\n *dst = value;\n }\n}\n\ntemplate <>\n__device__ __inline__ void apply_additive_mask<__half, 1>(__half* dst, const __half* additive_mask) {\n *dst += *additive_mask;\n}\n\ntemplate <>\n__device__ __inline__ void apply_additive_mask<__half, 4>(__half* dst, const __half* additive_mask) {\n *dst += *additive_mask;\n *(dst + 1) += *(additive_mask + 1);\n *(dst + 2) += *(additive_mask + 2);\n *(dst + 3) += *(additive_mask + 3);\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n// Warp Softmax forward\n////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n// WARP_BATCH number of batches.\n// WARP_ITERATOINS The number of iterations required for one warp to iterate\n// over all data. WARP_SIZE number of elements working on a single batch, has to\n// be a power of two. ELEMENTS_PER_LDG_STG has to be 1.\ntemplate \n__global__ void softmax_warp_forward(input_t* dst, const output_t* src, int batch_size, int stride, int element_count) {\n assert(ELEMENTS_PER_LDG_STG == 1);\n\n int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n\n // batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the\n // batch\n int local_idx = threadIdx.x;\n\n src += first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n dst += first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n\n // load data from global memory\n input_t elements_input[WARP_BATCH][WARP_ITERATIONS];\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements_input[i][it + element] = -std::numeric_limits::infinity();\n }\n\n if (element_index < batch_element_count) {\n copy_vector(&elements_input[i][it], src + i * element_count + it * WARP_SIZE);\n }\n }\n }\n\n // convert input_t to acc_t\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = elements_input[i][it];\n }\n }\n\n constexpr uint32_t FULL_MASK = 0xffffffff;\n\n // compute local max_value\n\n // take the max_value of the first element to avoid one max call\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n }\n\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n\n// reduction max_value\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n float val[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);\n }\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];\n }\n }\n\n // compute local sum\n acc_t sum[WARP_BATCH]{0.0f};\n\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n // elements[i][it] = expf(elements[i][it] - max_value[i]);\n elements[i][it] = std::exp(elements[i][it] - max_value[i]);\n sum[i] += elements[i][it];\n }\n }\n\n// reduction sum\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);\n }\n }\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // dst[i * element_count + it * WARP_SIZE] = elements[i][it] / sum[i];\n output_t out[ELEMENTS_PER_LDG_STG];\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = elements[i][it + element] / sum[i];\n }\n copy_vector(dst + i * element_count + it * WARP_SIZE, out);\n } else {\n break;\n }\n }\n }\n}\n\n// WARP_BATCH number of batches.\n// WARP_ITERATOINS The number of iterations required for one warp to iterate\n// over all data. WARP_SIZE number of elements working on a single batch, has to\n// be a power of two. ELEMENTS_PER_LDG_STG has to be 1.\ntemplate \nusing softmax_forward_func = void (*)(input_t* dst, const output_t* src, int batch_size, int stride, int element_count);\n\ntemplate \nbool warp_softmax_kernel(int log2_elements, int& warp_size, int& batches_per_warp,\n softmax_forward_func& kernel) {\n // determine size of a warp\n const int next_power_of_two = 1 << log2_elements;\n warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;\n\n // determine how many batches a warp should process.\n batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n switch (log2_elements) {\n case 0: // 1\n kernel = &softmax_warp_forward;\n break;\n case 1: // 2\n kernel = &softmax_warp_forward;\n break;\n case 2: // 4\n kernel = &softmax_warp_forward;\n break;\n case 3: // 8\n kernel = &softmax_warp_forward;\n break;\n case 4: // 16\n kernel = &softmax_warp_forward;\n break;\n case 5: // 32\n kernel = &softmax_warp_forward;\n break;\n case 6: // 64\n kernel = &softmax_warp_forward;\n break;\n case 7: // 128\n kernel = &softmax_warp_forward;\n break;\n case 8: // 256\n kernel = &softmax_warp_forward;\n break;\n case 9: // 512\n kernel = &softmax_warp_forward;\n break;\n case 10: // 1024\n kernel = &softmax_warp_forward;\n break;\n default:\n return false;\n }\n return true;\n}\n\ntemplate \nbool dispatch_softmax(output_t* dst, const input_t* src, int softmax_elements, int softmax_elements_stride,\n int batch_count) {\n if (softmax_elements == 0) {\n return true;\n } else if (softmax_elements <= 1024) {\n // compute function index. there's a function for each power of two size up\n // to 1024.\n int log2_elements = 0;\n while ((1 << log2_elements) < softmax_elements) ++log2_elements;\n\n softmax_forward_func kernel;\n int warp_size, batches_per_warp;\n if (!warp_softmax_kernel(log2_elements, warp_size, batches_per_warp, kernel)) {\n return false;\n }\n\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n\n // compute warps per block.\n int warps_per_block = (threads_per_block / warp_size);\n\n // compute launch size\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n\n // launch\n kernel<<>>(dst, src, batch_count, softmax_elements_stride,\n softmax_elements);\n return true;\n }\n return false;\n}\n\ntemplate \n__global__ void additive_masked_softmax_dropout_warp_forward_vec4(output_t* dst, uint8_t* dropout_mask,\n const input_t* src, const input_t* pad_mask,\n int batch_size, int stride, int element_count,\n int pad_batch_stride, at::PhiloxCudaState philox_args,\n float p) {\n assert(ELEMENTS_PER_LDG_STG == 4);\n int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n int tid = blockIdx.x * blockDim.x * blockDim.y + threadIdx.y * blockDim.x + threadIdx.x;\n acc_t pinv = acc_t(1) / p;\n // batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the\n // batch\n int local_idx = threadIdx.x;\n // vectorize if element_count is multiple of 4, else don't vectorize\n input_t elements_input[WARP_BATCH][WARP_ITERATIONS];\n\n int thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n src += thread_offset;\n dst += thread_offset;\n dropout_mask += thread_offset;\n\n // load data from global memory\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n int pad_thread_offset = ((first_batch + i) / pad_batch_stride) * stride + ELEMENTS_PER_LDG_STG * local_idx;\n const half* curr_mask = pad_mask + pad_thread_offset;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n // masking_value is a large negative value\n elements_input[i][it + element] = -10000;\n }\n\n if (element_index < batch_element_count) {\n int itr_jmp = it * WARP_SIZE;\n int itr_idx = i * element_count + itr_jmp;\n copy_vector(&elements_input[i][it], src + itr_idx);\n apply_additive_mask(\n &elements_input[i][it],\n curr_mask + itr_jmp); //(__half)-std::numeric_limits::infinity()\n }\n }\n }\n // convert input_t to acc_t\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = elements_input[i][it];\n }\n }\n\n constexpr uint32_t FULL_MASK = 0xffffffff;\n\n // compute local max_value\n\n // take the max_value of the first element to avoid one max call\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n }\n\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n\n// reduction max_value\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n float val[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);\n }\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];\n }\n }\n\n // compute local sum\n acc_t sum[WARP_BATCH]{0.0f};\n\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = std::exp(elements[i][it] - max_value[i]);\n sum[i] += elements[i][it];\n }\n }\n\n// reduction sum\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);\n }\n }\n auto seeds = at::cuda::philox::unpack(philox_args);\n Philox ph(std::get<0>(seeds), tid, std::get<1>(seeds));\n uint8_t rands[WARP_BATCH][WARP_ITERATIONS];\n float4 rand_num;\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n rand_num = uniform4(ph());\n rands[i][it] = (rand_num.x <= p) > 0.5;\n rands[i][it + 1] = (rand_num.y <= p) > 0.5;\n rands[i][it + 2] = (rand_num.z <= p) > 0.5;\n rands[i][it + 3] = (rand_num.w <= p) > 0.5;\n copy_vector(dropout_mask + i * element_count + it * WARP_SIZE, &rands[i][it]);\n }\n }\n }\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n output_t out[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = rands[i][it + element] * (pinv * (elements[i][it + element] / sum[i]));\n }\n copy_vector(dst + i * element_count + it * WARP_SIZE, out);\n\n } else {\n break;\n }\n }\n }\n}\n\ntemplate \n__global__ void additive_masked_softmax_dropout_warp_forward(output_t* dst, uint8_t* dropout_mask, const input_t* src,\n const input_t* pad_mask, int batch_size, int stride,\n int element_count, int pad_batch_stride,\n at::PhiloxCudaState philox_args, float p) {\n assert(ELEMENTS_PER_LDG_STG == 1);\n int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n int tid = blockIdx.x * blockDim.x * blockDim.y + threadIdx.y * blockDim.x + threadIdx.x;\n acc_t pinv = acc_t(1) / p;\n // batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the\n // batch\n int local_idx = threadIdx.x;\n // vectorize if element_count is multiple of 4, else don't vectorize\n input_t elements_input[WARP_BATCH][WARP_ITERATIONS];\n\n int thread_offset = first_batch * stride + local_idx;\n src += thread_offset;\n dst += thread_offset;\n dropout_mask += thread_offset;\n\n // load data from global memory\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n int pad_thread_offset = ((first_batch + i) / pad_batch_stride) * stride + local_idx;\n const half* curr_mask = pad_mask + pad_thread_offset;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += 1) {\n int element_index = local_idx + it * WARP_SIZE;\n#pragma unroll\n for (int element = 0; element < 1; ++element) {\n // masking_value is a large negative value\n elements_input[i][it + element] = -10000;\n }\n\n if (element_index < batch_element_count) {\n int itr_jmp = it * WARP_SIZE;\n int itr_idx = i * element_count + itr_jmp;\n copy_vector(&elements_input[i][it], src + itr_idx);\n apply_additive_mask(&elements_input[i][it], curr_mask + itr_jmp);\n }\n }\n }\n // convert input_t to acc_t\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = elements_input[i][it];\n }\n }\n\n constexpr uint32_t FULL_MASK = 0xffffffff;\n\n // compute local max_value\n\n // take the max_value of the first element to avoid one max call\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n }\n\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n\n// reduction max_value\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n float val[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);\n }\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];\n }\n }\n\n // compute local sum\n acc_t sum[WARP_BATCH]{0.0f};\n\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = std::exp(elements[i][it] - max_value[i]);\n sum[i] += elements[i][it];\n }\n }\n\n// reduction sum\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);\n }\n }\n curandStatePhilox4_32_10_t state;\n auto seeds = at::cuda::philox::unpack(philox_args);\n curand_init(std::get<0>(seeds), tid, std::get<1>(seeds), &state);\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += 1) {\n int element_index = local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n output_t out[1];\n acc_t softmax_out[1];\n uint8_t dropout_mask_temp[1];\n // generate a vector of random numbers here\n float rand = curand_uniform(&state);\n float* rand_ptr = (float*)(&rand);\n#pragma unroll\n for (int element = 0; element < 1; ++element) {\n softmax_out[element] = (elements[i][it + element] / sum[i]);\n rand_ptr[element] = rand_ptr[element] <= p;\n out[element] = rand_ptr[element] * pinv * softmax_out[element];\n dropout_mask_temp[element] = rand_ptr[element] > 0.5; // just to distinguish 0.0f and 1.0f\n }\n copy_vector(dst + i * element_count + it * WARP_SIZE, out);\n copy_vector(dropout_mask + i * element_count + it * WARP_SIZE, dropout_mask_temp);\n\n } else {\n break;\n }\n }\n }\n}\n\n// WARP_BATCH number of batches.\n// WARP_ITERATOINS The number of iterations required for one warp to iterate\n// over all data. WARP_SIZE number of elements working on a single batch, has to\n// be a power of two. ELEMENTS_PER_LDG_STG has to be 1.\ntemplate \nusing additive_masked_softmax_dropout_forward_func = void (*)(output_t* dst, uint8_t* dropout_mask, const input_t* src,\n const input_t* pad_mask, int batch_size, int stride,\n int element_count, int pad_batch_stride,\n at::PhiloxCudaState philox_args, float p);\n\ntemplate \nbool warp_additive_masked_softmax_dropout_kernel(\n int element_count, int log2_elements, int& warp_size, int& batches_per_warp,\n additive_masked_softmax_dropout_forward_func& kernel) {\n // determine size of a warp\n const int next_power_of_two = 1 << log2_elements;\n warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;\n\n // determine how many batches a warp should process.\n batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n bool flag_vec4 = (element_count % 4 == 0);\n switch (log2_elements) {\n case 0: // 1\n kernel = &additive_masked_softmax_dropout_warp_forward;\n break;\n case 1: // 2\n kernel = &additive_masked_softmax_dropout_warp_forward;\n break;\n case 2: // 4\n kernel = &additive_masked_softmax_dropout_warp_forward;\n break;\n case 3: // 8\n kernel = &additive_masked_softmax_dropout_warp_forward;\n break;\n case 4: // 16\n kernel = &additive_masked_softmax_dropout_warp_forward;\n break;\n case 5: // 32\n kernel = &additive_masked_softmax_dropout_warp_forward;\n break;\n case 6: // 64\n kernel = &additive_masked_softmax_dropout_warp_forward;\n break;\n case 7: // 128\n if (flag_vec4)\n kernel = &additive_masked_softmax_dropout_warp_forward_vec4;\n else\n kernel = &additive_masked_softmax_dropout_warp_forward;\n break;\n case 8: // 256\n if (flag_vec4)\n kernel = &additive_masked_softmax_dropout_warp_forward_vec4;\n else\n kernel = &additive_masked_softmax_dropout_warp_forward;\n break;\n case 9: // 512\n if (flag_vec4)\n kernel = &additive_masked_softmax_dropout_warp_forward_vec4;\n else\n kernel = &additive_masked_softmax_dropout_warp_forward;\n break;\n case 10: // 1024\n if (flag_vec4)\n kernel = &additive_masked_softmax_dropout_warp_forward_vec4;\n else\n kernel = &additive_masked_softmax_dropout_warp_forward;\n break;\n case 11: // 2048\n if (flag_vec4)\n kernel = &additive_masked_softmax_dropout_warp_forward_vec4;\n else\n kernel = &additive_masked_softmax_dropout_warp_forward;\n break;\n default:\n return false;\n }\n return true;\n}\n\ntemplate \nbool dispatch_additive_masked_softmax_dropout(output_t* dst, uint8_t* dropout_mask, const input_t* src,\n const input_t* pad_mask, int totalElements, int softmax_elements,\n int softmax_elements_stride, int batch_count, int pad_batch_stride,\n float p,\n cudaStream_t streamid) // p is the probability to keep, not drop\n{\n if (softmax_elements == 0) {\n return true;\n } else if (softmax_elements <= 2048) {\n // compute function index. there's a function for each power of two size up\n // to 1024.\n int log2_elements = 0;\n while ((1 << log2_elements) < softmax_elements) ++log2_elements;\n\n additive_masked_softmax_dropout_forward_func kernel;\n int warp_size, batches_per_warp;\n if (!warp_additive_masked_softmax_dropout_kernel(softmax_elements, log2_elements,\n warp_size, batches_per_warp, kernel)) {\n return false;\n }\n\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n // compute warps per block.\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n c10::optional gen_;\n auto gen = at::get_generator_or_default(gen_, at::cuda::detail::getDefaultCUDAGenerator());\n int64_t counter_offset = (totalElements / (blocks * threads_per_block) + 1);\n at::PhiloxCudaState rng_engine_inputs;\n {\n std::lock_guard lock(gen->mutex_);\n rng_engine_inputs = gen->philox_cuda_state(counter_offset);\n }\n\n // compute launch size\n dim3 threads(warp_size, warps_per_block, 1);\n\n // launch\n kernel<<>>(dst, dropout_mask, src, pad_mask, batch_count, softmax_elements_stride,\n softmax_elements, pad_batch_stride, rng_engine_inputs, p);\n return true;\n }\n return false;\n}\n\n// WARP_BATCH number of batches.\n// WARP_ITERATOINS The number of iterations required for one warp to iterate\n// over all data. WARP_SIZE number of elements working on a single batch, has to\n// be a power of two. ELEMENTS_PER_LDG_STG has to be 1.\ntemplate \n__global__ void additive_masked_softmax_warp_forward(input_t* dst, const output_t* src, const input_t* pad_mask,\n int batch_size, int stride, int element_count,\n int pad_batch_stride) {\n assert(ELEMENTS_PER_LDG_STG == 1);\n\n int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n\n // batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the\n // batch\n int local_idx = threadIdx.x;\n\n int thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n src += thread_offset;\n dst += thread_offset;\n\n // load data from global memory\n input_t elements_input[WARP_BATCH][WARP_ITERATIONS];\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n int pad_thread_offset = ((first_batch + i) / pad_batch_stride) * stride + ELEMENTS_PER_LDG_STG * local_idx;\n const half* curr_mask = pad_mask + pad_thread_offset;\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n // masking_value is a large negative value\n elements_input[i][it + element] = -10000;\n }\n\n if (element_index < batch_element_count) {\n int itr_jmp = it * WARP_SIZE;\n int itr_idx = i * element_count + itr_jmp;\n copy_vector(&elements_input[i][it], src + itr_idx);\n // apply_mask(&elements_input[i][it],\n // (__half)-std::numeric_limits::infinity(),\n // curr_mask + itr_jmp);\n elements_input[i][it] += *(curr_mask + itr_jmp);\n }\n }\n }\n\n // convert input_t to acc_t\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = elements_input[i][it];\n }\n }\n\n constexpr uint32_t FULL_MASK = 0xffffffff;\n\n // compute local max_value\n\n // take the max_value of the first element to avoid one max call\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n }\n\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n\n// reduction max_value\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n float val[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);\n }\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];\n }\n }\n\n // compute local sum\n acc_t sum[WARP_BATCH]{0.0f};\n\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n // elements[i][it] = expf(elements[i][it] - max_value[i]);\n elements[i][it] = std::exp(elements[i][it] - max_value[i]);\n sum[i] += elements[i][it];\n }\n }\n\n// reduction sum\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);\n }\n }\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // dst[i * element_count + it * WARP_SIZE] = elements[i][it] / sum[i];\n output_t out[ELEMENTS_PER_LDG_STG];\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = elements[i][it + element] / sum[i];\n }\n copy_vector(dst + i * element_count + it * WARP_SIZE, out);\n } else {\n break;\n }\n }\n }\n}\n\n// WARP_BATCH number of batches.\n// WARP_ITERATOINS The number of iterations required for one warp to iterate\n// over all data. WARP_SIZE number of elements working on a single batch, has to\n// be a power of two. ELEMENTS_PER_LDG_STG has to be 1.\ntemplate \nusing additive_masked_softmax_forward_func = void (*)(input_t* dst, const output_t* src, const half* pad_mask,\n int batch_size, int stride, int element_count,\n int pad_batch_stride);\n\ntemplate \nbool warp_additive_masked_softmax_kernel(int log2_elements, int& warp_size, int& batches_per_warp,\n additive_masked_softmax_forward_func& kernel) {\n // determine size of a warp\n const int next_power_of_two = 1 << log2_elements;\n warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;\n\n // determine how many batches a warp should process.\n batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n switch (log2_elements) {\n case 0: // 1\n kernel = &additive_masked_softmax_warp_forward;\n break;\n case 1: // 2\n kernel = &additive_masked_softmax_warp_forward;\n break;\n case 2: // 4\n kernel = &additive_masked_softmax_warp_forward;\n break;\n case 3: // 8\n kernel = &additive_masked_softmax_warp_forward;\n break;\n case 4: // 16\n kernel = &additive_masked_softmax_warp_forward;\n break;\n case 5: // 32\n kernel = &additive_masked_softmax_warp_forward;\n break;\n case 6: // 64\n kernel = &additive_masked_softmax_warp_forward;\n break;\n case 7: // 128\n kernel = &additive_masked_softmax_warp_forward;\n break;\n case 8: // 256\n kernel = &additive_masked_softmax_warp_forward;\n break;\n case 9: // 512\n kernel = &additive_masked_softmax_warp_forward;\n break;\n case 10: // 1024\n kernel = &additive_masked_softmax_warp_forward;\n break;\n default:\n return false;\n }\n return true;\n}\n\ntemplate \nbool dispatch_additive_masked_softmax(output_t* dst, const input_t* src, const input_t* pad_mask, int softmax_elements,\n int softmax_elements_stride, int batch_count, int pad_batch_stride) {\n if (softmax_elements == 0) {\n return true;\n } else if (softmax_elements <= 1024) {\n // compute function index. there's a function for each power of two size up\n // to 1024.\n int log2_elements = 0;\n while ((1 << log2_elements) < softmax_elements) ++log2_elements;\n\n additive_masked_softmax_forward_func kernel;\n int warp_size, batches_per_warp;\n if (!warp_additive_masked_softmax_kernel(log2_elements, warp_size, batches_per_warp,\n kernel)) {\n return false;\n }\n\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n\n // compute warps per block.\n int warps_per_block = (threads_per_block / warp_size);\n\n // compute launch size\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n\n // launch\n kernel<<>>(\n dst, src, pad_mask, batch_count, softmax_elements_stride, softmax_elements, pad_batch_stride);\n return true;\n }\n return false;\n}\n\ntemplate \nbool dispatch_additive_masked_softmax_stream(output_t* dst, const input_t* src, const input_t* pad_mask,\n int softmax_elements, int softmax_elements_stride, int batch_count,\n int pad_batch_stride, cudaStream_t streamid) {\n if (softmax_elements == 0) {\n return true;\n } else if (softmax_elements <= 1024) {\n // compute function index. there's a function for each power of two size up\n // to 1024.\n int log2_elements = 0;\n while ((1 << log2_elements) < softmax_elements) ++log2_elements;\n additive_masked_softmax_forward_func kernel;\n int warp_size, batches_per_warp;\n if (!warp_additive_masked_softmax_kernel(log2_elements, warp_size, batches_per_warp,\n kernel)) {\n return false;\n }\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n // compute warps per block.\n int warps_per_block = (threads_per_block / warp_size);\n // compute launch size\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n // launch\n kernel<<>>(dst, src, pad_mask, batch_count, softmax_elements_stride, softmax_elements,\n pad_batch_stride);\n return true;\n }\n return false;\n}\n\n// WARP_BATCH number of batches.\n// WARP_ITERATOINS The number of iterations required for one warp to iterate\n// over all data. WARP_SIZE number of elements working on a single batch, has to\n// be a power of two. ELEMENTS_PER_LDG_STG has to be 1.\ntemplate \n__global__ void masked_softmax_warp_forward(input_t* dst, const output_t* src, const uint8_t* pad_mask, int batch_size,\n int stride, int element_count, int pad_batch_stride) {\n assert(ELEMENTS_PER_LDG_STG == 1);\n\n int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n\n // batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the\n // batch\n int local_idx = threadIdx.x;\n\n int thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n src += thread_offset;\n dst += thread_offset;\n\n // load data from global memory\n input_t elements_input[WARP_BATCH][WARP_ITERATIONS];\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n int pad_thread_offset = ((first_batch + i) / pad_batch_stride) * stride + ELEMENTS_PER_LDG_STG * local_idx;\n const uint8_t* curr_mask = pad_mask + pad_thread_offset;\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements_input[i][it + element] = -std::numeric_limits::infinity();\n }\n\n if (element_index < batch_element_count) {\n int itr_jmp = it * WARP_SIZE;\n int itr_idx = i * element_count + itr_jmp;\n copy_vector(&elements_input[i][it], src + itr_idx);\n apply_mask(\n &elements_input[i][it], __float2half(-std::numeric_limits::infinity()), curr_mask + itr_jmp);\n }\n }\n }\n\n // convert input_t to acc_t\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = elements_input[i][it];\n }\n }\n\n constexpr uint32_t FULL_MASK = 0xffffffff;\n\n // compute local max_value\n\n // take the max_value of the first element to avoid one max call\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n }\n\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n\n// reduction max_value\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n float val[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);\n }\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];\n }\n }\n\n // compute local sum\n acc_t sum[WARP_BATCH]{0.0f};\n\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n // elements[i][it] = expf(elements[i][it] - max_value[i]);\n elements[i][it] = std::exp(elements[i][it] - max_value[i]);\n sum[i] += elements[i][it];\n }\n }\n\n// reduction sum\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);\n }\n }\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // dst[i * element_count + it * WARP_SIZE] = elements[i][it] / sum[i];\n output_t out[ELEMENTS_PER_LDG_STG];\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = elements[i][it + element] / sum[i];\n }\n copy_vector(dst + i * element_count + it * WARP_SIZE, out);\n } else {\n break;\n }\n }\n }\n}\n\n// WARP_BATCH number of batches.\n// WARP_ITERATOINS The number of iterations required for one warp to iterate\n// over all data. WARP_SIZE number of elements working on a single batch, has to\n// be a power of two. ELEMENTS_PER_LDG_STG has to be 1.\ntemplate \nusing masked_softmax_forward_func = void (*)(input_t* dst, const output_t* src, const uint8_t* pad_mask, int batch_size,\n int stride, int element_count, int pad_batch_stride);\n\ntemplate \nbool warp_masked_softmax_kernel(int log2_elements, int& warp_size, int& batches_per_warp,\n masked_softmax_forward_func& kernel) {\n // determine size of a warp\n const int next_power_of_two = 1 << log2_elements;\n warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;\n\n // determine how many batches a warp should process.\n batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n switch (log2_elements) {\n case 0: // 1\n kernel = &masked_softmax_warp_forward;\n break;\n case 1: // 2\n kernel = &masked_softmax_warp_forward;\n break;\n case 2: // 4\n kernel = &masked_softmax_warp_forward;\n break;\n case 3: // 8\n kernel = &masked_softmax_warp_forward;\n break;\n case 4: // 16\n kernel = &masked_softmax_warp_forward;\n break;\n case 5: // 32\n kernel = &masked_softmax_warp_forward;\n break;\n case 6: // 64\n kernel = &masked_softmax_warp_forward;\n break;\n case 7: // 128\n kernel = &masked_softmax_warp_forward;\n break;\n case 8: // 256\n kernel = &masked_softmax_warp_forward;\n break;\n case 9: // 512\n kernel = &masked_softmax_warp_forward;\n break;\n case 10: // 1024\n kernel = &masked_softmax_warp_forward;\n break;\n default:\n return false;\n }\n return true;\n}\n\ntemplate \nbool dispatch_masked_softmax(output_t* dst, const input_t* src, const uint8_t* pad_mask, int softmax_elements,\n int softmax_elements_stride, int batch_count, int pad_batch_stride) {\n if (softmax_elements == 0) {\n return true;\n } else if (softmax_elements <= 1024) {\n // compute function index. there's a function for each power of two size up\n // to 1024.\n int log2_elements = 0;\n while ((1 << log2_elements) < softmax_elements) ++log2_elements;\n\n masked_softmax_forward_func kernel;\n int warp_size, batches_per_warp;\n if (!warp_masked_softmax_kernel(log2_elements, warp_size, batches_per_warp, kernel)) {\n return false;\n }\n\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n\n // compute warps per block.\n int warps_per_block = (threads_per_block / warp_size);\n\n // compute launch size\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n\n // launch\n kernel<<>>(\n dst, src, pad_mask, batch_count, softmax_elements_stride, softmax_elements, pad_batch_stride);\n return true;\n }\n return false;\n}\n\n// WARP_BATCH number of batches.\n// WARP_ITERATOINS The number of iterations required for one warp to iterate\n// over all data. WARP_SIZE number of elements working on a single batch, has to\n// be a power of two. ELEMENTS_PER_LDG_STG has to be 1.\ntemplate \n__global__ void time_masked_softmax_warp_forward(input_t* dst, const output_t* src, const uint8_t* pad_mask,\n int batch_size, int stride, int element_count, int mod_seq_len) {\n assert(ELEMENTS_PER_LDG_STG == 1);\n\n int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n\n // batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the\n // batch\n int local_idx = threadIdx.x;\n\n int thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n src += thread_offset;\n dst += thread_offset;\n\n // load data from global memory\n input_t elements_input[WARP_BATCH][WARP_ITERATIONS];\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n int pad_thread_offset = ((first_batch + i) % mod_seq_len) * stride + ELEMENTS_PER_LDG_STG * local_idx;\n const uint8_t* curr_mask = pad_mask + pad_thread_offset;\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n elements_input[i][it + element] = -std::numeric_limits::infinity();\n }\n\n if (element_index < batch_element_count) {\n int itr_jmp = it * WARP_SIZE;\n int itr_idx = i * element_count + itr_jmp;\n copy_vector(&elements_input[i][it], src + itr_idx);\n apply_mask(\n &elements_input[i][it], __float2half(-std::numeric_limits::infinity()), curr_mask + itr_jmp);\n }\n }\n }\n\n // convert input_t to acc_t\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = elements_input[i][it];\n }\n }\n\n constexpr uint32_t FULL_MASK = 0xffffffff;\n\n // compute local max_value\n\n // take the max_value of the first element to avoid one max call\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n }\n\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n\n// reduction max_value\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n float val[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);\n }\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];\n }\n }\n\n // compute local sum\n acc_t sum[WARP_BATCH]{0.0f};\n\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n // elements[i][it] = expf(elements[i][it] - max_value[i]);\n elements[i][it] = std::exp(elements[i][it] - max_value[i]);\n sum[i] += elements[i][it];\n }\n }\n\n// reduction sum\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);\n }\n }\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // dst[i * element_count + it * WARP_SIZE] = elements[i][it] / sum[i];\n output_t out[ELEMENTS_PER_LDG_STG];\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = elements[i][it + element] / sum[i];\n }\n copy_vector(dst + i * element_count + it * WARP_SIZE, out);\n } else {\n break;\n }\n }\n }\n}\n\n// WARP_BATCH number of batches.\n// WARP_ITERATOINS The number of iterations required for one warp to iterate\n// over all data. WARP_SIZE number of elements working on a single batch, has to\n// be a power of two. ELEMENTS_PER_LDG_STG has to be 1.\ntemplate \nusing time_masked_softmax_forward_func = void (*)(input_t* dst, const output_t* src, const uint8_t* pad_mask,\n int batch_size, int stride, int element_count, int mod_seq_len);\n\ntemplate \nbool warp_time_masked_softmax_kernel(int log2_elements, int& warp_size, int& batches_per_warp,\n time_masked_softmax_forward_func& kernel) {\n // determine size of a warp\n const int next_power_of_two = 1 << log2_elements;\n warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;\n\n // determine how many batches a warp should process.\n batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n switch (log2_elements) {\n case 0: // 1\n kernel = &time_masked_softmax_warp_forward;\n break;\n case 1: // 2\n kernel = &time_masked_softmax_warp_forward;\n break;\n case 2: // 4\n kernel = &time_masked_softmax_warp_forward;\n break;\n case 3: // 8\n kernel = &time_masked_softmax_warp_forward;\n break;\n case 4: // 16\n kernel = &time_masked_softmax_warp_forward;\n break;\n case 5: // 32\n kernel = &time_masked_softmax_warp_forward;\n break;\n case 6: // 64\n kernel = &time_masked_softmax_warp_forward;\n break;\n case 7: // 128\n kernel = &time_masked_softmax_warp_forward;\n break;\n case 8: // 256\n kernel = &time_masked_softmax_warp_forward;\n break;\n case 9: // 512\n kernel = &time_masked_softmax_warp_forward;\n break;\n case 10: // 1024\n kernel = &time_masked_softmax_warp_forward;\n break;\n default:\n return false;\n }\n return true;\n}\n\ntemplate \nbool dispatch_time_masked_softmax(output_t* dst, const input_t* src, const uint8_t* pad_mask, int softmax_elements,\n int softmax_elements_stride, int batch_count, int mod_seq_len) {\n if (softmax_elements == 0) {\n return true;\n } else if (softmax_elements <= 1024) {\n // compute function index. there's a function for each power of two size up\n // to 1024.\n int log2_elements = 0;\n while ((1 << log2_elements) < softmax_elements) ++log2_elements;\n\n time_masked_softmax_forward_func kernel;\n int warp_size, batches_per_warp;\n if (!warp_time_masked_softmax_kernel(log2_elements, warp_size, batches_per_warp,\n kernel)) {\n return false;\n }\n\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n\n // compute warps per block.\n int warps_per_block = (threads_per_block / warp_size);\n\n // compute launch size\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n\n // launch\n kernel<<>>(\n dst, src, pad_mask, batch_count, softmax_elements_stride, softmax_elements, mod_seq_len);\n return true;\n }\n return false;\n}\n\nint log2_ceil_native(int value) {\n int log2_value = 0;\n while ((1 << log2_value) < value) ++log2_value;\n return log2_value;\n}\n\ntemplate \n__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize,\n unsigned int mask = 0xffffffff) {\n#if CUDA_VERSION >= 9000\n return __shfl_xor_sync(mask, value, laneMask, width);\n#else\n return __shfl_xor(value, laneMask, width);\n#endif\n}\n\ntemplate \n__device__ __forceinline__ void warp_reduce_sum(acc_t* sum) {\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);\n sum[i] = sum[i] + b;\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n// Warp softmax backward functions as fused variants of\n// at::softmax_backward_data function\n////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\n// softmax backward data function is taken from native pytorch, elementwise mul\n// is fused in the epolog, as well as masking and scaling for fusing dropout\n\ntemplate \n__global__ void masked_scale_softmax_warp_backward_masked_dgrad(output_t* gradInput, const input_t* grad,\n const input_t* output, const uint8_t* mask,\n const uint8_t* pad_mask, acc_t scale, int batch_size,\n int stride, int element_count, int heads) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_backward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n\n int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n\n // batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the\n // batch\n int local_idx = threadIdx.x % WARP_SIZE;\n\n // the first element to process by the current thread\n int thread_offset = first_batch * stride + local_idx;\n grad += thread_offset;\n output += thread_offset;\n gradInput += thread_offset;\n mask += thread_offset;\n\n // The nested loops over WARP_BATCH and then WARP_ITERATIONS can be simplified\n // to one loop, but I think doing so would obfuscate the logic of the\n // algorithm, thus I chose to keep the nested loops. This should have no\n // impact on performance because the loops are unrolled anyway.\n\n // load data from global memory\n acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS];\n acc_t output_reg[WARP_BATCH][WARP_ITERATIONS];\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n int element_index = local_idx + it * WARP_SIZE;\n if (element_index < batch_element_count) {\n grad_reg[i][it] = (input_t)((acc_t)mask[i * element_count + it * WARP_SIZE] *\n (acc_t)grad[i * element_count + it * WARP_SIZE] * (acc_t)scale) *\n output[i * element_count + it * WARP_SIZE];\n output_reg[i][it] = output[i * element_count + it * WARP_SIZE];\n } else {\n grad_reg[i][it] = acc_t(0);\n output_reg[i][it] = acc_t(0);\n }\n }\n }\n\n acc_t sum[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] = grad_reg[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n sum[i] += grad_reg[i][it];\n }\n }\n warp_reduce_sum(sum);\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n int element_index = local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // compute gradients\n int total_ind = thread_offset + i * element_count + it * WARP_SIZE;\n int pad_mask_ind =\n element_count * (total_ind / (heads * element_count * element_count)) + total_ind % element_count;\n uint8_t pad_mask_element = 1 - pad_mask[pad_mask_ind];\n if (pad_mask_element == 0)\n gradInput[i * element_count + it * WARP_SIZE] = 0;\n else {\n if (is_log_softmax) {\n gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - std::exp(output_reg[i][it]) * sum[i]);\n } else {\n gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - output_reg[i][it] * sum[i]);\n }\n }\n }\n }\n }\n}\ntemplate \nvoid dispatch_masked_scale_softmax_backward_masked_out(output_t* grad_input, const input_t* grad, const input_t* output,\n const uint8_t* mask, const uint8_t* pad_mask, acc_t scale,\n int softmax_elements, int softmax_elements_stride,\n int batch_count, int heads) {\n TORCH_INTERNAL_ASSERT(softmax_elements >= 0 && softmax_elements <= 1024);\n if (softmax_elements == 0) {\n return;\n } else {\n int log2_elements = log2_ceil_native(softmax_elements);\n const int next_power_of_two = 1 << log2_elements;\n\n // This value must match the WARP_SIZE constexpr value computed inside\n // softmax_warp_backward.\n int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;\n\n // This value must match the WARP_BATCH constexpr value computed inside\n // softmax_warp_backward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale,\n batch_count, softmax_elements_stride,\n softmax_elements, heads);\n break;\n case 1: // 2\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale,\n batch_count, softmax_elements_stride,\n softmax_elements, heads);\n break;\n case 2: // 4\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale,\n batch_count, softmax_elements_stride,\n softmax_elements, heads);\n break;\n case 3: // 8\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale,\n batch_count, softmax_elements_stride,\n softmax_elements, heads);\n break;\n case 4: // 16\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale,\n batch_count, softmax_elements_stride,\n softmax_elements, heads);\n break;\n case 5: // 32\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale,\n batch_count, softmax_elements_stride,\n softmax_elements, heads);\n break;\n case 6: // 64\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale,\n batch_count, softmax_elements_stride,\n softmax_elements, heads);\n break;\n case 7: // 128\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale,\n batch_count, softmax_elements_stride,\n softmax_elements, heads);\n break;\n case 8: // 256\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale,\n batch_count, softmax_elements_stride,\n softmax_elements, heads);\n break;\n case 9: // 512\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale,\n batch_count, softmax_elements_stride,\n softmax_elements, heads);\n break;\n case 10: // 1024\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale,\n batch_count, softmax_elements_stride,\n softmax_elements, heads);\n break;\n default:\n break;\n }\n }\n}\n\ntemplate \nvoid dispatch_masked_scale_softmax_backward_masked_out_stream(output_t* grad_input, const input_t* grad,\n const input_t* output, const uint8_t* mask,\n const uint8_t* pad_mask, acc_t scale,\n int softmax_elements, int softmax_elements_stride,\n int batch_count, int heads, cudaStream_t streamid) {\n TORCH_INTERNAL_ASSERT(softmax_elements >= 0 && softmax_elements <= 1024);\n if (softmax_elements == 0) {\n return;\n } else {\n int log2_elements = log2_ceil_native(softmax_elements);\n const int next_power_of_two = 1 << log2_elements;\n // This value must match the WARP_SIZE constexpr value computed inside\n // softmax_warp_backward.\n int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;\n // This value must match the WARP_BATCH constexpr value computed inside\n // softmax_warp_backward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale, batch_count,\n softmax_elements_stride, softmax_elements, heads);\n break;\n case 1: // 2\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale, batch_count,\n softmax_elements_stride, softmax_elements, heads);\n break;\n case 2: // 4\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale, batch_count,\n softmax_elements_stride, softmax_elements, heads);\n break;\n case 3: // 8\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale, batch_count,\n softmax_elements_stride, softmax_elements, heads);\n break;\n case 4: // 16\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale, batch_count,\n softmax_elements_stride, softmax_elements, heads);\n break;\n case 5: // 32\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale, batch_count,\n softmax_elements_stride, softmax_elements, heads);\n break;\n case 6: // 64\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale, batch_count,\n softmax_elements_stride, softmax_elements, heads);\n break;\n case 7: // 128\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale, batch_count,\n softmax_elements_stride, softmax_elements, heads);\n break;\n case 8: // 256\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale, batch_count,\n softmax_elements_stride, softmax_elements, heads);\n break;\n case 9: // 512\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale, batch_count,\n softmax_elements_stride, softmax_elements, heads);\n break;\n case 10: // 1024\n masked_scale_softmax_warp_backward_masked_dgrad\n <<>>(grad_input, grad, output, mask, pad_mask, scale, batch_count,\n softmax_elements_stride, softmax_elements, heads);\n break;\n default:\n break;\n }\n }\n}\n\ntemplate \n__global__ void masked_scale_softmax_warp_backward(output_t* gradInput, const input_t* grad, const input_t* output,\n const uint8_t* mask, acc_t scale, int batch_size, int stride,\n int element_count) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_backward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n\n int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n\n // batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the\n // batch\n int local_idx = threadIdx.x % WARP_SIZE;\n\n // the first element to process by the current thread\n int thread_offset = first_batch * stride + local_idx;\n grad += thread_offset;\n output += thread_offset;\n gradInput += thread_offset;\n mask += thread_offset;\n\n // The nested loops over WARP_BATCH and then WARP_ITERATIONS can be simplified\n // to one loop, but I think doing so would obfuscate the logic of the\n // algorithm, thus I chose to keep the nested loops. This should have no\n // impact on performance because the loops are unrolled anyway.\n\n // load data from global memory\n acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS];\n acc_t output_reg[WARP_BATCH][WARP_ITERATIONS];\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n int element_index = local_idx + it * WARP_SIZE;\n if (element_index < batch_element_count) {\n grad_reg[i][it] = (input_t)((acc_t)mask[i * element_count + it * WARP_SIZE] *\n (acc_t)grad[i * element_count + it * WARP_SIZE] * (acc_t)scale) *\n output[i * element_count + it * WARP_SIZE];\n output_reg[i][it] = output[i * element_count + it * WARP_SIZE];\n } else {\n grad_reg[i][it] = acc_t(0);\n output_reg[i][it] = acc_t(0);\n }\n }\n }\n\n acc_t sum[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] = grad_reg[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n sum[i] += grad_reg[i][it];\n }\n }\n warp_reduce_sum(sum);\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n int element_index = local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // compute gradients\n if (is_log_softmax) {\n gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - std::exp(output_reg[i][it]) * sum[i]);\n } else {\n gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - output_reg[i][it] * sum[i]);\n }\n }\n }\n }\n}\n\ntemplate \n__global__ void masked_scale_softmax_warp_backward_recompute(output_t* gradInput, const input_t* grad,\n const input_t* softmax_input, const input_t* pad_mask,\n const uint8_t* mask, acc_t scale, int batch_size,\n int stride, int pad_batch_stride, int element_count) {\n int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n\n // batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the\n // batch\n int local_idx = threadIdx.x % WARP_SIZE;\n // vectorize if a row length is multiple of 4\n int flag_vec4 = element_count & 3 == 0;\n acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS];\n input_t elements_input[WARP_BATCH][WARP_ITERATIONS];\n\n // the first element to process by the current thread\n int thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n\n grad += thread_offset;\n softmax_input += thread_offset;\n gradInput += thread_offset;\n mask += thread_offset;\n\n // The nested loops over WARP_BATCH and then WARP_ITERATIONS can be simplified\n // to one loop, but I think doing so would obfuscate the logic of the\n // algorithm, thus I chose to keep the nested loops. This should have no\n // impact on performance because the loops are unrolled anyway.\n\n // load data from global memory\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n int pad_thread_offset = ((first_batch + i) / pad_batch_stride) * stride + ELEMENTS_PER_LDG_STG * local_idx;\n const input_t* curr_mask = pad_mask + pad_thread_offset;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n // masking_value is a large negative value\n elements_input[i][it + element] = -10000;\n grad_reg[i][it + element] = acc_t(0);\n }\n\n if (element_index < batch_element_count) {\n int itr_jmp = it * WARP_SIZE;\n int itr_idx = i * element_count + itr_jmp;\n copy_vector(&elements_input[i][it], softmax_input + itr_idx);\n apply_additive_mask(\n &elements_input[i][it],\n curr_mask + itr_jmp); //(__half)-std::numeric_limits::infinity()\n uint8_t mask_temp[ELEMENTS_PER_LDG_STG];\n input_t grad_temp[ELEMENTS_PER_LDG_STG];\n copy_vector(&mask_temp[0], mask + itr_idx);\n copy_vector(&grad_temp[0], grad + itr_idx);\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n grad_reg[i][it + element] = ((acc_t)mask_temp[element] * (acc_t)grad_temp[element] * (acc_t)scale);\n }\n }\n }\n }\n // load data from global memory\n\n // convert input_t to acc_t\n // TODO : remove this, input is already acc_t type in register\n acc_t elements[WARP_BATCH][WARP_ITERATIONS];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n elements[i][it] = elements_input[i][it];\n }\n }\n\n constexpr uint32_t FULL_MASK = 0xffffffff;\n\n // compute local max_value\n\n // take the max_value of the first element to avoid one max call\n acc_t max_value[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = elements[i][0];\n }\n\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];\n }\n }\n\n// reduction max_value\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n float val[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);\n }\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];\n }\n }\n\n // compute local sum\n acc_t sum[WARP_BATCH]{0.0f};\n\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n // elements[i][it] = expf(elements[i][it] - max_value[i]);\n elements[i][it] = std::exp(elements[i][it] - max_value[i]);\n sum[i] += elements[i][it];\n }\n }\n\n// reduction sum\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);\n }\n }\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it++) {\n elements[i][it] = elements[i][it] / sum[i];\n grad_reg[i][it] = grad_reg[i][it] * elements[i][it];\n }\n }\n\n acc_t grad_sum[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n grad_sum[i] = grad_reg[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n grad_sum[i] += grad_reg[i][it];\n }\n }\n warp_reduce_sum(grad_sum);\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // compute gradients\n output_t grad_input_reg[ELEMENTS_PER_LDG_STG];\n#pragma unroll\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; element++) {\n if (is_log_softmax) {\n grad_input_reg[element] = (grad_reg[i][it + element] - std::exp(elements[i][it + element]) * grad_sum[i]);\n } else {\n grad_input_reg[element] = (grad_reg[i][it + element] - elements[i][it + element] * grad_sum[i]);\n }\n }\n copy_vector(gradInput + i * element_count + it * WARP_SIZE, grad_input_reg);\n }\n }\n }\n}\n\ntemplate \nusing masked_scale_softmax_warp_backward_recompute_func = void (*)(output_t* gradInput, const input_t* grad,\n const input_t* softmax_input,\n const input_t* pad_mask, const uint8_t* mask,\n acc_t scale, int batch_size, int stride,\n int pad_batch_stride, int element_count);\n\ntemplate \nbool masked_scale_softmax_warp_backward_recompute_kernel(\n int element_count, int log2_elements, int& warp_size, int& batches_per_warp,\n masked_scale_softmax_warp_backward_recompute_func& kernel) {\n // determine size of a warp\n const int next_power_of_two = 1 << log2_elements;\n warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;\n\n // determine how many batches a warp should process.\n batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n bool flag_vec4 = (element_count % 4 == 0);\n switch (log2_elements) {\n case 0: // 1\n kernel = &masked_scale_softmax_warp_backward_recompute;\n break;\n case 1: // 2\n kernel = &masked_scale_softmax_warp_backward_recompute;\n break;\n case 2: // 4\n kernel = &masked_scale_softmax_warp_backward_recompute;\n break;\n case 3: // 8\n kernel = &masked_scale_softmax_warp_backward_recompute;\n break;\n case 4: // 16\n kernel = &masked_scale_softmax_warp_backward_recompute;\n break;\n case 5: // 32\n kernel = &masked_scale_softmax_warp_backward_recompute;\n break;\n case 6: // 64\n kernel = &masked_scale_softmax_warp_backward_recompute;\n break;\n case 7: // 128\n kernel = &masked_scale_softmax_warp_backward_recompute;\n break;\n case 8: // 256\n if (flag_vec4)\n kernel = &masked_scale_softmax_warp_backward_recompute;\n else\n kernel = &masked_scale_softmax_warp_backward_recompute;\n break;\n case 9: // 512\n if (flag_vec4)\n kernel = &masked_scale_softmax_warp_backward_recompute;\n else\n kernel = &masked_scale_softmax_warp_backward_recompute;\n break;\n case 10: // 1024\n if (flag_vec4)\n kernel = &masked_scale_softmax_warp_backward_recompute;\n else\n kernel = &masked_scale_softmax_warp_backward_recompute;\n break;\n case 11: // 2048\n if (flag_vec4)\n kernel = &masked_scale_softmax_warp_backward_recompute;\n else\n kernel = &masked_scale_softmax_warp_backward_recompute;\n break;\n default:\n return false;\n }\n return true;\n}\n\ntemplate \nbool dispatch_masked_scale_softmax_backward_recompute(output_t* grad_input, const input_t* grad,\n const input_t* softmax_input, const input_t* pad_mask,\n const uint8_t* mask, acc_t scale, int softmax_elements,\n int softmax_elements_stride, int pad_batch_stride,\n int batch_count, cudaStream_t streamid) {\n if (softmax_elements == 0) {\n return true;\n } else if (softmax_elements <= 2048) {\n // compute function index. there's a function for each power of two size up\n // to 1024.\n int log2_elements = 0;\n while ((1 << log2_elements) < softmax_elements) ++log2_elements;\n\n masked_scale_softmax_warp_backward_recompute_func kernel;\n int warp_size, batches_per_warp;\n if (!masked_scale_softmax_warp_backward_recompute_kernel(\n softmax_elements, log2_elements, warp_size, batches_per_warp, kernel)) {\n return false;\n }\n\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n // compute warps per block.\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n\n // compute launch size\n dim3 threads(warp_size, warps_per_block, 1);\n\n // launch\n kernel<<>>(grad_input, grad, softmax_input, pad_mask, mask, scale, batch_count,\n softmax_elements_stride, pad_batch_stride, softmax_elements);\n return true;\n }\n return false;\n}\n\ntemplate \nvoid dispatch_masked_scale_softmax_backward_stream(output_t* grad_input, const input_t* grad, const input_t* output,\n const uint8_t* mask, acc_t scale, int softmax_elements,\n int softmax_elements_stride, int batch_count,\n cudaStream_t streamid) {\n TORCH_INTERNAL_ASSERT(softmax_elements >= 0 && softmax_elements <= 1024);\n if (softmax_elements == 0) {\n return;\n } else {\n int log2_elements = log2_ceil_native(softmax_elements);\n const int next_power_of_two = 1 << log2_elements;\n // This value must match the WARP_SIZE constexpr value computed inside\n // softmax_warp_backward.\n int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;\n // This value must match the WARP_BATCH constexpr value computed inside\n // softmax_warp_backward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n masked_scale_softmax_warp_backward\n <<>>(grad_input, grad, output, mask, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 1: // 2\n masked_scale_softmax_warp_backward\n <<>>(grad_input, grad, output, mask, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 2: // 4\n masked_scale_softmax_warp_backward\n <<>>(grad_input, grad, output, mask, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 3: // 8\n masked_scale_softmax_warp_backward\n <<>>(grad_input, grad, output, mask, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 4: // 16\n masked_scale_softmax_warp_backward\n <<>>(grad_input, grad, output, mask, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 5: // 32\n masked_scale_softmax_warp_backward\n <<>>(grad_input, grad, output, mask, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 6: // 64\n masked_scale_softmax_warp_backward\n <<>>(grad_input, grad, output, mask, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 7: // 128\n masked_scale_softmax_warp_backward\n <<>>(grad_input, grad, output, mask, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 8: // 256\n masked_scale_softmax_warp_backward\n <<>>(grad_input, grad, output, mask, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 9: // 512\n masked_scale_softmax_warp_backward\n <<>>(grad_input, grad, output, mask, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 10: // 1024\n masked_scale_softmax_warp_backward\n <<>>(grad_input, grad, output, mask, scale, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n default:\n break;\n }\n }\n}\n\n// elementwise multiplication called in at::softmax_backward_data is fused\n// inside softmax dgrad kernel as a result of fusion, intermediate\n// multiplication result is stored in fp32 in registers, instead of fp16\ntemplate \n__global__ void softmax_warp_backward_fused_native(output_t* gradInput, const input_t* grad, const input_t* output,\n int batch_size, int stride, int element_count) {\n // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and\n // warp_size of method warp_softmax_backward_kernel.\n constexpr int next_power_of_two = 1 << log2_elements;\n constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;\n constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;\n constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;\n\n int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n\n // batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the\n // batch\n int local_idx = threadIdx.x % WARP_SIZE;\n\n // the first element to process by the current thread\n int thread_offset = first_batch * stride + local_idx;\n grad += thread_offset;\n output += thread_offset;\n gradInput += thread_offset;\n\n // The nested loops over WARP_BATCH and then WARP_ITERATIONS can be simplified\n // to one loop, but I think doing so would obfuscate the logic of the\n // algorithm, thus I chose to keep the nested loops. This should have no\n // impact on performance because the loops are unrolled anyway.\n\n // load data from global memory\n acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS];\n acc_t output_reg[WARP_BATCH][WARP_ITERATIONS];\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n int element_index = local_idx + it * WARP_SIZE;\n if (element_index < batch_element_count) {\n grad_reg[i][it] = grad[i * element_count + it * WARP_SIZE] * output[i * element_count + it * WARP_SIZE];\n output_reg[i][it] = output[i * element_count + it * WARP_SIZE];\n } else {\n grad_reg[i][it] = acc_t(0);\n output_reg[i][it] = acc_t(0);\n }\n }\n }\n\n acc_t sum[WARP_BATCH];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] = grad_reg[i][0]; //* output_reg[i][0];\n#pragma unroll\n for (int it = 1; it < WARP_ITERATIONS; ++it) {\n sum[i] += grad_reg[i][it]; // * output_reg[i][it];\n }\n }\n warp_reduce_sum(sum);\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n int element_index = local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // compute gradients\n if (is_log_softmax) {\n gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - std::exp(output_reg[i][it]) * sum[i]);\n } else {\n gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - output_reg[i][it] * sum[i]);\n }\n }\n }\n }\n}\n\ntemplate \nvoid dispatch_softmax_backward_fused_native(output_t* grad_input, const input_t* grad, const input_t* output,\n int softmax_elements, int softmax_elements_stride, int batch_count) {\n TORCH_INTERNAL_ASSERT(softmax_elements >= 0 && softmax_elements <= 1024);\n if (softmax_elements == 0) {\n return;\n } else {\n int log2_elements = log2_ceil_native(softmax_elements);\n const int next_power_of_two = 1 << log2_elements;\n\n // This value must match the WARP_SIZE constexpr value computed inside\n // softmax_warp_backward.\n int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;\n\n // This value must match the WARP_BATCH constexpr value computed inside\n // softmax_warp_backward.\n int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n\n int warps_per_block = (threads_per_block / warp_size);\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n // Launch code would be more elegant if C++ supported FOR CONSTEXPR\n switch (log2_elements) {\n case 0: // 1\n softmax_warp_backward_fused_native\n <<>>(grad_input, grad, output, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 1: // 2\n softmax_warp_backward_fused_native\n <<>>(grad_input, grad, output, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 2: // 4\n softmax_warp_backward_fused_native\n <<>>(grad_input, grad, output, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 3: // 8\n softmax_warp_backward_fused_native\n <<>>(grad_input, grad, output, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 4: // 16\n softmax_warp_backward_fused_native\n <<>>(grad_input, grad, output, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 5: // 32\n softmax_warp_backward_fused_native\n <<>>(grad_input, grad, output, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 6: // 64\n softmax_warp_backward_fused_native\n <<>>(grad_input, grad, output, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 7: // 128\n softmax_warp_backward_fused_native\n <<>>(grad_input, grad, output, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 8: // 256\n softmax_warp_backward_fused_native\n <<>>(grad_input, grad, output, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 9: // 512\n softmax_warp_backward_fused_native\n <<>>(grad_input, grad, output, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n case 10: // 1024\n softmax_warp_backward_fused_native\n <<>>(grad_input, grad, output, batch_count,\n softmax_elements_stride, softmax_elements);\n break;\n default:\n break;\n }\n }\n}\n\n////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n// Warp softmax backward\n////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////\n\ntemplate \n__global__ void softmax_warp_backward(__half* gradInput, const __half* grad, const __half* output, int batch_size,\n int stride, int element_count) {\n int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n\n // batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the\n // batch\n int local_idx = threadIdx.x;\n\n // the first element to process by the current thread\n int thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n grad += thread_offset;\n output += thread_offset;\n gradInput += thread_offset;\n\n // load data from global memory\n input_t grad_reg_input[WARP_BATCH][WARP_ITERATIONS] = {0.0f};\n input_t output_reg_input[WARP_BATCH][WARP_ITERATIONS] = {0.0f};\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < batch_element_count) {\n copy_vector(&grad_reg_input[i][it], grad + i * element_count + it * WARP_SIZE);\n copy_vector(&output_reg_input[i][it],\n output + i * element_count + it * WARP_SIZE);\n }\n }\n }\n\n // convert half to floating point\n acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS];\n acc_t output_reg[WARP_BATCH][WARP_ITERATIONS];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n grad_reg[i][it] = grad_reg_input[i][it];\n output_reg[i][it] = output_reg_input[i][it];\n }\n }\n\n // compute thread local sum\n acc_t sum[WARP_BATCH] = {0};\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] += grad_reg[i][it] * output_reg[i][it];\n }\n }\n\n // reduction sum\n constexpr uint32_t FULL_MASK = 0xffffffff;\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);\n }\n }\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // compute gradients\n output_t out[ELEMENTS_PER_LDG_STG];\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = (output_reg[i][it + element] * (grad_reg[i][it + element] - sum[i]));\n }\n // store them in global memory\n copy_vector(gradInput + i * element_count + it * WARP_SIZE, out);\n }\n }\n }\n}\n\n// WARP_BATCH number of batches.\n// WARP_ITERATOINS The number of iterations required for one warp to iterate\n// over all data. WARP_SIZE number of elements working on a single batch, has to\n// be a power of two. ELEMENTS_PER_LDG_STG has to be 1.\ntemplate \nusing softmax_backward_func = void (*)(output_t* gradInput, const input_t* grad, const input_t* output, int batch_size,\n int stride, int element_count);\n\ntemplate \nbool warp_softmax_backward_kernel(int log2_elements, int& warp_size, int& batches_per_warp,\n softmax_backward_func& kernel) {\n // determine size of a warp\n const int next_power_of_two = 1 << log2_elements;\n warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;\n\n // determine how many batches a warp should process.\n batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n switch (log2_elements) {\n case 0: // 1\n kernel = &softmax_warp_backward;\n break;\n case 1: // 2\n kernel = &softmax_warp_backward;\n break;\n case 2: // 4\n kernel = &softmax_warp_backward;\n break;\n case 3: // 8\n kernel = &softmax_warp_backward;\n break;\n case 4: // 16\n kernel = &softmax_warp_backward;\n break;\n case 5: // 32\n kernel = &softmax_warp_backward;\n break;\n case 6: // 64\n kernel = &softmax_warp_backward;\n break;\n case 7: // 128\n kernel = &softmax_warp_backward;\n break;\n case 8: // 256\n kernel = &softmax_warp_backward;\n break;\n case 9: // 512\n kernel = &softmax_warp_backward;\n break;\n case 10: // 1024\n kernel = &softmax_warp_backward;\n break;\n default:\n return false;\n }\n return true;\n}\n\ntemplate \nbool dispatch_softmax_backward(output_t* grad_input, const input_t* grad, const input_t* output, int softmax_elements,\n int softmax_elements_stride, int batch_count) {\n if (softmax_elements == 0) {\n return true;\n } else if (softmax_elements <= 1024) {\n // compute function index. there's a function for each power of two size up\n // to 1024.\n int log2_elements = 0;\n while ((1 << log2_elements) < softmax_elements) ++log2_elements;\n\n softmax_backward_func kernel;\n int warp_size, batches_per_warp;\n if (!warp_softmax_backward_kernel(log2_elements, warp_size, batches_per_warp, kernel)) {\n return false;\n }\n\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n\n // compute warps per block.\n int warps_per_block = (threads_per_block / warp_size);\n\n // compute launch size\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n\n // launch\n kernel<<>>(grad_input, grad, output, batch_count,\n softmax_elements_stride, softmax_elements);\n return true;\n }\n return false;\n}\n\ntemplate \nbool dispatch_softmax_backward_stream(output_t* grad_input, const input_t* grad, const input_t* output,\n int softmax_elements, int softmax_elements_stride, int batch_count,\n cudaStream_t streamid) {\n if (softmax_elements == 0) {\n return true;\n } else if (softmax_elements <= 1024) {\n // compute function index. there's a function for each power of two size up\n // to 1024.\n int log2_elements = 0;\n while ((1 << log2_elements) < softmax_elements) ++log2_elements;\n softmax_backward_func kernel;\n int warp_size, batches_per_warp;\n if (!warp_softmax_backward_kernel(log2_elements, warp_size, batches_per_warp, kernel)) {\n return false;\n }\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n // compute warps per block.\n int warps_per_block = (threads_per_block / warp_size);\n // compute launch size\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n // launch\n kernel<<>>(grad_input, grad, output, batch_count, softmax_elements_stride,\n softmax_elements);\n return true;\n }\n return false;\n}\n\ntemplate \n__global__ void masked_softmax_warp_backward(__half* gradInput, const __half* grad, const __half* output,\n const uint8_t* pad_mask, int batch_size, int stride, int element_count,\n int pad_batch_stride) {\n int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;\n\n // batch_size might not be a multiple of WARP_BATCH. Check how\n // many batches have to computed within this WARP.\n int local_batches = batch_size - first_batch;\n if (local_batches > WARP_BATCH) local_batches = WARP_BATCH;\n\n // there might be multiple batches per warp. compute the index within the\n // batch\n int local_idx = threadIdx.x;\n\n // the first element to process by the current thread\n int thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;\n grad += thread_offset;\n output += thread_offset;\n gradInput += thread_offset;\n\n // load data from global memory\n input_t grad_reg_input[WARP_BATCH][WARP_ITERATIONS] = {0.0f};\n input_t output_reg_input[WARP_BATCH][WARP_ITERATIONS] = {0.0f};\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n int batch_element_count = (i >= local_batches) ? 0 : element_count;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < batch_element_count) {\n copy_vector(&grad_reg_input[i][it], grad + i * element_count + it * WARP_SIZE);\n copy_vector(&output_reg_input[i][it],\n output + i * element_count + it * WARP_SIZE);\n }\n }\n }\n\n // convert half to floating point\n acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS];\n acc_t output_reg[WARP_BATCH][WARP_ITERATIONS];\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n grad_reg[i][it] = grad_reg_input[i][it];\n output_reg[i][it] = output_reg_input[i][it];\n }\n }\n\n // compute thread local sum\n acc_t sum[WARP_BATCH] = {0};\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; ++it) {\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] += grad_reg[i][it] * output_reg[i][it];\n }\n }\n\n // reduction sum\n constexpr uint32_t FULL_MASK = 0xffffffff;\n#pragma unroll\n for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);\n }\n }\n\n// store result\n#pragma unroll\n for (int i = 0; i < WARP_BATCH; ++i) {\n if (i >= local_batches) break;\n int pad_thread_offset = ((first_batch + i) / pad_batch_stride) * stride + ELEMENTS_PER_LDG_STG * local_idx;\n const uint8_t* curr_mask = pad_mask + pad_thread_offset;\n#pragma unroll\n for (int it = 0; it < WARP_ITERATIONS; it += ELEMENTS_PER_LDG_STG) {\n int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;\n if (element_index < element_count) {\n // compute gradients\n output_t out[ELEMENTS_PER_LDG_STG];\n for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {\n out[element] = (output_reg[i][it + element] * (grad_reg[i][it + element] - sum[i]));\n }\n // store them in global memory\n int itr_jmp = it * WARP_SIZE;\n int itr_idx = i * element_count + itr_jmp;\n // It is kind of unfortunate this has to be here to zero something out\n // that is close to zero in the first place\n apply_mask(&out[0], 0.0, curr_mask + itr_jmp);\n copy_vector(gradInput + itr_idx, out);\n }\n }\n }\n}\n\n// WARP_BATCH number of batches.\n// WARP_ITERATOINS The number of iterations required for one warp to iterate\n// over all data. WARP_SIZE number of elements working on a single batch, has to\n// be a power of two. ELEMENTS_PER_LDG_STG has to be 1.\ntemplate \nusing masked_softmax_backward_func = void (*)(output_t* gradInput, const input_t* grad, const input_t* output,\n const uint8_t* pad_mask, int batch_size, int stride, int element_count,\n int pad_batch_stride);\n\ntemplate \nbool warp_masked_softmax_backward_kernel(int log2_elements, int& warp_size, int& batches_per_warp,\n masked_softmax_backward_func& kernel) {\n // determine size of a warp\n const int next_power_of_two = 1 << log2_elements;\n warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;\n\n // determine how many batches a warp should process.\n batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;\n\n switch (log2_elements) {\n case 0: // 1\n kernel = &masked_softmax_warp_backward;\n break;\n case 1: // 2\n kernel = &masked_softmax_warp_backward;\n break;\n case 2: // 4\n kernel = &masked_softmax_warp_backward;\n break;\n case 3: // 8\n kernel = &masked_softmax_warp_backward;\n break;\n case 4: // 16\n kernel = &masked_softmax_warp_backward;\n break;\n case 5: // 32\n kernel = &masked_softmax_warp_backward;\n break;\n case 6: // 64\n kernel = &masked_softmax_warp_backward;\n break;\n case 7: // 128\n kernel = &masked_softmax_warp_backward;\n break;\n case 8: // 256\n kernel = &masked_softmax_warp_backward;\n break;\n case 9: // 512\n kernel = &masked_softmax_warp_backward;\n break;\n case 10: // 1024\n kernel = &masked_softmax_warp_backward;\n break;\n default:\n return false;\n }\n return true;\n}\n\ntemplate \nbool dispatch_masked_softmax_backward(output_t* grad_input, const input_t* grad, const input_t* output,\n const uint8_t* pad_mask, int softmax_elements, int softmax_elements_stride,\n int batch_count, int pad_batch_stride) {\n if (softmax_elements == 0) {\n return true;\n } else if (softmax_elements <= 1024) {\n // compute function index. there's a function for each power of two size up\n // to 1024.\n int log2_elements = 0;\n while ((1 << log2_elements) < softmax_elements) ++log2_elements;\n\n masked_softmax_backward_func kernel;\n int warp_size, batches_per_warp;\n if (!warp_masked_softmax_backward_kernel(log2_elements, warp_size, batches_per_warp,\n kernel)) {\n return false;\n }\n\n // use 128 threads per block to maximize gpu utilization\n constexpr int threads_per_block = 128;\n\n // compute warps per block.\n int warps_per_block = (threads_per_block / warp_size);\n\n // compute launch size\n int batches_per_block = warps_per_block * batches_per_warp;\n int blocks = (batch_count + batches_per_block - 1) / batches_per_block;\n dim3 threads(warp_size, warps_per_block, 1);\n\n // launch\n kernel<<>>(\n grad_input, grad, output, pad_mask, batch_count, softmax_elements_stride, softmax_elements, pad_batch_stride);\n return true;\n }\n return false;\n}\n} // namespace\n" }, { "path": "apex/contrib/csrc/multihead_attn/strided_batched_gemm.cuh", "content": "#pragma once\n#include \n#include \n#include \n#include \n\n#include \n#include \n\n// #include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace {\ncublasOperation_t convertTransToCublasOperation(char trans) {\n if (trans == 't')\n return CUBLAS_OP_T;\n else if (trans == 'n')\n return CUBLAS_OP_N;\n else if (trans == 'c')\n return CUBLAS_OP_C;\n else {\n TORCH_CHECK(false, \"trans must be one of: t, n, c\");\n return CUBLAS_OP_T;\n }\n}\n\nvoid CublasStridedBatchedGemm(char transa, char transb, long m, long n, long k, float alpha, const half* a, long lda,\n long strideA, const half* b, long ldb, long strideB, float beta, half* c, long ldc,\n long strideC, long batchCount, cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP) {\n cublasOperation_t opa = convertTransToCublasOperation(transa);\n cublasOperation_t opb = convertTransToCublasOperation(transb);\n\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();\n cublasSetStream(handle, stream);\n float fAlpha = alpha;\n float fBeta = beta;\n TORCH_CUDABLAS_CHECK(cublasGemmStridedBatchedEx(\n handle, opa, opb, (int)m, (int)n, (int)k, (void*)&fAlpha, a, CUDA_R_16F, (int)lda, strideA, b, CUDA_R_16F,\n (int)ldb, strideB, (void*)&fBeta, c, CUDA_R_16F, (int)ldc, strideC, (int)batchCount, CUDA_R_32F, algo));\n}\n\n} // namespace\n\n// TODO(mkozuki): Make use of the int template parameters or discard them.\ntemplate \nvoid CutlassGemm_FP32Accum(cudaStream_t stream, long m, long n, long k, float alpha, const half* a, long lda,\n long long int batch_stride_A, const half* b, long ldb, long long int batch_stride_B,\n float beta, half* c, long ldc, long long int batch_stride_C, long batch_count) {\n using Gemm = cutlass::gemm::device::GemmBatched<\n /* Element type of A matrix */ half, /* Layout of A matrix */ LayoutA,\n /* Element type of B matrix */ half, /* Layout of B matrix */ LayoutB,\n /* Element type of C matrix */ half, /* Layout of C matrix */ cutlass::layout::ColumnMajor,\n /* Element Accumulator*/ float>;\n Gemm gemm_op;\n cutlass::Status status = gemm_op({{static_cast(m), static_cast(n), static_cast(k)},\n {a, lda},\n batch_stride_A,\n {b, ldb},\n batch_stride_B,\n {c, ldc},\n batch_stride_C,\n {c, ldc},\n batch_stride_C,\n {alpha, beta},\n static_cast(batch_count)},\n nullptr, stream);\n C10_CUDA_CHECK(status != cutlass::Status::kSuccess ? cudaErrorUnknown : cudaSuccess);\n}\n\nnamespace {\nvoid gemm_switch_fp32accum(char transa, char transb, long m, long n, long k, float alpha, const half* a, long lda,\n long strideA, const half* b, long ldb, long strideB, float beta, half* c, long ldc,\n long strideC, long batchCount) {\n auto stream = c10::cuda::getCurrentCUDAStream();\n // printf(\"GEMM -> %c%c M: %i N: %i K: %i Alpha: %f Beta: %f\\n\", (transa ==\n // 't' ? 'T' : 'N'), (transb =='t' ? 'T' : 'N'), m, n, k, alpha, beta);\n if ((transa == 't') && (transb == 'n')) {\n if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x7)) {\n CublasStridedBatchedGemm(transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC,\n batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP);\n } else if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else {\n CublasStridedBatchedGemm(transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC,\n batchCount);\n }\n } else if ((transa == 'n') && (transb == 'n')) {\n if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x7)) {\n CublasStridedBatchedGemm(transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC,\n batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP);\n } else if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else {\n CublasStridedBatchedGemm(transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC,\n batchCount);\n }\n } else if ((transa == 'n') && (transb == 't')) {\n if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x7)) {\n CublasStridedBatchedGemm(transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC,\n batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP);\n } else if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x7)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x3)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x1)) {\n CutlassGemm_FP32Accum(\n stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);\n } else {\n CublasStridedBatchedGemm(transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC,\n batchCount);\n }\n } else {\n TORCH_CHECK(false, \"TransA and TransB are invalid\");\n }\n}\n\nvoid adjustLdLevel3(char transa, char transb, int64_t m, int64_t n, int64_t k, int64_t* lda, int64_t* ldb,\n int64_t* ldc) {\n int transa_ = ((transa == 't') || (transa == 'T'));\n int transb_ = ((transb == 't') || (transb == 'T'));\n\n // Note: leading dimensions generally are checked that they are > 0 and at\n // least as big the result requires (even if the value won't be used).\n if (n <= 1) *ldc = std::max(m, 1);\n\n if (transa_) {\n if (m <= 1) *lda = std::max(k, 1);\n } else {\n if (k <= 1) *lda = std::max(m, 1);\n }\n\n if (transb_) {\n if (k <= 1) *ldb = std::max(n, 1);\n } else {\n if (n <= 1) *ldb = std::max(k, 1);\n }\n}\n\nvoid HgemmStridedBatched(char transa, char transb, long m, long n, long k, float alpha, const half* a, long lda,\n long strideA, const half* b, long ldb, long strideB, float beta, half* c, long ldc,\n long strideC, long batchCount) {\n if ((m >= INT_MAX) || (n >= INT_MAX) || (k >= INT_MAX) || (lda >= INT_MAX) || (ldb >= INT_MAX) || (ldc >= INT_MAX) ||\n (batchCount >= INT_MAX))\n\n {\n TORCH_CHECK(false,\n \"Cublas_SgemmStridedBatched only supports m, n, k, lda, ldb, ldc, \"\n \"batchCount\"\n \"with the bound [val] <= %d\",\n INT_MAX);\n }\n\n adjustLdLevel3(transa, transb, m, n, k, &lda, &ldb, &ldc);\n\n gemm_switch_fp32accum(transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC,\n batchCount);\n}\n\n} // namespace\n" }, { "path": "apex/contrib/csrc/nccl_allocator/NCCLAllocator.cpp", "content": "#include \n#include \n#include \n#include \n#include \n\n#define NCCL_CHECK(cmd) \\\n do { \\\n ncclResult_t result = cmd; \\\n if (result != ncclSuccess) { \\\n std::string err = \"NCCL error in: \" + std::string(__FILE__) + \":\" + std::to_string(__LINE__) + \", \" + \\\n std::string(ncclGetErrorString(result)); \\\n TORCH_CHECK(false, err); \\\n } \\\n } while (0)\n\nvoid* nccl_alloc_plug(size_t size, int device, void* stream) {\n void* ptr;\n NCCL_CHECK(ncclMemAlloc(&ptr, size));\n return ptr;\n}\n\nvoid nccl_free_plug(void* ptr, std::size_t size, int device, void* stream) { NCCL_CHECK(ncclMemFree(ptr)); }\n\nstd::shared_ptr nccl_allocator;\n\nvoid maybe_init() {\n if (!nccl_allocator) {\n nccl_allocator =\n std::make_shared(nccl_alloc_plug, nccl_free_plug);\n }\n}\n\nstd::shared_ptr get_nccl_allocator() {\n maybe_init();\n return nccl_allocator;\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"get_nccl_allocator\", []() { return get_nccl_allocator(); });\n};\n" }, { "path": "apex/contrib/csrc/nccl_p2p/nccl_p2p.cpp", "content": "/**\n * Copyright (c) 2018-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#include \"nccl_p2p_cuda.cuh\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"get_unique_nccl_id\", &apex::contrib::nccl_p2p::get_unique_nccl_id, \"get_unique_nccl_id\",\n py::call_guard());\n m.def(\"init_nccl_comm\", &apex::contrib::nccl_p2p::init_nccl_comm, \"init_nccl_comm\",\n py::call_guard());\n m.def(\"left_right_halo_exchange_inplace\", &apex::contrib::nccl_p2p::left_right_halo_exchange_inplace,\n \"left_right_halo_exchange_inplace\", py::call_guard());\n m.def(\"left_right_halo_exchange\", &apex::contrib::nccl_p2p::left_right_halo_exchange, \"left_right_halo_exchange\",\n py::call_guard());\n m.def(\"add_delay\", &apex::contrib::nccl_p2p::add_delay, \"add_delay\", py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/nccl_p2p/nccl_p2p_cuda.cu", "content": "#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n#include \"nccl.h\"\n\n/*\n * This file implements a crude but effective mechanism for copying data between tenors owned by different ranks\n * on the same machine using cudaMemcpyAsync peer-to-peer transfers.\n */\n\nnamespace {\n\n__global__ void AddDelay_kernel(const int delay, int* counter) {\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n // waste time while doing something compiler can't predict, thus preventing it from optimizing away this code.\n int new_counter = 0;\n double elapsed = 0;\n clock_t start = clock();\n do {\n clock_t now = clock();\n elapsed = (double)(now - start) * 1e9 / CLOCKS_PER_SEC;\n ++new_counter;\n } while (elapsed < (double)delay);\n *counter = new_counter;\n }\n}\n\nclass NcclCommWrapper {\n private:\n ncclComm_t comm;\n int rank, world_size;\n\n ncclDataType_t get_nccl_type(at::Tensor input) {\n switch (input.scalar_type()) {\n case at::ScalarType::Half:\n return ncclFloat16;\n case at::ScalarType::Float:\n return ncclFloat32;\n case at::ScalarType::Double:\n return ncclFloat64;\n case at::ScalarType::Byte:\n return ncclUint8;\n case at::ScalarType::Char:\n return ncclInt8;\n case at::ScalarType::Int:\n return ncclInt32;\n case at::ScalarType::Long:\n return ncclInt64;\n case at::ScalarType::BFloat16:\n return ncclBfloat16;\n default:\n assert(false);\n }\n }\n\n public:\n NcclCommWrapper() {\n memset(&comm, 0, sizeof(ncclComm_t));\n rank = 0;\n world_size = 0;\n }\n NcclCommWrapper(ncclUniqueId id, int my_rank, int num_ranks) {\n ncclCommInitRank(&comm, num_ranks, id, my_rank);\n rank = my_rank;\n world_size = num_ranks;\n }\n\n ~NcclCommWrapper() {\n printf(\"ncclCommDestroy()\\n\");\n ncclCommDestroy(comm);\n }\n\n void left_right_halo_exchange_inplace(int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo, at::Tensor left_input_halo,\n at::Tensor right_input_halo) {\n auto stream = at::cuda::getCurrentCUDAStream();\n ncclGroupStart();\n ncclDataType_t ncclType = get_nccl_type(left_output_halo);\n bool left_zero = (left_rank < 0);\n bool right_zero = (right_rank < 0);\n size_t left_n = torch::numel(left_output_halo);\n size_t right_n = torch::numel(right_output_halo);\n assert(left_n > 0 && left_n == right_n);\n if (left_zero) {\n left_input_halo.zero_();\n } else {\n AT_DISPATCH_ALL_TYPES_AND3(\n at::ScalarType::Bool, at::ScalarType::BFloat16, at::ScalarType::Half, left_output_halo.scalar_type(),\n \"left_halo_exch\", [&]() {\n // send left (to my_rank - 1)\n ncclSend(left_output_halo.data_ptr(), left_n, ncclType, left_rank, comm, stream);\n // receive left (from my_rank - 1)\n ncclRecv(left_input_halo.data_ptr(), right_n, ncclType, left_rank, comm, stream);\n });\n }\n if (right_zero) {\n right_input_halo.zero_();\n } else {\n AT_DISPATCH_ALL_TYPES_AND3(\n at::ScalarType::Bool, at::ScalarType::BFloat16, at::ScalarType::Half, right_output_halo.scalar_type(),\n \"right_halo_exch\", [&]() {\n // send right (to my_rank + 1 )\n ncclSend(right_output_halo.data_ptr(), right_n, ncclType, right_rank, comm, stream);\n // receive right (from my_rank + 1)\n ncclRecv(right_input_halo.data_ptr(), left_n, ncclType, right_rank, comm, stream);\n });\n }\n ncclGroupEnd();\n }\n\n std::vector left_right_halo_exchange(int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo) {\n // after halo exchange:\n // left_output_halo of rank+1 ends up in right_input_halo of rank\n // right_output_halo of rank-1 ends up in left_input_halo of rank\n auto right_input_halo = torch::empty_like(left_output_halo);\n auto left_input_halo = torch::empty_like(right_output_halo);\n left_right_halo_exchange_inplace(left_rank, right_rank, left_output_halo, right_output_halo, left_input_halo,\n right_input_halo);\n return {left_input_halo, right_input_halo};\n }\n};\n\nclass ManagedObjects {\n public:\n ManagedObjects() {}\n ~ManagedObjects() {\n for (auto it = _nccl_comms.begin(); it != _nccl_comms.end(); ++it) {\n delete *it;\n }\n }\n\n int add_comm(NcclCommWrapper* comm) {\n int handle = _nccl_comms.size();\n _nccl_comms.push_back(comm);\n return handle;\n }\n\n NcclCommWrapper& get_comm(int handle) {\n assert(handle >= 0 && handle < _nccl_comms.size());\n return *_nccl_comms[handle];\n }\n\n private:\n std::vector _nccl_comms;\n};\nclass ManagedObjects mo;\n\n} // end anonymous namespace\n\nnamespace apex {\nnamespace contrib {\nnamespace nccl_p2p {\n\nat::Tensor get_unique_nccl_id(int n) {\n ncclUniqueId id;\n ncclGetUniqueId(&id);\n auto id_tensor = torch::empty({n, (int)sizeof(ncclUniqueId)},\n torch::dtype(torch::kUInt8).device(torch::kCPU).requires_grad(false));\n auto id_ptr = id_tensor.data_ptr();\n size_t offset = 0;\n for (int i = 0; i < n; ++i) {\n ncclUniqueId id;\n ncclGetUniqueId(&id);\n memcpy(id_ptr + offset, &id, sizeof(ncclUniqueId));\n offset += sizeof(ncclUniqueId);\n }\n return id_tensor;\n}\n\nint init_nccl_comm(at::Tensor unique_nccl_id, int my_rank, int num_ranks) {\n ncclUniqueId id;\n auto unique_nccl_id_ptr = unique_nccl_id.data_ptr();\n memcpy(&id, unique_nccl_id_ptr, sizeof(ncclUniqueId));\n NcclCommWrapper* comm = new NcclCommWrapper(id, my_rank, num_ranks);\n int handle = mo.add_comm(comm);\n comm = 0L;\n return handle;\n}\n\nvoid left_right_halo_exchange_inplace(int handle, int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo, at::Tensor left_input_halo,\n at::Tensor right_input_halo) {\n class NcclCommWrapper& communicator = mo.get_comm(handle);\n return communicator.left_right_halo_exchange_inplace(left_rank, right_rank, left_output_halo, right_output_halo,\n left_input_halo, right_input_halo);\n}\n\nstd::vector left_right_halo_exchange(int handle, int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo) {\n class NcclCommWrapper& communicator = mo.get_comm(handle);\n return communicator.left_right_halo_exchange(left_rank, right_rank, left_output_halo, right_output_halo);\n}\n\nvoid add_delay(int delay) {\n auto stream = at::cuda::getCurrentCUDAStream();\n auto t = torch::empty({1}, torch::dtype(torch::kInt32).device(torch::kCUDA));\n AddDelay_kernel<<<1, 1, 0, stream>>>(delay, t.data_ptr());\n}\n\n} // namespace nccl_p2p\n} // namespace contrib\n} // namespace apex\n" }, { "path": "apex/contrib/csrc/nccl_p2p/nccl_p2p_cuda.cuh", "content": "/**\n * Copyright (c) 2018-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#pragma once\n#include \n#ifndef _nccl_p2p_h_\n#define _nccl_p2p_h_\n\nnamespace apex {\nnamespace contrib {\nnamespace nccl_p2p {\nat::Tensor get_unique_nccl_id(int n);\nint init_nccl_comm(at::Tensor unique_nccl_id, int my_rank, int num_ranks);\nvoid left_right_halo_exchange_inplace(int handle, int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo, at::Tensor left_input_halo,\n at::Tensor right_input_halo);\nstd::vector left_right_halo_exchange(int handle, int left_rank, int right_rank, at::Tensor left_output_halo,\n at::Tensor right_output_halo);\nvoid add_delay(int delay);\n} // namespace nccl_p2p\n} // namespace contrib\n} // namespace apex\n#endif\n" }, { "path": "apex/contrib/csrc/nccl_p2p/nccl_version.cpp", "content": "// Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n// This file is used to check the version of NCCL detected.\n#include \n\n#include \n\nstd::tuple get_nccl_version();\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def(\"get_nccl_version\", &get_nccl_version); }\n" }, { "path": "apex/contrib/csrc/nccl_p2p/nccl_version_check.cu", "content": "// Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.\n\n// This file is used to check the version of NCCL detected.\n#include \n\n#include \n\nstd::tuple get_nccl_version() { return {int(NCCL_MAJOR), int(NCCL_MINOR)}; }\n" }, { "path": "apex/contrib/csrc/optimizers/fused_adam_cuda.cpp", "content": "#include \n\n// CUDA forward declaration\nvoid fused_strided_check_finite(at::Tensor& overflow_flag, at::Tensor& p_copy, int stride, int clear_overflow_first);\n\nvoid fused_adam_cuda(at::Tensor& p, at::Tensor& p_copy, at::Tensor& m, at::Tensor& v, at::Tensor& g, float lr,\n float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction,\n float decay);\nvoid fused_reversible_adam_cuda(at::Tensor& p, at::Tensor& p_copy, at::Tensor& m, at::Tensor& v, at::Tensor& g,\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay);\nvoid fused_maybe_adam_undo_cuda(at::Tensor& overflow_flag, at::Tensor& p, at::Tensor& m, at::Tensor& v, at::Tensor& g,\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay);\n\nvoid fused_adam_cuda_mt(int chunk_size, at::Tensor overflow_flag, std::vector> tensor_lists,\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay);\n\nvoid maybe_cast_cuda(at::Tensor& overflow_flag, at::Tensor& p_in, at::Tensor& p_out);\nvoid maybe_cast_cuda_mt(int chunk_size, at::Tensor overflow_flag, std::vector> tensor_lists);\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n#define CHECK_INPUT(x) \\\n CHECK_CUDA(x); \\\n CHECK_CONTIGUOUS(x)\n\n// C++ interface\nvoid strided_check_finite(at::Tensor& overflow_flag, at::Tensor& p_copy, int stride, int clear_overflow_first) {\n CHECK_INPUT(p_copy);\n fused_strided_check_finite(overflow_flag, p_copy, stride, clear_overflow_first);\n}\nvoid adam(at::Tensor& p, at::Tensor& p_copy, at::Tensor& m, at::Tensor& v, at::Tensor& g, float lr, float beta1,\n float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay) {\n CHECK_INPUT(p);\n if (p_copy.numel() > 0) CHECK_INPUT(p_copy);\n CHECK_INPUT(m);\n CHECK_INPUT(v);\n CHECK_INPUT(g);\n int64_t num_elem = p.numel();\n TORCH_CHECK(m.numel() == num_elem, \"number of elements in m and p tensors should be equal\");\n TORCH_CHECK(v.numel() == num_elem, \"number of elements in v and p tensors should be equal\");\n TORCH_CHECK(g.numel() == num_elem, \"number of elements in g and p tensors should be equal\");\n TORCH_CHECK(p_copy.numel() == num_elem || p_copy.numel() == 0,\n \"number of elements in p_copy and p tensors should be equal, or p_copy should be empty\");\n\n fused_adam_cuda(p, p_copy, m, v, g, lr, beta1, beta2, eps, grad_scale, step, mode, bias_correction, decay);\n}\nvoid reversible_adam(at::Tensor& p, at::Tensor& p_copy, at::Tensor& m, at::Tensor& v, at::Tensor& g, float lr,\n float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction,\n float decay) {\n CHECK_INPUT(p);\n if (p_copy.numel() > 0) CHECK_INPUT(p_copy);\n CHECK_INPUT(m);\n CHECK_INPUT(v);\n CHECK_INPUT(g);\n int64_t num_elem = p.numel();\n TORCH_CHECK(m.numel() == num_elem, \"number of elements in m and p tensors should be equal\");\n TORCH_CHECK(v.numel() == num_elem, \"number of elements in v and p tensors should be equal\");\n TORCH_CHECK(g.numel() == num_elem, \"number of elements in g and p tensors should be equal\");\n TORCH_CHECK(p_copy.numel() == num_elem || p_copy.numel() == 0,\n \"number of elements in p_copy and p tensors should be equal, or p_copy should be empty\");\n\n fused_reversible_adam_cuda(p, p_copy, m, v, g, lr, beta1, beta2, eps, grad_scale, step, mode, bias_correction, decay);\n}\nvoid maybe_adam_undo(at::Tensor& overflow_flag, at::Tensor& p, at::Tensor& m, at::Tensor& v, at::Tensor& g, float lr,\n float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction,\n float decay) {\n CHECK_INPUT(p);\n CHECK_INPUT(m);\n CHECK_INPUT(v);\n CHECK_INPUT(g);\n int64_t num_elem = p.numel();\n TORCH_CHECK(m.numel() == num_elem, \"number of elements in m and p tensors should be equal\");\n TORCH_CHECK(v.numel() == num_elem, \"number of elements in v and p tensors should be equal\");\n TORCH_CHECK(g.numel() == num_elem, \"number of elements in g and p tensors should be equal\");\n\n fused_maybe_adam_undo_cuda(overflow_flag, p, m, v, g, lr, beta1, beta2, eps, grad_scale, step, mode, bias_correction,\n decay);\n}\nvoid maybe_cast(at::Tensor& overflow_flag, at::Tensor& p_in, at::Tensor& p_out) {\n CHECK_INPUT(p_in);\n CHECK_INPUT(p_out);\n int64_t num_elem = p_in.numel();\n TORCH_CHECK(p_out.numel() == num_elem, \"number of elements in p_in and p_out should be equal\");\n\n maybe_cast_cuda(overflow_flag, p_in, p_out);\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"strided_check_finite\", &strided_check_finite, \"Strided finite check.\",\n py::call_guard());\n m.def(\"adam\", &adam, \"Adam optimized CUDA implementation.\", py::call_guard());\n m.def(\"reversible_adam\", &reversible_adam, \"Reversible Adam optimized CUDA implementation.\",\n py::call_guard());\n m.def(\"adam_mt\", &fused_adam_cuda_mt, \"Multi tensor Adam optimized CUDA implementation.\",\n py::call_guard());\n m.def(\"maybe_adam_undo\", &maybe_adam_undo, \"Undo function for Adam optimized CUDA implementation.\",\n py::call_guard());\n m.def(\"maybe_cast\", &maybe_cast, \"Unpack byte tensor containing e5m2 floats.\",\n py::call_guard());\n m.def(\"maybe_cast_mt\", &maybe_cast_cuda_mt, \"Unpack byte tensor containing e5m2 floats.\",\n py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/optimizers/fused_adam_cuda_kernel.cu", "content": "#include \n#include \n#include \n\n#include \n\n#include \"ATen/ATen.h\"\n#include \"ATen/TensorUtils.h\"\n#include \"ATen/cuda/CUDAContext.h\"\n#include \"ATen/cuda/detail/IndexUtils.cuh\"\n// #include \"ATen/Type.h\"\n#include \"ATen/AccumulateType.h\"\n#include \"multi_tensor_apply.cuh\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\n#include \"type_shim.h\"\n\ntypedef enum {\n ADAM_MODE_0 = 0, // eps under square root\n ADAM_MODE_1 = 1 // eps outside square root\n} adamMode_t;\n\ntemplate \n__global__ void adam_cuda_kernel(T* __restrict__ p,\n GRAD_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed\n T* __restrict__ m, T* __restrict__ v, const GRAD_T* __restrict__ g, const float b1,\n const float b2, const float eps, const float grad_scale, const float step_size,\n const size_t tsize, adamMode_t mode, const float decay) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n for (int j = i; j < tsize; j += totThreads) {\n T scaled_grad = g[j] / grad_scale;\n m[j] = b1 * m[j] + (1 - b1) * scaled_grad;\n v[j] = b2 * v[j] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(v[j] + eps);\n else // Mode 1\n denom = sqrtf(v[j]) + eps;\n float update = (m[j] / denom) + (decay * p[j]);\n p[j] = p[j] - (step_size * update);\n if (p_copy != NULL) p_copy[j] = (GRAD_T)p[j];\n }\n}\n\ntemplate \nstruct AdamFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata& tl,\n const float b1, const float b2, const float eps, const float grad_scale,\n const float step_size, adamMode_t mode, const float decay) {\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n T* p = (T*)tl.addresses[0][tensor_loc];\n p += chunk_idx * chunk_size;\n T* m = (T*)tl.addresses[1][tensor_loc];\n m += chunk_idx * chunk_size;\n T* v = (T*)tl.addresses[2][tensor_loc];\n v += chunk_idx * chunk_size;\n GRAD_T* g = (GRAD_T*)tl.addresses[3][tensor_loc];\n g += chunk_idx * chunk_size;\n GRAD_T* p_copy = NULL;\n if (DEPTH == 5) {\n p_copy = (GRAD_T*)tl.addresses[4][tensor_loc];\n p_copy += chunk_idx * chunk_size;\n }\n\n n -= chunk_idx * chunk_size;\n\n T incoming_p[ILP];\n T incoming_m[ILP];\n T incoming_v[ILP];\n T incoming_g[ILP];\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(p) && is_aligned(m) && is_aligned(v) && is_aligned(g) &&\n is_aligned(p_copy)) {\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n GRAD_T tmp_g[ILP];\n load_store(incoming_p, p, 0, i_start);\n load_store(incoming_m, m, 0, i_start);\n load_store(incoming_v, v, 0, i_start);\n load_store(tmp_g, g, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n incoming_g[ii] = static_cast(tmp_g[ii]);\n T scaled_grad = incoming_g[ii] / grad_scale;\n incoming_m[ii] = b1 * incoming_m[ii] + (1 - b1) * scaled_grad;\n incoming_v[ii] = b2 * incoming_v[ii] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(incoming_v[ii] + eps);\n else // Mode 1\n denom = sqrtf(incoming_v[ii]) + eps;\n float update = (incoming_m[ii] / denom) + (decay * incoming_p[ii]);\n incoming_p[ii] = incoming_p[ii] - (step_size * update);\n if (DEPTH == 5) tmp_g[ii] = static_cast(incoming_p[ii]);\n }\n load_store(p, incoming_p, i_start, 0);\n load_store(m, incoming_m, i_start, 0);\n load_store(v, incoming_v, i_start, 0);\n if (DEPTH == 5) load_store(p_copy, tmp_g, i_start, 0);\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n incoming_p[ii] = 0;\n incoming_m[ii] = 0;\n incoming_v[ii] = 0;\n incoming_g[ii] = 0;\n\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n incoming_p[ii] = p[i];\n incoming_m[ii] = m[i];\n incoming_v[ii] = v[i];\n incoming_g[ii] = static_cast(g[i]);\n }\n }\n\n // note for clarification to future michael:\n // From a pure memory dependency perspective, there's likely no point unrolling\n // the write loop, since writes just fire off once their LDGs arrive.\n // Put another way, the STGs are dependent on the LDGs, but not on each other.\n // There is still compute ILP benefit from unrolling the loop though.\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = i_start + threadIdx.x + ii * blockDim.x;\n\n if (j < n && j < chunk_size) {\n T scaled_grad = incoming_g[ii] / grad_scale;\n m[j] = b1 * incoming_m[ii] + (1 - b1) * scaled_grad;\n v[j] = b2 * incoming_v[ii] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(v[j] + eps);\n else // Mode 1\n denom = sqrtf(v[j]) + eps;\n float update = (m[j] / denom) + (decay * incoming_p[ii]);\n p[j] = incoming_p[ii] - (step_size * update);\n if (DEPTH == 5) p_copy[j] = (GRAD_T)p[j];\n }\n }\n }\n }\n }\n};\n\nvoid fused_adam_cuda(at::Tensor& p, at::Tensor& p_copy, at::Tensor& m, at::Tensor& v, at::Tensor& g, float lr,\n float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction,\n float decay) {\n // using namespace at;\n\n // Get tensor size\n int tsize = p.numel();\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (g.scalar_type() == at::ScalarType::Half) {\n // all other values should be fp32 for half gradients\n TORCH_CHECK(p.scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dispatch is done on the gradient type\n using namespace at; // prevents \"toString is undefined\" errors\n DISPATCH_FLOAT_AND_HALF(g.scalar_type(), 0, \"adam_cuda_kernel\", using accscalar_t = at::acc_type;\n adam_cuda_kernel<<>>(\n p.data_ptr(), p_copy.numel() ? p_copy.data_ptr() : NULL,\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n } else {\n using namespace at;\n DISPATCH_DOUBLE_AND_FLOAT(g.scalar_type(), 0, \"adam_cuda_kernel\",\n adam_cuda_kernel<<>>(\n p.data_ptr(),\n NULL, // don't output p_copy for fp32, it's wasted write\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n }\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid fused_adam_cuda_mt(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists, // p, m, v, g, p_copy\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay) {\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n size_t tl_sz = tensor_lists.size();\n TORCH_CHECK(tl_sz == 4 || tl_sz == 5, \"expected tensor lists of size 4 or 5\");\n\n if (tensor_lists[3][0].scalar_type() == at::ScalarType::Half) {\n // alher values should be fp32 for half gradients\n TORCH_CHECK(tensor_lists[0][0].scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dich is done on the gradient type\n if (tl_sz == 5) {\n DISPATCH_FLOAT_AND_HALF(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n using accscalar_t = at::acc_type;\n multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<5, accscalar_t, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n } else {\n DISPATCH_FLOAT_AND_HALF(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n using accscalar_t = at::acc_type;\n multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<4, accscalar_t, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n }\n } else {\n if (tl_sz == 5) {\n DISPATCH_DOUBLE_AND_FLOAT(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<5, scalar_t_0, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n } else {\n DISPATCH_DOUBLE_AND_FLOAT(tensor_lists[3][0].scalar_type(), 0, \"adam_cuda_mt_kernel\",\n multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n AdamFunctor<4, scalar_t_0, scalar_t_0>(), beta1, beta2, eps,\n grad_scale, step_size, (adamMode_t)mode, decay););\n }\n }\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\ntemplate \n__device__ void convert(const FROM_T vi, TO_T& vo) {\n vo = static_cast(vi);\n}\n\ntemplate <>\n__device__ void convert(const float vi, uint8_t& vo) {\n union S {\n float as_float;\n int as_int;\n };\n S s;\n s.as_float = vi;\n s.as_int = s.as_int & 0xFF800000;\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_half = static_cast(vi + s.as_float / 8.0f);\n vo = t.as_byte[1];\n}\n\ntemplate <>\n__device__ void convert(const uint8_t vi, float& vo) {\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_byte[0] = 0;\n t.as_byte[1] = vi;\n vo = static_cast(t.as_half);\n}\n\ntemplate <>\n__device__ void convert(const at::Half vi, uint8_t& vo) {\n union S {\n float as_float;\n int as_int;\n };\n S s;\n s.as_float = static_cast(vi);\n s.as_int = s.as_int & 0xFF800000;\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_half = static_cast(vi + s.as_float / 8.0f);\n vo = t.as_byte[1];\n}\n\ntemplate <>\n__device__ void convert(const uint8_t vi, at::Half& vo) {\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_byte[0] = 0;\n t.as_byte[1] = vi;\n vo = t.as_half;\n}\n\ntemplate \n__global__ void strided_check_finite_cuda_kernel(volatile int* noop_gmem, GRAD_T* __restrict__ p_copy,\n const size_t tsize, int stride, int clear_overflow_first) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock) * stride;\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock * stride;\n\n if (clear_overflow_first) {\n if (i == 0) {\n *noop_gmem = 0;\n }\n __syncthreads();\n }\n\n for (int j = i; j < tsize; j += totThreads) {\n GRAD_T pi = p_copy[j];\n if (!isfinite(pi)) {\n *noop_gmem = 1;\n }\n }\n}\ntemplate <>\n__global__ void strided_check_finite_cuda_kernel(volatile int* noop_gmem, uint8_t* __restrict__ p_copy,\n const size_t tsize, int stride, int clear_overflow_first) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock) * stride;\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock * stride;\n\n if (clear_overflow_first) {\n if (i == 0) {\n *noop_gmem = 0;\n }\n __syncthreads();\n }\n\n for (int j = i; j < tsize; j += totThreads) {\n at::Half pi;\n convert(p_copy[j], pi);\n if (!isfinite(pi)) {\n *noop_gmem = 1;\n }\n }\n}\n\ntemplate \n__global__ void maybe_cast_kernel(volatile int* overflow_flag, const FROM_T* p_in, TO_T* p_out, const size_t tsize) {\n if (overflow_flag && *overflow_flag != 0) return;\n\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n FROM_T pi[ILP];\n TO_T po[ILP];\n\n for (int j_start = 0; j_start < tsize; j_start += totThreads * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n pi[ii] = 0;\n\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n pi[ii] = p_in[j];\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n convert(pi[ii], po[ii]);\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n p_out[j] = po[ii];\n }\n }\n }\n}\n\ntemplate \n__global__ void reversible_adam_cuda_kernel(\n T* __restrict__ p,\n REDU_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed\n T* __restrict__ m, T* __restrict__ v, const GRAD_T* __restrict__ g, const float b1, const float b2, const float eps,\n const float grad_scale, const float step_size, const size_t tsize, adamMode_t mode, const float decay) {\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n T mi[ILP];\n T vi[ILP];\n T pi[ILP];\n T gi[ILP];\n\n bool overflow = false;\n for (int j_start = 0; j_start < tsize; j_start += totThreads * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n mi[ii] = T(0);\n vi[ii] = T(0);\n pi[ii] = T(0);\n gi[ii] = GRAD_T(0);\n\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n pi[ii] = p[j];\n mi[ii] = m[j];\n vi[ii] = v[j];\n gi[ii] = static_cast(g[j]);\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n T scaled_grad = gi[ii] / grad_scale;\n if (isfinite(scaled_grad)) {\n mi[ii] = b1 * mi[ii] + (1 - b1) * scaled_grad;\n vi[ii] = b2 * vi[ii] + (1 - b2) * scaled_grad * scaled_grad;\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(vi[ii] + eps);\n else // Mode 1\n denom = sqrtf(vi[ii]) + eps;\n float update = (mi[ii] / denom) + (decay * pi[ii]);\n pi[ii] = pi[ii] - (step_size * update);\n } else {\n overflow = true;\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + i + totThreads * ii;\n if (j < tsize) {\n m[j] = mi[ii];\n v[j] = vi[ii];\n p[j] = pi[ii];\n if (p_copy != NULL) {\n convert(pi[ii], p_copy[j]);\n }\n }\n }\n }\n\n if (p_copy != NULL) {\n __syncthreads();\n if (overflow) {\n convert(float(INFINITY), p_copy[0]);\n }\n }\n}\n\ntemplate \n__global__ void maybe_adam_undo_cuda_kernel(volatile int* overflow_flag, T* __restrict__ p, T* __restrict__ m,\n T* __restrict__ v, const GRAD_T* __restrict__ g, const float b1,\n const float b2, const float eps, const float grad_scale,\n const float step_size, const size_t tsize, adamMode_t mode,\n const float decay) {\n // NB! Skip undo kernel when overflow flag is NOT set\n if (overflow_flag && *overflow_flag == 0) return;\n\n // Assuming 2D grids and 2D blocks\n const int blockId = gridDim.x * blockIdx.y + blockIdx.x;\n const int threadsPerBlock = blockDim.x * blockDim.y;\n const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;\n const int i = (blockId * threadsPerBlock + threadIdInBlock);\n const int totThreads = gridDim.x * gridDim.y * threadsPerBlock;\n\n T mi[ILP];\n T vi[ILP];\n T pi[ILP];\n T gi[ILP];\n\n for (int j_start = 0; j_start < tsize; j_start += totThreads * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n mi[ii] = T(0);\n vi[ii] = T(0);\n pi[ii] = T(0);\n gi[ii] = GRAD_T(0);\n\n int j = j_start + i * ILP;\n if (j < tsize) {\n pi[ii] = p[j];\n mi[ii] = m[j];\n vi[ii] = v[j];\n gi[ii] = static_cast(g[j]);\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n T scaled_grad = gi[ii] / grad_scale;\n if (isfinite(scaled_grad)) {\n float denom;\n if (mode == ADAM_MODE_0)\n denom = sqrtf(vi[ii] + eps);\n else // Mode 1\n denom = sqrtf(vi[ii]) + eps;\n pi[ii] = (pi[ii] + step_size * (mi[ii] / denom)) / (1.0f - step_size * decay);\n mi[ii] = (mi[ii] - (1 - b1) * scaled_grad) / b1;\n vi[ii] = (vi[ii] - (1 - b2) * scaled_grad * scaled_grad) / b2;\n // Make sure round off errors don't create (small) negative value.\n // This can happen if we have to revert the very first step.\n vi[ii] = vi[ii] >= 0.0f ? vi[ii] : 0.0f;\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + i * ILP;\n if (j < tsize) {\n m[j] = mi[ii];\n v[j] = vi[ii];\n p[j] = pi[ii];\n }\n }\n }\n}\n\ntemplate \nstruct MaybeCastFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* overflow_flag,\n TensorListMetadata& tl) {\n if (overflow_flag && *overflow_flag != 0) return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n FROM_T* p_in = (FROM_T*)tl.addresses[0][tensor_loc];\n p_in += chunk_idx * chunk_size;\n TO_T* p_out = (TO_T*)tl.addresses[1][tensor_loc];\n p_out += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n int dim = chunk_size < n ? chunk_size : n;\n\n FROM_T pi[ILP];\n TO_T po[ILP];\n\n for (int j_start = 0; j_start < dim; j_start += blockDim.x * ILP) {\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n pi[ii] = FROM_T(0);\n int j = j_start + threadIdx.x + ii * blockDim.x;\n if (j < dim) {\n pi[ii] = p_in[j];\n }\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n convert(pi[ii], po[ii]);\n }\n\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int j = j_start + threadIdx.x + ii * blockDim.x;\n if (j < dim) {\n p_out[j] = po[ii];\n }\n }\n }\n }\n};\n\nvoid fused_strided_check_finite(at::Tensor& overflow_flag, at::Tensor& p_copy, int stride, int clear_overflow_first) {\n // Get tensor size\n int tsize = p_copy.numel();\n int niter = (tsize + stride - 1) / stride;\n\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n // In order to avoid race condition, blocks must be 1 when clear_overflow_first flag is set.\n const dim3 blocks(clear_overflow_first ? 1 : (niter + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p_copy), \"parameter tensor is too large to be indexed with int32\");\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n using namespace at; // prevents \"toString is undefined\" errors\n DISPATCH_FLOAT_HALF_AND_BYTE(\n p_copy.scalar_type(), 0, \"check_finite_cuda_kernel\",\n strided_check_finite_cuda_kernel<<>>(\n overflow_flag.data_ptr(), p_copy.data_ptr(), tsize, stride, clear_overflow_first););\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid fused_reversible_adam_cuda(at::Tensor& p, at::Tensor& p_copy, at::Tensor& m, at::Tensor& v, at::Tensor& g,\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay) {\n // using namespace at;\n\n // Get tensor size\n int tsize = p.numel();\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (g.scalar_type() == at::ScalarType::Half) {\n // all other values should be fp32 for half gradients\n TORCH_CHECK(p.scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dispatch is done on the gradient type\n using namespace at; // prevents \"toString is undefined\" errors\n if (p_copy.numel() == 0 || p_copy.scalar_type() == g.scalar_type()) {\n DISPATCH_FLOAT_AND_HALF(\n g.scalar_type(), 0, \"adam_cuda_kernel\", using accscalar_t = at::acc_type;\n reversible_adam_cuda_kernel<<>>(\n p.data_ptr(), p_copy.numel() ? p_copy.data_ptr() : NULL,\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1, beta2, eps,\n grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n } else {\n TORCH_CHECK(p_copy.scalar_type() == at::ScalarType::Byte, \"expected parameter to be of byte type\");\n DISPATCH_FLOAT_AND_HALF(\n g.scalar_type(), 0, \"adam_cuda_e5m2_kernel\", using accscalar_t = at::acc_type;\n reversible_adam_cuda_kernel<<>>(\n p.data_ptr(), p_copy.data_ptr(), m.data_ptr(),\n v.data_ptr(), g.data_ptr(), beta1, beta2, eps, grad_scale, step_size, tsize,\n (adamMode_t)mode, decay););\n }\n } else {\n using namespace at;\n DISPATCH_DOUBLE_AND_FLOAT(g.scalar_type(), 0, \"adam_cuda_kernel\",\n reversible_adam_cuda_kernel\n <<>>(\n p.data_ptr(),\n NULL, // don't output p_copy for fp32, it's wasted write\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n }\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid maybe_cast_cuda(at::Tensor& overflow_flag, at::Tensor& p_in, at::Tensor& p_out) {\n // Get tensor size\n int tsize = p_in.numel();\n TORCH_CHECK(tsize == p_out.numel(), \"p_in.numel() must equal p_out.numel()\");\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p_in), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n DISPATCH_FLOAT_HALF_AND_BYTE(p_in.scalar_type(), 0,\n \"maybe_cast_cuda\" DISPATCH_FLOAT_HALF_AND_BYTE(\n p_out.scalar_type(), 1, \"maybe_cast_cuda\",\n maybe_cast_kernel<<>>(\n overflow_flag.numel() ? overflow_flag.data_ptr() : NULL,\n p_in.data_ptr(), p_out.data_ptr(), tsize);))\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid maybe_cast_cuda_mt(int chunk_size, at::Tensor overflow_flag,\n std::vector> tensor_lists) // p_in, p_out\n{\n // Constants\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n size_t tl_sz = tensor_lists.size();\n TORCH_CHECK(tl_sz == 2, \"expected tensor lists of size 2\");\n\n DISPATCH_FLOAT_HALF_AND_BYTE(\n tensor_lists[0][0].scalar_type(), 0, \"maybe_cast_cuda_mt_kernel\",\n DISPATCH_FLOAT_HALF_AND_BYTE(tensor_lists[1][0].scalar_type(), 1, \"maybe_cast_cuda_mt_kernel\",\n multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, overflow_flag, tensor_lists,\n MaybeCastFunctor<2, scalar_t_0, scalar_t_1>());))\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid fused_maybe_adam_undo_cuda(at::Tensor& overflow_flag, at::Tensor& p, at::Tensor& m, at::Tensor& v, at::Tensor& g,\n float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode,\n int bias_correction, float decay) {\n // Get tensor size\n int tsize = p.numel();\n // Determine #threads and #blocks\n const int threadsPerBlock = 512;\n const dim3 blocks((tsize + threadsPerBlock - 1) / threadsPerBlock);\n TORCH_CHECK(at::cuda::detail::canUse32BitIndexMath(p), \"parameter tensor is too large to be indexed with int32\");\n // Constants\n float step_size = 0;\n if (bias_correction == 1) {\n const float bias_correction1 = 1 - std::pow(beta1, step);\n const float bias_correction2 = 1 - std::pow(beta2, step);\n step_size = lr * std::sqrt(bias_correction2) / bias_correction1;\n } else {\n step_size = lr;\n }\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (g.scalar_type() == at::ScalarType::Half) {\n // all other values should be fp32 for half gradients\n TORCH_CHECK(p.scalar_type() == at::ScalarType::Float, \"expected parameter to be of float type\");\n // dispatch is done on the gradient type\n using namespace at; // prevents \"toString is undefined\" errors\n DISPATCH_FLOAT_AND_HALF(g.scalar_type(), 0, \"adam_cuda_kernel\", using accscalar_t = at::acc_type;\n maybe_adam_undo_cuda_kernel\n <<>>(\n overflow_flag.numel() ? overflow_flag.data_ptr() : NULL, p.data_ptr(),\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1,\n beta2, eps, grad_scale, step_size, tsize, (adamMode_t)mode, decay););\n } else {\n using namespace at;\n DISPATCH_DOUBLE_AND_FLOAT(\n g.scalar_type(), 0, \"adam_cuda_kernel\",\n maybe_adam_undo_cuda_kernel<<>>(\n overflow_flag.numel() ? overflow_flag.data_ptr() : NULL, p.data_ptr(),\n m.data_ptr(), v.data_ptr(), g.data_ptr(), beta1, beta2, eps, grad_scale,\n step_size, tsize, (adamMode_t)mode, decay););\n }\n C10_CUDA_CHECK(cudaGetLastError());\n}\n" }, { "path": "apex/contrib/csrc/optimizers/fused_lamb_cuda.cpp", "content": "#include \n\nvoid multi_tensor_lamb_cuda(int chunk_size, at::Tensor noop_flag, std::vector> tensor_lists,\n const float lr, const float beta1, const float beta2, const float epsilon, const int step,\n const int bias_correction, const float weight_decay, const int grad_averaging,\n const int mode, const float global_grad_norm, const float max_grad_norm);\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"lamb\", &multi_tensor_lamb_cuda, \"Computes and apply update for LAMB optimizer\",\n py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/optimizers/fused_lamb_cuda_kernel.cu", "content": "#include \n#include \n#include \n#include \n// Another possibility:\n// #include \n\n#include \n\n#include \"multi_tensor_apply.cuh\"\n#include \"type_shim.h\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntypedef enum {\n MOMENT_MODE_0 = 0, // L2 regularization mode\n MOMENT_MODE_1 = 1 // Decoupled weight decay mode\n} adamMode_t;\n\nstd::tuple multi_tensor_l2norm_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n at::optional per_tensor_python);\n\nusing MATH_T = float;\n\ntemplate \nstruct LAMBStage1Functor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<4>& tl,\n const float beta1, const float beta2, const float beta3,\n const float beta1_correction, const float beta2_correction,\n const float epsilon, adamMode_t mode, const float decay,\n const float global_grad_norm, const float max_global_grad_norm) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n float clipped_global_grad_norm =\n global_grad_norm > max_global_grad_norm ? global_grad_norm / max_global_grad_norm : 1.0f;\n\n T* g = (T*)tl.addresses[0][tensor_loc];\n g += chunk_idx * chunk_size;\n\n T* p = (T*)tl.addresses[1][tensor_loc];\n p += chunk_idx * chunk_size;\n\n T* m = (T*)tl.addresses[2][tensor_loc];\n m += chunk_idx * chunk_size;\n\n T* v = (T*)tl.addresses[3][tensor_loc];\n v += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n // see note in multi_tensor_scale_kernel.cu\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n MATH_T r_g[ILP];\n MATH_T r_p[ILP];\n MATH_T r_m[ILP];\n MATH_T r_v[ILP];\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n r_g[ii] = g[i];\n // special ?optimization? for lamb stage 1\n if (decay == 0) {\n r_p[ii] = MATH_T(0);\n } else {\n r_p[ii] = p[i];\n }\n r_m[ii] = m[i];\n r_v[ii] = v[i];\n } else {\n r_g[ii] = MATH_T(0);\n r_p[ii] = MATH_T(0);\n r_m[ii] = MATH_T(0);\n r_v[ii] = MATH_T(0);\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n if (mode == MOMENT_MODE_0) {\n MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;\n // L2 on scaled grad\n scaled_grad = scaled_grad + decay * r_p[ii];\n r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;\n r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;\n MATH_T next_m_unbiased = r_m[ii] / beta1_correction;\n MATH_T next_v_unbiased = r_v[ii] / beta2_correction;\n MATH_T denom = sqrtf(next_v_unbiased) + epsilon;\n r_p[ii] = next_m_unbiased / denom;\n } else {\n MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;\n r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;\n r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;\n MATH_T next_m_unbiased = r_m[ii] / beta1_correction;\n MATH_T next_v_unbiased = r_v[ii] / beta2_correction;\n MATH_T denom = sqrtf(next_v_unbiased) + epsilon;\n r_p[ii] = (next_m_unbiased / denom) + (decay * r_p[ii]);\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n g[i] = r_p[ii];\n m[i] = r_m[ii];\n v[i] = r_v[ii];\n }\n }\n }\n }\n};\n\n// Step 2 reads in 'update' value and per-tensor param_norm and update_norm.\n// It computes new parameter value.\ntemplate \nstruct LAMBStage2Functor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<2>& tl,\n const float* per_tensor_param_norm, const float* per_tensor_update_norm,\n const float learning_rate, const float decay) {\n // I'd like this kernel to propagate infs/nans.\n // if(*noop_gmem == 1)\n // return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int tensor_num = tl.start_tensor_this_launch + tensor_loc;\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n MATH_T ratio = learning_rate;\n // apply adaptive learning rate to parameters with non-zero weight decay\n if (decay != 0.0) {\n float param_norm = per_tensor_param_norm[tensor_num];\n float update_norm = per_tensor_update_norm[tensor_num];\n ratio = (update_norm != 0.0f && param_norm != 0.0f) ? learning_rate * (param_norm / update_norm) : learning_rate;\n }\n\n T* update = (T*)tl.addresses[0][tensor_loc];\n update += chunk_idx * chunk_size;\n\n T* p = (T*)tl.addresses[1][tensor_loc];\n p += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n MATH_T r_p[ILP];\n MATH_T r_update[ILP];\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n r_p[ii] = p[i];\n r_update[ii] = update[i];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_p[ii] = r_p[ii] - (ratio * r_update[ii]);\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n p[i] = r_p[ii];\n }\n }\n }\n }\n};\n\nvoid multi_tensor_lamb_cuda(int chunk_size, at::Tensor noop_flag, std::vector> tensor_lists,\n const float lr, const float beta1, const float beta2, const float epsilon, const int step,\n const int bias_correction, const float weight_decay, const int grad_averaging,\n const int mode, const float global_grad_norm, const float max_grad_norm) {\n using namespace at;\n // Master weight and 32bit momentum(potentially changing) is not handled by this\n // So we assume every tensor are all in the same type\n\n // Handle bias correction mode\n float bias_correction1 = 1.0f, bias_correction2 = 1.0f;\n if (bias_correction == 1) {\n bias_correction1 = 1 - std::pow(beta1, step);\n bias_correction2 = 1 - std::pow(beta2, step);\n }\n\n // Handle grad averaging mode\n float beta3 = 1.0f;\n if (grad_averaging == 1) beta3 = 1 - beta1;\n\n std::vector> grad_list(tensor_lists.begin(), tensor_lists.begin() + 1);\n std::vector> param_list(tensor_lists.begin() + 1, tensor_lists.begin() + 2);\n\n // Compute per tensor param norm\n auto param_norm_tuple = multi_tensor_l2norm_cuda(chunk_size, noop_flag, param_list, true);\n\n // We now in-place modify grad to store update before compute its norm\n // Generally this is not a issue since people modify grad in step() method all the time\n // We can also grab list of empty tensor to avoid this, but I'd like to save space/cpu code\n DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, \"lamb_stage_1\",\n multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n LAMBStage1Functor(), beta1, beta2,\n beta3, // 1-beta1 or 1 depends on averaging mode\n bias_correction1, bias_correction2, epsilon, (adamMode_t)mode,\n weight_decay, global_grad_norm, max_grad_norm);)\n\n // Compute update norms\n auto update_norm_tuple = multi_tensor_l2norm_cuda(chunk_size, noop_flag, grad_list, true);\n\n std::vector> grad_param_list(tensor_lists.begin(), tensor_lists.begin() + 2);\n\n DISPATCH_FLOAT_AND_HALF(\n tensor_lists[0][0].scalar_type(), 0, \"lamb_stage_2\",\n multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, noop_flag, grad_param_list, LAMBStage2Functor(),\n std::get<1>(param_norm_tuple).data_ptr(),\n std::get<1>(update_norm_tuple).data_ptr(), lr, weight_decay);)\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n" }, { "path": "apex/contrib/csrc/optimizers/multi_tensor_distopt_adam.cpp", "content": "#include \n\nvoid multi_tensor_fused_adam_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists, at::Tensor grad_scale, float lr,\n float beta1, float beta2, float eps, int step, int mode, int bias_correction,\n float weight_decay);\n\nvoid multi_tensor_fused_adam_capturable_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists, at::Tensor grad_scale,\n at::Tensor lr, float beta1, float beta2, float eps, at::Tensor step,\n int mode, int bias_correction, float weight_decay);\n\nvoid multi_tensor_fused_adam_with_param_remainders_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n at::Tensor grad_scale, float lr, float beta1, float beta2,\n float eps, int step, int mode, int bias_correction,\n float weight_decay);\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"multi_tensor_fused_adam\", &multi_tensor_fused_adam_cuda,\n \"CUDA kernels for multi-tensor Adam, \"\n \"with param copy\",\n py::call_guard());\n m.def(\"multi_tensor_fused_adam_capturable\", &multi_tensor_fused_adam_capturable_cuda,\n \"CUDA kernels for multi-tensor Adam, \"\n \"with param copy, capturable for CUDA graph\",\n py::call_guard());\n m.def(\"multi_tensor_fused_adam_with_param_remainders\", &multi_tensor_fused_adam_with_param_remainders_cuda,\n \"CUDA kernel for multi-tensor Adam, \"\n \"with stored param remainders and param copy\",\n py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/optimizers/multi_tensor_distopt_adam_kernel.cu", "content": "#include \n#include \n#include \n#include \n// Another possibility:\n// #include \n\n#include \n\n#include \n\n#include \"multi_tensor_apply.cuh\"\n#include \"type_shim.h\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(const T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, const T* src, int dst_offset = 0, int src_offset = 0) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((const LT*)src)[src_offset];\n}\n\n// (1-t)*x + t*y\n// Note: Named _lerp to avoid ambiguity with std::lerp under C++20.\n__device__ __forceinline__ float _lerp(float t, float x, float y) {\n // See https://developer.nvidia.com/blog/lerp-faster-cuda/\n return fma(t, y, fma(-t, x, x));\n}\n\ntypedef enum {\n ADAM_MODE_0 = 0, // L2 regularization mode\n ADAM_MODE_1 = 1 // Decoupled weight decay mode(AdamW)\n} adamMode_t;\n\n/* Multi-tensor Adam\n *\n * Updates params in-place and outputs a copy with a desired datatype.\n */\ntemplate \nstruct DistAdamFunctor {\n // Vectorized local compute\n __device__ __forceinline__ static void local_step(T p[ILP], T m[ILP], T v[ILP], const GRAD_T g[ILP],\n const float grad_scale, const float beta1, const float beta2,\n const float beta1_correction, const float beta2_correction,\n const float eps, const float lr, adamMode_t mode,\n const float weight_decay) {\n if (mode == ADAM_MODE_0) { // L2\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float scaled_grad = (g[ii] * grad_scale) + (weight_decay * p[ii]);\n float next_m = _lerp(beta1, scaled_grad, m[ii]);\n float next_v = _lerp(beta2, scaled_grad * scaled_grad, v[ii]);\n float next_m_unbiased = next_m / beta1_correction;\n float next_v_unbiased = next_v / beta2_correction;\n float denom = sqrtf(next_v_unbiased) + eps;\n float update = next_m_unbiased / denom;\n m[ii] = next_m;\n v[ii] = next_v;\n p[ii] -= lr * update;\n }\n } else { // weight decay\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float scaled_grad = g[ii] * grad_scale;\n float next_m = _lerp(beta1, scaled_grad, m[ii]);\n float next_v = _lerp(beta2, scaled_grad * scaled_grad, v[ii]);\n float next_m_unbiased = next_m / beta1_correction;\n float next_v_unbiased = next_v / beta2_correction;\n float denom = sqrtf(next_v_unbiased) + eps;\n float update = (next_m_unbiased / denom) + (weight_decay * p[ii]);\n m[ii] = next_m;\n v[ii] = next_v;\n p[ii] -= lr * update;\n }\n }\n }\n\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<5>& tl,\n const float* grad_scale_ptr, const float beta1, const float beta2,\n const float beta1_correction, const float beta2_correction,\n const float eps, const float lr, adamMode_t mode,\n const float weight_decay) const {\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n const float grad_scale = *grad_scale_ptr;\n\n T* p_in = (T*)tl.addresses[0][tensor_loc];\n p_in += chunk_idx * chunk_size;\n T* m = (T*)tl.addresses[1][tensor_loc];\n m += chunk_idx * chunk_size;\n T* v = (T*)tl.addresses[2][tensor_loc];\n v += chunk_idx * chunk_size;\n const GRAD_T* g = (GRAD_T*)tl.addresses[3][tensor_loc];\n g += chunk_idx * chunk_size;\n PARAM_OUT_T* p_out = (PARAM_OUT_T*)tl.addresses[4][tensor_loc];\n p_out += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n n = chunk_size < n ? chunk_size : n;\n\n const bool aligned =\n (n % ILP == 0 && is_aligned(p_in) && is_aligned(m) && is_aligned(v) && is_aligned(g) && is_aligned(p_out));\n\n for (int i_start = threadIdx.x * ILP; i_start < n; i_start += blockDim.x * ILP) {\n T local_p[ILP];\n T local_m[ILP];\n T local_v[ILP];\n GRAD_T local_g[ILP];\n PARAM_OUT_T local_p_out[ILP];\n\n // Load\n if (aligned) {\n load_store(local_p, p_in + i_start);\n load_store(local_m, m + i_start);\n load_store(local_v, v + i_start);\n load_store(local_g, g + i_start);\n } else {\n#pragma unroll\n for (int ii = 0, i = i_start; ii < ILP; ii++, i++) {\n if (i < n) {\n local_p[ii] = p_in[i];\n local_m[ii] = m[i];\n local_v[ii] = v[i];\n local_g[ii] = g[i];\n } else {\n local_p[ii] = 0;\n local_m[ii] = 0;\n local_v[ii] = 0;\n local_g[ii] = 0;\n }\n }\n }\n\n // Local compute\n local_step(local_p, local_m, local_v, local_g, grad_scale, beta1, beta2, beta1_correction, beta2_correction, eps,\n lr, mode, weight_decay);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n local_p_out[ii] = static_cast(local_p[ii]);\n }\n\n // Store\n if (aligned) {\n load_store(p_in + i_start, local_p);\n load_store(m + i_start, local_m);\n load_store(v + i_start, local_v);\n load_store(p_out + i_start, local_p_out);\n } else {\n#pragma unroll\n for (int ii = 0, i = i_start; ii < ILP; ii++, i++) {\n if (i < n) {\n p_in[i] = local_p[ii];\n m[i] = local_m[ii];\n v[i] = local_v[ii];\n p_out[i] = local_p_out[ii];\n }\n }\n }\n }\n }\n};\n\n/* Multi-tensor Adam with CUDA Graph Support\n *\n * Updates params in-place and outputs a copy with a desired datatype.\n */\ntemplate \nstruct DistAdamCapturableFunctor {\n // Vectorized local compute\n __device__ __forceinline__ static void local_step(T p[ILP], T m[ILP], T v[ILP], const GRAD_T g[ILP],\n const float grad_scale, const float beta1, const float beta2,\n const float beta1_correction, const float beta2_correction,\n const float eps, const float lr, adamMode_t mode,\n const float weight_decay) {\n if (mode == ADAM_MODE_0) { // L2\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float scaled_grad = (g[ii] * grad_scale) + (weight_decay * p[ii]);\n float next_m = _lerp(beta1, scaled_grad, m[ii]);\n float next_v = _lerp(beta2, scaled_grad * scaled_grad, v[ii]);\n float next_m_unbiased = next_m / beta1_correction;\n float next_v_unbiased = next_v / beta2_correction;\n float denom = sqrtf(next_v_unbiased) + eps;\n float update = next_m_unbiased / denom;\n m[ii] = next_m;\n v[ii] = next_v;\n p[ii] -= lr * update;\n }\n } else { // weight decay\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n float scaled_grad = g[ii] * grad_scale;\n float next_m = _lerp(beta1, scaled_grad, m[ii]);\n float next_v = _lerp(beta2, scaled_grad * scaled_grad, v[ii]);\n float next_m_unbiased = next_m / beta1_correction;\n float next_v_unbiased = next_v / beta2_correction;\n float denom = sqrtf(next_v_unbiased) + eps;\n float update = (next_m_unbiased / denom) + (weight_decay * p[ii]);\n m[ii] = next_m;\n v[ii] = next_v;\n p[ii] -= lr * update;\n }\n }\n }\n\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<5>& tl,\n const float* grad_scale_ptr, const float beta1, const float beta2,\n const int* step, const int bias_correction, const float eps,\n const float* lr, adamMode_t mode, const float weight_decay) const {\n assert(noop_gmem);\n assert(grad_scale_ptr);\n assert(step);\n assert(lr);\n\n if (*noop_gmem == 1) return;\n\n float beta1_correction = 1.0f, beta2_correction = 1.0f;\n if (bias_correction == 1) {\n beta1_correction = 1 - pow(beta1, *step);\n beta2_correction = 1 - pow(beta2, *step);\n }\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n const float grad_scale = *grad_scale_ptr;\n\n T* p_in = (T*)tl.addresses[0][tensor_loc];\n p_in += chunk_idx * chunk_size;\n T* m = (T*)tl.addresses[1][tensor_loc];\n m += chunk_idx * chunk_size;\n T* v = (T*)tl.addresses[2][tensor_loc];\n v += chunk_idx * chunk_size;\n const GRAD_T* g = (GRAD_T*)tl.addresses[3][tensor_loc];\n g += chunk_idx * chunk_size;\n PARAM_OUT_T* p_out = (PARAM_OUT_T*)tl.addresses[4][tensor_loc];\n p_out += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n n = chunk_size < n ? chunk_size : n;\n\n const bool aligned =\n (n % ILP == 0 && is_aligned(p_in) && is_aligned(m) && is_aligned(v) && is_aligned(g) && is_aligned(p_out));\n\n for (int i_start = threadIdx.x * ILP; i_start < n; i_start += blockDim.x * ILP) {\n T local_p[ILP];\n T local_m[ILP];\n T local_v[ILP];\n GRAD_T local_g[ILP];\n PARAM_OUT_T local_p_out[ILP];\n\n // Load\n if (aligned) {\n load_store(local_p, p_in + i_start);\n load_store(local_m, m + i_start);\n load_store(local_v, v + i_start);\n load_store(local_g, g + i_start);\n } else {\n#pragma unroll\n for (int ii = 0, i = i_start; ii < ILP; ii++, i++) {\n if (i < n) {\n local_p[ii] = p_in[i];\n local_m[ii] = m[i];\n local_v[ii] = v[i];\n local_g[ii] = g[i];\n } else {\n local_p[ii] = 0;\n local_m[ii] = 0;\n local_v[ii] = 0;\n local_g[ii] = 0;\n }\n }\n }\n\n // Local compute\n local_step(local_p, local_m, local_v, local_g, grad_scale, beta1, beta2, beta1_correction, beta2_correction, eps,\n *lr, mode, weight_decay);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n local_p_out[ii] = static_cast(local_p[ii]);\n }\n\n // Store\n if (aligned) {\n load_store(p_in + i_start, local_p);\n load_store(m + i_start, local_m);\n load_store(v + i_start, local_v);\n load_store(p_out + i_start, local_p_out);\n } else {\n#pragma unroll\n for (int ii = 0, i = i_start; ii < ILP; ii++, i++) {\n if (i < n) {\n p_in[i] = local_p[ii];\n m[i] = local_m[ii];\n v[i] = local_v[ii];\n p_out[i] = local_p_out[ii];\n }\n }\n }\n }\n }\n};\n\n/* Functor for multi-tensor Adam with implicit main params\n *\n * If params are BF16 and optimizer state is FP32, it is not necessary\n * to store FP32 main params. Instead, store 16-bit param remainder\n * and combine with BF16 param to reconstruct the FP32 main param.\n */\ntemplate \nstruct DistAdamWithParamRemaindersFunctor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<6>& tl,\n const float* grad_scale_ptr, const float beta1, const float beta2,\n const float beta1_correction, const float beta2_correction,\n const float eps, const float lr, adamMode_t mode,\n const float weight_decay) const {\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n const float grad_scale = *grad_scale_ptr;\n\n int16_t* p_in = (int16_t*)tl.addresses[0][tensor_loc];\n p_in += chunk_idx * chunk_size;\n int16_t* p_rem = (int16_t*)tl.addresses[1][tensor_loc];\n p_rem += chunk_idx * chunk_size;\n float* m = (float*)tl.addresses[2][tensor_loc];\n m += chunk_idx * chunk_size;\n float* v = (float*)tl.addresses[3][tensor_loc];\n v += chunk_idx * chunk_size;\n const GRAD_T* g = (GRAD_T*)tl.addresses[4][tensor_loc];\n g += chunk_idx * chunk_size;\n int16_t* p_out = (int16_t*)tl.addresses[5][tensor_loc];\n p_out += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n n = chunk_size < n ? chunk_size : n;\n\n const bool aligned = (n % ILP == 0 && is_aligned(p_in) && is_aligned(p_rem) && is_aligned(m) && is_aligned(v) &&\n is_aligned(g) && is_aligned(p_out));\n\n for (int i_start = threadIdx.x * ILP; i_start < n; i_start += blockDim.x * ILP) {\n union fp32_or_int162 {\n float fp32;\n int16_t int16[2];\n };\n fp32_or_int162 local_p[ILP];\n int16_t local_p_bf16[ILP];\n int16_t local_p_rem[ILP];\n float local_m[ILP];\n float local_v[ILP];\n GRAD_T local_g[ILP];\n\n // Load\n if (aligned) {\n load_store(local_p_bf16, p_in + i_start);\n load_store(local_p_rem, p_rem + i_start);\n load_store(local_m, m + i_start);\n load_store(local_v, v + i_start);\n load_store(local_g, g + i_start);\n } else {\n#pragma unroll\n for (int ii = 0, i = i_start; ii < ILP; ii++, i++) {\n if (i < n) {\n local_p_bf16[ii] = p_in[i];\n local_p_rem[ii] = p_rem[i];\n local_m[ii] = m[i];\n local_v[ii] = v[i];\n local_g[ii] = g[i];\n } else {\n local_p_bf16[ii] = 0;\n local_p_rem[ii] = 0;\n local_m[ii] = 0;\n local_v[ii] = 0;\n local_g[ii] = 0;\n }\n }\n }\n\n // Reconstruct FP32 params\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n if (local_p_rem[ii] < 0) local_p_bf16[ii]--; // Undo rounding\n local_p[ii].int16[1] = local_p_bf16[ii];\n local_p[ii].int16[0] = local_p_rem[ii];\n }\n\n // Local compute\n using LocalFunctor = DistAdamFunctor;\n LocalFunctor::local_step(reinterpret_cast(local_p), local_m, local_v, local_g, grad_scale, beta1, beta2,\n beta1_correction, beta2_correction, eps, lr, mode, weight_decay);\n\n // Split into BF16 params (rounded-to-nearest) and remainders\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n local_p_bf16[ii] = local_p[ii].int16[1];\n local_p_rem[ii] = local_p[ii].int16[0];\n if (local_p_rem[ii] < 0) local_p_bf16[ii]++; // Round up\n }\n\n // Store\n if (aligned) {\n load_store(p_rem + i_start, local_p_rem);\n load_store(m + i_start, local_m);\n load_store(v + i_start, local_v);\n load_store(p_out + i_start, local_p_bf16);\n } else {\n#pragma unroll\n for (int ii = 0, i = i_start; ii < ILP; ii++, i++) {\n if (i < n) {\n p_rem[i] = local_p_rem[ii];\n m[i] = local_m[ii];\n v[i] = local_v[ii];\n p_out[i] = local_p_bf16[ii];\n }\n }\n }\n }\n }\n};\n\nvoid multi_tensor_fused_adam_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists, // p_in, m, v, g, p_out\n at::Tensor grad_scale, float lr, float beta1, float beta2, float eps, int step,\n int mode, int bias_correction, float weight_decay) {\n using namespace at;\n\n // Expect p_in, m, v, g, p_out\n size_t tl_sz = tensor_lists.size();\n TORCH_CHECK(tl_sz == 5, \"expected tensor lists of size 5\");\n const auto p_in_type = tensor_lists[0][0].scalar_type();\n const auto g_type = tensor_lists[3][0].scalar_type();\n const auto p_out_type = tensor_lists[4][0].scalar_type();\n\n float beta1_correction = 1.0f, beta2_correction = 1.0f;\n if (bias_correction == 1) {\n beta1_correction = 1 - std::pow(beta1, step);\n beta2_correction = 1 - std::pow(beta2, step);\n }\n\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n p_in_type, 0, \"dist_adam_cuda_kernel\",\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n g_type, 1, \"dist_adam_cuda_kernel\",\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n p_out_type, 2, \"dist_adam_cuda_kernel\",\n multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n DistAdamFunctor(), grad_scale.data_ptr(),\n beta1, beta2, beta1_correction, beta2_correction, eps, lr, (adamMode_t)mode,\n weight_decay);)));\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid multi_tensor_fused_adam_capturable_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists, // p_in, m, v, g, p_out\n at::Tensor grad_scale, at::Tensor lr, float beta1, float beta2, float eps,\n at::Tensor step, int mode, int bias_correction, float weight_decay) {\n using namespace at;\n\n // Expect p_in, m, v, g, p_out\n size_t tl_sz = tensor_lists.size();\n TORCH_CHECK(tl_sz == 5, \"expected tensor lists of size 5\");\n const auto p_in_type = tensor_lists[0][0].scalar_type();\n const auto g_type = tensor_lists[3][0].scalar_type();\n const auto p_out_type = tensor_lists[4][0].scalar_type();\n\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n p_in_type, 0, \"dist_adam_capturable_cuda_kernel\",\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n g_type, 1, \"dist_adam_capturable_cuda_kernel\",\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n p_out_type, 2, \"dist_adam_capturable_cuda_kernel\",\n multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n DistAdamCapturableFunctor(),\n grad_scale.data_ptr(), beta1, beta2, step.data_ptr(), bias_correction,\n eps, lr.data_ptr(), (adamMode_t)mode, weight_decay);)));\n C10_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid multi_tensor_fused_adam_with_param_remainders_cuda(\n int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists, // p_in, p_rem, m, v, g, p_out\n at::Tensor grad_scale, float lr, float beta1, float beta2, float eps, int step, int mode, int bias_correction,\n float weight_decay) {\n using namespace at;\n\n // Expect p_in, p_rem, m, v, g, p_out\n size_t tl_sz = tensor_lists.size();\n TORCH_CHECK(tl_sz == 6, \"expected tensor lists of size 6\");\n const auto g_type = tensor_lists[4][0].scalar_type();\n\n float beta1_correction = 1.0f, beta2_correction = 1.0f;\n if (bias_correction == 1) {\n beta1_correction = 1 - std::pow(beta1, step);\n beta2_correction = 1 - std::pow(beta2, step);\n }\n\n DISPATCH_FLOAT_HALF_AND_BFLOAT(\n g_type, 0, \"dist_adam_with_param_remainders_cuda_kernel\",\n multi_tensor_apply<6>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n DistAdamWithParamRemaindersFunctor(), grad_scale.data_ptr(), beta1,\n beta2, beta1_correction, beta2_correction, eps, lr, (adamMode_t)mode, weight_decay););\n C10_CUDA_CHECK(cudaGetLastError());\n}\n" }, { "path": "apex/contrib/csrc/optimizers/multi_tensor_distopt_lamb.cpp", "content": "#include \n\nvoid multi_tensor_lamb_compute_update_term_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n at::Tensor per_tensor_beta1, at::Tensor per_tensor_beta2,\n at::Tensor per_tensor_beta3, at::Tensor per_tensor_bias_correction,\n at::Tensor step, at::Tensor per_tensor_epsilon, const int mode,\n at::Tensor per_tensor_decay, at::Tensor global_scale,\n at::Tensor global_grad_norm, const float max_grad_norm);\n\nvoid multi_tensor_lamb_update_weights_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n at::Tensor per_tensor_param_norm, at::Tensor per_tensor_update_norm,\n at::Tensor update_norm_offset, at::Tensor learning_rate,\n at::Tensor per_tensor_decay, at::Tensor global_grad_norm, bool use_nvlamb);\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"multi_tensor_lamb_compute_update_term\", &multi_tensor_lamb_compute_update_term_cuda,\n \"Computes update term for LAMB optimizer\", py::call_guard());\n m.def(\"multi_tensor_lamb_update_weights\", &multi_tensor_lamb_update_weights_cuda,\n \"Applies update term for LAMB optimizer\", py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/optimizers/multi_tensor_distopt_lamb_kernel.cu", "content": "#include \n#include \n#include \n#include \n// Another possibility:\n// #include \n\n#include \n\n#include \"multi_tensor_apply.cuh\"\n#include \"type_shim.h\"\n\n#define BLOCK_SIZE 512\n#define ILP 4\n\ntemplate \n__device__ __forceinline__ bool is_aligned(T* p) {\n return ((uint64_t)p) % (ILP * sizeof(T)) == 0;\n}\n\ntemplate \n__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) {\n typedef typename std::aligned_storage::type LT;\n ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];\n}\n\ntemplate \n__device__ void convert(const FROM_T vi, TO_T& vo) {\n vo = static_cast(vi);\n}\n\ntemplate <>\n__device__ void convert(const float vi, uint8_t& vo) {\n union S {\n float as_float;\n int as_int;\n };\n S s;\n s.as_float = vi;\n s.as_int = s.as_int & 0xFF800000;\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_half = static_cast(vi + s.as_float / 8.0f);\n vo = t.as_byte[1];\n}\n\ntemplate <>\n__device__ void convert(const uint8_t vi, float& vo) {\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_byte[0] = 0;\n t.as_byte[1] = vi;\n vo = static_cast(t.as_half);\n}\n\ntemplate <>\n__device__ void convert(const at::Half vi, uint8_t& vo) {\n union S {\n float as_float;\n int as_int;\n };\n S s;\n s.as_float = static_cast(vi);\n s.as_int = s.as_int & 0xFF800000;\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_half = static_cast(vi + s.as_float / 8.0f);\n vo = t.as_byte[1];\n}\n\ntemplate <>\n__device__ void convert(const uint8_t vi, at::Half& vo) {\n union T {\n at::Half as_half;\n uint8_t as_byte[2];\n };\n T t;\n t.as_byte[0] = 0;\n t.as_byte[1] = vi;\n vo = t.as_half;\n}\n\ntypedef enum {\n MOMENT_MODE_0 = 0, // L2 regularization mode\n MOMENT_MODE_1 = 1 // Decoupled weight decay mode\n} adamMode_t;\n\ntemplate \nstruct DistOptLAMBStage1Functor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<5>& tl,\n const MATH_T* per_tensor_beta1, const MATH_T* per_tensor_beta2,\n const MATH_T* per_tensor_beta3, const int* per_tensor_bias_correction,\n const int* step, const MATH_T* per_tensor_epsilon, adamMode_t mode,\n const MATH_T* per_tensor_decay, const MATH_T* global_scale,\n const MATH_T* global_grad_norm, const float max_grad_norm) {\n // I'd like this kernel to propagate infs/nans.\n if (*noop_gmem == 1) return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int tensor_num = tl.start_tensor_this_launch + tensor_loc;\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n float combined_scale = *global_scale;\n if (max_grad_norm > 0) {\n combined_scale = max_grad_norm / (*global_grad_norm / *global_scale + 1e-6);\n combined_scale = *global_scale / std::min((float)1.0, combined_scale);\n }\n\n MATH_T beta1 = per_tensor_beta1[tensor_num];\n MATH_T beta2 = per_tensor_beta2[tensor_num];\n MATH_T beta3 = 1 - beta1;\n MATH_T beta1_correction, beta2_correction;\n if (per_tensor_bias_correction[tensor_num] == 1) {\n beta1_correction = 1 - pow(beta1, *step);\n beta2_correction = 1 - pow(beta2, *step);\n } else {\n beta1_correction = (MATH_T)1.0;\n beta2_correction = (MATH_T)1.0;\n }\n MATH_T epsilon = per_tensor_epsilon[tensor_num];\n MATH_T decay = per_tensor_decay[tensor_num];\n\n GRAD_T* g = (GRAD_T*)tl.addresses[0][tensor_loc];\n g += chunk_idx * chunk_size;\n\n T* p = (T*)tl.addresses[1][tensor_loc];\n p += chunk_idx * chunk_size;\n\n T* m = (T*)tl.addresses[2][tensor_loc];\n m += chunk_idx * chunk_size;\n\n T* v = (T*)tl.addresses[3][tensor_loc];\n v += chunk_idx * chunk_size;\n\n MATH_T* u = (MATH_T*)tl.addresses[4][tensor_loc];\n u += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n MATH_T r_g[ILP];\n MATH_T r_p[ILP];\n MATH_T r_m[ILP];\n MATH_T r_v[ILP];\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(g) && is_aligned(p) && is_aligned(m) && is_aligned(v)) {\n GRAD_T l_g[ILP];\n T l_p[ILP];\n T l_m[ILP];\n T l_v[ILP];\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(l_g, g, 0, i_start);\n if (decay != 0) load_store(l_p, p, 0, i_start);\n load_store(l_m, m, 0, i_start);\n load_store(l_v, v, 0, i_start);\n // unpack\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_g[ii] = l_g[ii];\n if (decay == 0) {\n r_p[ii] = MATH_T(0);\n } else {\n r_p[ii] = l_p[ii];\n }\n r_m[ii] = l_m[ii];\n r_v[ii] = l_v[ii];\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n if (mode == MOMENT_MODE_0) {\n MATH_T scaled_grad = r_g[ii] / combined_scale;\n // L2 on scaled grad\n scaled_grad = scaled_grad + decay * r_p[ii];\n r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;\n r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;\n MATH_T next_m_unbiased = r_m[ii] / beta1_correction;\n MATH_T next_v_unbiased = r_v[ii] / beta2_correction;\n MATH_T denom = sqrtf(next_v_unbiased) + epsilon;\n r_p[ii] = next_m_unbiased / denom;\n } else {\n MATH_T scaled_grad = r_g[ii] / combined_scale;\n r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;\n r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;\n MATH_T next_m_unbiased = r_m[ii] / beta1_correction;\n MATH_T next_v_unbiased = r_v[ii] / beta2_correction;\n MATH_T denom = sqrtf(next_v_unbiased) + epsilon;\n r_p[ii] = (next_m_unbiased / denom) + (decay * r_p[ii]);\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n l_m[ii] = r_m[ii];\n l_v[ii] = r_v[ii];\n }\n // store\n load_store(u, r_p, i_start, 0);\n load_store(m, l_m, i_start, 0);\n load_store(v, l_v, i_start, 0);\n }\n } else {\n // see note in multi_tensor_scale_kernel.cu\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n MATH_T r_g[ILP];\n MATH_T r_p[ILP];\n MATH_T r_m[ILP];\n MATH_T r_v[ILP];\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n r_g[ii] = g[i];\n // special ?optimization? for lamb stage 1\n if (decay == 0) {\n r_p[ii] = MATH_T(0);\n } else {\n r_p[ii] = p[i];\n }\n r_m[ii] = m[i];\n r_v[ii] = v[i];\n } else {\n r_g[ii] = MATH_T(0);\n r_p[ii] = MATH_T(0);\n r_m[ii] = MATH_T(0);\n r_v[ii] = MATH_T(0);\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n if (mode == MOMENT_MODE_0) {\n MATH_T scaled_grad = r_g[ii] / combined_scale;\n // L2 on scaled grad\n scaled_grad = scaled_grad + decay * r_p[ii];\n r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;\n r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;\n MATH_T next_m_unbiased = r_m[ii] / beta1_correction;\n MATH_T next_v_unbiased = r_v[ii] / beta2_correction;\n MATH_T denom = sqrtf(next_v_unbiased) + epsilon;\n r_p[ii] = next_m_unbiased / denom;\n } else {\n MATH_T scaled_grad = r_g[ii] / combined_scale;\n r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;\n r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;\n MATH_T next_m_unbiased = r_m[ii] / beta1_correction;\n MATH_T next_v_unbiased = r_v[ii] / beta2_correction;\n MATH_T denom = sqrtf(next_v_unbiased) + epsilon;\n r_p[ii] = (next_m_unbiased / denom) + (decay * r_p[ii]);\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n u[i] = r_p[ii];\n m[i] = r_m[ii];\n v[i] = r_v[ii];\n }\n }\n }\n }\n }\n};\n\n// Step 2 reads in 'update' value and per-tensor param_norm and update_norm.\n// It computes new parameter value.\ntemplate \nstruct DistOptLAMBStage2Functor {\n __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<3>& tl,\n const MATH_T* per_tensor_param_norm, const MATH_T* per_tensor_update_norm,\n const long* update_norm_offset, const MATH_T* learning_rate,\n const MATH_T* per_tensor_decay, const MATH_T* global_grad_norm,\n bool use_nvlamb) {\n // I'd like this kernel to propagate infs/nans.\n if (*noop_gmem == 1) return;\n\n int tensor_loc = tl.block_to_tensor[blockIdx.x];\n int tensor_num = tl.start_tensor_this_launch + tensor_loc;\n int chunk_idx = tl.block_to_chunk[blockIdx.x];\n int n = tl.sizes[tensor_loc];\n\n MATH_T decay = per_tensor_decay[tensor_num];\n\n MATH_T ratio = *learning_rate;\n // nvlamb: apply adaptive learning rate to all parameters\n // otherwise, only apply to those with non-zero weight decay\n if (use_nvlamb || (decay != (MATH_T)0.0)) {\n MATH_T param_norm = per_tensor_param_norm[tensor_num];\n MATH_T update_norm = per_tensor_update_norm[update_norm_offset[tensor_num]];\n ratio =\n (update_norm != 0.0 && param_norm != 0.0) ? (*learning_rate) * (param_norm / update_norm) : (*learning_rate);\n }\n\n MATH_T* update = (MATH_T*)tl.addresses[0][tensor_loc];\n update += chunk_idx * chunk_size;\n\n T* p = (T*)tl.addresses[1][tensor_loc];\n p += chunk_idx * chunk_size;\n\n GRAD_T* p_copy = (GRAD_T*)tl.addresses[2][tensor_loc];\n p_copy += chunk_idx * chunk_size;\n\n n -= chunk_idx * chunk_size;\n\n // to make things simple, we put aligned case in a different code path\n if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(p) && is_aligned(update)) {\n T r_p[ILP];\n MATH_T r_update[ILP];\n GRAD_T r_p_copy[ILP];\n for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) {\n // load\n load_store(r_p, p, 0, i_start);\n load_store(r_update, update, 0, i_start);\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_p[ii] = static_cast(r_p[ii]) - (ratio * r_update[ii]);\n convert(r_p[ii], r_p_copy[ii]);\n }\n load_store(p, r_p, i_start, 0);\n load_store(p_copy, r_p_copy, i_start, 0);\n }\n } else {\n for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {\n MATH_T r_p[ILP];\n MATH_T r_update[ILP];\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n r_p[ii] = p[i];\n r_update[ii] = update[i];\n }\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n r_p[ii] = r_p[ii] - (ratio * r_update[ii]);\n }\n#pragma unroll\n for (int ii = 0; ii < ILP; ii++) {\n int i = i_start + threadIdx.x + ii * blockDim.x;\n if (i < n && i < chunk_size) {\n p[i] = r_p[ii];\n convert(r_p[ii], p_copy[i]);\n }\n }\n }\n }\n }\n};\n\nvoid multi_tensor_lamb_compute_update_term_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n at::Tensor per_tensor_beta1, at::Tensor per_tensor_beta2,\n at::Tensor per_tensor_beta3, at::Tensor per_tensor_bias_correction,\n at::Tensor step, at::Tensor per_tensor_epsilon, const int mode,\n at::Tensor per_tensor_decay, at::Tensor global_scale,\n at::Tensor global_grad_norm, const float max_grad_norm) {\n using namespace at;\n\n DISPATCH_FLOAT_AND_HALF(\n tensor_lists[1][0].scalar_type(), 0, \"lamb_stage_1\",\n DISPATCH_FLOAT_AND_HALF(\n tensor_lists[0][0].scalar_type(), 1, \"lamb_stage_1\",\n DISPATCH_FLOAT_AND_HALF(\n tensor_lists[4][0].scalar_type(), 2, \"lamb_stage_1\",\n multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n DistOptLAMBStage1Functor(),\n per_tensor_beta1.data_ptr(), per_tensor_beta2.data_ptr(),\n per_tensor_beta3.data_ptr(), per_tensor_bias_correction.data_ptr(),\n step.data_ptr(), per_tensor_epsilon.data_ptr(), (adamMode_t)mode,\n per_tensor_decay.data_ptr(), global_scale.data_ptr(),\n global_grad_norm.data_ptr(), max_grad_norm);)))\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n\nvoid multi_tensor_lamb_update_weights_cuda(int chunk_size, at::Tensor noop_flag,\n std::vector> tensor_lists,\n at::Tensor per_tensor_param_norm, at::Tensor per_tensor_update_norm,\n at::Tensor update_norm_offset, at::Tensor learning_rate,\n at::Tensor per_tensor_decay, at::Tensor global_grad_norm, bool use_nvlamb) {\n using namespace at;\n\n DISPATCH_FLOAT_AND_HALF(\n tensor_lists[1][0].scalar_type(), 0, \"lamb_stage_2\",\n DISPATCH_FLOAT_HALF_AND_BYTE(\n tensor_lists[2][0].scalar_type(), 1, \"lamb_stage_2\",\n DISPATCH_FLOAT_AND_HALF(\n tensor_lists[0][0].scalar_type(), 2, \"lamb_stage_2\",\n multi_tensor_apply<3>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,\n DistOptLAMBStage2Functor(),\n per_tensor_param_norm.data_ptr(),\n per_tensor_update_norm.data_ptr(), update_norm_offset.data_ptr(),\n learning_rate.data_ptr(), per_tensor_decay.data_ptr(),\n global_grad_norm.data_ptr(), use_nvlamb);)))\n\n AT_CUDA_CHECK(cudaGetLastError());\n}\n" }, { "path": "apex/contrib/csrc/peer_memory/peer_memory.cpp", "content": "/**\n * Copyright (c) 2018-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#include \"peer_memory_cuda.cuh\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"allocate_raw\", &apex::contrib::peer_memory::allocate_raw, \"allocate_raw\",\n py::call_guard());\n m.def(\"free_raw\", &apex::contrib::peer_memory::free_raw, \"free_raw\", py::call_guard());\n m.def(\"zero\", &apex::contrib::peer_memory::zero, \"zero\", py::call_guard());\n m.def(\"get_raw_ipc_address\", &apex::contrib::peer_memory::get_raw_ipc_address, \"get_raw_ipc_address\",\n py::call_guard());\n m.def(\"get_raw_peers\", &apex::contrib::peer_memory::get_raw_peers, \"get_raw_peers\",\n py::call_guard());\n m.def(\"blob_view_half\", &apex::contrib::peer_memory::blob_view_half, \"blob_view_half\",\n py::call_guard());\n m.def(\"blob_view_float\", &apex::contrib::peer_memory::blob_view_float, \"blob_view_float\",\n py::call_guard());\n m.def(\"blob_view_int\", &apex::contrib::peer_memory::blob_view_int, \"blob_view_int\",\n py::call_guard());\n m.def(\"push_pull_halos_1d\", &apex::contrib::peer_memory::push_pull_halos_1d, \"push_pull_halos_1d\",\n py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/peer_memory/peer_memory_cuda.cu", "content": "#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n#include \"nccl.h\"\n\n#define CUDACHECK(cmd) \\\n do { \\\n cudaError_t err = cmd; \\\n if (err != cudaSuccess) { \\\n char hostname[1024]; \\\n gethostname(hostname, 1024); \\\n printf(\"%s: CUDA failure %s:%d '%s'\\n\", hostname, __FILE__, __LINE__, cudaGetErrorString(err)); \\\n } \\\n } while (0)\n\nnamespace {\n\nconstexpr int THREADS_PER_CTA = 128;\n\n/* Basic deleter function for from_blob function.\nvoid deleter(void* ptr)\n{\n printf(\"deleter(ptr=%p)\\n\",ptr);\n cudaFree(ptr);\n}\n*/\n\ntemplate \nat::Tensor blob_view(T* raw_ptr, std::vector shape, const at::TensorOptions& options, bool channels_last) {\n size_t size = 1;\n std::vector strides(shape.size());\n if (channels_last) {\n assert(shape.size() == 4);\n strides[0] = shape[1] * shape[2] * shape[3];\n strides[1] = 1;\n strides[2] = shape[1] * shape[3];\n strides[3] = shape[1];\n } else {\n int idx = strides.size();\n for (auto it = shape.rbegin(); it != shape.rend(); ++it) {\n strides[--idx] = size;\n size *= *it;\n }\n }\n size *= sizeof(T);\n // TODO: Implement dynamic reuse of pooled peer memory.\n // We provide no deleter function because all peer memory allocations are static in this implementation.\n return torch::from_blob((void*)raw_ptr, shape, strides, 0L, options);\n}\n\nvoid tensor_shape(at::Tensor t, bool explicit_nhwc, int& N, int& C, int& H, int& W) {\n if (t.dim() == 3) {\n N = 1;\n if (explicit_nhwc) {\n C = t.size(2);\n H = t.size(0);\n W = t.size(1);\n } else {\n C = t.size(0);\n H = t.size(1);\n W = t.size(2);\n }\n } else if (t.dim() == 4) {\n if (explicit_nhwc) {\n N = t.size(0);\n C = t.size(3);\n H = t.size(1);\n W = t.size(2);\n } else {\n N = t.size(0);\n C = t.size(1);\n H = t.size(2);\n W = t.size(3);\n }\n } else {\n printf(\"%s;%d - t.dim() must be either 3 or 4 (was %d)\\n\", __FILE__, __LINE__, int(t.dim()));\n assert(t.dim() == 3 || t.dim() == 4);\n }\n}\n\nvoid tensor_strides(at::Tensor t, bool explicit_nhwc, int& stride_N, int& stride_C, int& stride_H, int& stride_W) {\n if (t.dim() == 3) {\n if (explicit_nhwc) {\n stride_C = t.stride(2);\n stride_H = t.stride(0);\n stride_W = t.stride(1);\n } else {\n stride_C = t.stride(0);\n stride_H = t.stride(1);\n stride_W = t.stride(2);\n }\n stride_N = t.size(0) * t.size(1) * t.size(2);\n } else if (t.dim() == 4) {\n if (explicit_nhwc) {\n stride_N = t.stride(0);\n stride_C = t.stride(3);\n stride_H = t.stride(1);\n stride_W = t.stride(2);\n } else {\n stride_N = t.stride(0);\n stride_C = t.stride(1);\n stride_H = t.stride(2);\n stride_W = t.stride(3);\n }\n } else {\n printf(\"%s;%d - t.dim() must be either 3 or 4 (was %d)\\n\", __FILE__, __LINE__, t.dim());\n assert(t.dim() == 3 || t.dim() == 4);\n }\n}\n\ntemplate \ninline __device__ void __zero(T* dst) {\n *dst = T(0);\n}\n\ninline __device__ void __zero(int2* dst) { *dst = {0, 0}; }\n\ntemplate \ninline __device__ void zero_tensor(const int dim0, const int dim1, const int dim2, T* __restrict__ data,\n const int data_stride0, const int data_stride1, const int data_stride2,\n const int thread_id, const int block_id, const int num_blocks) {\n const int global_id = thread_id + block_id * THREADS_PER_CTA;\n const int num_threads = num_blocks * THREADS_PER_CTA;\n const int count = dim0 * dim1 * dim2;\n for (int i = global_id; i < count; i += num_threads) {\n int offset;\n if (contiguous) {\n offset = i;\n } else {\n const int j2 = i % dim2;\n const int k = i / dim2;\n const int j1 = k % dim1;\n const int j0 = k / dim1;\n offset = j0 * data_stride0 + j1 * data_stride1 + j2 * data_stride2;\n }\n __zero(data + offset);\n }\n}\n\ntemplate \ninline __device__ void push_pull_tensor(const int dim0, const int dim1, const int dim2, const T* __restrict__ data_in,\n const int data_in_stride0, const int data_in_stride1, const int data_in_stride2,\n T* __restrict__ data_out, const int data_out_stride0,\n const int data_out_stride1, const int data_out_stride2, int4* local_peer,\n int4* remote_peer, const int thread_id, const int block_id,\n const int num_blocks) {\n // 128b=16B NVLink flit\n // Note: Use last 4B as a semaphore\n static_assert(sizeof(T) <= 12);\n union Flit {\n T payload;\n uint uints[4];\n };\n // Communication bit indicates whether flit has been received from\n // a remote GPU\n constexpr uint communication_mask = 1 << 0;\n // Status bit is used to choose the active peer buffer in an\n // alternating double buffer scheme. We use buffer 1 if the bits\n // match, use buffer 2 if the bits differ, and invert the bit\n // after finishing with a buffer.\n constexpr uint status_mask = 1 << 1;\n\n // Split peer memory into two sets of buffers\n // Note: Each block owns a THREADS_PER_CTA*2*16B chunk of peer\n // memory\n const int peer_offset1 = block_id * THREADS_PER_CTA * 2 + thread_id;\n const int peer_offset2 = peer_offset1 + THREADS_PER_CTA;\n volatile int* local_peer1 = reinterpret_cast(local_peer + peer_offset1);\n volatile int* local_peer2 = reinterpret_cast(local_peer + peer_offset2);\n volatile int* remote_peer1 = reinterpret_cast(remote_peer + peer_offset1);\n volatile int* remote_peer2 = reinterpret_cast(remote_peer + peer_offset2);\n\n // Iterate through tensor entries\n const int num_threads = num_blocks * THREADS_PER_CTA;\n const int count = dim0 * dim1 * dim2;\n for (int i0 = block_id * THREADS_PER_CTA; i0 < count; i0 += num_threads) {\n const int i = i0 + thread_id;\n const bool has_data = i < count;\n\n // Calculate buffer positions\n int data_in_offset, data_out_offset;\n if (contiguous) {\n data_in_offset = i;\n data_out_offset = i;\n } else {\n const int j2 = i % dim2;\n const int k = i / dim2;\n const int j1 = k % dim1;\n const int j0 = k / dim1;\n data_in_offset = j0 * data_in_stride0 + j1 * data_in_stride1 + j2 * data_in_stride2;\n data_out_offset = j0 * data_out_stride0 + j1 * data_out_stride1 + j2 * data_out_stride2;\n }\n\n // Determine which peer memory buffer to use\n // Note: The status bit is not affected by asynchronous\n // communication from the remote GPU.\n Flit local_message1, local_message2;\n asm volatile(\"ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(local_message1.uints[0]), \"=r\"(local_message1.uints[1]), \"=r\"(local_message1.uints[2]),\n \"=r\"(local_message1.uints[3])\n : \"l\"(local_peer1)\n : \"memory\");\n asm volatile(\"ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(local_message2.uints[0]), \"=r\"(local_message2.uints[1]), \"=r\"(local_message2.uints[2]),\n \"=r\"(local_message2.uints[3])\n : \"l\"(local_peer2)\n : \"memory\");\n const uint status1 = local_message1.uints[3] & status_mask;\n const uint status2 = local_message2.uints[3] & status_mask;\n const bool peer1_is_active = (status1 ^ status2) == 0;\n volatile int* ox = peer1_is_active ? remote_peer1 : remote_peer2;\n volatile int* ix = peer1_is_active ? local_peer1 : local_peer2;\n const uint status = peer1_is_active ? status1 : status2;\n Flit recv_message = peer1_is_active ? local_message1 : local_message2;\n\n // Send flit to remote GPU\n // Note: Set communication bit and keep status bit\n Flit send_message;\n if (has_data) {\n send_message.payload = data_in[data_in_offset];\n }\n send_message.uints[3] = communication_mask | status;\n asm volatile(\"st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};\" ::\"l\"(ox), \"r\"(send_message.uints[0]),\n \"r\"(send_message.uints[1]), \"r\"(send_message.uints[2]), \"r\"(send_message.uints[3])\n : \"memory\");\n\n // Recieve flit from peer\n while ((recv_message.uints[3] & communication_mask) == 0) {\n asm volatile(\"ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];\"\n : \"=r\"(recv_message.uints[0]), \"=r\"(recv_message.uints[1]), \"=r\"(recv_message.uints[2]),\n \"=r\"(recv_message.uints[3])\n : \"l\"(ix)\n : \"memory\");\n }\n if (has_data) {\n data_out[data_out_offset] = recv_message.payload;\n }\n\n // Reset semaphore\n // Note: Clear communication bit and invert status bit\n uint flag = ~status & status_mask;\n asm volatile(\"st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};\" ::\"l\"(ix), \"n\"(0), \"n\"(0), \"n\"(0), \"r\"(flag)\n : \"memory\");\n if (i0 + num_threads < count) {\n __threadfence_system();\n }\n }\n}\n\ntemplate \n#if __CUDA_ARCH__ >= 700\n__launch_bounds__(THREADS_PER_CTA)\n#endif\n __global__ void push_pull_halos_1d_kernel(\n // top halo,\n T* toh, int toh_stride0, int toh_stride1, int toh_stride2, // top output halo (local)\n const T* tih, int tih_stride0, int tih_stride1, int tih_stride2, // top input halo (local)\n int4* tox, // top output transfer buffer (remote peer)\n int4* tix, // top input transfer buffer (local peer)\n // btm halo\n T* boh, int boh_stride0, int boh_stride1, int boh_stride2, // btm output halo (local)\n const T* bih, int bih_stride0, int bih_stride1, int bih_stride2, // btm input halo (local)\n int4* box, // btm output transfer buffer (remote peer)\n int4* bix, // btm input transfer buffer (local peer)\n // dimensions\n int dim0, int dim1, int dim2,\n bool top_first // whether to launch communicate top halo first\n ) {\n const int num_blocks_side = gridDim.x / 2;\n const int block_id_side = (blockIdx.x < num_blocks_side ? blockIdx.x : blockIdx.x - num_blocks_side);\n const bool in_top_block = top_first == (blockIdx.x < num_blocks_side);\n if (in_top_block) {\n if (top_zero) {\n zero_tensor(dim0, dim1, dim2, toh, toh_stride0, toh_stride1, toh_stride2, threadIdx.x,\n block_id_side, num_blocks_side);\n } else {\n push_pull_tensor(dim0, dim1, dim2, tih, tih_stride0, tih_stride1, tih_stride2, toh, toh_stride0,\n toh_stride1, toh_stride2, tix, tox, threadIdx.x, block_id_side, num_blocks_side);\n }\n } else {\n if (btm_zero) {\n zero_tensor(dim0, dim1, dim2, boh, boh_stride0, boh_stride1, boh_stride2, threadIdx.x,\n block_id_side, num_blocks_side);\n } else {\n push_pull_tensor(dim0, dim1, dim2, bih, bih_stride0, bih_stride1, bih_stride2, boh, boh_stride0,\n boh_stride1, boh_stride2, bix, box, threadIdx.x, block_id_side, num_blocks_side);\n }\n }\n}\n\n__global__ void delay_kernel(int delay_nanoseconds, int* counter) {\n if (blockIdx.x == 0 && threadIdx.x == 0) {\n // waste time while doing something compiler can't predict, thus preventing it from optimizing away this code.\n int new_counter = 0;\n double elapsed = 0;\n clock_t start = clock();\n do {\n clock_t now = clock();\n elapsed = (double)(now - start) * 1e9 / CLOCKS_PER_SEC;\n ++new_counter;\n } while (elapsed < (double)delay_nanoseconds);\n *counter = new_counter;\n }\n}\n\n} // namespace\n\nnamespace apex {\nnamespace contrib {\nnamespace peer_memory {\n\nint64_t allocate_raw(int64_t size) {\n float* ptr = 0L;\n cudaMalloc(&ptr, size);\n cudaMemset(ptr, 0, size);\n return (int64_t)ptr;\n}\n\nvoid free_raw(int64_t raw) { cudaFree((void*)raw); }\n\nvoid zero(int64_t raw, int64_t size) { cudaMemset((void*)raw, 0, size); }\n\nat::Tensor get_raw_ipc_address(int64_t raw) {\n cudaIpcMemHandle_t mem_handle;\n CUDACHECK(cudaIpcGetMemHandle(&mem_handle, (void*)raw));\n const int n = sizeof(cudaIpcMemHandle_t);\n auto address_tensor = torch::empty({n}, torch::dtype(torch::kUInt8));\n auto address_tensor_p = address_tensor.data_ptr();\n memcpy(address_tensor_p, (uint8_t*)&mem_handle, n);\n return address_tensor;\n}\n\nstd::vector get_raw_peers(at::Tensor ipc_addresses, int peer_rank, int64_t raw) {\n int peer_group_size = ipc_addresses.size(0);\n std::vector results(peer_group_size);\n for (int i = 0; i < peer_group_size; ++i) {\n if (i != peer_rank) {\n cudaIpcMemHandle_t mem_handle;\n memcpy(&mem_handle, ipc_addresses.index({i}).data_ptr(), sizeof(cudaIpcMemHandle_t));\n void* p = 0L;\n CUDACHECK(cudaIpcOpenMemHandle((void**)&p, mem_handle, cudaIpcMemLazyEnablePeerAccess));\n results[i] = (int64_t)p;\n } else {\n results[i] = (int64_t)raw;\n }\n }\n return results;\n}\n\nat::Tensor blob_view_half(int64_t raw, std::vector shape, bool channels_last) {\n return blob_view((at::Half*)raw, shape, torch::dtype(torch::kFloat16).device(torch::kCUDA), channels_last);\n}\n\nat::Tensor blob_view_float(int64_t raw, std::vector shape, bool channels_last) {\n return blob_view((float*)raw, shape, torch::dtype(torch::kFloat32).device(torch::kCUDA), channels_last);\n}\n\nat::Tensor blob_view_int(int64_t raw, std::vector shape, bool channels_last) {\n return blob_view((int*)raw, shape, torch::dtype(torch::kInt32).device(torch::kCUDA), channels_last);\n}\n\nvoid push_pull_halos_1d(\n bool diagnostics, bool explicit_nhwc,\n int numSM, // number of SMs to use (zero corresponds to all SMs)\n int rank, // rank in spatial parallel group\n bool top_zero, // if top halo should be zeroed\n at::Tensor top_in_halo, // top input halo buffer (in local device memory, sent to top neighbor)\n at::Tensor top_in_transfer, // top input transfer buffer (in local peer memory)\n at::Tensor top_out_transfer, // top output transfer buffer (in top neighbor peer memory)\n at::Tensor top_out_halo, // top output halo buffer (in local device memory, received from top neighbor)\n bool btm_zero, // if btm halo should be zeroed\n at::Tensor btm_in_halo, // btm input halo buffer (in local device memory, sent to btm neighbor)\n at::Tensor btm_in_transfer, // btm input transfer buffer (in local peer memory)\n at::Tensor btm_out_transfer, // btm output transfer buffer (in btm neighbor peer memory)\n at::Tensor btm_out_halo // btm output halo buffer (in local device memory, received from btm neighbor)\n) {\n // basic checks of inputs\n TORCH_CHECK(!(top_zero && btm_zero));\n TORCH_CHECK(top_in_halo.is_cuda());\n TORCH_CHECK(top_out_transfer.is_cuda());\n TORCH_CHECK(top_in_transfer.is_cuda());\n TORCH_CHECK(top_out_halo.is_cuda());\n TORCH_CHECK(btm_in_halo.is_cuda());\n TORCH_CHECK(btm_out_transfer.is_cuda());\n TORCH_CHECK(btm_in_transfer.is_cuda());\n TORCH_CHECK(btm_out_halo.is_cuda());\n\n // tensor shapes\n int tih_N, tih_C, tih_H, tih_W;\n tensor_shape(top_in_halo, explicit_nhwc, tih_N, tih_C, tih_H, tih_W);\n int toh_N, toh_C, toh_H, toh_W;\n tensor_shape(top_out_halo, explicit_nhwc, toh_N, toh_C, toh_H, toh_W);\n int bih_N, bih_C, bih_H, bih_W;\n tensor_shape(btm_in_halo, explicit_nhwc, bih_N, bih_C, bih_H, bih_W);\n int boh_N, boh_C, boh_H, boh_W;\n tensor_shape(btm_out_halo, explicit_nhwc, boh_N, boh_C, boh_H, boh_W);\n TORCH_CHECK(toh_N == tih_N && tih_N == boh_N && boh_N == bih_N && toh_C == tih_C && tih_C == boh_C &&\n boh_C == bih_C && toh_H == tih_H && tih_H == boh_H && boh_H == bih_H && toh_W == tih_W &&\n tih_W == boh_W && boh_W == bih_W);\n int NN = toh_N, NC = toh_C, NH = toh_H, NW = toh_W;\n if (diagnostics) {\n printf(\"rank %d: NN=%d, NC=%d, NH=%d, NW=%d\\n\", rank, NN, NC, NH, NW);\n }\n TORCH_CHECK(NN == 1);\n\n // tensor strides\n int tih_stride_N, tih_stride_C, tih_stride_H, tih_stride_W;\n tensor_strides(top_in_halo, explicit_nhwc, tih_stride_N, tih_stride_C, tih_stride_H, tih_stride_W);\n int toh_stride_N, toh_stride_C, toh_stride_H, toh_stride_W;\n tensor_strides(top_out_halo, explicit_nhwc, toh_stride_N, toh_stride_C, toh_stride_H, toh_stride_W);\n int bih_stride_N, bih_stride_C, bih_stride_H, bih_stride_W;\n tensor_strides(btm_in_halo, explicit_nhwc, bih_stride_N, bih_stride_C, bih_stride_H, bih_stride_W);\n int boh_stride_N, boh_stride_C, boh_stride_H, boh_stride_W;\n tensor_strides(btm_out_halo, explicit_nhwc, boh_stride_N, boh_stride_C, boh_stride_H, boh_stride_W);\n if (diagnostics) {\n printf(\"rank %d: tih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, tih_stride_N, tih_stride_C, tih_stride_H,\n tih_stride_W);\n printf(\"rank %d: toh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, toh_stride_N, toh_stride_C, toh_stride_H,\n toh_stride_W);\n printf(\"rank %d: bih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, bih_stride_N, bih_stride_C, bih_stride_H,\n bih_stride_W);\n printf(\"rank %d: boh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, boh_stride_N, boh_stride_C, boh_stride_H,\n boh_stride_W);\n }\n\n // determine if nhwc\n bool is_nhwc = (toh_stride_C == 1);\n if (diagnostics) {\n printf(\"rank %d: is_nhwc = %s\\n\", rank, is_nhwc ? \"true\" : \"false\");\n }\n\n // determine if contiguous\n bool contiguous = true;\n if ((NN - 1) * toh_stride_N + (NC - 1) * toh_stride_C + (NH - 1) * toh_stride_H + (NW - 1) * toh_stride_W !=\n NN * NC * NH * NW - 1) {\n contiguous = false;\n }\n if ((NN - 1) * boh_stride_N + (NC - 1) * boh_stride_C + (NH - 1) * boh_stride_H + (NW - 1) * boh_stride_W !=\n NN * NC * NH * NW - 1) {\n contiguous = false;\n }\n if (!top_zero) {\n if (toh_stride_N != tih_stride_N || toh_stride_C != tih_stride_C || toh_stride_H != tih_stride_H ||\n toh_stride_W != tih_stride_W) {\n contiguous = false;\n }\n }\n if (!btm_zero) {\n if (boh_stride_N != bih_stride_N || boh_stride_C != bih_stride_C || boh_stride_H != bih_stride_H ||\n boh_stride_W != bih_stride_W) {\n contiguous = false;\n }\n }\n if (diagnostics) {\n printf(\"rank %d: contiguous = %s\\n\", rank, contiguous ? \"true\" : \"false\");\n }\n\n // determine whether to communicate top halo first\n bool top_first = rank % 2 != 0;\n if (diagnostics) {\n printf(\"rank %d: top_first = %s\\n\", rank, top_first ? \"true\" : \"false\");\n }\n\n // peer memory buffers\n int tox_size = top_out_transfer.numel() * top_out_transfer.element_size();\n int tix_size = top_in_transfer.numel() * top_in_transfer.element_size();\n int box_size = btm_out_transfer.numel() * btm_out_transfer.element_size();\n int bix_size = btm_in_transfer.numel() * btm_in_transfer.element_size();\n if (!top_zero) {\n TORCH_CHECK(top_out_transfer.is_contiguous());\n TORCH_CHECK(top_in_transfer.is_contiguous());\n TORCH_CHECK(tox_size == tix_size);\n }\n if (!btm_zero) {\n TORCH_CHECK(btm_out_transfer.is_contiguous());\n TORCH_CHECK(btm_in_transfer.is_contiguous());\n TORCH_CHECK(box_size == bix_size);\n }\n\n // figure out launch parameters\n int device;\n cudaGetDevice(&device);\n cudaDeviceProp prop;\n cudaGetDeviceProperties(&prop, device);\n if (numSM <= 0 || numSM > prop.multiProcessorCount) {\n numSM = prop.multiProcessorCount;\n }\n auto current_stream = at::cuda::getCurrentCUDAStream();\n dim3 block(THREADS_PER_CTA, 1, 1);\n\n // helper macros to launch templated kernel\n#define LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, TOP_ZERO, BTM_ZERO, KERNEL_ARGS, NUM_ELEMENTS) \\\n do { \\\n /* kernel configuration */ \\\n int numBlocksPerSm; \\\n cudaOccupancyMaxActiveBlocksPerMultiprocessor( \\\n &numBlocksPerSm, push_pull_halos_1d_kernel, THREADS_PER_CTA, 0); \\\n dim3 grid(numSM * numBlocksPerSm, 1, 1); \\\n if (grid.x % 2 != 0) { \\\n /* require even number of blocks (half for top, half for bottom) */ \\\n grid.x -= 1; \\\n } \\\n if ((grid.x / 2) * THREADS_PER_CTA > NUM_ELEMENTS) { \\\n /* only need enough blocks to cover top and bottom halo elements */ \\\n grid.x = 2 * ((NUM_ELEMENTS + THREADS_PER_CTA - 1) / THREADS_PER_CTA); \\\n } \\\n if (!TOP_ZERO) { \\\n /* require 2*128b=32B peer memory per thread */ \\\n if ((grid.x / 2) * THREADS_PER_CTA * 32 > tox_size) { \\\n grid.x = 2 * (tox_size / (THREADS_PER_CTA * 32)); \\\n } \\\n } \\\n if (!BTM_ZERO) { \\\n /* require 2*128b=32B peer memory per thread */ \\\n if ((grid.x / 2) * THREADS_PER_CTA * 32 > box_size) { \\\n grid.x = 2 * (box_size / (THREADS_PER_CTA * 32)); \\\n } \\\n } \\\n TORCH_CHECK(grid.x >= 2); \\\n \\\n /* launch kernel */ \\\n cudaLaunchCooperativeKernel((void*)push_pull_halos_1d_kernel, grid, block, \\\n KERNEL_ARGS, 0, current_stream); \\\n } while (false)\n#define LAUNCH_PUSH_PULL_HALO_KERNEL(T, CONTIGUOUS, KERNEL_ARGS, NUM_ELEMENTS) \\\n do { \\\n if (top_zero) { \\\n LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, true, false, KERNEL_ARGS, NUM_ELEMENTS); \\\n } else if (btm_zero) { \\\n LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, false, true, KERNEL_ARGS, NUM_ELEMENTS); \\\n } else { \\\n LAUNCH_PUSH_PULL_HALO_KERNEL_BASE(T, CONTIGUOUS, false, false, KERNEL_ARGS, NUM_ELEMENTS); \\\n } \\\n } while (false)\n\n AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, top_out_halo.scalar_type(), \"push_pull_halos_1d_kernel\", [&] {\n if (diagnostics) {\n printf(\"rank %d: size(scalar_t) = %ld\\n\", rank, sizeof(scalar_t));\n }\n scalar_t* toh_p = top_out_halo.data_ptr();\n scalar_t* tih_p = top_in_halo.data_ptr();\n int4* tox_p = reinterpret_cast(top_out_transfer.data_ptr());\n int4* tix_p = reinterpret_cast(top_in_transfer.data_ptr());\n scalar_t* boh_p = btm_out_halo.data_ptr();\n scalar_t* bih_p = btm_in_halo.data_ptr();\n int4* box_p = reinterpret_cast(btm_out_transfer.data_ptr());\n int4* bix_p = reinterpret_cast(btm_in_transfer.data_ptr());\n if (diagnostics) printf(\"rank %d: choosing halo exchange kernel\\n\", rank);\n\n // do int2 vector loads if channel count permits\n if (contiguous && (NN * NH * NW * NC * sizeof(scalar_t)) % sizeof(int2) == 0) {\n // can do contiguous int2 transfers\n if (diagnostics) {\n }\n toh_stride_N = toh_stride_H = toh_stride_W = toh_stride_C = 1;\n tih_stride_N = tih_stride_H = tih_stride_W = tih_stride_C = 1;\n boh_stride_N = boh_stride_H = boh_stride_W = boh_stride_C = 1;\n bih_stride_N = bih_stride_H = bih_stride_W = bih_stride_C = 1;\n NC = (NN * NH * NW * NC * sizeof(scalar_t)) / sizeof(int2);\n NN = NH = NW = 1;\n if (diagnostics) {\n printf(\"rank %d: launching contiguous int2 halo exchange kernel\\n\", rank);\n printf(\"rank %d: NC=%d, NH=%d, NW=%d\\n\", rank, NC, NH, NW);\n }\n void* kernel_args[] = {(int2**)&toh_p,\n &toh_stride_H,\n &toh_stride_W,\n &toh_stride_C,\n (int2**)&tih_p,\n &tih_stride_H,\n &tih_stride_W,\n &tih_stride_C,\n &tox_p,\n &tix_p,\n (int2**)&boh_p,\n &boh_stride_H,\n &boh_stride_W,\n &boh_stride_C,\n (int2**)&bih_p,\n &bih_stride_H,\n &bih_stride_W,\n &bih_stride_C,\n &box_p,\n &bix_p,\n &NH,\n &NW,\n &NC,\n &top_first};\n int num_elem = NN * NH * NW * NC;\n LAUNCH_PUSH_PULL_HALO_KERNEL(int2, true, kernel_args, num_elem);\n } else if (is_nhwc && (NC * sizeof(scalar_t)) % sizeof(int2) == 0) {\n // can do strided int2 transfers\n int divisor = sizeof(int2) / sizeof(scalar_t);\n if (diagnostics) {\n printf(\"rank %d: launching strided int2 halo exchange kernel\\n\", rank);\n }\n toh_stride_N /= divisor;\n toh_stride_H /= divisor;\n toh_stride_W /= divisor;\n tih_stride_N /= divisor;\n tih_stride_H /= divisor;\n tih_stride_W /= divisor;\n boh_stride_N /= divisor;\n boh_stride_H /= divisor;\n boh_stride_W /= divisor;\n bih_stride_N /= divisor;\n bih_stride_H /= divisor;\n bih_stride_W /= divisor;\n NC /= divisor;\n if (diagnostics) {\n printf(\"rank %d: divisor=%d\\n\", rank, divisor);\n printf(\"rank %d: tih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, tih_stride_N, tih_stride_C, tih_stride_H,\n tih_stride_W);\n printf(\"rank %d: toh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, toh_stride_N, toh_stride_C, toh_stride_H,\n toh_stride_W);\n printf(\"rank %d: bih_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, bih_stride_N, bih_stride_C, bih_stride_H,\n bih_stride_W);\n printf(\"rank %d: boh_stride :: N=%d, C=%d, H=%d, W=%d\\n\", rank, boh_stride_N, boh_stride_C, boh_stride_H,\n boh_stride_W);\n printf(\"rank %d: NC=%d, NH=%d, NW=%d\\n\", rank, NC, NH, NW);\n }\n void* kernel_args[] = {(int2**)&toh_p,\n &toh_stride_H,\n &toh_stride_W,\n &toh_stride_C,\n (int2**)&tih_p,\n &tih_stride_H,\n &tih_stride_W,\n &tih_stride_C,\n &tox_p,\n &tix_p,\n (int2**)&boh_p,\n &boh_stride_H,\n &boh_stride_W,\n &boh_stride_C,\n (int2**)&bih_p,\n &bih_stride_H,\n &bih_stride_W,\n &bih_stride_C,\n &box_p,\n &bix_p,\n &NH,\n &NW,\n &NC,\n &top_first};\n int num_elem = NH * NW * NC;\n LAUNCH_PUSH_PULL_HALO_KERNEL(int2, false, kernel_args, num_elem);\n } else {\n // cannot do int2 transfers\n if (diagnostics) {\n printf(\"rank %d: launching non-int2 halo exchange kernel\\n\", rank);\n }\n int num_elem = NC * NH * NW;\n if (is_nhwc) {\n void* kernel_args[] = {&toh_p, &toh_stride_H, &toh_stride_W, &toh_stride_C, &tih_p, &tih_stride_H,\n &tih_stride_W, &tih_stride_C, &tox_p, &tix_p, &boh_p, &boh_stride_H,\n &boh_stride_W, &boh_stride_C, &bih_p, &bih_stride_H, &bih_stride_W, &bih_stride_C,\n &box_p, &bix_p, &NH, &NW, &NC, &top_first};\n LAUNCH_PUSH_PULL_HALO_KERNEL(scalar_t, false, kernel_args, num_elem);\n } else {\n void* kernel_args[] = {&toh_p, &toh_stride_C, &toh_stride_H, &toh_stride_W, &tih_p, &tih_stride_C,\n &tih_stride_H, &tih_stride_W, &tox_p, &tix_p, &boh_p, &boh_stride_C,\n &boh_stride_H, &boh_stride_W, &bih_p, &bih_stride_C, &bih_stride_H, &bih_stride_W,\n &box_p, &bix_p, &NC, &NH, &NW, &top_first};\n LAUNCH_PUSH_PULL_HALO_KERNEL(scalar_t, false, kernel_args, num_elem);\n }\n }\n });\n\n#undef LAUNCH_PUSH_PULL_HALO_KERNEL_BASE\n#undef LAUNCH_PUSH_PULL_HALO_KERNEL\n}\n\n} // namespace peer_memory\n} // namespace contrib\n} // namespace apex\n" }, { "path": "apex/contrib/csrc/peer_memory/peer_memory_cuda.cuh", "content": "/**\n * Copyright (c) 2018-2021, NVIDIA CORPORATION. All rights reserved.\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\n#pragma once\n#include \n#ifndef _peer_memory_h_\n#define _peer_memory_h_\n\nnamespace apex {\nnamespace contrib {\nnamespace peer_memory {\nint64_t allocate_raw(int64_t size);\nvoid free_raw(int64_t raw);\nvoid zero(int64_t raw, int64_t size);\nat::Tensor get_raw_ipc_address(int64_t raw);\nstd::vector get_raw_peers(at::Tensor ipc_addresses, int peer_rank, int64_t raw);\nat::Tensor blob_view_half(int64_t raw, std::vector shape, bool channels_last);\nat::Tensor blob_view_float(int64_t raw, std::vector shape, bool channels_last);\nat::Tensor blob_view_int(int64_t raw, std::vector shape, bool channels_last);\nvoid push_pull_halos_1d(\n bool diagnostics, bool explicit_nhwc,\n int numSM, // number of SMs to use\n int peer_rank, // rank in spatial parallel group\n bool top_zero, // if top halo should be zeroed\n at::Tensor top_out_halo, // top output halo buffer (in local device memory, received from top neighbor)\n at::Tensor top_inp_transfer, // top input transfer buffer (in local peer memory)\n at::Tensor top_out_transfer, // top output transfer buffer (in top neighbor peer memory)\n at::Tensor top_inp_halo, // top input halo buffer (in local device memory, sent to top neighbor)\n bool btm_zero, // if btm halo should be zeroed\n at::Tensor btm_out_halo, // btm output halo buffer (in local device memory, received from btm neighbor)\n at::Tensor btm_inp_transfer, // btm input transfer buffer (in local peer memory)\n at::Tensor btm_out_transfer, // btm output transfer buffer (in btm neighbor peer memory)\n at::Tensor btm_inp_halo // btm input halo buffer (in local device memory, sent to btm neighbor)\n);\n} // namespace peer_memory\n} // namespace contrib\n} // namespace apex\n#endif\n" }, { "path": "apex/contrib/csrc/transducer/transducer_joint.cpp", "content": "#include \r\n#include \r\n\r\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\r\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\r\n#define CHECK_INPUT(x) \\\r\n CHECK_CUDA(x); \\\r\n CHECK_CONTIGUOUS(x)\r\n\r\nstd::vector transducer_joint_cuda_forward(torch::Tensor f, torch::Tensor g, torch::Tensor fLen,\r\n torch::Tensor gLen, torch::Tensor batchOffset,\r\n int64_t packedBatch, int opt, bool packOutput, bool relu,\r\n bool dropout, float dropoutProb, int tileSize);\r\n\r\nstd::vector transducer_joint_cuda_backward(std::vector in, torch::Tensor fLen,\r\n torch::Tensor gLen, torch::Tensor batchOffset, int maxFLen,\r\n int maxGLen, bool packOutput, float scale);\r\n\r\nstd::vector transducer_joint_forward(torch::Tensor f, torch::Tensor g, torch::Tensor fLen,\r\n torch::Tensor gLen, torch::Tensor batchOffset, int64_t packedBatch,\r\n int opt, bool packOutput, bool relu, bool dropout,\r\n float dropoutProb, int tileSize) {\r\n CHECK_INPUT(f);\r\n CHECK_INPUT(g);\r\n CHECK_INPUT(fLen);\r\n CHECK_INPUT(gLen);\r\n if (packOutput) CHECK_INPUT(batchOffset);\r\n return transducer_joint_cuda_forward(f, g, fLen, gLen, batchOffset, packedBatch, opt, packOutput, relu, dropout,\r\n dropoutProb, tileSize);\r\n}\r\n\r\nstd::vector transducer_joint_backward(std::vector in, torch::Tensor fLen,\r\n torch::Tensor gLen, torch::Tensor batchOffset, int maxFLen,\r\n int maxGLen, bool packOutput, float scale) {\r\n for (auto t : in) {\r\n CHECK_INPUT(t);\r\n }\r\n CHECK_INPUT(fLen);\r\n CHECK_INPUT(gLen);\r\n if (packOutput) CHECK_INPUT(batchOffset);\r\n return transducer_joint_cuda_backward(in, fLen, gLen, batchOffset, maxFLen, maxGLen, packOutput, scale);\r\n}\r\n\r\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\r\n m.def(\"forward\", &transducer_joint_forward, \"transducer joint forward (CUDA)\",\r\n py::call_guard());\r\n m.def(\"backward\", &transducer_joint_backward, \"transducer joint backward (CUDA)\",\r\n py::call_guard());\r\n}\r\n" }, { "path": "apex/contrib/csrc/transducer/transducer_joint_kernel.cu", "content": "#include \r\n#include \r\n#include \r\n#include \r\n#include \r\n\r\n#ifdef OLD_GENERATOR_PATH\r\n#include \r\n#else\r\n#include \r\n#endif\r\n\r\n#include \r\n#include \r\n\r\n#include \r\n\r\n#include \"philox.cuh\"\r\n\r\n// Warp reduce kernels to reduce N groups of data into N numbers, where N = warpSize / width.\r\n// width should be a power of 2 and should be less than warpSize.\r\ntemplate \r\n__device__ __forceinline__ scalar_t warpReduce(scalar_t x, int width = C10_WARP_SIZE) {\r\n for (unsigned offset = width / 2; offset > 0; offset /= 2) {\r\n x += __shfl_down_sync(0xffffffff, x, offset, width);\r\n }\r\n return x;\r\n}\r\n\r\ninline int largestPowerOfTwo(int x) {\r\n int y = 1;\r\n while (y <= x) y <<= 1;\r\n return y >> 1;\r\n}\r\n\r\n/*\r\nFigure out vectorization type for masks.\r\nSimilar to how PyTorch figures out acc_t here:\r\naten/src/ATen/AccumulateType.h\r\n*/\r\ntemplate \r\nstruct MaskVecType {};\r\n\r\ntemplate <>\r\nstruct MaskVecType<1> {\r\n using type = uint8_t;\r\n};\r\ntemplate <>\r\nstruct MaskVecType<2> {\r\n using type = uint16_t;\r\n};\r\ntemplate <>\r\nstruct MaskVecType<4> {\r\n using type = uint32_t;\r\n};\r\n\r\ntemplate \r\nusing mvec_type = typename MaskVecType::type;\r\n\r\n// Helper class to calculate pointer offset that can be shared by different flavors of kernels.\r\n// For fwd, batch offset and stride are different for packing and non-packing mode.\r\nstruct OffsetCalFwd {\r\n __device__ __forceinline__ OffsetCalFwd(int64_t batch, const int64_t* batchOffset, int64_t maxFLen, int64_t maxGLen,\r\n int64_t gLen, int64_t hiddenSize, bool packOutput)\r\n : batch(batch),\r\n batchOffset(batchOffset),\r\n maxFLen(maxFLen),\r\n maxGLen(maxGLen),\r\n gLen(gLen),\r\n hiddenSize(hiddenSize),\r\n packOutput(packOutput) {}\r\n\r\n int64_t batch;\r\n const int64_t* batchOffset;\r\n int64_t maxFLen;\r\n int64_t maxGLen;\r\n int64_t gLen;\r\n int64_t hiddenSize;\r\n bool packOutput;\r\n\r\n __device__ __forceinline__ int64_t getBatchOffset() {\r\n return packOutput ? ((batch == 0) ? 0 : batchOffset[batch - 1]) * hiddenSize\r\n : batch * maxFLen * maxGLen * hiddenSize;\r\n }\r\n\r\n __device__ __forceinline__ int64_t getStrideF() { return packOutput ? gLen * hiddenSize : maxGLen * hiddenSize; }\r\n};\r\n\r\n// Helper class to calculate pointer offset that can be shared by different flavors of kernels\r\n// For bwd, batch offset and stride are different for packing and non-packing mode.\r\n// The reducion is done for two input tensors. Therefore, generating two sets of offsets\r\n// according to bwdFasterDim can lead to a unified implementation in the actual kernel.\r\nstruct OffsetCalBwd {\r\n __device__ __forceinline__ OffsetCalBwd(int64_t batch, const int64_t* batchOffset, const int* fLen, const int* gLen,\r\n int64_t maxFLen, int64_t maxGLen, int64_t hiddenSize, bool packOutput,\r\n bool bwdFasterDim)\r\n : batch(batch),\r\n batchOffset(batchOffset),\r\n maxFLen(maxFLen),\r\n maxGLen(maxGLen),\r\n fLen(fLen),\r\n gLen(gLen),\r\n hiddenSize(hiddenSize),\r\n packOutput(packOutput),\r\n bwdFasterDim(bwdFasterDim) {}\r\n\r\n int64_t batch;\r\n const int64_t* batchOffset;\r\n const int* fLen;\r\n const int* gLen;\r\n int64_t maxFLen;\r\n int64_t maxGLen;\r\n int64_t hiddenSize;\r\n bool packOutput;\r\n bool bwdFasterDim; // whether doing bwd on the faster moving dimension\r\n\r\n __device__ __forceinline__ int64_t getBatchOffset() {\r\n return packOutput ? ((batch == 0) ? 0 : batchOffset[batch - 1]) * hiddenSize\r\n : batch * maxFLen * maxGLen * hiddenSize;\r\n }\r\n\r\n __device__ __forceinline__ int64_t getMaxXLen() { return bwdFasterDim ? maxGLen : maxFLen; }\r\n\r\n __device__ __forceinline__ auto getMyXLen() -> decltype(gLen[batch]) {\r\n return bwdFasterDim ? gLen[batch] : fLen[batch];\r\n }\r\n\r\n __device__ __forceinline__ auto getMyYLen() -> decltype(gLen[batch]) {\r\n return bwdFasterDim ? fLen[batch] : gLen[batch];\r\n }\r\n\r\n __device__ __forceinline__ int64_t getStrideX() {\r\n return bwdFasterDim ? hiddenSize : ((packOutput ? gLen[batch] : maxGLen) * hiddenSize);\r\n }\r\n\r\n __device__ __forceinline__ int64_t getStrideY() {\r\n return bwdFasterDim ? ((packOutput ? gLen[batch] : maxGLen) * hiddenSize) : hiddenSize;\r\n }\r\n};\r\n\r\n// Vanila transducer joint forward kernel\r\n// Detail of this joint function can be found in:\r\n// [1] Sequence Transduction with Recurrent Neural Networks.\r\n\r\n// f is a tensor of shape [batch, T, H]\r\n// g is a tensor of shape [batch, U, H]\r\n// the transducer joint does\r\n// sum = f.unsqueeze(dim=2) + g.unsqueeze(dim=1)\r\n// The resultant tensor is of shape [batch, T, U, H]\r\n// Each thread block is working on one \"batch\" of data in the output tensor, [batch, t, u, :]\r\n\r\n// This joint function can optionally pack the output where the output tensor with a shape of\r\n// [B, T, U, H] is packed into [B_packed, H].\r\n// Don't-care region (t > fLen) or (u > gLen) is removed.\r\n// To enable packing, the starting offset for each batch need to be specified with batchOffset.\r\ntemplate \r\n__global__ void transducer_joint_forward(const scalar_t* f, const scalar_t* g, const int* fLen, const int* gLen,\r\n const int64_t* batchOffset, int64_t maxFLen, int64_t maxGLen,\r\n int64_t hiddenSize, bool packOutput, scalar_t* sum) {\r\n const int batch = blockIdx.z;\r\n const int t = blockIdx.y;\r\n const int u = blockIdx.x;\r\n const auto myFLen = fLen[batch];\r\n const auto myGLen = gLen[batch];\r\n\r\n OffsetCal offsetCal(batch, batchOffset, maxFLen, maxGLen, myGLen, hiddenSize, packOutput);\r\n const auto myBatchOffset = offsetCal.getBatchOffset();\r\n const auto strideF = offsetCal.getStrideF();\r\n scalar_t const* myF = f + batch * maxFLen * hiddenSize + t * hiddenSize;\r\n scalar_t const* myG = g + batch * maxGLen * hiddenSize + u * hiddenSize;\r\n scalar_t* mySum = sum + myBatchOffset + t * strideF + u * hiddenSize;\r\n\r\n if (t < myFLen and u < myGLen) {\r\n#pragma unroll\r\n for (int h = threadIdx.x; h < hiddenSize; h += blockDim.x) {\r\n if (h < hiddenSize) {\r\n mySum[h] = myF[h] + myG[h];\r\n }\r\n }\r\n } else if (packOutput == false and t < maxFLen and u < maxGLen) {\r\n// Need to write finite data to don't-care region because we instantiate the result tensor\r\n// with torch::empty for performance reasons. Even though it is don't-care region, the\r\n// contents need to be finite, otherwise could lead to NaN in WGRAD.\r\n// In packing mode, this write is no longer necessary as we remove the don't-care region\r\n// from the output.\r\n// Picking -1 (over 0) here for ease of testing.\r\n#pragma unroll\r\n for (int h = threadIdx.x; h < hiddenSize; h += blockDim.x) {\r\n if (h < hiddenSize) {\r\n mySum[h] = -1;\r\n }\r\n }\r\n }\r\n}\r\n\r\n/*\r\nTiled version of the joint forward kernel\r\nDetail of this joint function can be found in:\r\n[1] Sequence Transduction with Recurrent Neural Networks.\r\n\r\nf is a tensor of shape [batch, T, H]\r\ng is a tensor of shape [batch, U, H]\r\nthe transducer joint does\r\nsum = f.unsqueeze(dim=2) + g.unsqueeze(dim=1)\r\nThe resultant tensor is of shape [batch, T, U, H]\r\nEach thread is working on a tile of the shape of tileF x tileG in the result tensor.\r\nThe input for the tile is first loaded in the register and is reused tileG and tileF times.\r\n\r\nThis joint function can optionally pack the output where the output tensor with a shape of\r\n[B, T, U, H] is packed into [B_packed, H].\r\nDon't-care region (t > fLen) or (u > gLen) is removed.\r\nTo enable packing, the starting offset for each batch need to be specified with batchOffset.\r\n\r\nOptionally this joint function performs ReLU and/or dropout on the joint output, which is\r\ncontrolled by arguments relu and dropout, respectively. philoxArgs is argument used for generating\r\npseudorandom number. When at least one of operations in ReLU and dropout is activated, the joint\r\nfunction is a masked operation, which is controlled by the template argument masked. In this case,\r\nmasks are saved to backward.\r\n*/\r\ntemplate \r\n__global__ void transducer_joint_tiled_forward(const scalar_t* f, const scalar_t* g, const int* fLen, const int* gLen,\r\n const int64_t* batchOffset, int64_t maxFLen, int64_t maxGLen,\r\n int64_t hiddenSize, int64_t hiddenPerBlock, bool packOutput, bool relu,\r\n bool dropout, float p, at::PhiloxCudaState philoxArgs, scalar_t* sum,\r\n uint8_t* mask) {\r\n static_assert(U == 4, \"U has to be 4, as random numbers are generated in batch of 4\");\r\n\r\n const int batch = blockIdx.z;\r\n const int t = blockIdx.y * tileF;\r\n const int hiddenBlock = (hiddenSize + hiddenPerBlock - 1) / hiddenPerBlock;\r\n const int u = blockIdx.x / hiddenBlock * tileG;\r\n const int hOffset = (blockIdx.x % hiddenBlock) * hiddenPerBlock;\r\n const int h = threadIdx.x;\r\n const auto myFLen = fLen[batch];\r\n const auto myGLen = gLen[batch];\r\n\r\n OffsetCal offsetCal(batch, batchOffset, maxFLen, maxGLen, myGLen, hiddenSize, packOutput);\r\n const auto myBatchOffset = offsetCal.getBatchOffset();\r\n const auto strideF = offsetCal.getStrideF();\r\n\r\n scalar_t const* myF = f + batch * maxFLen * hiddenSize + t * hiddenSize + hOffset;\r\n scalar_t const* myG = g + batch * maxGLen * hiddenSize + u * hiddenSize + hOffset;\r\n scalar_t* mySum = sum + myBatchOffset + t * strideF + u * hiddenSize + hOffset;\r\n uint8_t* myMask = mask + myBatchOffset + t * strideF + u * hiddenSize + hOffset;\r\n\r\n // The following code is only needed for dropout. We try to bypass them as much as possible.\r\n auto seeds = masked ? at::cuda::philox::unpack(philoxArgs)\r\n : std::make_tuple(static_cast(0), static_cast(0));\r\n uint64_t tid =\r\n masked ? (static_cast(blockIdx.z) * gridDim.y * gridDim.x + blockIdx.y * gridDim.x + blockIdx.x) *\r\n blockDim.x +\r\n threadIdx.x\r\n : 0;\r\n Philox ph(std::get<0>(seeds), tid, std::get<1>(seeds));\r\n scalar_t scale = masked ? ((p == 0) ? 0 : 1 / p) : 0;\r\n bool dropoutMask[U];\r\n\r\n if (t < myFLen and u < myGLen and hOffset + h < hiddenSize) {\r\n // register buffers for tiled input reuse\r\n scalar_t fBuffer[tileF], gBuffer[tileG];\r\n for (int i = 0; i < tileF; ++i) {\r\n if (t + i < myFLen) fBuffer[i] = myF[i * hiddenSize + h];\r\n }\r\n for (int j = 0; j < tileG; ++j) {\r\n if (u + j < myGLen) gBuffer[j] = myG[j * hiddenSize + h];\r\n }\r\n#pragma unroll\r\n for (int i = 0; i < tileF; ++i) {\r\n if (t + i < myFLen) {\r\n#pragma unroll\r\n for (int j = 0; j < tileG; ++j) {\r\n int idx = i * tileG + j;\r\n if (masked and dropout and idx % U == 0) {\r\n // For performance, generate 4 random numbers in one shot\r\n // auto rand4 = curand_uniform4(&state);\r\n auto rand4 = uniform4(ph());\r\n dropoutMask[0] = rand4.x < p;\r\n dropoutMask[1] = rand4.y < p;\r\n dropoutMask[2] = rand4.z < p;\r\n dropoutMask[3] = rand4.w < p;\r\n }\r\n\r\n if (u + j < myGLen) {\r\n scalar_t out = fBuffer[i] + gBuffer[j];\r\n if (masked) {\r\n // Apply ReLU here when relu is True\r\n bool localMask = relu ? (out > 0) : 1;\r\n localMask = dropout ? localMask & dropoutMask[idx % U] : localMask;\r\n out = dropout ? out * localMask * scale : out * localMask;\r\n myMask[i * strideF + j * hiddenSize + h] = static_cast(localMask);\r\n }\r\n mySum[i * strideF + j * hiddenSize + h] = out;\r\n } else if (packOutput == false and u + j < maxGLen)\r\n mySum[i * strideF + j * hiddenSize + h] = -1;\r\n }\r\n } else if (packOutput == false and t + i < maxFLen) {\r\n// Again need to write finite data to don't-care region\r\n#pragma unroll\r\n for (int j = 0; j < tileG; ++j) {\r\n if (u + j < maxGLen) mySum[i * strideF + j * hiddenSize + h] = -1;\r\n }\r\n }\r\n }\r\n } else if (packOutput == false and t < maxFLen and u < maxGLen and hOffset + h < hiddenSize) {\r\n// Only need to ensure the finity in normal mode\r\n#pragma unroll\r\n for (int i = 0; i < tileF; ++i) {\r\n if (t + i < maxFLen) {\r\n#pragma unroll\r\n for (int j = 0; j < tileG; ++j) {\r\n if (u + j < maxGLen) mySum[i * strideF + j * hiddenSize + h] = -1;\r\n }\r\n }\r\n }\r\n }\r\n}\r\n\r\n/*\r\nBwd operation (reduction) on one input tensor. Since the operation performed for the two input\r\ntensors are exactly the same, only one kernel is needed, and the different indexing offsets\r\nand strides are handled by OffsetCalBwd.\r\n\r\nWhen packing is enabled in the fwd op, unpacking is needed to restore the gradients in a\r\nnon-packed form.\r\n\r\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\r\nand mask contains the mask information.\r\n*/\r\ntemplate \r\n__device__ void transducer_joint_single_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\r\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\r\n int64_t maxGLen, int64_t hiddenSize, bool packOutput,\r\n bool bwdFasterDim, // whether bwd on the faster moving dimension (u)\r\n float scale, scalar_t* inGrad, int yBlockOffset = 0) {\r\n const int batch = blockIdx.z;\r\n // For the second input tensor, this offset need to be subtracted because the first yBlockOffset\r\n // sets of thread blocks are for the first input tensor.\r\n const int x = blockIdx.y - yBlockOffset;\r\n const int hOffset = blockIdx.x * C10_WARP_SIZE;\r\n const int wid = threadIdx.y;\r\n const int lid = threadIdx.x;\r\n const int numWarp = blockDim.y;\r\n extern __shared__ char smem8[];\r\n auto smem = reinterpret_cast(smem8);\r\n\r\n OffsetCal offsetCal(batch, batchOffset, fLen, gLen, maxFLen, maxGLen, hiddenSize, packOutput, bwdFasterDim);\r\n const auto maxXLen = offsetCal.getMaxXLen();\r\n const auto myXLen = offsetCal.getMyXLen();\r\n const auto myYLen = offsetCal.getMyYLen();\r\n scalar_t* myInGrad = inGrad + batch * maxXLen * hiddenSize + x * hiddenSize + hOffset;\r\n\r\n if (x < myXLen) {\r\n const auto myBatchOffset = offsetCal.getBatchOffset();\r\n const auto strideX = offsetCal.getStrideX();\r\n const auto strideY = offsetCal.getStrideY();\r\n const scalar_t* myGrad = grad + myBatchOffset + x * strideX + hOffset;\r\n const uint8_t* myMask = masked ? mask + myBatchOffset + x * strideX + hOffset : nullptr;\r\n\r\n // Each warp reduces numYPerWarp \"y\" first\r\n acc_t warpSum = 0;\r\n auto numYPerWarp = (myYLen + numWarp - 1) / numWarp;\r\n#pragma unroll\r\n for (int warpY = 0; warpY < numYPerWarp; ++warpY) {\r\n auto y = wid * numYPerWarp + warpY;\r\n if (y < myYLen and (hOffset + lid) < hiddenSize)\r\n if (masked)\r\n warpSum += static_cast(myGrad[y * strideY + lid]) * myMask[y * strideY + lid] * scale;\r\n else\r\n warpSum += myGrad[y * strideY + lid];\r\n }\r\n\r\n // transpose partial sum in SMEM and reduce further using warpReduce\r\n smem[lid * numWarp + wid] = warpSum;\r\n __syncthreads();\r\n auto sum = smem[wid * C10_WARP_SIZE + lid];\r\n sum = warpReduce(sum, numWarp);\r\n\r\n // a a b b c c d d\r\n // a a b b c c d d\r\n // a a b b c c d d\r\n // a a b b c c d d\r\n // example of 4 warps (a, b, c, d) with 8 threads per warp\r\n // Each warp need 8 / 4 = 2 threads to write the results.\r\n if (hOffset + wid * C10_WARP_SIZE / numWarp + lid / numWarp < hiddenSize) {\r\n if (lid % numWarp == 0) {\r\n myInGrad[wid * C10_WARP_SIZE / numWarp + lid / numWarp] = sum;\r\n }\r\n }\r\n } else if (wid == 0 and hOffset + lid < hiddenSize) {\r\n // Need to ensure the grad is zero for don't care region\r\n myInGrad[lid] = 0;\r\n }\r\n}\r\n\r\n/*\r\nActual bwd (reduction) kernel get launched.\r\nCall transducer_joint_single_backward twice on two input tensors.\r\nThe two bwd ops are launched together, the first op uses blockIdx.y < maxFLen, and the second op\r\nuses the rest.\r\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\r\nand mask contains the mask information.\r\n*/\r\ntemplate \r\n__global__ void transducer_joint_combined_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\r\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\r\n int64_t maxGLen, int64_t hiddenSize, bool packOutput, float scale,\r\n scalar_t* fGrad, scalar_t* gGrad) {\r\n if (blockIdx.y < maxFLen) {\r\n transducer_joint_single_backward(\r\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, false, scale, fGrad);\r\n } else {\r\n transducer_joint_single_backward(\r\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, true, scale, gGrad, maxFLen);\r\n }\r\n}\r\n\r\n/*\r\nVectorized version of transducer_joint_single_backward\r\nDoing exact same operation as transducer_joint_single_backward except the load and store are\r\nvectorized.\r\nWhen packing is enabled in the fwd op, unpacking is needed to restore the gradients in a\r\nnon-packed form.\r\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\r\nand mask contains the mask information.\r\n*/\r\ntemplate \r\n__device__ void transducer_joint_single_vec_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\r\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\r\n int64_t maxGLen, int64_t hiddenSize, bool packOutput,\r\n bool bwdFasterDim, float scale, scalar_t* inGrad,\r\n int yBlockOffset = 0) {\r\n const int batch = blockIdx.z;\r\n const int x = blockIdx.y - yBlockOffset;\r\n const int hOffset = blockIdx.x * C10_WARP_SIZE * V;\r\n const int wid = threadIdx.y;\r\n const int lid = threadIdx.x;\r\n const int numWarp = blockDim.y;\r\n\r\n // Figure out the vectorization type for mask\r\n using mvec_t = mvec_type;\r\n\r\n OffsetCal offsetCal(batch, batchOffset, fLen, gLen, maxFLen, maxGLen, hiddenSize, packOutput, bwdFasterDim);\r\n const auto maxXLen = offsetCal.getMaxXLen();\r\n const auto myXLen = offsetCal.getMyXLen();\r\n const auto myYLen = offsetCal.getMyYLen();\r\n scalar_t* myInGrad = inGrad + batch * maxXLen * hiddenSize + x * hiddenSize + hOffset;\r\n extern __shared__ char smem8[];\r\n auto smem = reinterpret_cast(smem8);\r\n\r\n acc_t warpSum[V];\r\n scalar_t inBuffer[V];\r\n uint8_t maskBuffer[V];\r\n scalar_t outBuffer[V];\r\n auto myInGradVec = reinterpret_cast(myInGrad);\r\n auto outBufferVec = reinterpret_cast(outBuffer);\r\n\r\n if (x < myXLen) {\r\n const auto myBatchOffset = offsetCal.getBatchOffset();\r\n const auto strideX = offsetCal.getStrideX();\r\n const auto strideY = offsetCal.getStrideY();\r\n const scalar_t* myGrad = grad + myBatchOffset + x * strideX + hOffset;\r\n const uint8_t* myMask = masked ? mask + myBatchOffset + x * strideX + hOffset : nullptr;\r\n\r\n for (int i = 0; i < V; ++i) warpSum[i] = 0;\r\n\r\n // Each warp reduces numYPerWarp \"y\" first\r\n auto numYPerWarp = (myYLen + numWarp - 1) / numWarp;\r\n for (int warpY = 0; warpY < numYPerWarp; ++warpY) {\r\n auto y = wid * numYPerWarp + warpY;\r\n auto myGradVec = reinterpret_cast(myGrad + y * strideY);\r\n auto myMaskVec = masked ? reinterpret_cast(myMask + y * strideY) : nullptr;\r\n auto inBufferVec = reinterpret_cast(inBuffer);\r\n auto maskBufferVec = reinterpret_cast(maskBuffer);\r\n if (hOffset + lid * V < hiddenSize and y < myYLen) {\r\n *inBufferVec = myGradVec[lid]; // vectorized load\r\n if (masked) {\r\n *maskBufferVec = myMaskVec[lid];\r\n#pragma unroll\r\n for (int i = 0; i < V; ++i) warpSum[i] += static_cast(inBuffer[i]) * maskBuffer[i] * scale;\r\n } else {\r\n#pragma unroll\r\n for (int i = 0; i < V; ++i) warpSum[i] += inBuffer[i];\r\n }\r\n }\r\n }\r\n\r\n // transpose partial sum in SMEM and reduce further using warpReduce\r\n for (int i = 0; i < V; ++i) {\r\n smem[lid * numWarp + wid] = warpSum[i];\r\n __syncthreads();\r\n auto sum = smem[wid * C10_WARP_SIZE + lid];\r\n\r\n if (hOffset + (wid * C10_WARP_SIZE / numWarp) * V < hiddenSize) {\r\n sum = warpReduce(sum, numWarp);\r\n if (lid % numWarp == 0) {\r\n outBuffer[i] = sum;\r\n }\r\n }\r\n __syncthreads();\r\n }\r\n\r\n // a a b b c c d d\r\n // a a b b c c d d\r\n // a a b b c c d d\r\n // a a b b c c d d\r\n // example of 4 warps (a, b, c, d) with 8 threads per warp\r\n // Each warp need 8 / 4 = 2 threads to write the results.\r\n if (lid % numWarp == 0 and hOffset + (wid * C10_WARP_SIZE / numWarp + lid / numWarp) * V < hiddenSize)\r\n myInGradVec[wid * C10_WARP_SIZE / numWarp + lid / numWarp] = *outBufferVec;\r\n } else if (wid == 0 and hOffset + lid * V < hiddenSize) {\r\n // Need to ensure the grad is zero for don't care region\r\n myInGradVec[lid] = 0;\r\n }\r\n}\r\n\r\n/*\r\nVecotrized version of transducer_joint_combined_backward\r\nCall transducer_joint_single_vec_backward twice on two input tensors.\r\nThe two bwd ops are launched together, the first op uses blockIdx.y < maxFLen, and the second op\r\nuses the rest.\r\nWhen ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,\r\nand mask contains the mask information.\r\n*/\r\ntemplate \r\n__global__ void transducer_joint_combined_vec_backward(const scalar_t* grad, const uint8_t* mask, const int* fLen,\r\n const int* gLen, const int64_t* batchOffset, int64_t maxFLen,\r\n int64_t maxGLen, int64_t hiddenSize, bool packOutput,\r\n float scale, scalar_t* fGrad, scalar_t* gGrad) {\r\n if (blockIdx.y < maxFLen) {\r\n transducer_joint_single_vec_backward(\r\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, false, scale, fGrad);\r\n } else {\r\n transducer_joint_single_vec_backward(\r\n grad, mask, fLen, gLen, batchOffset, maxFLen, maxGLen, hiddenSize, packOutput, true, scale, gGrad, maxFLen);\r\n }\r\n}\r\n\r\nstd::vector transducer_joint_cuda_forward(torch::Tensor f, torch::Tensor g, torch::Tensor fLen,\r\n torch::Tensor gLen, torch::Tensor batchOffset,\r\n int64_t packedBatch, int opt, bool packOutput, bool relu,\r\n bool dropout, float dropoutProb, int tileSize) {\r\n auto tensorOpt = f.options();\r\n auto dtype = f.scalar_type();\r\n const auto batchSize = f.size(0);\r\n const auto maxFLen = f.size(1);\r\n const auto maxGLen = g.size(1);\r\n const auto hiddenSize = f.size(2);\r\n bool masked = dropout or relu;\r\n\r\n int64_t* batchOffsetPtr = nullptr;\r\n torch::Tensor sum, mask;\r\n auto maskOpt = tensorOpt.dtype(torch::kUInt8);\r\n if (!packOutput) {\r\n sum = torch::empty({batchSize, maxFLen, maxGLen, hiddenSize}, tensorOpt);\r\n batchOffsetPtr = nullptr;\r\n if (masked) mask = torch::empty({batchSize, maxFLen, maxGLen, hiddenSize}, maskOpt);\r\n } else {\r\n sum = torch::empty({packedBatch, hiddenSize}, tensorOpt);\r\n batchOffsetPtr = batchOffset.data_ptr();\r\n if (masked) mask = torch::empty({packedBatch, hiddenSize}, maskOpt);\r\n }\r\n uint8_t* maskPtr = masked ? mask.data_ptr() : nullptr;\r\n\r\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\r\n\r\n TORCH_CHECK(opt == 0 or opt == 1, \"Got an invalid optimization level \", opt);\r\n // Simple heuristics\r\n const int numThread =\r\n std::min(128, (static_cast(hiddenSize) + C10_WARP_SIZE - 1) / C10_WARP_SIZE * C10_WARP_SIZE);\r\n\r\n if (opt == 0) {\r\n // vanilla kernel\r\n const int threads = numThread;\r\n const dim3 blocks(maxGLen, maxFLen, batchSize);\r\n\r\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\r\n dtype, \"transducer_joint_forward\", ([&] {\r\n transducer_joint_forward<<>>(\r\n f.data_ptr(), g.data_ptr(), fLen.data_ptr(), gLen.data_ptr(),\r\n batchOffsetPtr, maxFLen, maxGLen, hiddenSize, packOutput, sum.data_ptr());\r\n }));\r\n }\r\n if (opt == 1) {\r\n // tiled version. For simplicity, assume tileF == tileG, even though the kernel can\r\n // support more general cases.\r\n const int threads = numThread;\r\n const int hiddenPerBlock = numThread;\r\n const int hiddenBlock = (hiddenSize + hiddenPerBlock - 1) / hiddenPerBlock;\r\n const dim3 blocks((maxGLen + tileSize - 1) / tileSize * hiddenBlock, (maxFLen + tileSize - 1) / tileSize,\r\n batchSize);\r\n\r\n TORCH_CHECK(tileSize == 1 or tileSize == 2 or tileSize == 4, \"Expected tileSize to be in [1, 2, 4], but got \",\r\n tileSize);\r\n\r\n at::PhiloxCudaState rng_engine_inputs;\r\n if (masked) {\r\n // set up PRG when the input is masked. rng_engine_inputs will be used as a space filler\r\n // for non-masked calls.\r\n // Therefore no need to initialize.\r\n c10::optional gen_;\r\n auto gen = at::get_generator_or_default(gen_, at::cuda::detail::getDefaultCUDAGenerator());\r\n // counterOffset records how many cuRAND calls each thread makes. For a tiled kernel,\r\n // each thread processes tileF * tileG output elements.\r\n int64_t counterOffset = tileSize * tileSize;\r\n {\r\n std::lock_guard lock(gen->mutex_);\r\n rng_engine_inputs = gen->philox_cuda_state(counterOffset);\r\n }\r\n }\r\n\r\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\r\n dtype, \"transducer_joint_forward\", ([&] {\r\n void (*kernel)(const scalar_t*, const scalar_t*, const int*, const int*, const int64_t*, int64_t, int64_t,\r\n int64_t, int64_t, bool, bool, bool, float, at::PhiloxCudaState, scalar_t*, uint8_t*);\r\n if (masked) {\r\n switch (tileSize) {\r\n case 2:\r\n kernel = &transducer_joint_tiled_forward;\r\n break;\r\n case 4:\r\n kernel = &transducer_joint_tiled_forward;\r\n break;\r\n }\r\n } else {\r\n switch (tileSize) {\r\n case 1:\r\n kernel = &transducer_joint_tiled_forward;\r\n break;\r\n case 2:\r\n kernel = &transducer_joint_tiled_forward;\r\n break;\r\n case 4:\r\n kernel = &transducer_joint_tiled_forward;\r\n break;\r\n }\r\n }\r\n\r\n kernel<<>>(f.data_ptr(), g.data_ptr(), fLen.data_ptr(),\r\n gLen.data_ptr(), batchOffsetPtr, maxFLen, maxGLen, hiddenSize,\r\n hiddenPerBlock, packOutput, relu, dropout, 1.0f - dropoutProb,\r\n rng_engine_inputs, sum.data_ptr(), maskPtr);\r\n }));\r\n }\r\n\r\n C10_CUDA_CHECK(cudaGetLastError());\r\n if (masked)\r\n return {sum, mask};\r\n else\r\n return {sum};\r\n}\r\n\r\nstd::vector transducer_joint_cuda_backward(std::vector in, torch::Tensor fLen,\r\n torch::Tensor gLen, torch::Tensor batchOffset, int maxFLen,\r\n int maxGLen, bool packOutput, float scale) {\r\n auto grad = in[0];\r\n bool masked = (in.size() == 2);\r\n uint8_t* maskPtr = masked ? in[1].data_ptr() : nullptr;\r\n\r\n auto tensorOpt = grad.options();\r\n auto dtype = grad.scalar_type();\r\n const int batchSize = fLen.size(0);\r\n const int hiddenSize = grad.size(-1);\r\n\r\n const auto deviceProperties = at::cuda::getCurrentDeviceProperties();\r\n const int maxNumWarp = deviceProperties->maxThreadsPerBlock / C10_WARP_SIZE;\r\n\r\n torch::Tensor fGrad = torch::empty({batchSize, maxFLen, hiddenSize}, tensorOpt);\r\n torch::Tensor gGrad = torch::empty({batchSize, maxGLen, hiddenSize}, tensorOpt);\r\n\r\n int64_t* batchOffsetPtr = (!packOutput) ? nullptr : batchOffset.data_ptr();\r\n\r\n // The number \"y\" I would like each thread to work on\r\n const int workPerThread = 32;\r\n // Since the bwd for f and g have the same thread block size, we need to use the max of the two.\r\n int numWarp = largestPowerOfTwo((std::max(maxFLen, maxGLen) + workPerThread - 1) / workPerThread);\r\n // Would like to have at least 2 warps\r\n numWarp = std::max(2, numWarp);\r\n // cap on the maximum number of warps allowed\r\n numWarp = std::min(maxNumWarp, numWarp);\r\n\r\n // Need smem for transposing the partial sum. The partial sum is in a matrix of the shape\r\n // numWarp x warpSize\r\n const int smemSize = numWarp * C10_WARP_SIZE;\r\n const dim3 threads(C10_WARP_SIZE, numWarp, 1);\r\n\r\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\r\n dtype, \"transducer_joint_cuda_backward_kernel\", ([&] {\r\n auto gradPtr = grad.data_ptr();\r\n auto fLenPtr = fLen.data_ptr();\r\n auto gLenPtr = gLen.data_ptr();\r\n auto fGradPtr = fGrad.data_ptr();\r\n auto gGradPtr = gGrad.data_ptr();\r\n\r\n // resolve the acc_t type\r\n using acc_t = at::acc_type;\r\n using vec_t = uint64_t;\r\n\r\n constexpr int vectFactor = sizeof(vec_t) / sizeof(scalar_t);\r\n constexpr int vecAlignment = std::alignment_of::value;\r\n\r\n // if all input and output tensors meet the alignment requirement\r\n bool memAlign = (reinterpret_cast(gradPtr) % vecAlignment == 0) and\r\n (reinterpret_cast(fGradPtr) % vecAlignment == 0) and\r\n (reinterpret_cast(gGradPtr) % vecAlignment == 0);\r\n\r\n if (vectFactor > 1 and hiddenSize % vectFactor == 0 and memAlign) {\r\n // If vectorization helps and the alignment requirement is met, use the vectorized\r\n // kernel. For simplicity, hiddenSize needs to be a multiple vecFactor.\r\n const dim3 blocks((hiddenSize + C10_WARP_SIZE * vectFactor - 1) / (C10_WARP_SIZE * vectFactor),\r\n maxFLen + maxGLen, batchSize);\r\n if (masked) {\r\n transducer_joint_combined_vec_backward\r\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\r\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\r\n fGradPtr, gGradPtr);\r\n } else {\r\n transducer_joint_combined_vec_backward\r\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\r\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\r\n fGradPtr, gGradPtr);\r\n }\r\n } else {\r\n const dim3 blocks((hiddenSize + C10_WARP_SIZE - 1) / C10_WARP_SIZE, maxFLen + maxGLen, batchSize);\r\n if (masked) {\r\n transducer_joint_combined_backward\r\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\r\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\r\n fGradPtr, gGradPtr);\r\n } else {\r\n transducer_joint_combined_backward\r\n <<>>(gradPtr, maskPtr, fLenPtr, gLenPtr, batchOffsetPtr,\r\n maxFLen, maxGLen, hiddenSize, packOutput, scale,\r\n fGradPtr, gGradPtr);\r\n }\r\n }\r\n }));\r\n\r\n return {fGrad, gGrad};\r\n}\r\n" }, { "path": "apex/contrib/csrc/transducer/transducer_loss.cpp", "content": "#include \n\n#include \n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n#define CHECK_INPUT(x) \\\n CHECK_CUDA(x); \\\n CHECK_CONTIGUOUS(x)\n\nstd::vector transducer_loss_cuda_forward(torch::Tensor x, torch::Tensor label, torch::Tensor audLen,\n torch::Tensor txtLen, torch::Tensor batchOffset, int maxFLen,\n int blankIdx, int opt, bool packedInput);\n\ntorch::Tensor transducer_loss_cuda_backward(torch::Tensor x, torch::Tensor lossGrad, torch::Tensor alpha,\n torch::Tensor beta, torch::Tensor audLen, torch::Tensor txtLen,\n torch::Tensor label, torch::Tensor batchOffset, int maxFLen, int blankIdx,\n int opt, bool fuseSoftmaxBackward, bool packedInput);\n\nstd::vector transducer_loss_forward(torch::Tensor x, torch::Tensor label, torch::Tensor fLen,\n torch::Tensor yLen, torch::Tensor batchOffset, int maxFLen,\n int blankIdx, int opt, bool packedInput) {\n CHECK_INPUT(x);\n CHECK_INPUT(label);\n CHECK_INPUT(fLen);\n CHECK_INPUT(yLen);\n if (packedInput) CHECK_INPUT(batchOffset);\n return transducer_loss_cuda_forward(x, label, fLen, yLen, batchOffset, maxFLen, blankIdx, opt, packedInput);\n}\n\ntorch::Tensor transducer_loss_backward(torch::Tensor x, torch::Tensor lossGrad, torch::Tensor alpha, torch::Tensor beta,\n torch::Tensor fLen, torch::Tensor yLen, torch::Tensor label,\n torch::Tensor batchOffset, int maxFLen, int blankIdx, int opt,\n bool fuseSoftmaxBackward, bool packedInput) {\n CHECK_INPUT(x);\n CHECK_INPUT(label);\n CHECK_INPUT(lossGrad);\n CHECK_INPUT(alpha);\n CHECK_INPUT(beta);\n CHECK_INPUT(fLen);\n CHECK_INPUT(yLen);\n if (packedInput) CHECK_INPUT(batchOffset);\n\n return transducer_loss_cuda_backward(x, lossGrad, alpha, beta, fLen, yLen, label, batchOffset, maxFLen, blankIdx, opt,\n fuseSoftmaxBackward, packedInput);\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"forward\", &transducer_loss_forward, \"transducer loss forward (CUDA)\",\n py::call_guard());\n m.def(\"backward\", &transducer_loss_backward, \"transducer loss backward (CUDA)\",\n py::call_guard());\n}\n" }, { "path": "apex/contrib/csrc/transducer/transducer_loss_kernel.cu", "content": "#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n\ntemplate \n__device__ __forceinline__ scalar_t logSumExp(scalar_t a, scalar_t b) {\n // standard log-sum-exp trick is used here to provide better numerical stability\n return (a >= b) ? a + std::log1p(exp(b - a)) : b + std::log1p(exp(a - b));\n}\n\n// Vanilla transducer loss function (i.e. forward-backward algorithm)\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n\n// Forward (alpha) and backward (beta) path are launched together. Input is assumed to be converted\n// into log scale by the preceding log_softmax layer\n// Diagonal wavefront advancing usually used in dynamic programming is leveraged here.\n// alpha and beta are of acc_t type, as they are essentially accumulators.\n\n// This loss function supports packed input where a tensor of shape [B, T, U, H] is packed into\n// [B_packed, H].\n// Don't-care region (t > audLen) or (u > txtLen) is removed.\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_forward(const scalar_t* x, const int* label, const int* audLen, const int* txtLen,\n const int64_t* batchOffset,\n int64_t dictSize, // 64-bit indexing for data tensor\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n acc_t* alpha, acc_t* beta, scalar_t* loss) {\n const int batch = blockIdx.y;\n const int tid = threadIdx.x;\n const auto myFLen = audLen[batch];\n // Note that start of the sentence is added as 1 here\n const auto myGLen = txtLen[batch] + 1;\n const auto myLabel = label + batch * (maxGLen - 1);\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n const scalar_t* myX = x + myBatchOffset * dictSize;\n int u = tid;\n\n if (blockIdx.x == 0) {\n // alpha path\n acc_t* myAlpha = alpha + batch * maxFLen * maxGLen;\n if (u == 0) myAlpha[0] = 0;\n __syncthreads();\n\n for (int64_t step = 1; step < myFLen + myGLen - 1; ++step) {\n // Move along the diagonal wavefront to leverage available parallelism\n for (u = tid; u < myGLen; u += blockDim.x) {\n int64_t t = step - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(16) in [1]\n if (u == 0) {\n // alpha(t, u) = alpha(t-1, u) * null(t-1, u)\n myAlpha[t * maxGLen + u] = myAlpha[(t - 1) * maxGLen] + myX[((t - 1) * myStrideT) * dictSize + blankIdx];\n } else if (t == 0) {\n // alpha(t, u-1) = alpha(t, u-1) * y(t, u-1)\n myAlpha[u] = myAlpha[u - 1] + myX[(u - 1) * dictSize + myLabel[u - 1]];\n } else {\n // alpha(t, u) = alpha(t-1, u) * null(t-1, u) + alpha(t, u-1) * y(t, u-1)\n acc_t current = myAlpha[(t - 1) * maxGLen + u] + myX[((t - 1) * myStrideT + u) * dictSize + blankIdx];\n acc_t next = myAlpha[t * maxGLen + u - 1] + myX[(t * myStrideT + u - 1) * dictSize + myLabel[u - 1]];\n myAlpha[t * maxGLen + u] = logSumExp(next, current);\n }\n }\n }\n __syncthreads();\n }\n } else if (blockIdx.x == 1) {\n // beta path\n acc_t* myBeta = beta + batch * maxFLen * maxGLen;\n if (u == 0) {\n myBeta[(myFLen - 1) * maxGLen + myGLen - 1] = myX[((myFLen - 1) * myStrideT + myGLen - 1) * dictSize + blankIdx];\n }\n __syncthreads();\n\n for (int64_t step = myFLen + myGLen - 3; step >= 0; --step) {\n for (u = tid; u < myGLen; u += blockDim.x) {\n int64_t t = step - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(18) in [1]\n if (u == myGLen - 1) {\n // beta(t, u) = beta(t+1, u) * null(t, u)\n myBeta[t * maxGLen + u] = myBeta[(t + 1) * maxGLen + u] + myX[(t * myStrideT + u) * dictSize + blankIdx];\n } else if (t == myFLen - 1) {\n // beta(t, u) = beta(t, u+1) * y(t, u)\n myBeta[t * maxGLen + u] = myBeta[t * maxGLen + u + 1] + myX[(t * myStrideT + u) * dictSize + myLabel[u]];\n } else {\n // beta(t, u) = beta(t+1, u)*null(t, u) + beta(t, u+1)*y(t, u)\n acc_t current = myBeta[(t + 1) * maxGLen + u] + myX[(t * myStrideT + u) * dictSize + blankIdx];\n acc_t next = myBeta[t * maxGLen + u + 1] + myX[(t * myStrideT + u) * dictSize + myLabel[u]];\n myBeta[t * maxGLen + u] = logSumExp(next, current);\n }\n }\n }\n __syncthreads();\n }\n if (tid == 0) loss[batch] = -myBeta[0];\n }\n}\n\n// transudcer loss function (i.e. forward-backward algorithm) with batch loading optimization.\n// Compared to the vanilla version, there are two optimizations:\n// 1. load x in batch through loop unrolling to reduce the latency.\n// 2. Use registers and shared memory to hold alpha and beta values passed from one step the next.\n// For simplicity, this kernel currently only supports U <= maxThread, which should be the common\n// case. For cases where U > maxThread, the vanilla kernel is used as a fallback option.\n\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// Forward (alpha) and backward (beta) path are launched together. Input is assumed to be converted\n// into log scale by the preceding log_softmax layer\n// Diagonal wavefront advancing usually used in dynamic programming is leveraged here.\n// alpha and beta are of acc_t type, as they are essentially accumulators.\n\n// This loss function supports packed input where a tensor of shape [B, T, U, H] is packed into\n// [B_packed, H].\n// Don't-care region (t > audLen) or (u > txtLen) is removed.\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_batch_load_forward(const scalar_t* x, const int* label, const int* audLen,\n const int* txtLen, const int64_t* batchOffset, int64_t dictSize,\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n acc_t* alpha, acc_t* beta, scalar_t* loss) {\n const int batch = blockIdx.y;\n int u = threadIdx.x;\n const auto myFLen = audLen[batch];\n const auto myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n const scalar_t* myX = x + myBatchOffset * dictSize;\n scalar_t next[batchLdSize], current[batchLdSize];\n extern __shared__ char smem8[];\n auto smem = reinterpret_cast(smem8);\n\n if (blockIdx.x == 0) {\n // alpha path\n acc_t* myAlpha = alpha + batch * maxFLen * maxGLen;\n // two SMEM regions for double buffering read and write data to avoid data race\n acc_t* const sharedAlpha[2] = {smem, smem + maxGLen};\n\n sharedAlpha[0][u] = 0;\n __syncthreads();\n\n if (u == 0) myAlpha[0] = 0;\n\n auto myAlphaLabel = (u == 0) ? 0 : label[batch * (maxGLen - 1) + u - 1];\n // register used to pass value to the next step for the same thread\n acc_t prvStepAlpha = 0;\n for (int64_t step = 1; step < myFLen + myGLen - 1 + batchLdSize; step += batchLdSize) {\n// Move along the diagonal wavefront to leverage available parallelism\n// Batch loading X through loop unrolling\n#pragma unroll\n for (int i = 0; i < batchLdSize; ++i) {\n if (step + i < myFLen + myGLen - 1) {\n // index computing\n int64_t t = step + i - u;\n int64_t currentId = ((t - 1) * myStrideT + u) * dictSize + blankIdx;\n int64_t nextId = (t * myStrideT + u - 1) * dictSize + myAlphaLabel;\n // main loading loop\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n if (u == 0) {\n current[i] = myX[currentId];\n } else if (t == 0) {\n next[i] = myX[nextId];\n } else {\n current[i] = myX[currentId];\n next[i] = myX[nextId];\n }\n }\n }\n }\n // main computing loop\n for (int i = 0; i < batchLdSize; ++i) {\n // swap the pointer for double buffering\n auto sharedAlphaRd = sharedAlpha[(step + i - 1) % 2];\n auto sharedAlphaWr = sharedAlpha[(step + i) % 2];\n if (step + i < myFLen + myGLen - 1) {\n int64_t t = step + i - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(16) in [1]\n if (u == 0)\n prvStepAlpha = prvStepAlpha + current[i];\n else if (t == 0)\n prvStepAlpha = sharedAlphaRd[u - 1] + next[i];\n else\n prvStepAlpha = logSumExp(prvStepAlpha + current[i], sharedAlphaRd[u - 1] + next[i]);\n sharedAlphaWr[u] = prvStepAlpha;\n myAlpha[t * maxGLen + u] = prvStepAlpha;\n }\n }\n __syncthreads();\n }\n }\n } else if (blockIdx.x == 1) {\n // beta path\n acc_t* myBeta = beta + batch * maxFLen * maxGLen;\n // two SMEM regions for double buffering read and write data to avoid data race\n acc_t* const sharedBeta[2] = {smem, smem + maxGLen};\n sharedBeta[0][u] = myX[((myFLen - 1) * myStrideT + myGLen - 1) * dictSize + blankIdx];\n __syncthreads();\n\n auto myBetaLabel = (u == maxGLen - 1) ? 0 : label[batch * (maxGLen - 1) + u];\n // register used to pass value to the next step for the same thread\n acc_t prvStepBeta = myX[((myFLen - 1) * myStrideT + myGLen - 1) * dictSize + blankIdx];\n if (u == 0) myBeta[(myFLen - 1) * maxGLen + myGLen - 1] = prvStepBeta;\n\n for (int64_t step = 1; step < myFLen + myGLen - 1; step += batchLdSize) {\n// Move along the diagonal wavefront to leverage available parallelism\n// Batch loading X\n#pragma unroll\n for (int i = 0; i < batchLdSize; ++i) {\n if (step + i < myFLen + myGLen - 1) {\n // index computing\n int64_t t = myFLen + myGLen - (step + i) - 2 - u;\n int64_t currentId = (t * myStrideT + u) * dictSize + blankIdx;\n int64_t nextId = (t * myStrideT + u) * dictSize + myBetaLabel;\n // main loading loop\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n if (u == myGLen - 1) {\n current[i] = myX[currentId];\n } else if (t == myFLen - 1) {\n next[i] = myX[nextId];\n } else {\n current[i] = myX[currentId];\n next[i] = myX[nextId];\n }\n }\n }\n }\n // main computing loop\n for (int i = 0; i < batchLdSize; ++i) {\n // swap the pointer for double buffering\n auto sharedBetaRd = sharedBeta[(step + i - 1) % 2];\n auto sharedBetaWr = sharedBeta[(step + i) % 2];\n if (step + i < myFLen + myGLen - 1) {\n int64_t t = myFLen + myGLen - (step + i) - 2 - u;\n if (t >= 0 and t < myFLen and u >= 0 and u < myGLen) {\n // Eq(18) in [1]\n if (u == myGLen - 1)\n prvStepBeta = prvStepBeta + current[i];\n else if (t == myFLen - 1)\n prvStepBeta = sharedBetaRd[u + 1] + next[i];\n else\n prvStepBeta = logSumExp(prvStepBeta + current[i], sharedBetaRd[u + 1] + next[i]);\n sharedBetaWr[u] = prvStepBeta;\n myBeta[t * maxGLen + u] = prvStepBeta;\n }\n }\n __syncthreads();\n }\n }\n if (u == 0) loss[batch] = -prvStepBeta;\n }\n}\n\n// Vanilla transudcer loss backward operation.\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// For this backward kernel, bwd op for the preceding softmax is assumed to be handled elsewhere,\n// hence only Eq(20) in [1] is implemented in this kernel.\n\n// Each thread block works on [batch, t, :, :] of data. Each thread works on a specific u at a time\n// Since only gradients for the correct token and null token need to be updated, gradients at other\n// locations are initialized to 0.\n\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_backward(const scalar_t* x, const scalar_t* lossGrad, const int* audLen,\n const int* txtLen, const int* label, const acc_t* alpha, const acc_t* beta,\n const int64_t* batchOffset, int64_t dictSize, int64_t blankIdx,\n int64_t maxFLen, int64_t maxGLen, bool packedInput, scalar_t* xGrad) {\n const int tid = threadIdx.x;\n const int t = blockIdx.x;\n const int batch = blockIdx.y;\n const int64_t myFLen = audLen[batch];\n const int64_t myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n auto myX = x + (myBatchOffset + t * myStrideT) * dictSize;\n auto myAlpha = alpha + batch * maxFLen * maxGLen;\n auto myBeta = beta + batch * maxFLen * maxGLen;\n auto myXGrad = xGrad + (myBatchOffset + t * myStrideT) * dictSize;\n auto myLabel = label + batch * (maxGLen - 1);\n\n int64_t u = tid;\n while (t < myFLen and u < myGLen) {\n // Do the update\n // loss = -ln(Pr(y*|x))\n acc_t grad = std::log(lossGrad[batch]) + myAlpha[t * maxGLen + u] - myBeta[0];\n if (u != myGLen - 1)\n myXGrad[u * dictSize + myLabel[u]] =\n -std::exp(grad + myBeta[t * maxGLen + u + 1] + myX[u * dictSize + myLabel[u]]);\n if (t == myFLen - 1 and u == myGLen - 1)\n myXGrad[u * dictSize + blankIdx] = -std::exp(grad + myX[u * dictSize + blankIdx]);\n else if (t != myFLen - 1)\n myXGrad[u * dictSize + blankIdx] = -std::exp(grad + myBeta[(t + 1) * maxGLen + u] + myX[u * dictSize + blankIdx]);\n\n u += blockDim.x;\n }\n}\n\n// Fused transudcer loss backward operation.\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// The bwd op of the preceding softmax layer is fused in this kernel.\n// Each thread block works on [batch, t, u, :] of data. Each thread works on a specific h at a time\n\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_fused_backward(const scalar_t* x, const scalar_t* lossGrad, const int* audLen,\n const int* txtLen, const int* label, const acc_t* alpha,\n const acc_t* beta, const int64_t* batchOffset, int64_t dictSize,\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n scalar_t* xGrad) {\n const int tid = threadIdx.x;\n const int u = blockIdx.x;\n const int t = blockIdx.y;\n const int batch = blockIdx.z;\n const int64_t myFLen = audLen[batch];\n const int64_t myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n\n __shared__ acc_t commonFactor, myBetaTU, myBetaTUp1, myBetaTp1U, myLabelShared;\n auto myXGrad = xGrad + (myBatchOffset + t * myStrideT + u) * dictSize;\n\n if (t < myFLen and u < myGLen) {\n auto myX = x + (myBatchOffset + t * myStrideT + u) * dictSize;\n auto myAlpha = alpha + batch * maxFLen * maxGLen;\n auto myBeta = beta + batch * maxFLen * maxGLen;\n auto myLabel = label + batch * (maxGLen - 1);\n\n // load and store shared variables in SMEM\n if (tid == 0) {\n commonFactor = std::log(lossGrad[batch]) + myAlpha[t * maxGLen + u] - myBeta[0];\n myBetaTU = myBeta[t * maxGLen + u];\n myBetaTUp1 = myBeta[t * maxGLen + u + 1];\n myBetaTp1U = myBeta[(t + 1) * maxGLen + u];\n myLabelShared = myLabel[u];\n }\n\n __syncthreads();\n\n for (int64_t h = tid; h < dictSize; h += blockDim.x) {\n // Do the update\n acc_t grad = commonFactor + myX[h]; // loss = -ln(Pr(y*|x))\n acc_t myGrad = std::exp(grad + myBetaTU);\n if (u != myGLen - 1 and h == myLabelShared) {\n myGrad -= std::exp(grad + myBetaTUp1);\n } else if (h == blankIdx) {\n if (t == myFLen - 1 and u == myGLen - 1)\n myGrad -= std::exp(grad);\n else if (t != myFLen - 1)\n myGrad -= std::exp(grad + myBetaTp1U);\n }\n myXGrad[h] = myGrad;\n }\n } else if (!packedInput) {\n // In non-pack mode, need to make sure the gradients for don't-care regions are zero.\n for (int64_t h = tid; h < dictSize; h += blockDim.x) {\n myXGrad[h] = 0;\n }\n }\n}\n\n// Vectorized version of fused transudcer loss backward operation.\n// Detail of this loss function can be found in:\n// [1] Sequence Transduction with Recurrent Neural Networks.\n// The bwd op of the preceding softmax layer is fused in this kernel.\n// Each thread block works on [batch, t, u, :] of data. Each thread works on a specific h at a time\n\n// To support the packed input, the starting offsets for each batch need to be specified with\n// batchOffset.\ntemplate \n__global__ void transducer_loss_fused_vec_backward(const scalar_t* x, const scalar_t* lossGrad, const int* audLen,\n const int* txtLen, const int* label, const acc_t* alpha,\n const acc_t* beta, const int64_t* batchOffset, int64_t dictSize,\n int64_t blankIdx, int64_t maxFLen, int64_t maxGLen, bool packedInput,\n scalar_t* xGrad) {\n const int tid = threadIdx.x;\n const int u = blockIdx.x;\n const int t = blockIdx.y;\n const int batch = blockIdx.z;\n const int64_t myFLen = audLen[batch];\n const int64_t myGLen = txtLen[batch] + 1;\n const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch - 1]) : batch * maxFLen * maxGLen;\n const int64_t myStrideT = packedInput ? myGLen : maxGLen;\n\n __shared__ acc_t commonFactor, myBetaTU, myBetaTUp1, myBetaTp1U, myLabelShared;\n auto myXGrad = xGrad + (myBatchOffset + t * myStrideT + u) * dictSize;\n auto myX = x + (myBatchOffset + t * myStrideT + u) * dictSize;\n auto myAlpha = alpha + batch * maxFLen * maxGLen;\n auto myBeta = beta + batch * maxFLen * maxGLen;\n auto myLabel = label + batch * (maxGLen - 1);\n\n // Variabels for vectorization\n scalar_t myXBuffer[V], myXGradBuffer[V];\n auto myXVec = reinterpret_cast(myX);\n auto myXGradVec = reinterpret_cast(myXGrad);\n auto myXBufferVec = reinterpret_cast(myXBuffer);\n auto myXGradBufferVec = reinterpret_cast(myXGradBuffer);\n if (t < myFLen and u < myGLen) {\n // load and store shared variables in SMEM\n if (tid == 0) {\n commonFactor = std::log(lossGrad[batch]) + myAlpha[t * maxGLen + u] - myBeta[0];\n myBetaTU = myBeta[t * maxGLen + u];\n if (t != myFLen - 1) myBetaTp1U = myBeta[(t + 1) * maxGLen + u];\n if (u != myGLen - 1) {\n myBetaTUp1 = myBeta[t * maxGLen + u + 1];\n myLabelShared = myLabel[u];\n }\n }\n\n __syncthreads();\n\n#pragma unroll\n for (int64_t h0 = tid * V; h0 < dictSize; h0 += blockDim.x * V) {\n // Load myX in a vector form\n *myXBufferVec = myXVec[h0 / V];\n// Do the update for a vector of input\n#pragma unroll\n for (int i = 0; i < V; ++i) {\n auto h = h0 + i;\n acc_t grad = commonFactor + myXBuffer[i]; // loss = -ln(Pr(y*|x))\n acc_t myGrad = std::exp(grad + myBetaTU);\n if (u != myGLen - 1 and h == myLabelShared) {\n myGrad -= std::exp(grad + myBetaTUp1);\n } else if (h == blankIdx) {\n if (t == myFLen - 1 and u == myGLen - 1)\n myGrad -= std::exp(grad);\n else if (t != myFLen - 1)\n myGrad -= std::exp(grad + myBetaTp1U);\n }\n myXGradBuffer[i] = myGrad;\n }\n\n // Store myXGrad in a vector form\n myXGradVec[h0 / V] = *myXGradBufferVec;\n }\n } else if (!packedInput) {\n // In non-pack mode, need to make sure the gradients for don't-care regions are zero.\n for (int64_t h0 = tid * V; h0 < dictSize; h0 += blockDim.x * V) {\n myXGradVec[h0 / V] = 0;\n }\n }\n}\n\nstd::vector transducer_loss_cuda_forward(torch::Tensor x, torch::Tensor label, torch::Tensor audLen,\n torch::Tensor txtLen, torch::Tensor batchOffset, int maxFLen,\n int blankIdx, int opt, bool packedInput) {\n auto scalarType = x.scalar_type();\n auto tensorOpt = x.options();\n const int batchSize = label.size(0);\n const int maxGLen = label.size(1) + 1;\n const int dictSize = x.size(-1);\n\n TORCH_CHECK(blankIdx >= 0 and blankIdx < dictSize, \"Expected blank index to be in the range of 0 to \", dictSize - 1,\n \", but got \", blankIdx);\n TORCH_CHECK(opt == -1 or opt == 0 or opt == 1, \"Got an invalid optimization level \", opt);\n\n // The data type of alpha and beta will be resolved at dispatch time,\n // hence defined here and assigned later\n torch::Tensor alpha;\n torch::Tensor beta;\n torch::Tensor loss = torch::empty({batchSize}, tensorOpt);\n const auto deviceProperties = at::cuda::getCurrentDeviceProperties();\n const auto maxThreadPerBlock = deviceProperties->maxThreadsPerBlock;\n const auto maxSmemPerBlock = deviceProperties->sharedMemPerBlock;\n const auto batchOffsetPtr = packedInput ? batchOffset.data_ptr() : nullptr;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n scalarType, \"transducer_loss_cuda_forward\", ([&] {\n // resolve accumulation type\n using acc_t = at::acc_type;\n auto accType = c10::CppTypeToScalarType::value;\n auto accTensorOpt = tensorOpt.dtype(accType);\n alpha = torch::empty({batchSize, maxFLen, maxGLen}, accTensorOpt);\n beta = torch::empty({batchSize, maxFLen, maxGLen}, accTensorOpt);\n\n // decide what kernel to launch based on the problem size\n // if the required SMEM size or number threads exceeds the limit, fall back to the vanilla\n // kernel.\n const auto smemSize = 2 * maxGLen * sizeof(acc_t);\n const auto optFallBack = (maxGLen > maxThreadPerBlock or smemSize > maxSmemPerBlock) ? 0\n : (opt == -1) ? 1\n : opt;\n const int threads = std::min(maxThreadPerBlock, maxGLen);\n const dim3 blocks(2, batchSize, 1);\n\n if (optFallBack == 0)\n transducer_loss_forward<<>>(\n x.data_ptr(), label.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n batchOffsetPtr, dictSize, blankIdx, maxFLen, maxGLen, packedInput, alpha.data_ptr(),\n beta.data_ptr(), loss.data_ptr());\n else if (optFallBack == 1)\n transducer_loss_batch_load_forward<<>>(\n x.data_ptr(), label.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n batchOffsetPtr, dictSize, blankIdx, maxFLen, maxGLen, packedInput, alpha.data_ptr(),\n beta.data_ptr(), loss.data_ptr());\n }));\n C10_CUDA_CHECK(cudaGetLastError());\n\n return {alpha, beta, loss};\n}\n\ntorch::Tensor transducer_loss_cuda_backward(torch::Tensor x, torch::Tensor lossGrad, torch::Tensor alpha,\n torch::Tensor beta, torch::Tensor audLen, torch::Tensor txtLen,\n torch::Tensor label, torch::Tensor batchOffset, int maxFLen, int blankIdx,\n int opt, bool fuseSoftmaxBackward, bool packedInput) {\n auto dtype = x.scalar_type();\n torch::Tensor xGrad;\n const int batchSize = label.size(0);\n const int maxGLen = label.size(1) + 1;\n const int dictSize = x.size(-1);\n const auto deviceProperties = at::cuda::getCurrentDeviceProperties();\n const int maxThreadPerBlock = deviceProperties->maxThreadsPerBlock;\n const int warpSize = deviceProperties->warpSize;\n const auto batchOffsetPtr = packedInput ? batchOffset.data_ptr() : nullptr;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (fuseSoftmaxBackward) {\n // alloc empty tensors for performance, hence need to ensure zeros are writtern to\n // don't-care region in the kernel.\n xGrad = torch::empty_like(x);\n\n // Would like each thread to work on 4 hidden units\n const int workPerThread = 4;\n // Don't want to have more than 128 threads per thread block\n const int maxThreadPerElmt = std::min(128, maxThreadPerBlock);\n const int threads = std::min(maxThreadPerElmt, std::max(warpSize, (dictSize + workPerThread - 1) / workPerThread));\n const dim3 blocks(maxGLen, maxFLen, batchSize);\n\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(\n dtype, \"transducer_loss_cuda_backward\", ([&] {\n using vec_t = uint64_t;\n using acc_t = at::acc_type;\n constexpr int vectFactor = sizeof(vec_t) / sizeof(scalar_t);\n constexpr int vecAlignment = std::alignment_of::value;\n // if all input and output tensors meet the alignment requirement\n bool memAlign = reinterpret_cast(x.data_ptr()) % vecAlignment == 0 and\n reinterpret_cast(xGrad.data_ptr()) % vecAlignment == 0;\n\n if (vectFactor > 1 and dictSize % vectFactor == 0 and memAlign) {\n transducer_loss_fused_vec_backward<<>>(\n x.data_ptr(), lossGrad.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n label.data_ptr(), alpha.data_ptr(), beta.data_ptr(), batchOffsetPtr, dictSize,\n blankIdx, maxFLen, maxGLen, packedInput, xGrad.data_ptr());\n } else {\n transducer_loss_fused_backward<<>>(\n x.data_ptr(), lossGrad.data_ptr(), audLen.data_ptr(), txtLen.data_ptr(),\n label.data_ptr(), alpha.data_ptr(), beta.data_ptr(), batchOffsetPtr, dictSize,\n blankIdx, maxFLen, maxGLen, packedInput, xGrad.data_ptr());\n }\n }));\n } else {\n // for non-fused kernel, the gradients need to be writtern are very sparse, hence initialize\n // the tensor with all zeros.\n xGrad = torch::zeros_like(x);\n // don't launch more threads than needed.\n const int threads = std::min(maxThreadPerBlock, maxGLen);\n const dim3 blocks(maxFLen, batchSize);\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(dtype, \"transducer_loss_cuda_backward\", ([&] {\n using acc_t = at::acc_type;\n transducer_loss_backward<<>>(\n x.data_ptr(), lossGrad.data_ptr(),\n audLen.data_ptr(), txtLen.data_ptr(), label.data_ptr(),\n alpha.data_ptr(), beta.data_ptr(), batchOffsetPtr, dictSize,\n blankIdx, maxFLen, maxGLen, packedInput, xGrad.data_ptr());\n }));\n }\n C10_CUDA_CHECK(cudaGetLastError());\n\n return xGrad;\n}\n" }, { "path": "apex/contrib/csrc/xentropy/interface.cpp", "content": "#include \n\n#include \n\n// CUDA forward declarations\n\nstd::vector softmax_xentropy_cuda(const at::Tensor& input, const at::Tensor& labels, const float smoothing,\n const bool half_to_float);\n\nat::Tensor softmax_xentropy_backward_cuda(const at::Tensor& grad_loss, const at::Tensor& logits,\n const at::Tensor& max_log_sum_exp, const at::Tensor& labels,\n const float smoothing);\n\n// C++ interface\n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor\")\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\n#define CHECK_INPUT(x) \\\n CHECK_CUDA(x); \\\n CHECK_CONTIGUOUS(x)\n\nstd::vector softmax_xentropy_forward(const at::Tensor& input, const at::Tensor& labels,\n const float smoothing, const bool half_to_float) {\n CHECK_CUDA(input);\n CHECK_INPUT(labels);\n\n return softmax_xentropy_cuda(input, labels, smoothing, half_to_float);\n}\n\nat::Tensor softmax_xentropy_backward(const at::Tensor& grad_loss, const at::Tensor& logits,\n const at::Tensor& max_log_sum_exp, const at::Tensor& labels,\n const float smoothing) {\n CHECK_CUDA(grad_loss);\n CHECK_CUDA(logits);\n CHECK_INPUT(max_log_sum_exp);\n CHECK_INPUT(labels);\n\n return softmax_xentropy_backward_cuda(grad_loss, logits, max_log_sum_exp, labels, smoothing);\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"forward\", &softmax_xentropy_forward, \"Softmax cross entropy loss with label smoothing forward (CUDA)\",\n py::call_guard());\n m.def(\"backward\", &softmax_xentropy_backward, \"Softmax cross entropy loss with label smoothing backward (CUDA)\",\n py::call_guard());\n // ref: https://pybind11.readthedocs.io/en/stable/basics.html#exporting-variables\n py::object version = py::cast(\n#ifdef XENTROPY_VER\n XENTROPY_VER\n#else\n std::string{}\n#endif\n );\n m.attr(\"__version__\") = version;\n}\n" }, { "path": "apex/contrib/csrc/xentropy/xentropy_kernel.cu", "content": "/**\n * From PyTorch:\n *\n * Copyright (c) 2016- Facebook, Inc (Adam Paszke)\n * Copyright (c) 2014- Facebook, Inc (Soumith Chintala)\n * Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)\n * Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu)\n * Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)\n * Copyright (c) 2011-2013 NYU (Clement Farabet)\n * Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)\n * Copyright (c) 2006 Idiap Research Institute (Samy Bengio)\n * Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)\n *\n * From Caffe2:\n *\n * Copyright (c) 2016-present, Facebook Inc. All rights reserved.\n *\n * All contributions by Facebook:\n * Copyright (c) 2016 Facebook Inc.\n *\n * All contributions by Google:\n * Copyright (c) 2015 Google Inc.\n * All rights reserved.\n *\n * All contributions by Yangqing Jia:\n * Copyright (c) 2015 Yangqing Jia\n * All rights reserved.\n *\n * All contributions from Caffe:\n * Copyright(c) 2013, 2014, 2015, the respective contributors\n * All rights reserved.\n *\n * All other contributions:\n * Copyright(c) 2015, 2016 the respective contributors\n * All rights reserved.\n *\n * Caffe2 uses a copyright model similar to Caffe: each contributor holds\n * copyright over their contributions to Caffe2. The project versioning records\n * all such contribution and copyright details. If a contributor wants to further\n * mark their specific copyright on a particular contribution, they should\n * indicate their copyright solely in the commit message of the change when it is\n * committed.\n *\n * All rights reserved.\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n *\n * 1. Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n *\n * 2. Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n *\n * 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America\n * and IDIAP Research Institute nor the names of its contributors may be\n * used to endorse or promote products derived from this software without\n * specific prior written permission.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE\n * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n * POSSIBILITY OF SUCH DAMAGE.\n */\n#include \n#include \n#include \n\n#include \n\n#include \"type_shim.h\"\n\n#define ALIGN_BYTES 16\n\nusing Tensor = at::Tensor;\nusing TensorList = at::TensorList;\nusing ScalarType = at::ScalarType;\nusing at::acc_type;\n\ntemplate \nstruct LogSoftMaxForwardEpilogue {\n __device__ __forceinline__ LogSoftMaxForwardEpilogue(AccumT max_input, AccumT sum)\n : logsum(max_input + std::log(sum)) {}\n\n __device__ __forceinline__ LogSoftMaxForwardEpilogue(AccumT max_log_sum_exp) : logsum(max_log_sum_exp) {}\n\n __device__ __forceinline__ OutT operator()(T input) const { return static_cast(input - logsum); }\n\n const AccumT logsum;\n};\n\ntemplate \nstruct LogSoftMaxBackwardEpilogue {\n __device__ __forceinline__ LogSoftMaxBackwardEpilogue(AccumT sum) : sum(sum) {}\n\n __device__ __forceinline__ T operator()(OutT gradOutput, OutT output) const {\n return static_cast(gradOutput - std::exp(static_cast(output)) * sum);\n }\n\n const AccumT sum;\n};\n\nconst int max_threads = 1024;\n\ninline dim3 SoftMax_getBlockSize(int ILP, uint64_t dim_size) {\n uint64_t block_size = 1;\n uint64_t max_block_size = std::min(dim_size / ILP, static_cast(max_threads));\n while (block_size < (max_block_size / 2)) block_size *= 2;\n // Launch at least a single warp - the kernel assumes that.\n block_size = std::max(block_size, static_cast(32));\n return dim3(block_size);\n}\n\ntemplate \nstruct Add {\n __device__ __forceinline__ T operator()(T a, T b) const { return a + b; }\n};\n\ntemplate \nstruct Max {\n __device__ __forceinline__ T operator()(T a, T b) const { return a < b ? b : a; }\n};\n\n////////////////////////////////////////////////////////////////////////////////\n// Regular kernel (fast when dim_size is large; requires inner_size == 1)\n////////////////////////////////////////////////////////////////////////////////\n\ntemplate \nstruct MaxFloat {\n __device__ __forceinline__ AccumT operator()(AccumT max, T v) const { return ::max(max, (AccumT)v); }\n};\n\ntemplate \nstruct AddFloat {\n __device__ __forceinline__ AccumT operator()(AccumT sum, T v) const { return sum + v; }\n};\n\ntemplate \nstruct SumExpFloat {\n __device__ __forceinline__ SumExpFloat(AccumT v) : max_k(v) {}\n\n __device__ __forceinline__ AccumT operator()(AccumT sum, T v) const { return sum + std::exp(v - max_k); }\n\n const AccumT max_k;\n};\n\ntemplate