[
{
"path": ".idea/CCSR.iml",
"content": "\n\n \n \n \n \n \n \n \n \n \n"
},
{
"path": ".idea/inspectionProfiles/Project_Default.xml",
"content": "\n \n \n \n \n \n \n"
},
{
"path": ".idea/inspectionProfiles/profiles_settings.xml",
"content": "\n \n \n \n \n"
},
{
"path": ".idea/modules.xml",
"content": "\n\n \n \n \n \n \n"
},
{
"path": ".idea/vcs.xml",
"content": "\n\n \n \n \n"
},
{
"path": ".idea/workspace.xml",
"content": "\n\n \n
\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n 1734539270044\n \n \n 1734539270044\n \n \n \n"
},
{
"path": "ADD/dnnlib/__init__.py",
"content": "# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\nfrom .util import EasyDict, make_cache_dir_path\n"
},
{
"path": "ADD/dnnlib/util.py",
"content": "# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n\"\"\"Miscellaneous utility classes and functions.\"\"\"\n\nimport ctypes\nimport fnmatch\nimport importlib\nimport inspect\nimport os\nimport sys\nimport types\nimport io\nimport pickle\nimport re\nimport requests\nimport html\nimport hashlib\nimport glob\nimport tempfile\nimport urllib\nimport urllib.request\nimport uuid\nfrom typing import Any, List, Tuple, Union, Optional\nfrom distutils.util import strtobool\nimport shutil\n\nimport numpy as np\n\n\n# Util classes\n# ------------------------------------------------------------------------------------------\n\nclass EasyDict(dict):\n \"\"\"Convenience class that behaves like a dict but allows access with the attribute syntax.\"\"\"\n\n def __getattr__(self, name: str) -> Any:\n try:\n return self[name]\n except KeyError:\n raise AttributeError(name)\n\n def __setattr__(self, name: str, value: Any) -> None:\n self[name] = value\n\n def __delattr__(self, name: str) -> None:\n del self[name]\n\n\nclass Logger(object):\n \"\"\"Redirect stderr to stdout, optionally print stdout to a file, and optionally force flushing on both stdout and the file.\"\"\"\n\n def __init__(self, file_name: Optional[str] = None, file_mode: str = \"w\", should_flush: bool = True):\n self.file = None\n\n if file_name is not None:\n self.file = open(file_name, file_mode)\n\n self.should_flush = should_flush\n self.stdout = sys.stdout\n self.stderr = sys.stderr\n\n sys.stdout = self\n sys.stderr = self\n\n def __enter__(self) -> \"Logger\":\n return self\n\n def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:\n self.close()\n\n def write(self, text: Union[str, bytes]) -> None:\n \"\"\"Write text to stdout (and a file) and optionally flush.\"\"\"\n if isinstance(text, bytes):\n text = text.decode()\n if len(text) == 0: # workaround for a bug in VSCode debugger: sys.stdout.write(''); sys.stdout.flush() => crash\n return\n\n if self.file is not None:\n self.file.write(text)\n\n self.stdout.write(text)\n\n if self.should_flush:\n self.flush()\n\n def flush(self) -> None:\n \"\"\"Flush written text to both stdout and a file, if open.\"\"\"\n if self.file is not None:\n self.file.flush()\n\n self.stdout.flush()\n\n def close(self) -> None:\n \"\"\"Flush, close possible files, and remove stdout/stderr mirroring.\"\"\"\n self.flush()\n\n # if using multiple loggers, prevent closing in wrong order\n if sys.stdout is self:\n sys.stdout = self.stdout\n if sys.stderr is self:\n sys.stderr = self.stderr\n\n if self.file is not None:\n self.file.close()\n self.file = None\n\n\n# Cache directories\n# ------------------------------------------------------------------------------------------\n\n_dnnlib_cache_dir = None\n\ndef set_cache_dir(path: str) -> None:\n global _dnnlib_cache_dir\n _dnnlib_cache_dir = path\n\n\ndef make_cache_dir_path(*paths: str) -> str:\n if _dnnlib_cache_dir is not None:\n return os.path.join(_dnnlib_cache_dir, *paths)\n if 'DNNLIB_CACHE_DIR' in os.environ:\n return os.path.join(os.environ['DNNLIB_CACHE_DIR'], *paths)\n if 'HOME' in os.environ:\n return os.path.join(os.environ['HOME'], '.cache', 'dnnlib', *paths)\n if 'USERPROFILE' in os.environ:\n return os.path.join(os.environ['USERPROFILE'], '.cache', 'dnnlib', *paths)\n return os.path.join(tempfile.gettempdir(), '.cache', 'dnnlib', *paths)\n\n\n# Small util functions\n# ------------------------------------------------------------------------------------------\n\ndef format_time(seconds: Union[int, float]) -> str:\n \"\"\"Convert the seconds to human readable string with days, hours, minutes and seconds.\"\"\"\n s = int(np.rint(seconds))\n\n if s < 60:\n return \"{0}s\".format(s)\n elif s < 60 * 60:\n return \"{0}m {1:02}s\".format(s // 60, s % 60)\n elif s < 24 * 60 * 60:\n return \"{0}h {1:02}m {2:02}s\".format(s // (60 * 60), (s // 60) % 60, s % 60)\n else:\n return \"{0}d {1:02}h {2:02}m\".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24, (s // 60) % 60)\n\n\ndef format_time_brief(seconds: Union[int, float]) -> str:\n \"\"\"Convert the seconds to human readable string with days, hours, minutes and seconds.\"\"\"\n s = int(np.rint(seconds))\n\n if s < 60:\n return \"{0}s\".format(s)\n elif s < 60 * 60:\n return \"{0}m {1:02}s\".format(s // 60, s % 60)\n elif s < 24 * 60 * 60:\n return \"{0}h {1:02}m\".format(s // (60 * 60), (s // 60) % 60)\n else:\n return \"{0}d {1:02}h\".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24)\n\n\ndef ask_yes_no(question: str) -> bool:\n \"\"\"Ask the user the question until the user inputs a valid answer.\"\"\"\n while True:\n try:\n print(\"{0} [y/n]\".format(question))\n return strtobool(input().lower())\n except ValueError:\n pass\n\n\ndef tuple_product(t: Tuple) -> Any:\n \"\"\"Calculate the product of the tuple elements.\"\"\"\n result = 1\n\n for v in t:\n result *= v\n\n return result\n\n\n_str_to_ctype = {\n \"uint8\": ctypes.c_ubyte,\n \"uint16\": ctypes.c_uint16,\n \"uint32\": ctypes.c_uint32,\n \"uint64\": ctypes.c_uint64,\n \"int8\": ctypes.c_byte,\n \"int16\": ctypes.c_int16,\n \"int32\": ctypes.c_int32,\n \"int64\": ctypes.c_int64,\n \"float32\": ctypes.c_float,\n \"float64\": ctypes.c_double\n}\n\n\ndef get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:\n \"\"\"Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes.\"\"\"\n type_str = None\n\n if isinstance(type_obj, str):\n type_str = type_obj\n elif hasattr(type_obj, \"__name__\"):\n type_str = type_obj.__name__\n elif hasattr(type_obj, \"name\"):\n type_str = type_obj.name\n else:\n raise RuntimeError(\"Cannot infer type name from input\")\n\n assert type_str in _str_to_ctype.keys()\n\n my_dtype = np.dtype(type_str)\n my_ctype = _str_to_ctype[type_str]\n\n assert my_dtype.itemsize == ctypes.sizeof(my_ctype)\n\n return my_dtype, my_ctype\n\n\ndef is_pickleable(obj: Any) -> bool:\n try:\n with io.BytesIO() as stream:\n pickle.dump(obj, stream)\n return True\n except:\n return False\n\n\n# Functionality to import modules/objects by name, and call functions by name\n# ------------------------------------------------------------------------------------------\n\ndef get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, str]:\n \"\"\"Searches for the underlying module behind the name to some python object.\n Returns the module and the object name (original name with module part removed).\"\"\"\n\n # allow convenience shorthands, substitute them by full names\n obj_name = re.sub(\"^np.\", \"numpy.\", obj_name)\n obj_name = re.sub(\"^tf.\", \"tensorflow.\", obj_name)\n\n # list alternatives for (module_name, local_obj_name)\n parts = obj_name.split(\".\")\n name_pairs = [(\".\".join(parts[:i]), \".\".join(parts[i:])) for i in range(len(parts), 0, -1)]\n\n # try each alternative in turn\n for module_name, local_obj_name in name_pairs:\n try:\n module = importlib.import_module(module_name) # may raise ImportError\n get_obj_from_module(module, local_obj_name) # may raise AttributeError\n return module, local_obj_name\n except:\n pass\n\n # maybe some of the modules themselves contain errors?\n for module_name, _local_obj_name in name_pairs:\n try:\n importlib.import_module(module_name) # may raise ImportError\n except ImportError:\n if not str(sys.exc_info()[1]).startswith(\"No module named '\" + module_name + \"'\"):\n raise\n\n # maybe the requested attribute is missing?\n for module_name, local_obj_name in name_pairs:\n try:\n module = importlib.import_module(module_name) # may raise ImportError\n get_obj_from_module(module, local_obj_name) # may raise AttributeError\n except ImportError:\n pass\n\n # we are out of luck, but we have no idea why\n raise ImportError(obj_name)\n\n\ndef get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:\n \"\"\"Traverses the object name and returns the last (rightmost) python object.\"\"\"\n if obj_name == '':\n return module\n obj = module\n for part in obj_name.split(\".\"):\n obj = getattr(obj, part)\n return obj\n\n\ndef get_obj_by_name(name: str) -> Any:\n \"\"\"Finds the python object with the given name.\"\"\"\n module, obj_name = get_module_from_obj_name(name)\n return get_obj_from_module(module, obj_name)\n\n\ndef call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:\n \"\"\"Finds the python object with the given name and calls it as a function.\"\"\"\n assert func_name is not None\n func_obj = get_obj_by_name(func_name)\n assert callable(func_obj)\n return func_obj(*args, **kwargs)\n\n\ndef construct_class_by_name(*args, class_name: str = None, **kwargs) -> Any:\n \"\"\"Finds the python class with the given name and constructs it with the given arguments.\"\"\"\n return call_func_by_name(*args, func_name=class_name, **kwargs)\n\n\ndef get_module_dir_by_obj_name(obj_name: str) -> str:\n \"\"\"Get the directory path of the module containing the given object name.\"\"\"\n module, _ = get_module_from_obj_name(obj_name)\n return os.path.dirname(inspect.getfile(module))\n\n\ndef is_top_level_function(obj: Any) -> bool:\n \"\"\"Determine whether the given object is a top-level function, i.e., defined at module scope using 'def'.\"\"\"\n return callable(obj) and obj.__name__ in sys.modules[obj.__module__].__dict__\n\n\ndef get_top_level_function_name(obj: Any) -> str:\n \"\"\"Return the fully-qualified name of a top-level function.\"\"\"\n assert is_top_level_function(obj)\n module = obj.__module__\n if module == '__main__':\n module = os.path.splitext(os.path.basename(sys.modules[module].__file__))[0]\n return module + \".\" + obj.__name__\n\n\n# File system helpers\n# ------------------------------------------------------------------------------------------\n\ndef list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False) -> List[Tuple[str, str]]:\n \"\"\"List all files recursively in a given directory while ignoring given file and directory names.\n Returns list of tuples containing both absolute and relative paths.\"\"\"\n assert os.path.isdir(dir_path)\n base_name = os.path.basename(os.path.normpath(dir_path))\n\n if ignores is None:\n ignores = []\n\n result = []\n\n for root, dirs, files in os.walk(dir_path, topdown=True):\n for ignore_ in ignores:\n dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)]\n\n # dirs need to be edited in-place\n for d in dirs_to_remove:\n dirs.remove(d)\n\n files = [f for f in files if not fnmatch.fnmatch(f, ignore_)]\n\n absolute_paths = [os.path.join(root, f) for f in files]\n relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths]\n\n if add_base_to_relative:\n relative_paths = [os.path.join(base_name, p) for p in relative_paths]\n\n assert len(absolute_paths) == len(relative_paths)\n result += zip(absolute_paths, relative_paths)\n\n return result\n\n\ndef copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:\n \"\"\"Takes in a list of tuples of (src, dst) paths and copies files.\n Will create all necessary directories.\"\"\"\n for file in files:\n target_dir_name = os.path.dirname(file[1])\n\n # will create all intermediate-level directories\n if not os.path.exists(target_dir_name):\n os.makedirs(target_dir_name)\n\n shutil.copyfile(file[0], file[1])\n\n\n# URL helpers\n# ------------------------------------------------------------------------------------------\n\ndef is_url(obj: Any, allow_file_urls: bool = False) -> bool:\n \"\"\"Determine whether the given object is a valid URL string.\"\"\"\n if not isinstance(obj, str) or not \"://\" in obj:\n return False\n if allow_file_urls and obj.startswith('file://'):\n return True\n try:\n res = requests.compat.urlparse(obj)\n if not res.scheme or not res.netloc or not \".\" in res.netloc:\n return False\n res = requests.compat.urlparse(requests.compat.urljoin(obj, \"/\"))\n if not res.scheme or not res.netloc or not \".\" in res.netloc:\n return False\n except:\n return False\n return True\n\n\ndef open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False, cache: bool = True) -> Any:\n \"\"\"Download the given URL and return a binary-mode file object to access the data.\"\"\"\n assert num_attempts >= 1\n assert not (return_filename and (not cache))\n\n # Doesn't look like an URL scheme so interpret it as a local filename.\n if not re.match('^[a-z]+://', url):\n return url if return_filename else open(url, \"rb\")\n\n # Handle file URLs. This code handles unusual file:// patterns that\n # arise on Windows:\n #\n # file:///c:/foo.txt\n #\n # which would translate to a local '/c:/foo.txt' filename that's\n # invalid. Drop the forward slash for such pathnames.\n #\n # If you touch this code path, you should test it on both Linux and\n # Windows.\n #\n # Some internet resources suggest using urllib.request.url2pathname() but\n # but that converts forward slashes to backslashes and this causes\n # its own set of problems.\n if url.startswith('file://'):\n filename = urllib.parse.urlparse(url).path\n if re.match(r'^/[a-zA-Z]:', filename):\n filename = filename[1:]\n return filename if return_filename else open(filename, \"rb\")\n\n assert is_url(url)\n\n # Lookup from cache.\n if cache_dir is None:\n cache_dir = make_cache_dir_path('downloads')\n\n url_md5 = hashlib.md5(url.encode(\"utf-8\")).hexdigest()\n if cache:\n cache_files = glob.glob(os.path.join(cache_dir, url_md5 + \"_*\"))\n if len(cache_files) == 1:\n filename = cache_files[0]\n return filename if return_filename else open(filename, \"rb\")\n\n # Download.\n url_name = None\n url_data = None\n with requests.Session() as session:\n if verbose:\n print(\"Downloading %s ...\" % url, end=\"\", flush=True)\n for attempts_left in reversed(range(num_attempts)):\n try:\n with session.get(url) as res:\n res.raise_for_status()\n if len(res.content) == 0:\n raise IOError(\"No data received\")\n\n if len(res.content) < 8192:\n content_str = res.content.decode(\"utf-8\")\n if \"download_warning\" in res.headers.get(\"Set-Cookie\", \"\"):\n links = [html.unescape(link) for link in content_str.split('\"') if \"export=download\" in link]\n if len(links) == 1:\n url = requests.compat.urljoin(url, links[0])\n raise IOError(\"Google Drive virus checker nag\")\n if \"Google Drive - Quota exceeded\" in content_str:\n raise IOError(\"Google Drive download quota exceeded -- please try again later\")\n\n match = re.search(r'filename=\"([^\"]*)\"', res.headers.get(\"Content-Disposition\", \"\"))\n url_name = match[1] if match else url\n url_data = res.content\n if verbose:\n print(\" done\")\n break\n except KeyboardInterrupt:\n raise\n except:\n if not attempts_left:\n if verbose:\n print(\" failed\")\n raise\n if verbose:\n print(\".\", end=\"\", flush=True)\n\n # Save to cache.\n if cache:\n safe_name = re.sub(r\"[^0-9a-zA-Z-._]\", \"_\", url_name)\n safe_name = safe_name[:min(len(safe_name), 128)]\n cache_file = os.path.join(cache_dir, url_md5 + \"_\" + safe_name)\n temp_file = os.path.join(cache_dir, \"tmp_\" + uuid.uuid4().hex + \"_\" + url_md5 + \"_\" + safe_name)\n os.makedirs(cache_dir, exist_ok=True)\n with open(temp_file, \"wb\") as f:\n f.write(url_data)\n os.replace(temp_file, cache_file) # atomic\n if return_filename:\n return cache_file\n\n # Return data as file object.\n assert not return_filename\n return io.BytesIO(url_data)\n"
},
{
"path": "ADD/layers/__init__.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the Apache License, Version 2.0\n# found in the LICENSE file in the root directory of this source tree.\n\nfrom .dino_head import DINOHead\nfrom .mlp import Mlp\nfrom .patch_embed import PatchEmbed\nfrom .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused\nfrom .block import NestedTensorBlock\nfrom .attention import MemEffAttention\n"
},
{
"path": "ADD/layers/attention.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the Apache License, Version 2.0\n# found in the LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py\n\nimport logging\nimport os\nimport warnings\n\nfrom torch import Tensor\nfrom torch import nn\n\n\nlogger = logging.getLogger(\"dinov2\")\n\n\nXFORMERS_ENABLED = os.environ.get(\"XFORMERS_DISABLED\") is None\ntry:\n if XFORMERS_ENABLED:\n from xformers.ops import memory_efficient_attention, unbind\n\n XFORMERS_AVAILABLE = True\n warnings.warn(\"xFormers is available (Attention)\")\n else:\n warnings.warn(\"xFormers is disabled (Attention)\")\n raise ImportError\nexcept ImportError:\n XFORMERS_AVAILABLE = False\n warnings.warn(\"xFormers is not available (Attention)\")\n\n\nclass Attention(nn.Module):\n def __init__(\n self,\n dim: int,\n num_heads: int = 8,\n qkv_bias: bool = False,\n proj_bias: bool = True,\n attn_drop: float = 0.0,\n proj_drop: float = 0.0,\n ) -> None:\n super().__init__()\n self.num_heads = num_heads\n head_dim = dim // num_heads\n self.scale = head_dim**-0.5\n\n self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n self.attn_drop = nn.Dropout(attn_drop)\n self.proj = nn.Linear(dim, dim, bias=proj_bias)\n self.proj_drop = nn.Dropout(proj_drop)\n\n def forward(self, x: Tensor) -> Tensor:\n B, N, C = x.shape\n qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n\n q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]\n attn = q @ k.transpose(-2, -1)\n\n attn = attn.softmax(dim=-1)\n attn = self.attn_drop(attn)\n\n x = (attn @ v).transpose(1, 2).reshape(B, N, C)\n x = self.proj(x)\n x = self.proj_drop(x)\n return x\n\n\nclass MemEffAttention(Attention):\n def forward(self, x: Tensor, attn_bias=None) -> Tensor:\n if not XFORMERS_AVAILABLE:\n if attn_bias is not None:\n raise AssertionError(\"xFormers is required for using nested tensors\")\n return super().forward(x)\n\n B, N, C = x.shape\n qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)\n\n q, k, v = unbind(qkv, 2)\n\n x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)\n x = x.reshape([B, N, C])\n\n x = self.proj(x)\n x = self.proj_drop(x)\n return x\n"
},
{
"path": "ADD/layers/block.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the Apache License, Version 2.0\n# found in the LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py\n\nimport logging\nimport os\nfrom typing import Callable, List, Any, Tuple, Dict\nimport warnings\n\nimport torch\nfrom torch import nn, Tensor\n\nfrom .attention import Attention, MemEffAttention\nfrom .drop_path import DropPath\nfrom .layer_scale import LayerScale\nfrom .mlp import Mlp\n\n\nlogger = logging.getLogger(\"dinov2\")\n\n\nXFORMERS_ENABLED = os.environ.get(\"XFORMERS_DISABLED\") is None\ntry:\n if XFORMERS_ENABLED:\n from xformers.ops import fmha, scaled_index_add, index_select_cat\n\n XFORMERS_AVAILABLE = True\n warnings.warn(\"xFormers is available (Block)\")\n else:\n warnings.warn(\"xFormers is disabled (Block)\")\n raise ImportError\nexcept ImportError:\n XFORMERS_AVAILABLE = False\n\n warnings.warn(\"xFormers is not available (Block)\")\n\n\nclass Block(nn.Module):\n def __init__(\n self,\n dim: int,\n num_heads: int,\n mlp_ratio: float = 4.0,\n qkv_bias: bool = False,\n proj_bias: bool = True,\n ffn_bias: bool = True,\n drop: float = 0.0,\n attn_drop: float = 0.0,\n init_values=None,\n drop_path: float = 0.0,\n act_layer: Callable[..., nn.Module] = nn.GELU,\n norm_layer: Callable[..., nn.Module] = nn.LayerNorm,\n attn_class: Callable[..., nn.Module] = Attention,\n ffn_layer: Callable[..., nn.Module] = Mlp,\n ) -> None:\n super().__init__()\n # print(f\"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}\")\n self.norm1 = norm_layer(dim)\n self.attn = attn_class(\n dim,\n num_heads=num_heads,\n qkv_bias=qkv_bias,\n proj_bias=proj_bias,\n attn_drop=attn_drop,\n proj_drop=drop,\n )\n self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()\n self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()\n\n self.norm2 = norm_layer(dim)\n mlp_hidden_dim = int(dim * mlp_ratio)\n self.mlp = ffn_layer(\n in_features=dim,\n hidden_features=mlp_hidden_dim,\n act_layer=act_layer,\n drop=drop,\n bias=ffn_bias,\n )\n self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()\n self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()\n\n self.sample_drop_ratio = drop_path\n\n def forward(self, x: Tensor) -> Tensor:\n def attn_residual_func(x: Tensor) -> Tensor:\n return self.ls1(self.attn(self.norm1(x)))\n\n def ffn_residual_func(x: Tensor) -> Tensor:\n return self.ls2(self.mlp(self.norm2(x)))\n\n if self.training and self.sample_drop_ratio > 0.1:\n # the overhead is compensated only for a drop path rate larger than 0.1\n x = drop_add_residual_stochastic_depth(\n x,\n residual_func=attn_residual_func,\n sample_drop_ratio=self.sample_drop_ratio,\n )\n x = drop_add_residual_stochastic_depth(\n x,\n residual_func=ffn_residual_func,\n sample_drop_ratio=self.sample_drop_ratio,\n )\n elif self.training and self.sample_drop_ratio > 0.0:\n x = x + self.drop_path1(attn_residual_func(x))\n x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2\n else:\n x = x + attn_residual_func(x)\n x = x + ffn_residual_func(x)\n return x\n\n\ndef drop_add_residual_stochastic_depth(\n x: Tensor,\n residual_func: Callable[[Tensor], Tensor],\n sample_drop_ratio: float = 0.0,\n) -> Tensor:\n # 1) extract subset using permutation\n b, n, d = x.shape\n sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)\n brange = (torch.randperm(b, device=x.device))[:sample_subset_size]\n x_subset = x[brange]\n\n # 2) apply residual_func to get residual\n residual = residual_func(x_subset)\n\n x_flat = x.flatten(1)\n residual = residual.flatten(1)\n\n residual_scale_factor = b / sample_subset_size\n\n # 3) add the residual\n x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)\n return x_plus_residual.view_as(x)\n\n\ndef get_branges_scales(x, sample_drop_ratio=0.0):\n b, n, d = x.shape\n sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)\n brange = (torch.randperm(b, device=x.device))[:sample_subset_size]\n residual_scale_factor = b / sample_subset_size\n return brange, residual_scale_factor\n\n\ndef add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):\n if scaling_vector is None:\n x_flat = x.flatten(1)\n residual = residual.flatten(1)\n x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)\n else:\n x_plus_residual = scaled_index_add(\n x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor\n )\n return x_plus_residual\n\n\nattn_bias_cache: Dict[Tuple, Any] = {}\n\n\ndef get_attn_bias_and_cat(x_list, branges=None):\n \"\"\"\n this will perform the index select, cat the tensors, and provide the attn_bias from cache\n \"\"\"\n batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]\n all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))\n if all_shapes not in attn_bias_cache.keys():\n seqlens = []\n for b, x in zip(batch_sizes, x_list):\n for _ in range(b):\n seqlens.append(x.shape[1])\n attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)\n attn_bias._batch_sizes = batch_sizes\n attn_bias_cache[all_shapes] = attn_bias\n\n if branges is not None:\n cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])\n else:\n tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)\n cat_tensors = torch.cat(tensors_bs1, dim=1)\n\n return attn_bias_cache[all_shapes], cat_tensors\n\n\ndef drop_add_residual_stochastic_depth_list(\n x_list: List[Tensor],\n residual_func: Callable[[Tensor, Any], Tensor],\n sample_drop_ratio: float = 0.0,\n scaling_vector=None,\n) -> Tensor:\n # 1) generate random set of indices for dropping samples in the batch\n branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]\n branges = [s[0] for s in branges_scales]\n residual_scale_factors = [s[1] for s in branges_scales]\n\n # 2) get attention bias and index+concat the tensors\n attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)\n\n # 3) apply residual_func to get residual, and split the result\n residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore\n\n outputs = []\n for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):\n outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))\n return outputs\n\n\nclass NestedTensorBlock(Block):\n def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:\n \"\"\"\n x_list contains a list of tensors to nest together and run\n \"\"\"\n assert isinstance(self.attn, MemEffAttention)\n\n if self.training and self.sample_drop_ratio > 0.0:\n\n def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:\n return self.attn(self.norm1(x), attn_bias=attn_bias)\n\n def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:\n return self.mlp(self.norm2(x))\n\n x_list = drop_add_residual_stochastic_depth_list(\n x_list,\n residual_func=attn_residual_func,\n sample_drop_ratio=self.sample_drop_ratio,\n scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,\n )\n x_list = drop_add_residual_stochastic_depth_list(\n x_list,\n residual_func=ffn_residual_func,\n sample_drop_ratio=self.sample_drop_ratio,\n scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,\n )\n return x_list\n else:\n\n def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:\n return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))\n\n def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:\n return self.ls2(self.mlp(self.norm2(x)))\n\n attn_bias, x = get_attn_bias_and_cat(x_list)\n x = x + attn_residual_func(x, attn_bias=attn_bias)\n x = x + ffn_residual_func(x)\n return attn_bias.split(x)\n\n def forward(self, x_or_x_list):\n if isinstance(x_or_x_list, Tensor):\n return super().forward(x_or_x_list)\n elif isinstance(x_or_x_list, list):\n if not XFORMERS_AVAILABLE:\n raise AssertionError(\"xFormers is required for using nested tensors\")\n return self.forward_nested(x_or_x_list)\n else:\n raise AssertionError\n"
},
{
"path": "ADD/layers/dino_head.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the Apache License, Version 2.0\n# found in the LICENSE file in the root directory of this source tree.\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn.init import trunc_normal_\nfrom torch.nn.utils import weight_norm\n\n\nclass DINOHead(nn.Module):\n def __init__(\n self,\n in_dim,\n out_dim,\n use_bn=False,\n nlayers=3,\n hidden_dim=2048,\n bottleneck_dim=256,\n mlp_bias=True,\n ):\n super().__init__()\n nlayers = max(nlayers, 1)\n self.mlp = _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=hidden_dim, use_bn=use_bn, bias=mlp_bias)\n self.apply(self._init_weights)\n self.last_layer = weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))\n self.last_layer.weight_g.data.fill_(1)\n\n def _init_weights(self, m):\n if isinstance(m, nn.Linear):\n trunc_normal_(m.weight, std=0.02)\n if isinstance(m, nn.Linear) and m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def forward(self, x):\n x = self.mlp(x)\n eps = 1e-6 if x.dtype == torch.float16 else 1e-12\n x = nn.functional.normalize(x, dim=-1, p=2, eps=eps)\n x = self.last_layer(x)\n return x\n\n\ndef _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=None, use_bn=False, bias=True):\n if nlayers == 1:\n return nn.Linear(in_dim, bottleneck_dim, bias=bias)\n else:\n layers = [nn.Linear(in_dim, hidden_dim, bias=bias)]\n if use_bn:\n layers.append(nn.BatchNorm1d(hidden_dim))\n layers.append(nn.GELU())\n for _ in range(nlayers - 2):\n layers.append(nn.Linear(hidden_dim, hidden_dim, bias=bias))\n if use_bn:\n layers.append(nn.BatchNorm1d(hidden_dim))\n layers.append(nn.GELU())\n layers.append(nn.Linear(hidden_dim, bottleneck_dim, bias=bias))\n return nn.Sequential(*layers)\n"
},
{
"path": "ADD/layers/drop_path.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the Apache License, Version 2.0\n# found in the LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py\n\n\nfrom torch import nn\n\n\ndef drop_path(x, drop_prob: float = 0.0, training: bool = False):\n if drop_prob == 0.0 or not training:\n return x\n keep_prob = 1 - drop_prob\n shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets\n random_tensor = x.new_empty(shape).bernoulli_(keep_prob)\n if keep_prob > 0.0:\n random_tensor.div_(keep_prob)\n output = x * random_tensor\n return output\n\n\nclass DropPath(nn.Module):\n \"\"\"Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).\"\"\"\n\n def __init__(self, drop_prob=None):\n super(DropPath, self).__init__()\n self.drop_prob = drop_prob\n\n def forward(self, x):\n return drop_path(x, self.drop_prob, self.training)\n"
},
{
"path": "ADD/layers/layer_scale.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the Apache License, Version 2.0\n# found in the LICENSE file in the root directory of this source tree.\n\n# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110\n\nfrom typing import Union\n\nimport torch\nfrom torch import Tensor\nfrom torch import nn\n\n\nclass LayerScale(nn.Module):\n def __init__(\n self,\n dim: int,\n init_values: Union[float, Tensor] = 1e-5,\n inplace: bool = False,\n ) -> None:\n super().__init__()\n self.inplace = inplace\n self.gamma = nn.Parameter(init_values * torch.ones(dim))\n\n def forward(self, x: Tensor) -> Tensor:\n return x.mul_(self.gamma) if self.inplace else x * self.gamma\n"
},
{
"path": "ADD/layers/mlp.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the Apache License, Version 2.0\n# found in the LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py\n\n\nfrom typing import Callable, Optional\n\nfrom torch import Tensor, nn\n\n\nclass Mlp(nn.Module):\n def __init__(\n self,\n in_features: int,\n hidden_features: Optional[int] = None,\n out_features: Optional[int] = None,\n act_layer: Callable[..., nn.Module] = nn.GELU,\n drop: float = 0.0,\n bias: bool = True,\n ) -> None:\n super().__init__()\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)\n self.act = act_layer()\n self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)\n self.drop = nn.Dropout(drop)\n\n def forward(self, x: Tensor) -> Tensor:\n x = self.fc1(x)\n x = self.act(x)\n x = self.drop(x)\n x = self.fc2(x)\n x = self.drop(x)\n return x\n"
},
{
"path": "ADD/layers/patch_embed.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the Apache License, Version 2.0\n# found in the LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py\n\nfrom typing import Callable, Optional, Tuple, Union\n\nfrom torch import Tensor\nimport torch.nn as nn\n\n\ndef make_2tuple(x):\n if isinstance(x, tuple):\n assert len(x) == 2\n return x\n\n assert isinstance(x, int)\n return (x, x)\n\n\nclass PatchEmbed(nn.Module):\n \"\"\"\n 2D image to patch embedding: (B,C,H,W) -> (B,N,D)\n\n Args:\n img_size: Image size.\n patch_size: Patch token size.\n in_chans: Number of input image channels.\n embed_dim: Number of linear projection output channels.\n norm_layer: Normalization layer.\n \"\"\"\n\n def __init__(\n self,\n img_size: Union[int, Tuple[int, int]] = 224,\n patch_size: Union[int, Tuple[int, int]] = 16,\n in_chans: int = 3,\n embed_dim: int = 768,\n norm_layer: Optional[Callable] = None,\n flatten_embedding: bool = True,\n ) -> None:\n super().__init__()\n\n image_HW = make_2tuple(img_size)\n patch_HW = make_2tuple(patch_size)\n patch_grid_size = (\n image_HW[0] // patch_HW[0],\n image_HW[1] // patch_HW[1],\n )\n\n self.img_size = image_HW\n self.patch_size = patch_HW\n self.patches_resolution = patch_grid_size\n self.num_patches = patch_grid_size[0] * patch_grid_size[1]\n\n #self.in_chans = in_chans\n self.embed_dim = embed_dim\n\n self.flatten_embedding = flatten_embedding\n\n self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)\n self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()\n\n def forward(self, x: Tensor) -> Tensor:\n _, _, H, W = x.shape\n patch_H, patch_W = self.patch_size\n\n assert H % patch_H == 0, f\"Input image height {H} is not a multiple of patch height {patch_H}\"\n assert W % patch_W == 0, f\"Input image width {W} is not a multiple of patch width: {patch_W}\"\n\n x = self.proj(x) # B C H W\n H, W = x.size(2), x.size(3)\n x = x.flatten(2).transpose(1, 2) # B HW C\n x = self.norm(x)\n if not self.flatten_embedding:\n x = x.reshape(-1, H, W, self.embed_dim) # B H W C\n return x\n\n #def flops(self) -> float:\n #Ho, Wo = self.patches_resolution\n #flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])\n #if self.norm is not None:\n # flops += Ho * Wo * self.embed_dim\n #return flops\n"
},
{
"path": "ADD/layers/swiglu_ffn.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the Apache License, Version 2.0\n# found in the LICENSE file in the root directory of this source tree.\n\nimport os\nfrom typing import Callable, Optional\nimport warnings\n\nfrom torch import Tensor, nn\nimport torch.nn.functional as F\n\n\nclass SwiGLUFFN(nn.Module):\n def __init__(\n self,\n in_features: int,\n hidden_features: Optional[int] = None,\n out_features: Optional[int] = None,\n act_layer: Callable[..., nn.Module] = None,\n drop: float = 0.0,\n bias: bool = True,\n ) -> None:\n super().__init__()\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)\n self.w3 = nn.Linear(hidden_features, out_features, bias=bias)\n\n def forward(self, x: Tensor) -> Tensor:\n x12 = self.w12(x)\n x1, x2 = x12.chunk(2, dim=-1)\n hidden = F.silu(x1) * x2\n return self.w3(hidden)\n\n\nXFORMERS_ENABLED = os.environ.get(\"XFORMERS_DISABLED\") is None\ntry:\n if XFORMERS_ENABLED:\n from xformers.ops import SwiGLU\n\n XFORMERS_AVAILABLE = True\n warnings.warn(\"xFormers is available (SwiGLU)\")\n else:\n warnings.warn(\"xFormers is disabled (SwiGLU)\")\n raise ImportError\nexcept ImportError:\n SwiGLU = SwiGLUFFN\n XFORMERS_AVAILABLE = False\n\n warnings.warn(\"xFormers is not available (SwiGLU)\")\n\n\nclass SwiGLUFFNFused(SwiGLU):\n def __init__(\n self,\n in_features: int,\n hidden_features: Optional[int] = None,\n out_features: Optional[int] = None,\n act_layer: Callable[..., nn.Module] = None,\n drop: float = 0.0,\n bias: bool = True,\n ) -> None:\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8\n super().__init__(\n in_features=in_features,\n hidden_features=hidden_features,\n out_features=out_features,\n bias=bias,\n )\n"
},
{
"path": "ADD/models/discriminator.py",
"content": "# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n\"\"\"\nProjected discriminator architecture from\n\"StyleGAN-T: Unlocking the Power of GANs for Fast Large-Scale Text-to-Image Synthesis\".\n\"\"\"\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.nn.utils.spectral_norm import SpectralNorm\nfrom torchvision.transforms import RandomCrop, Normalize\nimport timm\nfrom timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD\n\nfrom ADD.th_utils import misc\nfrom models.shared import ResidualBlock, FullyConnectedLayer\nfrom models.vit_utils import make_vit_backbone, forward_vit, make_sd_backbone\nfrom models.DiffAugment import DiffAugment\nfrom ADD.utils.util_net import reload_model_\n\nfrom functools import partial\n\nclass SpectralConv1d(nn.Conv1d):\n def __init__(self, *args, **kwargs):\n super().__init__(*args, **kwargs)\n SpectralNorm.apply(self, name='weight', n_power_iterations=1, dim=0, eps=1e-12)\n\n\nclass BatchNormLocal(nn.Module):\n def __init__(self, num_features: int, affine: bool = True, virtual_bs: int = 3, eps: float = 1e-5):\n super().__init__()\n self.virtual_bs = virtual_bs\n self.eps = eps\n self.affine = affine\n\n if self.affine:\n self.weight = nn.Parameter(torch.ones(num_features))\n self.bias = nn.Parameter(torch.zeros(num_features))\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n shape = x.size()\n\n # Reshape batch into groups.\n G = np.ceil(x.size(0)/self.virtual_bs).astype(int)\n x = x.view(G, -1, x.size(-2), x.size(-1))\n\n # Calculate stats.\n mean = x.mean([1, 3], keepdim=True)\n var = x.var([1, 3], keepdim=True, unbiased=False)\n x = (x - mean) / (torch.sqrt(var + self.eps))\n\n if self.affine:\n x = x * self.weight[None, :, None] + self.bias[None, :, None]\n\n return x.view(shape)\n\n\ndef make_block(channels: int, kernel_size: int) -> nn.Module:\n return nn.Sequential(\n SpectralConv1d(\n channels,\n channels,\n kernel_size = kernel_size,\n padding = kernel_size//2,\n padding_mode = 'circular',\n ),\n #BatchNormLocal(channels),\n nn.GroupNorm(4, channels),\n nn.LeakyReLU(0.2, True),\n )\n\nclass DiscHead(nn.Module):\n def __init__(self, channels: int, c_dim: int, cmap_dim: int = 64):\n super().__init__()\n self.channels = channels\n self.c_dim = c_dim\n self.cmap_dim = cmap_dim\n\n self.main = nn.Sequential(\n make_block(channels, kernel_size=1),\n ResidualBlock(make_block(channels, kernel_size=9))\n )\n\n if self.c_dim > 0:\n self.cmapper = FullyConnectedLayer(self.c_dim, cmap_dim)\n self.cls = SpectralConv1d(channels, cmap_dim, kernel_size=1, padding=0)\n else:\n self.cls = SpectralConv1d(channels, 1, kernel_size=1, padding=0)\n\n def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:\n h = self.main(x)\n out = self.cls(h)\n\n if self.c_dim > 0:\n cmap = self.cmapper(c).unsqueeze(-1)\n out = (out * cmap).sum(1, keepdim=True) * (1 / np.sqrt(self.cmap_dim))\n\n return out\n\nclass DINO(torch.nn.Module):\n def __init__(self, hooks: list[int] = [2,5,8,11], hook_patch: bool = True):\n super().__init__()\n self.n_hooks = len(hooks) + int(hook_patch)\n\n self.model = make_vit_backbone(\n timm.create_model('vit_small_patch16_224.dino', pretrained=False),\n patch_size=[16,16], hooks=hooks, hook_patch=hook_patch,\n )\n reload_model_(self.model, torch.load('preset/models/dino/dino_deitsmall16_pretrain.pth'))\n self.model = self.model.eval().requires_grad_(False)\n\n\n self.img_resolution = self.model.model.patch_embed.img_size[0]\n self.embed_dim = self.model.model.embed_dim\n self.norm = Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n ''' input: x in [0, 1]; output: dict of activations '''\n x = F.interpolate(x, self.img_resolution, mode='area')\n x = self.norm(x)\n features = forward_vit(self.model, x)\n return features\n\n\nclass ProjectedDiscriminator(nn.Module):\n def __init__(self, c_dim: int, diffaug: bool = True, p_crop: float = 0.5):\n super().__init__()\n self.c_dim = c_dim\n self.diffaug = diffaug\n self.p_crop = p_crop\n\n self.dino = DINO()\n\n heads = []\n for i in range(self.dino.n_hooks):\n heads += [str(i), DiscHead(self.dino.embed_dim, c_dim)],\n self.heads = nn.ModuleDict(heads)\n\n def train(self, mode: bool = True):\n self.dino = self.dino.train(False)\n self.heads = self.heads.train(mode)\n return self\n\n def eval(self):\n return self.train(False)\n\n def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:\n # Apply augmentation (x in [-1, 1]).\n if self.diffaug:\n x = DiffAugment(x, policy='translation,cutout')\n\n # Transform to [0, 1].\n x = x.add(1).div(2)\n\n # Take crops with probablity p_crop if the image is larger.\n if x.size(-1) > self.dino.img_resolution and np.random.random() < self.p_crop:\n x = RandomCrop(self.dino.img_resolution)(x)\n\n # Forward pass through DINO ViT.\n features = self.dino(x)\n\n # Apply discriminator heads.\n logits = []\n for k, head in self.heads.items():\n features[k].requires_grad_(True)\n logits.append(head(features[k], c).view(x.size(0), -1))\n #logits = torch.cat(logits, dim=1)\n\n return logits, features\n\n"
},
{
"path": "ADD/models/vit.py",
"content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the Apache License, Version 2.0\n# found in the LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py\n\nfrom functools import partial\nimport math\nimport logging\nfrom typing import Sequence, Tuple, Union, Callable\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.checkpoint\nfrom torch.nn.init import trunc_normal_\n\nfrom ADD.layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block\n\n\nlogger = logging.getLogger(\"dinov2\")\n\n\ndef named_apply(fn: Callable, module: nn.Module, name=\"\", depth_first=True, include_root=False) -> nn.Module:\n if not depth_first and include_root:\n fn(module=module, name=name)\n for child_name, child_module in module.named_children():\n child_name = \".\".join((name, child_name)) if name else child_name\n named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)\n if depth_first and include_root:\n fn(module=module, name=name)\n return module\n\n\nclass BlockChunk(nn.ModuleList):\n def forward(self, x):\n for b in self:\n x = b(x)\n return x\n\n\nclass DinoVisionTransformer(nn.Module):\n def __init__(\n self,\n img_size=224,\n patch_size=16,\n in_chans=3,\n embed_dim=768,\n depth=12,\n num_heads=12,\n mlp_ratio=4.0,\n qkv_bias=True,\n ffn_bias=True,\n proj_bias=True,\n drop_path_rate=0.0,\n drop_path_uniform=False,\n init_values=None, # for layerscale: None or 0 => no layerscale\n embed_layer=PatchEmbed,\n act_layer=nn.GELU,\n block_fn=Block,\n ffn_layer=\"mlp\",\n block_chunks=1,\n num_register_tokens=0,\n interpolate_antialias=False,\n interpolate_offset=0.1,\n ):\n \"\"\"\n Args:\n img_size (int, tuple): input image size\n patch_size (int, tuple): patch size\n in_chans (int): number of input channels\n embed_dim (int): embedding dimension\n depth (int): depth of transformer\n num_heads (int): number of attention heads\n mlp_ratio (int): ratio of mlp hidden dim to embedding dim\n qkv_bias (bool): enable bias for qkv if True\n proj_bias (bool): enable bias for proj in attn if True\n ffn_bias (bool): enable bias for ffn if True\n drop_path_rate (float): stochastic depth rate\n drop_path_uniform (bool): apply uniform drop rate across blocks\n weight_init (str): weight init scheme\n init_values (float): layer-scale init values\n embed_layer (nn.Module): patch embedding layer\n act_layer (nn.Module): MLP activation layer\n block_fn (nn.Module): transformer block class\n ffn_layer (str): \"mlp\", \"swiglu\", \"swiglufused\" or \"identity\"\n block_chunks: (int) split block sequence into block_chunks units for FSDP wrap\n num_register_tokens: (int) number of extra cls tokens (so-called \"registers\")\n interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings\n interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings\n \"\"\"\n super().__init__()\n norm_layer = partial(nn.LayerNorm, eps=1e-6)\n\n self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models\n self.num_tokens = 1\n self.n_blocks = depth\n self.num_heads = num_heads\n self.patch_size = patch_size\n self.num_register_tokens = num_register_tokens\n self.interpolate_antialias = interpolate_antialias\n self.interpolate_offset = interpolate_offset\n\n self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)\n num_patches = self.patch_embed.num_patches\n\n self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))\n assert num_register_tokens >= 0\n self.register_tokens = (\n nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None\n )\n\n if drop_path_uniform is True:\n dpr = [drop_path_rate] * depth\n else:\n dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule\n\n if ffn_layer == \"mlp\":\n logger.info(\"using MLP layer as FFN\")\n ffn_layer = Mlp\n elif ffn_layer == \"swiglufused\" or ffn_layer == \"swiglu\":\n logger.info(\"using SwiGLU layer as FFN\")\n ffn_layer = SwiGLUFFNFused\n elif ffn_layer == \"identity\":\n logger.info(\"using Identity layer as FFN\")\n\n def f(*args, **kwargs):\n return nn.Identity()\n\n ffn_layer = f\n else:\n raise NotImplementedError\n\n blocks_list = [\n block_fn(\n dim=embed_dim,\n num_heads=num_heads,\n mlp_ratio=mlp_ratio,\n qkv_bias=qkv_bias,\n proj_bias=proj_bias,\n ffn_bias=ffn_bias,\n drop_path=dpr[i],\n norm_layer=norm_layer,\n act_layer=act_layer,\n ffn_layer=ffn_layer,\n init_values=init_values,\n )\n for i in range(depth)\n ]\n if block_chunks > 0:\n self.chunked_blocks = True\n chunked_blocks = []\n chunksize = depth // block_chunks\n for i in range(0, depth, chunksize):\n # this is to keep the block index consistent if we chunk the block list\n chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])\n self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])\n else:\n self.chunked_blocks = False\n self.blocks = nn.ModuleList(blocks_list)\n\n self.norm = norm_layer(embed_dim)\n self.head = nn.Identity()\n\n self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))\n\n self.init_weights()\n\n def init_weights(self):\n trunc_normal_(self.pos_embed, std=0.02)\n nn.init.normal_(self.cls_token, std=1e-6)\n if self.register_tokens is not None:\n nn.init.normal_(self.register_tokens, std=1e-6)\n named_apply(init_weights_vit_timm, self)\n\n def interpolate_pos_encoding(self, x, w, h):\n previous_dtype = x.dtype\n npatch = x.shape[1] - 1\n N = self.pos_embed.shape[1] - 1\n if npatch == N and w == h:\n return self.pos_embed\n pos_embed = self.pos_embed.float()\n class_pos_embed = pos_embed[:, 0]\n patch_pos_embed = pos_embed[:, 1:]\n dim = x.shape[-1]\n w0 = w // self.patch_size\n h0 = h // self.patch_size\n # we add a small number to avoid floating point error in the interpolation\n # see discussion at https://github.com/facebookresearch/dino/issues/8\n w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset\n\n sqrt_N = math.sqrt(N)\n sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N\n patch_pos_embed = nn.functional.interpolate(\n patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),\n scale_factor=(sx, sy),\n mode=\"bicubic\",\n antialias=self.interpolate_antialias,\n )\n\n assert int(w0) == patch_pos_embed.shape[-2]\n assert int(h0) == patch_pos_embed.shape[-1]\n patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)\n return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)\n\n def prepare_tokens_with_masks(self, x, masks=None):\n B, nc, w, h = x.shape\n x = self.patch_embed(x)\n if masks is not None:\n x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)\n\n x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)\n x = x + self.interpolate_pos_encoding(x, w, h)\n\n if self.register_tokens is not None:\n x = torch.cat(\n (\n x[:, :1],\n self.register_tokens.expand(x.shape[0], -1, -1),\n x[:, 1:],\n ),\n dim=1,\n )\n\n return x\n\n def forward_features_list(self, x_list, masks_list):\n x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]\n for blk in self.blocks:\n x = blk(x)\n\n all_x = x\n output = []\n for x, masks in zip(all_x, masks_list):\n x_norm = self.norm(x)\n output.append(\n {\n \"x_norm_clstoken\": x_norm[:, 0],\n \"x_norm_regtokens\": x_norm[:, 1 : self.num_register_tokens + 1],\n \"x_norm_patchtokens\": x_norm[:, self.num_register_tokens + 1 :],\n \"x_prenorm\": x,\n \"masks\": masks,\n }\n )\n return output\n\n def forward_features(self, x, masks=None):\n fea_list = []\n counter = 0\n if isinstance(x, list):\n return self.forward_features_list(x, masks)\n\n x = self.prepare_tokens_with_masks(x, masks)\n fea_list.append(x[:, self.num_register_tokens + 1 :].permute(0, 2, 1))\n\n for blk in self.blocks:\n x = blk(x)\n counter += 1\n if counter % 3 == 0:\n fea_list.append(x[:, self.num_register_tokens + 1 :].permute(0, 2, 1))\n\n x_norm = self.norm(x)\n return fea_list, x_norm[:, 0]\n\n def _get_intermediate_layers_not_chunked(self, x, n=1):\n x = self.prepare_tokens_with_masks(x)\n # If n is an int, take the n last blocks. If it's a list, take them\n output, total_block_len = [], len(self.blocks)\n blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n\n for i, blk in enumerate(self.blocks):\n x = blk(x)\n if i in blocks_to_take:\n output.append(x)\n assert len(output) == len(blocks_to_take), f\"only {len(output)} / {len(blocks_to_take)} blocks found\"\n return output\n\n def _get_intermediate_layers_chunked(self, x, n=1):\n x = self.prepare_tokens_with_masks(x)\n output, i, total_block_len = [], 0, len(self.blocks[-1])\n # If n is an int, take the n last blocks. If it's a list, take them\n blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n\n for block_chunk in self.blocks:\n for blk in block_chunk[i:]: # Passing the nn.Identity()\n x = blk(x)\n if i in blocks_to_take:\n output.append(x)\n i += 1\n assert len(output) == len(blocks_to_take), f\"only {len(output)} / {len(blocks_to_take)} blocks found\"\n return output\n\n def get_intermediate_layers(\n self,\n x: torch.Tensor,\n n: Union[int, Sequence] = 1, # Layers or n last layers to take\n reshape: bool = False,\n return_class_token: bool = False,\n norm=True,\n ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:\n if self.chunked_blocks:\n outputs = self._get_intermediate_layers_chunked(x, n)\n else:\n outputs = self._get_intermediate_layers_not_chunked(x, n)\n if norm:\n outputs = [self.norm(out) for out in outputs]\n class_tokens = [out[:, 0] for out in outputs]\n outputs = [out[:, 1 + self.num_register_tokens:] for out in outputs]\n if reshape:\n B, _, w, h = x.shape\n outputs = [\n out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()\n for out in outputs\n ]\n if return_class_token:\n return tuple(zip(outputs, class_tokens))\n return tuple(outputs)\n\n def forward(self, *args, is_training=False, **kwargs):\n ret = self.forward_features(*args, **kwargs)\n if is_training:\n return ret\n else:\n return ret#self.head(ret[\"x_norm_clstoken\"])\n\n\ndef init_weights_vit_timm(module: nn.Module, name: str = \"\"):\n \"\"\"ViT weight initialization, original timm impl (for reproducibility)\"\"\"\n if isinstance(module, nn.Linear):\n trunc_normal_(module.weight, std=0.02)\n if module.bias is not None:\n nn.init.zeros_(module.bias)\n\n\ndef vit_small(patch_size=16, num_register_tokens=0, **kwargs):\n model = DinoVisionTransformer(\n patch_size=patch_size,\n embed_dim=384,\n depth=12,\n num_heads=6,\n mlp_ratio=4,\n block_fn=partial(Block, attn_class=MemEffAttention),\n num_register_tokens=num_register_tokens,\n **kwargs,\n )\n return model\n\ndef vit_large(patch_size=16, num_register_tokens=0, **kwargs):\n model = DinoVisionTransformer(\n patch_size=patch_size,\n embed_dim=1024,\n depth=24,\n num_heads=16,\n mlp_ratio=4,\n block_fn=partial(Block, attn_class=MemEffAttention),\n num_register_tokens=num_register_tokens,\n **kwargs,\n )\n return model\n\n\n# net = vit_small(patch_size=14, img_size=518, block_chunks=0, init_values=1.0)\n# prefile = torch.load('../weights/dinov2_vits14_pretrain.pth')\n# net.load_state_dict(prefile, True)\n# out = net(torch.rand(1, 3, 518, 518))\n# print(out.shape)\n\n# net = vit_large(patch_size=14, img_size=526, block_chunks=0, init_values=1.0, num_register_tokens=4)\n# prefile = torch.load('../weights/dinov2_vitl14_reg4_pretrain.pth')\n# net.load_state_dict(prefile, True)\n# out = net(torch.rand(1, 3, 70, 70))\n# print(out.shape)\n"
},
{
"path": "ADD/th_utils/__init__.py",
"content": "# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n# empty\n"
},
{
"path": "ADD/th_utils/custom_ops.py",
"content": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\nimport glob\nimport hashlib\nimport importlib\nimport os\nimport re\nimport shutil\nimport uuid\n\nimport torch\nimport torch.utils.cpp_extension\nfrom torch.utils.file_baton import FileBaton\n\n#----------------------------------------------------------------------------\n# Global options.\n\nverbosity = 'brief' # Verbosity level: 'none', 'brief', 'full'\n\n#----------------------------------------------------------------------------\n# Internal helper funcs.\n\ndef _find_compiler_bindir():\n patterns = [\n 'C:/Program Files*/Microsoft Visual Studio/*/Professional/VC/Tools/MSVC/*/bin/Hostx64/x64',\n 'C:/Program Files*/Microsoft Visual Studio/*/BuildTools/VC/Tools/MSVC/*/bin/Hostx64/x64',\n 'C:/Program Files*/Microsoft Visual Studio/*/Community/VC/Tools/MSVC/*/bin/Hostx64/x64',\n 'C:/Program Files*/Microsoft Visual Studio */vc/bin',\n ]\n for pattern in patterns:\n matches = sorted(glob.glob(pattern))\n if len(matches):\n return matches[-1]\n return None\n\n#----------------------------------------------------------------------------\n\ndef _get_mangled_gpu_name():\n name = torch.cuda.get_device_name().lower()\n out = []\n for c in name:\n if re.match('[a-z0-9_-]+', c):\n out.append(c)\n else:\n out.append('-')\n return ''.join(out)\n\n#----------------------------------------------------------------------------\n# Main entry point for compiling and loading C++/CUDA plugins.\n\n_cached_plugins = dict()\n\ndef get_plugin(module_name, sources, headers=None, source_dir=None, **build_kwargs):\n assert verbosity in ['none', 'brief', 'full']\n if headers is None:\n headers = []\n if source_dir is not None:\n sources = [os.path.join(source_dir, fname) for fname in sources]\n headers = [os.path.join(source_dir, fname) for fname in headers]\n\n # Already cached?\n if module_name in _cached_plugins:\n return _cached_plugins[module_name]\n\n # Print status.\n if verbosity == 'full':\n print(f'Setting up PyTorch plugin \"{module_name}\"...')\n elif verbosity == 'brief':\n print(f'Setting up PyTorch plugin \"{module_name}\"... ', end='', flush=True)\n verbose_build = (verbosity == 'full')\n\n # Compile and load.\n try: # pylint: disable=too-many-nested-blocks\n # Make sure we can find the necessary compiler binaries.\n if os.name == 'nt' and os.system(\"where cl.exe >nul 2>nul\") != 0:\n compiler_bindir = _find_compiler_bindir()\n if compiler_bindir is None:\n raise RuntimeError(f'Could not find MSVC/GCC/CLANG installation on this computer. Check _find_compiler_bindir() in \"{__file__}\".')\n os.environ['PATH'] += ';' + compiler_bindir\n\n # Some containers set TORCH_CUDA_ARCH_LIST to a list that can either\n # break the build or unnecessarily restrict what's available to nvcc.\n # Unset it to let nvcc decide based on what's available on the\n # machine.\n os.environ['TORCH_CUDA_ARCH_LIST'] = ''\n\n # Incremental build md5sum trickery. Copies all the input source files\n # into a cached build directory under a combined md5 digest of the input\n # source files. Copying is done only if the combined digest has changed.\n # This keeps input file timestamps and filenames the same as in previous\n # extension builds, allowing for fast incremental rebuilds.\n #\n # This optimization is done only in case all the source files reside in\n # a single directory (just for simplicity) and if the TORCH_EXTENSIONS_DIR\n # environment variable is set (we take this as a signal that the user\n # actually cares about this.)\n #\n # EDIT: We now do it regardless of TORCH_EXTENSIOS_DIR, in order to work\n # around the *.cu dependency bug in ninja config.\n #\n all_source_files = sorted(sources + headers)\n all_source_dirs = set(os.path.dirname(fname) for fname in all_source_files)\n if len(all_source_dirs) == 1: # and ('TORCH_EXTENSIONS_DIR' in os.environ):\n\n # Compute combined hash digest for all source files.\n hash_md5 = hashlib.md5()\n for src in all_source_files:\n with open(src, 'rb') as f:\n hash_md5.update(f.read())\n\n # Select cached build directory name.\n source_digest = hash_md5.hexdigest()\n build_top_dir = torch.utils.cpp_extension._get_build_directory(module_name, verbose=verbose_build) # pylint: disable=protected-access\n cached_build_dir = os.path.join(build_top_dir, f'{source_digest}-{_get_mangled_gpu_name()}')\n\n if not os.path.isdir(cached_build_dir):\n tmpdir = f'{build_top_dir}/srctmp-{uuid.uuid4().hex}'\n os.makedirs(tmpdir)\n for src in all_source_files:\n shutil.copyfile(src, os.path.join(tmpdir, os.path.basename(src)))\n try:\n os.replace(tmpdir, cached_build_dir) # atomic\n except OSError:\n # source directory already exists, delete tmpdir and its contents.\n shutil.rmtree(tmpdir)\n if not os.path.isdir(cached_build_dir): raise\n\n # Compile.\n cached_sources = [os.path.join(cached_build_dir, os.path.basename(fname)) for fname in sources]\n torch.utils.cpp_extension.load(name=module_name, build_directory=cached_build_dir,\n verbose=verbose_build, sources=cached_sources, **build_kwargs)\n else:\n torch.utils.cpp_extension.load(name=module_name, verbose=verbose_build, sources=sources, **build_kwargs)\n\n # Load.\n module = importlib.import_module(module_name)\n\n except:\n if verbosity == 'brief':\n print('Failed!')\n raise\n\n # Print status and add to cache dict.\n if verbosity == 'full':\n print(f'Done setting up PyTorch plugin \"{module_name}\".')\n elif verbosity == 'brief':\n print('Done.')\n _cached_plugins[module_name] = module\n return module\n\n#----------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/misc.py",
"content": "# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\nimport re\nimport contextlib\nimport numpy as np\nimport torch\nimport warnings\nimport ADD.dnnlib as dnnlib\n\n#----------------------------------------------------------------------------\n# Cached construction of constant tensors. Avoids CPU=>GPU copy when the\n# same constant is used multiple times.\n\n_constant_cache = dict()\n\ndef constant(value, shape=None, dtype=None, device=None, memory_format=None):\n value = np.asarray(value)\n if shape is not None:\n shape = tuple(shape)\n if dtype is None:\n dtype = torch.get_default_dtype()\n if device is None:\n device = torch.device('cpu')\n if memory_format is None:\n memory_format = torch.contiguous_format\n\n key = (value.shape, value.dtype, value.tobytes(), shape, dtype, device, memory_format)\n tensor = _constant_cache.get(key, None)\n if tensor is None:\n tensor = torch.as_tensor(value.copy(), dtype=dtype, device=device)\n if shape is not None:\n tensor, _ = torch.broadcast_tensors(tensor, torch.empty(shape))\n tensor = tensor.contiguous(memory_format=memory_format)\n _constant_cache[key] = tensor\n return tensor\n\n#----------------------------------------------------------------------------\n# Replace NaN/Inf with specified numerical values.\n\ntry:\n nan_to_num = torch.nan_to_num # 1.8.0a0\nexcept AttributeError:\n def nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None): # pylint: disable=redefined-builtin\n assert isinstance(input, torch.Tensor)\n if posinf is None:\n posinf = torch.finfo(input.dtype).max\n if neginf is None:\n neginf = torch.finfo(input.dtype).min\n assert nan == 0\n return torch.clamp(input.unsqueeze(0).nansum(0), min=neginf, max=posinf, out=out)\n\n#----------------------------------------------------------------------------\n# Symbolic assert.\n\ntry:\n symbolic_assert = torch._assert # 1.8.0a0 # pylint: disable=protected-access\nexcept AttributeError:\n symbolic_assert = torch.Assert # 1.7.0\n\n#----------------------------------------------------------------------------\n# Context manager to temporarily suppress known warnings in torch.jit.trace().\n# Note: Cannot use catch_warnings because of https://bugs.python.org/issue29672\n\n@contextlib.contextmanager\ndef suppress_tracer_warnings():\n flt = ('ignore', None, torch.jit.TracerWarning, None, 0)\n warnings.filters.insert(0, flt)\n yield\n warnings.filters.remove(flt)\n\n#----------------------------------------------------------------------------\n# Assert that the shape of a tensor matches the given list of integers.\n# None indicates that the size of a dimension is allowed to vary.\n# Performs symbolic assertion when used in torch.jit.trace().\n\ndef assert_shape(tensor, ref_shape):\n if tensor.ndim != len(ref_shape):\n raise AssertionError(f'Wrong number of dimensions: got {tensor.ndim}, expected {len(ref_shape)}')\n for idx, (size, ref_size) in enumerate(zip(tensor.shape, ref_shape)):\n if ref_size is None:\n pass\n elif isinstance(ref_size, torch.Tensor):\n with suppress_tracer_warnings(): # as_tensor results are registered as constants\n symbolic_assert(torch.equal(torch.as_tensor(size), ref_size), f'Wrong size for dimension {idx}')\n elif isinstance(size, torch.Tensor):\n with suppress_tracer_warnings(): # as_tensor results are registered as constants\n symbolic_assert(torch.equal(size, torch.as_tensor(ref_size)), f'Wrong size for dimension {idx}: expected {ref_size}')\n elif size != ref_size:\n raise AssertionError(f'Wrong size for dimension {idx}: got {size}, expected {ref_size}')\n\n#----------------------------------------------------------------------------\n# Function decorator that calls torch.autograd.profiler.record_function().\n\ndef profiled_function(fn):\n def decorator(*args, **kwargs):\n with torch.autograd.profiler.record_function(fn.__name__):\n return fn(*args, **kwargs)\n decorator.__name__ = fn.__name__\n return decorator\n\n#----------------------------------------------------------------------------\n# Sampler for torch.utils.data.DataLoader that loops over the dataset\n# indefinitely, shuffling items as it goes.\n\nclass InfiniteSampler(torch.utils.data.Sampler):\n def __init__(self, dataset, rank=0, num_replicas=1, shuffle=True, seed=0, window_size=0.5):\n assert len(dataset) > 0\n assert num_replicas > 0\n assert 0 <= rank < num_replicas\n assert 0 <= window_size <= 1\n super().__init__(dataset)\n self.dataset = dataset\n self.rank = rank\n self.num_replicas = num_replicas\n self.shuffle = shuffle\n self.seed = seed\n self.window_size = window_size\n\n def __iter__(self):\n order = np.arange(len(self.dataset))\n rnd = None\n window = 0\n if self.shuffle:\n rnd = np.random.RandomState(self.seed)\n rnd.shuffle(order)\n window = int(np.rint(order.size * self.window_size))\n\n idx = 0\n while True:\n i = idx % order.size\n if idx % self.num_replicas == self.rank:\n yield order[i]\n if window >= 2:\n j = (i - rnd.randint(window)) % order.size\n order[i], order[j] = order[j], order[i]\n idx += 1\n\n#----------------------------------------------------------------------------\n# Utilities for operating with torch.nn.Module parameters and buffers.\ndef spectral_to_cpu(model: torch.nn.Module):\n def wrapped_in_spectral(m): return hasattr(m, 'weight_v')\n children = get_children(model)\n for child in children:\n if wrapped_in_spectral(child):\n child.weight = child.weight.cpu()\n return model\n\ndef get_children(model: torch.nn.Module):\n children = list(model.children())\n flatt_children = []\n if children == []:\n return model\n else:\n for child in children:\n try:\n flatt_children.extend(get_children(child))\n except TypeError:\n flatt_children.append(get_children(child))\n return flatt_children\n\ndef params_and_buffers(module):\n assert isinstance(module, torch.nn.Module)\n return list(module.parameters()) + list(module.buffers())\n\ndef named_params_and_buffers(module):\n assert isinstance(module, torch.nn.Module)\n return list(module.named_parameters()) + list(module.named_buffers())\n\ndef copy_params_and_buffers(src_module, dst_module, require_all=False):\n assert isinstance(src_module, torch.nn.Module)\n assert isinstance(dst_module, torch.nn.Module)\n src_tensors = dict(named_params_and_buffers(src_module))\n for name, tensor in named_params_and_buffers(dst_module):\n assert (name in src_tensors) or (not require_all)\n if name in src_tensors:\n tensor.copy_(src_tensors[name].detach()).requires_grad_(tensor.requires_grad)\n\n#----------------------------------------------------------------------------\n# Context manager for easily enabling/disabling DistributedDataParallel\n# synchronization.\n\n@contextlib.contextmanager\ndef ddp_sync(module, sync):\n assert isinstance(module, torch.nn.Module)\n if sync or not isinstance(module, torch.nn.parallel.DistributedDataParallel):\n yield\n else:\n with module.no_sync():\n yield\n\n#----------------------------------------------------------------------------\n# Check DistributedDataParallel consistency across processes.\n\ndef check_ddp_consistency(module, ignore_regex=None):\n assert isinstance(module, torch.nn.Module)\n for name, tensor in named_params_and_buffers(module):\n fullname = type(module).__name__ + '.' + name\n if ignore_regex is not None and re.fullmatch(ignore_regex, fullname):\n continue\n tensor = tensor.detach()\n if tensor.is_floating_point():\n tensor = nan_to_num(tensor)\n other = tensor.clone()\n torch.distributed.broadcast(tensor=other, src=0)\n assert (tensor == other).all(), fullname\n\n#----------------------------------------------------------------------------\n# Print summary table of module hierarchy.\n\ndef print_module_summary(module, inputs, max_nesting=3, skip_redundant=True):\n assert isinstance(module, torch.nn.Module)\n assert not isinstance(module, torch.jit.ScriptModule)\n assert isinstance(inputs, (tuple, list))\n\n # Register hooks.\n entries = []\n nesting = [0]\n def pre_hook(_mod, _inputs):\n nesting[0] += 1\n def post_hook(mod, _inputs, outputs):\n nesting[0] -= 1\n if nesting[0] <= max_nesting:\n outputs = list(outputs) if isinstance(outputs, (tuple, list)) else [outputs]\n outputs = [t for t in outputs if isinstance(t, torch.Tensor)]\n entries.append(dnnlib.EasyDict(mod=mod, outputs=outputs))\n hooks = [mod.register_forward_pre_hook(pre_hook) for mod in module.modules()]\n hooks += [mod.register_forward_hook(post_hook) for mod in module.modules()]\n\n # Run module.\n outputs = module(*inputs)\n for hook in hooks:\n hook.remove()\n\n # Identify unique outputs, parameters, and buffers.\n tensors_seen = set()\n for e in entries:\n e.unique_params = [t for t in e.mod.parameters() if id(t) not in tensors_seen]\n e.unique_buffers = [t for t in e.mod.buffers() if id(t) not in tensors_seen]\n e.unique_outputs = [t for t in e.outputs if id(t) not in tensors_seen]\n tensors_seen |= {id(t) for t in e.unique_params + e.unique_buffers + e.unique_outputs}\n\n # Filter out redundant entries.\n if skip_redundant:\n entries = [e for e in entries if len(e.unique_params) or len(e.unique_buffers) or len(e.unique_outputs)]\n\n # Construct table.\n rows = [[type(module).__name__, 'Parameters', 'Buffers', 'Output shape', 'Datatype']]\n rows += [['---'] * len(rows[0])]\n param_total = 0\n buffer_total = 0\n submodule_names = {mod: name for name, mod in module.named_modules()}\n for e in entries:\n name = '' if e.mod is module else submodule_names[e.mod]\n param_size = sum(t.numel() for t in e.unique_params)\n buffer_size = sum(t.numel() for t in e.unique_buffers)\n output_shapes = [str(list(t.shape)) for t in e.outputs]\n output_dtypes = [str(t.dtype).split('.')[-1] for t in e.outputs]\n rows += [[\n name + (':0' if len(e.outputs) >= 2 else ''),\n str(param_size) if param_size else '-',\n str(buffer_size) if buffer_size else '-',\n (output_shapes + ['-'])[0],\n (output_dtypes + ['-'])[0],\n ]]\n for idx in range(1, len(e.outputs)):\n rows += [[name + f':{idx}', '-', '-', output_shapes[idx], output_dtypes[idx]]]\n param_total += param_size\n buffer_total += buffer_size\n rows += [['---'] * len(rows[0])]\n rows += [['Total', str(param_total), str(buffer_total), '-', '-']]\n\n # Print table.\n widths = [max(len(cell) for cell in column) for column in zip(*rows)]\n print()\n for row in rows:\n print(' '.join(cell + ' ' * (width - len(cell)) for cell, width in zip(row, widths)))\n print()\n return outputs\n"
},
{
"path": "ADD/th_utils/ops/__init__.py",
"content": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n# empty\n"
},
{
"path": "ADD/th_utils/ops/bias_act.cpp",
"content": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors retain all intellectual property\n// and proprietary rights in and to this software, related documentation\n// and any modifications thereto. Any use, reproduction, disclosure or\n// distribution of this software and related documentation without an express\n// license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n#include \n#include \n#include \n#include \"bias_act.h\"\n\n//------------------------------------------------------------------------\n\nstatic bool has_same_layout(torch::Tensor x, torch::Tensor y)\n{\n if (x.dim() != y.dim())\n return false;\n for (int64_t i = 0; i < x.dim(); i++)\n {\n if (x.size(i) != y.size(i))\n return false;\n if (x.size(i) >= 2 && x.stride(i) != y.stride(i))\n return false;\n }\n return true;\n}\n\n//------------------------------------------------------------------------\n\nstatic torch::Tensor bias_act(torch::Tensor x, torch::Tensor b, torch::Tensor xref, torch::Tensor yref, torch::Tensor dy, int grad, int dim, int act, float alpha, float gain, float clamp)\n{\n // Validate arguments.\n TORCH_CHECK(x.is_cuda(), \"x must reside on CUDA device\");\n TORCH_CHECK(b.numel() == 0 || (b.dtype() == x.dtype() && b.device() == x.device()), \"b must have the same dtype and device as x\");\n TORCH_CHECK(xref.numel() == 0 || (xref.sizes() == x.sizes() && xref.dtype() == x.dtype() && xref.device() == x.device()), \"xref must have the same shape, dtype, and device as x\");\n TORCH_CHECK(yref.numel() == 0 || (yref.sizes() == x.sizes() && yref.dtype() == x.dtype() && yref.device() == x.device()), \"yref must have the same shape, dtype, and device as x\");\n TORCH_CHECK(dy.numel() == 0 || (dy.sizes() == x.sizes() && dy.dtype() == x.dtype() && dy.device() == x.device()), \"dy must have the same dtype and device as x\");\n TORCH_CHECK(x.numel() <= INT_MAX, \"x is too large\");\n TORCH_CHECK(b.dim() == 1, \"b must have rank 1\");\n TORCH_CHECK(b.numel() == 0 || (dim >= 0 && dim < x.dim()), \"dim is out of bounds\");\n TORCH_CHECK(b.numel() == 0 || b.numel() == x.size(dim), \"b has wrong number of elements\");\n TORCH_CHECK(grad >= 0, \"grad must be non-negative\");\n\n // Validate layout.\n TORCH_CHECK(x.is_non_overlapping_and_dense(), \"x must be non-overlapping and dense\");\n TORCH_CHECK(b.is_contiguous(), \"b must be contiguous\");\n TORCH_CHECK(xref.numel() == 0 || has_same_layout(xref, x), \"xref must have the same layout as x\");\n TORCH_CHECK(yref.numel() == 0 || has_same_layout(yref, x), \"yref must have the same layout as x\");\n TORCH_CHECK(dy.numel() == 0 || has_same_layout(dy, x), \"dy must have the same layout as x\");\n\n // Create output tensor.\n const at::cuda::OptionalCUDAGuard device_guard(device_of(x));\n torch::Tensor y = torch::empty_like(x);\n TORCH_CHECK(has_same_layout(y, x), \"y must have the same layout as x\");\n\n // Initialize CUDA kernel parameters.\n bias_act_kernel_params p;\n p.x = x.data_ptr();\n p.b = (b.numel()) ? b.data_ptr() : NULL;\n p.xref = (xref.numel()) ? xref.data_ptr() : NULL;\n p.yref = (yref.numel()) ? yref.data_ptr() : NULL;\n p.dy = (dy.numel()) ? dy.data_ptr() : NULL;\n p.y = y.data_ptr();\n p.grad = grad;\n p.act = act;\n p.alpha = alpha;\n p.gain = gain;\n p.clamp = clamp;\n p.sizeX = (int)x.numel();\n p.sizeB = (int)b.numel();\n p.stepB = (b.numel()) ? (int)x.stride(dim) : 1;\n\n // Choose CUDA kernel.\n void* kernel;\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), \"upfirdn2d_cuda\", [&]\n {\n kernel = choose_bias_act_kernel(p);\n });\n TORCH_CHECK(kernel, \"no CUDA kernel found for the specified activation func\");\n\n // Launch CUDA kernel.\n p.loopX = 4;\n int blockSize = 4 * 32;\n int gridSize = (p.sizeX - 1) / (p.loopX * blockSize) + 1;\n void* args[] = {&p};\n AT_CUDA_CHECK(cudaLaunchKernel(kernel, gridSize, blockSize, args, 0, at::cuda::getCurrentCUDAStream()));\n return y;\n}\n\n//------------------------------------------------------------------------\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m)\n{\n m.def(\"bias_act\", &bias_act);\n}\n\n//------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/bias_act.cu",
"content": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors retain all intellectual property\n// and proprietary rights in and to this software, related documentation\n// and any modifications thereto. Any use, reproduction, disclosure or\n// distribution of this software and related documentation without an express\n// license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n#include \n#include \"bias_act.h\"\n\n//------------------------------------------------------------------------\n// Helpers.\n\ntemplate struct InternalType;\ntemplate <> struct InternalType { typedef double scalar_t; };\ntemplate <> struct InternalType { typedef float scalar_t; };\ntemplate <> struct InternalType { typedef float scalar_t; };\n\n//------------------------------------------------------------------------\n// CUDA kernel.\n\ntemplate \n__global__ void bias_act_kernel(bias_act_kernel_params p)\n{\n typedef typename InternalType::scalar_t scalar_t;\n int G = p.grad;\n scalar_t alpha = (scalar_t)p.alpha;\n scalar_t gain = (scalar_t)p.gain;\n scalar_t clamp = (scalar_t)p.clamp;\n scalar_t one = (scalar_t)1;\n scalar_t two = (scalar_t)2;\n scalar_t expRange = (scalar_t)80;\n scalar_t halfExpRange = (scalar_t)40;\n scalar_t seluScale = (scalar_t)1.0507009873554804934193349852946;\n scalar_t seluAlpha = (scalar_t)1.6732632423543772848170429916717;\n\n // Loop over elements.\n int xi = blockIdx.x * p.loopX * blockDim.x + threadIdx.x;\n for (int loopIdx = 0; loopIdx < p.loopX && xi < p.sizeX; loopIdx++, xi += blockDim.x)\n {\n // Load.\n scalar_t x = (scalar_t)((const T*)p.x)[xi];\n scalar_t b = (p.b) ? (scalar_t)((const T*)p.b)[(xi / p.stepB) % p.sizeB] : 0;\n scalar_t xref = (p.xref) ? (scalar_t)((const T*)p.xref)[xi] : 0;\n scalar_t yref = (p.yref) ? (scalar_t)((const T*)p.yref)[xi] : 0;\n scalar_t dy = (p.dy) ? (scalar_t)((const T*)p.dy)[xi] : one;\n scalar_t yy = (gain != 0) ? yref / gain : 0;\n scalar_t y = 0;\n\n // Apply bias.\n ((G == 0) ? x : xref) += b;\n\n // linear\n if (A == 1)\n {\n if (G == 0) y = x;\n if (G == 1) y = x;\n }\n\n // relu\n if (A == 2)\n {\n if (G == 0) y = (x > 0) ? x : 0;\n if (G == 1) y = (yy > 0) ? x : 0;\n }\n\n // lrelu\n if (A == 3)\n {\n if (G == 0) y = (x > 0) ? x : x * alpha;\n if (G == 1) y = (yy > 0) ? x : x * alpha;\n }\n\n // tanh\n if (A == 4)\n {\n if (G == 0) { scalar_t c = exp(x); scalar_t d = one / c; y = (x < -expRange) ? -one : (x > expRange) ? one : (c - d) / (c + d); }\n if (G == 1) y = x * (one - yy * yy);\n if (G == 2) y = x * (one - yy * yy) * (-two * yy);\n }\n\n // sigmoid\n if (A == 5)\n {\n if (G == 0) y = (x < -expRange) ? 0 : one / (exp(-x) + one);\n if (G == 1) y = x * yy * (one - yy);\n if (G == 2) y = x * yy * (one - yy) * (one - two * yy);\n }\n\n // elu\n if (A == 6)\n {\n if (G == 0) y = (x >= 0) ? x : exp(x) - one;\n if (G == 1) y = (yy >= 0) ? x : x * (yy + one);\n if (G == 2) y = (yy >= 0) ? 0 : x * (yy + one);\n }\n\n // selu\n if (A == 7)\n {\n if (G == 0) y = (x >= 0) ? seluScale * x : (seluScale * seluAlpha) * (exp(x) - one);\n if (G == 1) y = (yy >= 0) ? x * seluScale : x * (yy + seluScale * seluAlpha);\n if (G == 2) y = (yy >= 0) ? 0 : x * (yy + seluScale * seluAlpha);\n }\n\n // softplus\n if (A == 8)\n {\n if (G == 0) y = (x > expRange) ? x : log(exp(x) + one);\n if (G == 1) y = x * (one - exp(-yy));\n if (G == 2) { scalar_t c = exp(-yy); y = x * c * (one - c); }\n }\n\n // swish\n if (A == 9)\n {\n if (G == 0)\n y = (x < -expRange) ? 0 : x / (exp(-x) + one);\n else\n {\n scalar_t c = exp(xref);\n scalar_t d = c + one;\n if (G == 1)\n y = (xref > halfExpRange) ? x : x * c * (xref + d) / (d * d);\n else\n y = (xref > halfExpRange) ? 0 : x * c * (xref * (two - d) + two * d) / (d * d * d);\n yref = (xref < -expRange) ? 0 : xref / (exp(-xref) + one) * gain;\n }\n }\n\n // Apply gain.\n y *= gain * dy;\n\n // Clamp.\n if (clamp >= 0)\n {\n if (G == 0)\n y = (y > -clamp & y < clamp) ? y : (y >= 0) ? clamp : -clamp;\n else\n y = (yref > -clamp & yref < clamp) ? y : 0;\n }\n\n // Store.\n ((T*)p.y)[xi] = (T)y;\n }\n}\n\n//------------------------------------------------------------------------\n// CUDA kernel selection.\n\ntemplate void* choose_bias_act_kernel(const bias_act_kernel_params& p)\n{\n if (p.act == 1) return (void*)bias_act_kernel;\n if (p.act == 2) return (void*)bias_act_kernel;\n if (p.act == 3) return (void*)bias_act_kernel;\n if (p.act == 4) return (void*)bias_act_kernel;\n if (p.act == 5) return (void*)bias_act_kernel;\n if (p.act == 6) return (void*)bias_act_kernel;\n if (p.act == 7) return (void*)bias_act_kernel;\n if (p.act == 8) return (void*)bias_act_kernel;\n if (p.act == 9) return (void*)bias_act_kernel;\n return NULL;\n}\n\n//------------------------------------------------------------------------\n// Template specializations.\n\ntemplate void* choose_bias_act_kernel (const bias_act_kernel_params& p);\ntemplate void* choose_bias_act_kernel (const bias_act_kernel_params& p);\ntemplate void* choose_bias_act_kernel (const bias_act_kernel_params& p);\n\n//------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/bias_act.h",
"content": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors retain all intellectual property\n// and proprietary rights in and to this software, related documentation\n// and any modifications thereto. Any use, reproduction, disclosure or\n// distribution of this software and related documentation without an express\n// license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n//------------------------------------------------------------------------\n// CUDA kernel parameters.\n\nstruct bias_act_kernel_params\n{\n const void* x; // [sizeX]\n const void* b; // [sizeB] or NULL\n const void* xref; // [sizeX] or NULL\n const void* yref; // [sizeX] or NULL\n const void* dy; // [sizeX] or NULL\n void* y; // [sizeX]\n\n int grad;\n int act;\n float alpha;\n float gain;\n float clamp;\n\n int sizeX;\n int sizeB;\n int stepB;\n int loopX;\n};\n\n//------------------------------------------------------------------------\n// CUDA kernel selection.\n\ntemplate void* choose_bias_act_kernel(const bias_act_kernel_params& p);\n\n//------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/bias_act.py",
"content": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n\"\"\"Custom PyTorch ops for efficient bias and activation.\"\"\"\n\nimport os\nimport numpy as np\nimport torch\nimport ADD.dnnlib as dnnlib\n\nfrom .. import custom_ops\nfrom .. import misc\n\n#----------------------------------------------------------------------------\n\nactivation_funcs = {\n 'linear': dnnlib.EasyDict(func=lambda x, **_: x, def_alpha=0, def_gain=1, cuda_idx=1, ref='', has_2nd_grad=False),\n 'relu': dnnlib.EasyDict(func=lambda x, **_: torch.nn.functional.relu(x), def_alpha=0, def_gain=np.sqrt(2), cuda_idx=2, ref='y', has_2nd_grad=False),\n 'lrelu': dnnlib.EasyDict(func=lambda x, alpha, **_: torch.nn.functional.leaky_relu(x, alpha), def_alpha=0.2, def_gain=np.sqrt(2), cuda_idx=3, ref='y', has_2nd_grad=False),\n 'tanh': dnnlib.EasyDict(func=lambda x, **_: torch.tanh(x), def_alpha=0, def_gain=1, cuda_idx=4, ref='y', has_2nd_grad=True),\n 'sigmoid': dnnlib.EasyDict(func=lambda x, **_: torch.sigmoid(x), def_alpha=0, def_gain=1, cuda_idx=5, ref='y', has_2nd_grad=True),\n 'elu': dnnlib.EasyDict(func=lambda x, **_: torch.nn.functional.elu(x), def_alpha=0, def_gain=1, cuda_idx=6, ref='y', has_2nd_grad=True),\n 'selu': dnnlib.EasyDict(func=lambda x, **_: torch.nn.functional.selu(x), def_alpha=0, def_gain=1, cuda_idx=7, ref='y', has_2nd_grad=True),\n 'softplus': dnnlib.EasyDict(func=lambda x, **_: torch.nn.functional.softplus(x), def_alpha=0, def_gain=1, cuda_idx=8, ref='y', has_2nd_grad=True),\n 'swish': dnnlib.EasyDict(func=lambda x, **_: torch.sigmoid(x) * x, def_alpha=0, def_gain=np.sqrt(2), cuda_idx=9, ref='x', has_2nd_grad=True),\n}\n\n#----------------------------------------------------------------------------\n\n_plugin = None\n_null_tensor = torch.empty([0])\n\ndef _init():\n global _plugin\n if _plugin is None:\n _plugin = custom_ops.get_plugin(\n module_name='bias_act_plugin',\n sources=['bias_act.cpp', 'bias_act.cu'],\n headers=['bias_act.h'],\n source_dir=os.path.dirname(__file__),\n extra_cuda_cflags=['--use_fast_math', '--allow-unsupported-compiler'],\n )\n return True\n\n#----------------------------------------------------------------------------\n\ndef bias_act(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None, impl='cuda'):\n r\"\"\"Fused bias and activation function.\n\n Adds bias `b` to activation tensor `x`, evaluates activation function `act`,\n and scales the result by `gain`. Each of the steps is optional. In most cases,\n the fused op is considerably more efficient than performing the same calculation\n using standard PyTorch ops. It supports first and second order gradients,\n but not third order gradients.\n\n Args:\n x: Input activation tensor. Can be of any shape.\n b: Bias vector, or `None` to disable. Must be a 1D tensor of the same type\n as `x`. The shape must be known, and it must match the dimension of `x`\n corresponding to `dim`.\n dim: The dimension in `x` corresponding to the elements of `b`.\n The value of `dim` is ignored if `b` is not specified.\n act: Name of the activation function to evaluate, or `\"linear\"` to disable.\n Can be e.g. `\"relu\"`, `\"lrelu\"`, `\"tanh\"`, `\"sigmoid\"`, `\"swish\"`, etc.\n See `activation_funcs` for a full list. `None` is not allowed.\n alpha: Shape parameter for the activation function, or `None` to use the default.\n gain: Scaling factor for the output tensor, or `None` to use default.\n See `activation_funcs` for the default scaling of each activation function.\n If unsure, consider specifying 1.\n clamp: Clamp the output values to `[-clamp, +clamp]`, or `None` to disable\n the clamping (default).\n impl: Name of the implementation to use. Can be `\"ref\"` or `\"cuda\"` (default).\n\n Returns:\n Tensor of the same shape and datatype as `x`.\n \"\"\"\n assert isinstance(x, torch.Tensor)\n assert impl in ['ref', 'cuda']\n if impl == 'cuda' and x.device.type == 'cuda' and _init():\n return _bias_act_cuda(dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp).apply(x, b)\n return _bias_act_ref(x=x, b=b, dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp)\n\n#----------------------------------------------------------------------------\n\n@misc.profiled_function\ndef _bias_act_ref(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None):\n \"\"\"Slow reference implementation of `bias_act()` using standard TensorFlow ops.\n \"\"\"\n assert isinstance(x, torch.Tensor)\n assert clamp is None or clamp >= 0\n spec = activation_funcs[act]\n alpha = float(alpha if alpha is not None else spec.def_alpha)\n gain = float(gain if gain is not None else spec.def_gain)\n clamp = float(clamp if clamp is not None else -1)\n\n # Add bias.\n if b is not None:\n assert isinstance(b, torch.Tensor) and b.ndim == 1\n assert 0 <= dim < x.ndim\n assert b.shape[0] == x.shape[dim]\n x = x + b.reshape([-1 if i == dim else 1 for i in range(x.ndim)])\n\n # Evaluate activation function.\n alpha = float(alpha)\n x = spec.func(x, alpha=alpha)\n\n # Scale by gain.\n gain = float(gain)\n if gain != 1:\n x = x * gain\n\n # Clamp.\n if clamp >= 0:\n x = x.clamp(-clamp, clamp) # pylint: disable=invalid-unary-operand-type\n return x\n\n#----------------------------------------------------------------------------\n\n_bias_act_cuda_cache = dict()\n\ndef _bias_act_cuda(dim=1, act='linear', alpha=None, gain=None, clamp=None):\n \"\"\"Fast CUDA implementation of `bias_act()` using custom ops.\n \"\"\"\n # Parse arguments.\n assert clamp is None or clamp >= 0\n spec = activation_funcs[act]\n alpha = float(alpha if alpha is not None else spec.def_alpha)\n gain = float(gain if gain is not None else spec.def_gain)\n clamp = float(clamp if clamp is not None else -1)\n\n # Lookup from cache.\n key = (dim, act, alpha, gain, clamp)\n if key in _bias_act_cuda_cache:\n return _bias_act_cuda_cache[key]\n\n # Forward op.\n class BiasActCuda(torch.autograd.Function):\n @staticmethod\n def forward(ctx, x, b): # pylint: disable=arguments-differ\n ctx.memory_format = torch.channels_last if x.ndim > 2 and x.stride(1) == 1 else torch.contiguous_format\n x = x.contiguous(memory_format=ctx.memory_format)\n b = b.contiguous() if b is not None else _null_tensor\n y = x\n if act != 'linear' or gain != 1 or clamp >= 0 or b is not _null_tensor:\n y = _plugin.bias_act(x, b, _null_tensor, _null_tensor, _null_tensor, 0, dim, spec.cuda_idx, alpha, gain, clamp)\n ctx.save_for_backward(\n x if 'x' in spec.ref or spec.has_2nd_grad else _null_tensor,\n b if 'x' in spec.ref or spec.has_2nd_grad else _null_tensor,\n y if 'y' in spec.ref else _null_tensor)\n return y\n\n @staticmethod\n def backward(ctx, dy): # pylint: disable=arguments-differ\n dy = dy.contiguous(memory_format=ctx.memory_format)\n x, b, y = ctx.saved_tensors\n dx = None\n db = None\n\n if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]:\n dx = dy\n if act != 'linear' or gain != 1 or clamp >= 0:\n dx = BiasActCudaGrad.apply(dy, x, b, y)\n\n if ctx.needs_input_grad[1]:\n db = dx.sum([i for i in range(dx.ndim) if i != dim])\n\n return dx, db\n\n # Backward op.\n class BiasActCudaGrad(torch.autograd.Function):\n @staticmethod\n def forward(ctx, dy, x, b, y): # pylint: disable=arguments-differ\n ctx.memory_format = torch.channels_last if dy.ndim > 2 and dy.stride(1) == 1 else torch.contiguous_format\n dx = _plugin.bias_act(dy, b, x, y, _null_tensor, 1, dim, spec.cuda_idx, alpha, gain, clamp)\n ctx.save_for_backward(\n dy if spec.has_2nd_grad else _null_tensor,\n x, b, y)\n return dx\n\n @staticmethod\n def backward(ctx, d_dx): # pylint: disable=arguments-differ\n d_dx = d_dx.contiguous(memory_format=ctx.memory_format)\n dy, x, b, y = ctx.saved_tensors\n d_dy = None\n d_x = None\n d_b = None\n d_y = None\n\n if ctx.needs_input_grad[0]:\n d_dy = BiasActCudaGrad.apply(d_dx, x, b, y)\n\n if spec.has_2nd_grad and (ctx.needs_input_grad[1] or ctx.needs_input_grad[2]):\n d_x = _plugin.bias_act(d_dx, b, x, y, dy, 2, dim, spec.cuda_idx, alpha, gain, clamp)\n\n if spec.has_2nd_grad and ctx.needs_input_grad[2]:\n d_b = d_x.sum([i for i in range(d_x.ndim) if i != dim])\n\n return d_dy, d_x, d_b, d_y\n\n # Add to cache.\n _bias_act_cuda_cache[key] = BiasActCuda\n return BiasActCuda\n\n#----------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/conv2d_gradfix.py",
"content": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n\"\"\"Custom replacement for `torch.nn.functional.conv2d` that supports\narbitrarily high order gradients with zero performance penalty.\"\"\"\n\nimport contextlib\nimport torch\nfrom pkg_resources import parse_version\n\n# pylint: disable=redefined-builtin\n# pylint: disable=arguments-differ\n# pylint: disable=protected-access\n\n#----------------------------------------------------------------------------\n\nenabled = False # Enable the custom op by setting this to true.\nweight_gradients_disabled = False # Forcefully disable computation of gradients with respect to the weights.\n_use_pytorch_1_11_api = parse_version(torch.__version__) >= parse_version('1.11.0a') # Allow prerelease builds of 1.11\n\n@contextlib.contextmanager\ndef no_weight_gradients(disable=True):\n global weight_gradients_disabled\n old = weight_gradients_disabled\n if disable:\n weight_gradients_disabled = True\n yield\n weight_gradients_disabled = old\n\n#----------------------------------------------------------------------------\n\ndef conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):\n if _should_use_custom_op(input):\n return _conv2d_gradfix(transpose=False, weight_shape=weight.shape, stride=stride, padding=padding, output_padding=0, dilation=dilation, groups=groups).apply(input, weight, bias)\n return torch.nn.functional.conv2d(input=input, weight=weight, bias=bias, stride=stride, padding=padding, dilation=dilation, groups=groups)\n\ndef conv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1):\n if _should_use_custom_op(input):\n return _conv2d_gradfix(transpose=True, weight_shape=weight.shape, stride=stride, padding=padding, output_padding=output_padding, groups=groups, dilation=dilation).apply(input, weight, bias)\n return torch.nn.functional.conv_transpose2d(input=input, weight=weight, bias=bias, stride=stride, padding=padding, output_padding=output_padding, groups=groups, dilation=dilation)\n\n#----------------------------------------------------------------------------\n\ndef _should_use_custom_op(input):\n assert isinstance(input, torch.Tensor)\n if (not enabled) or (not torch.backends.cudnn.enabled):\n return False\n if _use_pytorch_1_11_api:\n # The work-around code doesn't work on PyTorch 1.11.0 onwards\n return False\n if input.device.type != 'cuda':\n return False\n return True\n\ndef _tuple_of_ints(xs, ndim):\n xs = tuple(xs) if isinstance(xs, (tuple, list)) else (xs,) * ndim\n assert len(xs) == ndim\n assert all(isinstance(x, int) for x in xs)\n return xs\n\n#----------------------------------------------------------------------------\n\n_conv2d_gradfix_cache = dict()\n_null_tensor = torch.empty([0])\n\ndef _conv2d_gradfix(transpose, weight_shape, stride, padding, output_padding, dilation, groups):\n # Parse arguments.\n ndim = 2\n weight_shape = tuple(weight_shape)\n stride = _tuple_of_ints(stride, ndim)\n padding = _tuple_of_ints(padding, ndim)\n output_padding = _tuple_of_ints(output_padding, ndim)\n dilation = _tuple_of_ints(dilation, ndim)\n\n # Lookup from cache.\n key = (transpose, weight_shape, stride, padding, output_padding, dilation, groups)\n if key in _conv2d_gradfix_cache:\n return _conv2d_gradfix_cache[key]\n\n # Validate arguments.\n assert groups >= 1\n assert len(weight_shape) == ndim + 2\n assert all(stride[i] >= 1 for i in range(ndim))\n assert all(padding[i] >= 0 for i in range(ndim))\n assert all(dilation[i] >= 0 for i in range(ndim))\n if not transpose:\n assert all(output_padding[i] == 0 for i in range(ndim))\n else: # transpose\n assert all(0 <= output_padding[i] < max(stride[i], dilation[i]) for i in range(ndim))\n\n # Helpers.\n common_kwargs = dict(stride=stride, padding=padding, dilation=dilation, groups=groups)\n def calc_output_padding(input_shape, output_shape):\n if transpose:\n return [0, 0]\n return [\n input_shape[i + 2]\n - (output_shape[i + 2] - 1) * stride[i]\n - (1 - 2 * padding[i])\n - dilation[i] * (weight_shape[i + 2] - 1)\n for i in range(ndim)\n ]\n\n # Forward & backward.\n class Conv2d(torch.autograd.Function):\n @staticmethod\n def forward(ctx, input, weight, bias):\n assert weight.shape == weight_shape\n ctx.save_for_backward(\n input if weight.requires_grad else _null_tensor,\n weight if input.requires_grad else _null_tensor,\n )\n ctx.input_shape = input.shape\n\n # Simple 1x1 convolution => cuBLAS (only on Volta, not on Ampere).\n if weight_shape[2:] == stride == dilation == (1, 1) and padding == (0, 0) and torch.cuda.get_device_capability(input.device) < (8, 0):\n a = weight.reshape(groups, weight_shape[0] // groups, weight_shape[1])\n b = input.reshape(input.shape[0], groups, input.shape[1] // groups, -1)\n c = (a.transpose(1, 2) if transpose else a) @ b.permute(1, 2, 0, 3).flatten(2)\n c = c.reshape(-1, input.shape[0], *input.shape[2:]).transpose(0, 1)\n c = c if bias is None else c + bias.unsqueeze(0).unsqueeze(2).unsqueeze(3)\n return c.contiguous(memory_format=(torch.channels_last if input.stride(1) == 1 else torch.contiguous_format))\n\n # General case => cuDNN.\n if transpose:\n return torch.nn.functional.conv_transpose2d(input=input, weight=weight, bias=bias, output_padding=output_padding, **common_kwargs)\n return torch.nn.functional.conv2d(input=input, weight=weight, bias=bias, **common_kwargs)\n\n @staticmethod\n def backward(ctx, grad_output):\n input, weight = ctx.saved_tensors\n input_shape = ctx.input_shape\n grad_input = None\n grad_weight = None\n grad_bias = None\n\n if ctx.needs_input_grad[0]:\n p = calc_output_padding(input_shape=input_shape, output_shape=grad_output.shape)\n op = _conv2d_gradfix(transpose=(not transpose), weight_shape=weight_shape, output_padding=p, **common_kwargs)\n grad_input = op.apply(grad_output, weight, None)\n assert grad_input.shape == input_shape\n\n if ctx.needs_input_grad[1] and not weight_gradients_disabled:\n grad_weight = Conv2dGradWeight.apply(grad_output, input)\n assert grad_weight.shape == weight_shape\n\n if ctx.needs_input_grad[2]:\n grad_bias = grad_output.sum([0, 2, 3])\n\n return grad_input, grad_weight, grad_bias\n\n # Gradient with respect to the weights.\n class Conv2dGradWeight(torch.autograd.Function):\n @staticmethod\n def forward(ctx, grad_output, input):\n ctx.save_for_backward(\n grad_output if input.requires_grad else _null_tensor,\n input if grad_output.requires_grad else _null_tensor,\n )\n ctx.grad_output_shape = grad_output.shape\n ctx.input_shape = input.shape\n\n # Simple 1x1 convolution => cuBLAS (on both Volta and Ampere).\n if weight_shape[2:] == stride == dilation == (1, 1) and padding == (0, 0):\n a = grad_output.reshape(grad_output.shape[0], groups, grad_output.shape[1] // groups, -1).permute(1, 2, 0, 3).flatten(2)\n b = input.reshape(input.shape[0], groups, input.shape[1] // groups, -1).permute(1, 2, 0, 3).flatten(2)\n c = (b @ a.transpose(1, 2) if transpose else a @ b.transpose(1, 2)).reshape(weight_shape)\n return c.contiguous(memory_format=(torch.channels_last if input.stride(1) == 1 else torch.contiguous_format))\n\n # General case => cuDNN.\n name = 'aten::cudnn_convolution_transpose_backward_weight' if transpose else 'aten::cudnn_convolution_backward_weight'\n flags = [torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic, torch.backends.cudnn.allow_tf32]\n return torch._C._jit_get_operation(name)(weight_shape, grad_output, input, padding, stride, dilation, groups, *flags)\n\n @staticmethod\n def backward(ctx, grad2_grad_weight):\n grad_output, input = ctx.saved_tensors\n grad_output_shape = ctx.grad_output_shape\n input_shape = ctx.input_shape\n grad2_grad_output = None\n grad2_input = None\n\n if ctx.needs_input_grad[0]:\n grad2_grad_output = Conv2d.apply(input, grad2_grad_weight, None)\n assert grad2_grad_output.shape == grad_output_shape\n\n if ctx.needs_input_grad[1]:\n p = calc_output_padding(input_shape=input_shape, output_shape=grad_output_shape)\n op = _conv2d_gradfix(transpose=(not transpose), weight_shape=weight_shape, output_padding=p, **common_kwargs)\n grad2_input = op.apply(grad_output, grad2_grad_weight, None)\n assert grad2_input.shape == input_shape\n\n return grad2_grad_output, grad2_input\n\n _conv2d_gradfix_cache[key] = Conv2d\n return Conv2d\n\n#----------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/conv2d_resample.py",
"content": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n\"\"\"2D convolution with optional up/downsampling.\"\"\"\n\nimport torch\n\nfrom .. import misc\nfrom . import conv2d_gradfix\nfrom . import upfirdn2d\nfrom .upfirdn2d import _parse_padding\nfrom .upfirdn2d import _get_filter_size\n\n#----------------------------------------------------------------------------\n\ndef _get_weight_shape(w):\n with misc.suppress_tracer_warnings(): # this value will be treated as a constant\n shape = [int(sz) for sz in w.shape]\n misc.assert_shape(w, shape)\n return shape\n\n#----------------------------------------------------------------------------\n\ndef _conv2d_wrapper(x, w, stride=1, padding=0, groups=1, transpose=False, flip_weight=True):\n \"\"\"Wrapper for the underlying `conv2d()` and `conv_transpose2d()` implementations.\n \"\"\"\n _out_channels, _in_channels_per_group, kh, kw = _get_weight_shape(w)\n\n # Flip weight if requested.\n # Note: conv2d() actually performs correlation (flip_weight=True) not convolution (flip_weight=False).\n if not flip_weight and (kw > 1 or kh > 1):\n w = w.flip([2, 3])\n\n # Execute using conv2d_gradfix.\n op = conv2d_gradfix.conv_transpose2d if transpose else conv2d_gradfix.conv2d\n return op(x, w, stride=stride, padding=padding, groups=groups)\n\n#----------------------------------------------------------------------------\n\n@misc.profiled_function\ndef conv2d_resample(x, w, f=None, up=1, down=1, padding=0, groups=1, flip_weight=True, flip_filter=False):\n r\"\"\"2D convolution with optional up/downsampling.\n\n Padding is performed only once at the beginning, not between the operations.\n\n Args:\n x: Input tensor of shape\n `[batch_size, in_channels, in_height, in_width]`.\n w: Weight tensor of shape\n `[out_channels, in_channels//groups, kernel_height, kernel_width]`.\n f: Low-pass filter for up/downsampling. Must be prepared beforehand by\n calling upfirdn2d.setup_filter(). None = identity (default).\n up: Integer upsampling factor (default: 1).\n down: Integer downsampling factor (default: 1).\n padding: Padding with respect to the upsampled image. Can be a single number\n or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`\n (default: 0).\n groups: Split input channels into N groups (default: 1).\n flip_weight: False = convolution, True = correlation (default: True).\n flip_filter: False = convolution, True = correlation (default: False).\n\n Returns:\n Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.\n \"\"\"\n # Validate arguments.\n assert isinstance(x, torch.Tensor) and (x.ndim == 4)\n assert isinstance(w, torch.Tensor) and (w.ndim == 4) and (w.dtype == x.dtype)\n assert f is None or (isinstance(f, torch.Tensor) and f.ndim in [1, 2] and f.dtype == torch.float32)\n assert isinstance(up, int) and (up >= 1)\n assert isinstance(down, int) and (down >= 1)\n assert isinstance(groups, int) and (groups >= 1)\n out_channels, in_channels_per_group, kh, kw = _get_weight_shape(w)\n fw, fh = _get_filter_size(f)\n px0, px1, py0, py1 = _parse_padding(padding)\n\n # Adjust padding to account for up/downsampling.\n if up > 1:\n px0 += (fw + up - 1) // 2\n px1 += (fw - up) // 2\n py0 += (fh + up - 1) // 2\n py1 += (fh - up) // 2\n if down > 1:\n px0 += (fw - down + 1) // 2\n px1 += (fw - down) // 2\n py0 += (fh - down + 1) // 2\n py1 += (fh - down) // 2\n\n # Fast path: 1x1 convolution with downsampling only => downsample first, then convolve.\n if kw == 1 and kh == 1 and (down > 1 and up == 1):\n x = upfirdn2d.upfirdn2d(x=x, f=f, down=down, padding=[px0,px1,py0,py1], flip_filter=flip_filter)\n x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight)\n return x\n\n # Fast path: 1x1 convolution with upsampling only => convolve first, then upsample.\n if kw == 1 and kh == 1 and (up > 1 and down == 1):\n x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight)\n x = upfirdn2d.upfirdn2d(x=x, f=f, up=up, padding=[px0,px1,py0,py1], gain=up**2, flip_filter=flip_filter)\n return x\n\n # Fast path: downsampling only => use strided convolution.\n if down > 1 and up == 1:\n x = upfirdn2d.upfirdn2d(x=x, f=f, padding=[px0,px1,py0,py1], flip_filter=flip_filter)\n x = _conv2d_wrapper(x=x, w=w, stride=down, groups=groups, flip_weight=flip_weight)\n return x\n\n # Fast path: upsampling with optional downsampling => use transpose strided convolution.\n if up > 1:\n if groups == 1:\n w = w.transpose(0, 1)\n else:\n w = w.reshape(groups, out_channels // groups, in_channels_per_group, kh, kw)\n w = w.transpose(1, 2)\n w = w.reshape(groups * in_channels_per_group, out_channels // groups, kh, kw)\n px0 -= kw - 1\n px1 -= kw - up\n py0 -= kh - 1\n py1 -= kh - up\n pxt = max(min(-px0, -px1), 0)\n pyt = max(min(-py0, -py1), 0)\n x = _conv2d_wrapper(x=x, w=w, stride=up, padding=[pyt,pxt], groups=groups, transpose=True, flip_weight=(not flip_weight))\n x = upfirdn2d.upfirdn2d(x=x, f=f, padding=[px0+pxt,px1+pxt,py0+pyt,py1+pyt], gain=up**2, flip_filter=flip_filter)\n if down > 1:\n x = upfirdn2d.upfirdn2d(x=x, f=f, down=down, flip_filter=flip_filter)\n return x\n\n # Fast path: no up/downsampling, padding supported by the underlying implementation => use plain conv2d.\n if up == 1 and down == 1:\n if px0 == px1 and py0 == py1 and px0 >= 0 and py0 >= 0:\n return _conv2d_wrapper(x=x, w=w, padding=[py0,px0], groups=groups, flip_weight=flip_weight)\n\n # Fallback: Generic reference implementation.\n x = upfirdn2d.upfirdn2d(x=x, f=(f if up > 1 else None), up=up, padding=[px0,px1,py0,py1], gain=up**2, flip_filter=flip_filter)\n x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight)\n if down > 1:\n x = upfirdn2d.upfirdn2d(x=x, f=f, down=down, flip_filter=flip_filter)\n return x\n\n#----------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/filtered_lrelu.cpp",
"content": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors retain all intellectual property\n// and proprietary rights in and to this software, related documentation\n// and any modifications thereto. Any use, reproduction, disclosure or\n// distribution of this software and related documentation without an express\n// license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n#include \n#include \n#include \n#include \"filtered_lrelu.h\"\n\n//------------------------------------------------------------------------\n\nstatic std::tuple filtered_lrelu(\n torch::Tensor x, torch::Tensor fu, torch::Tensor fd, torch::Tensor b, torch::Tensor si,\n int up, int down, int px0, int px1, int py0, int py1, int sx, int sy, float gain, float slope, float clamp, bool flip_filters, bool writeSigns)\n{\n // Set CUDA device.\n TORCH_CHECK(x.is_cuda(), \"x must reside on CUDA device\");\n const at::cuda::OptionalCUDAGuard device_guard(device_of(x));\n\n // Validate arguments.\n TORCH_CHECK(fu.device() == x.device() && fd.device() == x.device() && b.device() == x.device(), \"all input tensors must reside on the same device\");\n TORCH_CHECK(fu.dtype() == torch::kFloat && fd.dtype() == torch::kFloat, \"fu and fd must be float32\");\n TORCH_CHECK(b.dtype() == x.dtype(), \"x and b must have the same dtype\");\n TORCH_CHECK(x.dtype() == torch::kHalf || x.dtype() == torch::kFloat, \"x and b must be float16 or float32\");\n TORCH_CHECK(x.dim() == 4, \"x must be rank 4\");\n TORCH_CHECK(x.size(0) * x.size(1) <= INT_MAX && x.size(2) <= INT_MAX && x.size(3) <= INT_MAX, \"x is too large\");\n TORCH_CHECK(x.numel() > 0, \"x is empty\");\n TORCH_CHECK((fu.dim() == 1 || fu.dim() == 2) && (fd.dim() == 1 || fd.dim() == 2), \"fu and fd must be rank 1 or 2\");\n TORCH_CHECK(fu.size(0) <= INT_MAX && fu.size(-1) <= INT_MAX, \"fu is too large\");\n TORCH_CHECK(fd.size(0) <= INT_MAX && fd.size(-1) <= INT_MAX, \"fd is too large\");\n TORCH_CHECK(fu.numel() > 0, \"fu is empty\");\n TORCH_CHECK(fd.numel() > 0, \"fd is empty\");\n TORCH_CHECK(b.dim() == 1 && b.size(0) == x.size(1), \"b must be a vector with the same number of channels as x\");\n TORCH_CHECK(up >= 1 && down >= 1, \"up and down must be at least 1\");\n\n // Figure out how much shared memory is available on the device.\n int maxSharedBytes = 0;\n AT_CUDA_CHECK(cudaDeviceGetAttribute(&maxSharedBytes, cudaDevAttrMaxSharedMemoryPerBlockOptin, x.device().index()));\n int sharedKB = maxSharedBytes >> 10;\n\n // Populate enough launch parameters to check if a CUDA kernel exists.\n filtered_lrelu_kernel_params p;\n p.up = up;\n p.down = down;\n p.fuShape = make_int2((int)fu.size(-1), fu.dim() == 2 ? (int)fu.size(0) : 0); // shape [n, 0] indicates separable filter.\n p.fdShape = make_int2((int)fd.size(-1), fd.dim() == 2 ? (int)fd.size(0) : 0);\n filtered_lrelu_kernel_spec test_spec = choose_filtered_lrelu_kernel(p, sharedKB);\n if (!test_spec.exec)\n {\n // No kernel found - return empty tensors and indicate missing kernel with return code of -1.\n return std::make_tuple(torch::Tensor(), torch::Tensor(), -1);\n }\n\n // Input/output element size.\n int64_t sz = (x.dtype() == torch::kHalf) ? 2 : 4;\n\n // Input sizes.\n int64_t xw = (int)x.size(3);\n int64_t xh = (int)x.size(2);\n int64_t fut_w = (int)fu.size(-1) - 1;\n int64_t fut_h = (int)fu.size(0) - 1;\n int64_t fdt_w = (int)fd.size(-1) - 1;\n int64_t fdt_h = (int)fd.size(0) - 1;\n\n // Logical size of upsampled buffer.\n int64_t cw = xw * up + (px0 + px1) - fut_w;\n int64_t ch = xh * up + (py0 + py1) - fut_h;\n TORCH_CHECK(cw > fdt_w && ch > fdt_h, \"upsampled buffer must be at least the size of downsampling filter\");\n TORCH_CHECK(cw <= INT_MAX && ch <= INT_MAX, \"upsampled buffer is too large\");\n\n // Compute output size and allocate.\n int64_t yw = (cw - fdt_w + (down - 1)) / down;\n int64_t yh = (ch - fdt_h + (down - 1)) / down;\n TORCH_CHECK(yw > 0 && yh > 0, \"output must be at least 1x1\");\n TORCH_CHECK(yw <= INT_MAX && yh <= INT_MAX, \"output is too large\");\n torch::Tensor y = torch::empty({x.size(0), x.size(1), yh, yw}, x.options(), x.suggest_memory_format());\n\n // Allocate sign tensor.\n torch::Tensor so;\n torch::Tensor s = si;\n bool readSigns = !!s.numel();\n int64_t sw_active = 0; // Active width of sign tensor.\n if (writeSigns)\n {\n sw_active = yw * down - (down - 1) + fdt_w; // Active width in elements.\n int64_t sh = yh * down - (down - 1) + fdt_h; // Height = active height.\n int64_t sw = (sw_active + 15) & ~15; // Width = active width in elements, rounded up to multiple of 16.\n TORCH_CHECK(sh <= INT_MAX && (sw >> 2) <= INT_MAX, \"signs is too large\");\n s = so = torch::empty({x.size(0), x.size(1), sh, sw >> 2}, x.options().dtype(torch::kUInt8), at::MemoryFormat::Contiguous);\n }\n else if (readSigns)\n sw_active = s.size(3) << 2;\n\n // Validate sign tensor if in use.\n if (readSigns || writeSigns)\n {\n TORCH_CHECK(s.is_contiguous(), \"signs must be contiguous\");\n TORCH_CHECK(s.dtype() == torch::kUInt8, \"signs must be uint8\");\n TORCH_CHECK(s.device() == x.device(), \"signs must reside on the same device as x\");\n TORCH_CHECK(s.dim() == 4, \"signs must be rank 4\");\n TORCH_CHECK(s.size(0) == x.size(0) && s.size(1) == x.size(1), \"signs must have same batch & channels as x\");\n TORCH_CHECK(s.size(2) <= INT_MAX && s.size(3) <= INT_MAX, \"signs is too large\");\n }\n\n // Populate rest of CUDA kernel parameters.\n p.x = x.data_ptr();\n p.y = y.data_ptr();\n p.b = b.data_ptr();\n p.s = (readSigns || writeSigns) ? s.data_ptr() : 0;\n p.fu = fu.data_ptr();\n p.fd = fd.data_ptr();\n p.pad0 = make_int2(px0, py0);\n p.gain = gain;\n p.slope = slope;\n p.clamp = clamp;\n p.flip = (flip_filters) ? 1 : 0;\n p.xShape = make_int4((int)x.size(3), (int)x.size(2), (int)x.size(1), (int)x.size(0));\n p.yShape = make_int4((int)y.size(3), (int)y.size(2), (int)y.size(1), (int)y.size(0));\n p.sShape = (readSigns || writeSigns) ? make_int2((int)s.size(3), (int)s.size(2)) : make_int2(0, 0); // Width is in bytes. Contiguous.\n p.sOfs = make_int2(sx, sy);\n p.swLimit = (sw_active + 3) >> 2; // Rounded up to bytes.\n\n // x, y, b strides are in bytes.\n p.xStride = make_longlong4(sz * x.stride(3), sz * x.stride(2), sz * x.stride(1), sz * x.stride(0));\n p.yStride = make_longlong4(sz * y.stride(3), sz * y.stride(2), sz * y.stride(1), sz * y.stride(0));\n p.bStride = sz * b.stride(0);\n\n // fu, fd strides are in elements.\n p.fuStride = make_longlong3(fu.stride(-1), fu.dim() == 2 ? fu.stride(0) : 0, 0);\n p.fdStride = make_longlong3(fd.stride(-1), fd.dim() == 2 ? fd.stride(0) : 0, 0);\n\n // Determine if indices don't fit in int32. Support negative strides although Torch currently never produces those.\n bool index64b = false;\n if (std::abs(p.bStride * x.size(1)) > INT_MAX) index64b = true;\n if (std::min(x.size(0) * p.xStride.w, 0ll) + std::min(x.size(1) * p.xStride.z, 0ll) + std::min(x.size(2) * p.xStride.y, 0ll) + std::min(x.size(3) * p.xStride.x, 0ll) < -INT_MAX) index64b = true;\n if (std::max(x.size(0) * p.xStride.w, 0ll) + std::max(x.size(1) * p.xStride.z, 0ll) + std::max(x.size(2) * p.xStride.y, 0ll) + std::max(x.size(3) * p.xStride.x, 0ll) > INT_MAX) index64b = true;\n if (std::min(y.size(0) * p.yStride.w, 0ll) + std::min(y.size(1) * p.yStride.z, 0ll) + std::min(y.size(2) * p.yStride.y, 0ll) + std::min(y.size(3) * p.yStride.x, 0ll) < -INT_MAX) index64b = true;\n if (std::max(y.size(0) * p.yStride.w, 0ll) + std::max(y.size(1) * p.yStride.z, 0ll) + std::max(y.size(2) * p.yStride.y, 0ll) + std::max(y.size(3) * p.yStride.x, 0ll) > INT_MAX) index64b = true;\n if (s.numel() > INT_MAX) index64b = true;\n\n // Choose CUDA kernel.\n filtered_lrelu_kernel_spec spec = { 0 };\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), \"filtered_lrelu_cuda\", [&]\n {\n if constexpr (sizeof(scalar_t) <= 4) // Exclude doubles. constexpr prevents template instantiation.\n {\n // Choose kernel based on index type, datatype and sign read/write modes.\n if (!index64b && writeSigns && !readSigns) spec = choose_filtered_lrelu_kernel(p, sharedKB);\n else if (!index64b && !writeSigns && readSigns) spec = choose_filtered_lrelu_kernel(p, sharedKB);\n else if (!index64b && !writeSigns && !readSigns) spec = choose_filtered_lrelu_kernel(p, sharedKB);\n else if ( index64b && writeSigns && !readSigns) spec = choose_filtered_lrelu_kernel(p, sharedKB);\n else if ( index64b && !writeSigns && readSigns) spec = choose_filtered_lrelu_kernel(p, sharedKB);\n else if ( index64b && !writeSigns && !readSigns) spec = choose_filtered_lrelu_kernel(p, sharedKB);\n }\n });\n TORCH_CHECK(spec.exec, \"internal error - CUDA kernel not found\") // This should not happen because we tested earlier that kernel exists.\n\n // Launch CUDA kernel.\n void* args[] = {&p};\n int bx = spec.numWarps * 32;\n int gx = (p.yShape.x - 1) / spec.tileOut.x + 1;\n int gy = (p.yShape.y - 1) / spec.tileOut.y + 1;\n int gz = p.yShape.z * p.yShape.w;\n\n // Repeat multiple horizontal tiles in a CTA?\n if (spec.xrep)\n {\n p.tilesXrep = spec.xrep;\n p.tilesXdim = gx;\n\n gx = (gx + p.tilesXrep - 1) / p.tilesXrep;\n std::swap(gx, gy);\n }\n else\n {\n p.tilesXrep = 0;\n p.tilesXdim = 0;\n }\n\n // Launch filter setup kernel.\n AT_CUDA_CHECK(cudaLaunchKernel(spec.setup, 1, 1024, args, 0, at::cuda::getCurrentCUDAStream()));\n\n // Copy kernels to constant memory.\n if ( writeSigns && !readSigns) AT_CUDA_CHECK((copy_filters(at::cuda::getCurrentCUDAStream())));\n else if (!writeSigns && readSigns) AT_CUDA_CHECK((copy_filters(at::cuda::getCurrentCUDAStream())));\n else if (!writeSigns && !readSigns) AT_CUDA_CHECK((copy_filters(at::cuda::getCurrentCUDAStream())));\n\n // Set cache and shared memory configurations for main kernel.\n AT_CUDA_CHECK(cudaFuncSetCacheConfig(spec.exec, cudaFuncCachePreferShared));\n if (spec.dynamicSharedKB) // Need dynamically allocated shared memory?\n AT_CUDA_CHECK(cudaFuncSetAttribute(spec.exec, cudaFuncAttributeMaxDynamicSharedMemorySize, spec.dynamicSharedKB << 10));\n AT_CUDA_CHECK(cudaFuncSetSharedMemConfig(spec.exec, cudaSharedMemBankSizeFourByte));\n\n // Launch main kernel.\n const int maxSubGz = 65535; // CUDA maximum for block z dimension.\n for (int zofs=0; zofs < gz; zofs += maxSubGz) // Do multiple launches if gz is too big.\n {\n p.blockZofs = zofs;\n int subGz = std::min(maxSubGz, gz - zofs);\n AT_CUDA_CHECK(cudaLaunchKernel(spec.exec, dim3(gx, gy, subGz), bx, args, spec.dynamicSharedKB << 10, at::cuda::getCurrentCUDAStream()));\n }\n\n // Done.\n return std::make_tuple(y, so, 0);\n}\n\n//------------------------------------------------------------------------\n\nstatic torch::Tensor filtered_lrelu_act(torch::Tensor x, torch::Tensor si, int sx, int sy, float gain, float slope, float clamp, bool writeSigns)\n{\n // Set CUDA device.\n TORCH_CHECK(x.is_cuda(), \"x must reside on CUDA device\");\n const at::cuda::OptionalCUDAGuard device_guard(device_of(x));\n\n // Validate arguments.\n TORCH_CHECK(x.dim() == 4, \"x must be rank 4\");\n TORCH_CHECK(x.size(0) * x.size(1) <= INT_MAX && x.size(2) <= INT_MAX && x.size(3) <= INT_MAX, \"x is too large\");\n TORCH_CHECK(x.numel() > 0, \"x is empty\");\n TORCH_CHECK(x.dtype() == torch::kHalf || x.dtype() == torch::kFloat || x.dtype() == torch::kDouble, \"x must be float16, float32 or float64\");\n\n // Output signs if we don't have sign input.\n torch::Tensor so;\n torch::Tensor s = si;\n bool readSigns = !!s.numel();\n if (writeSigns)\n {\n int64_t sw = x.size(3);\n sw = (sw + 15) & ~15; // Round to a multiple of 16 for coalescing.\n s = so = torch::empty({x.size(0), x.size(1), x.size(2), sw >> 2}, x.options().dtype(torch::kUInt8), at::MemoryFormat::Contiguous);\n }\n\n // Validate sign tensor if in use.\n if (readSigns || writeSigns)\n {\n TORCH_CHECK(s.is_contiguous(), \"signs must be contiguous\");\n TORCH_CHECK(s.dtype() == torch::kUInt8, \"signs must be uint8\");\n TORCH_CHECK(s.device() == x.device(), \"signs must reside on the same device as x\");\n TORCH_CHECK(s.dim() == 4, \"signs must be rank 4\");\n TORCH_CHECK(s.size(0) == x.size(0) && s.size(1) == x.size(1), \"signs must have same batch & channels as x\");\n TORCH_CHECK(s.size(2) <= INT_MAX && (s.size(3) << 2) <= INT_MAX, \"signs tensor is too large\");\n }\n\n // Initialize CUDA kernel parameters.\n filtered_lrelu_act_kernel_params p;\n p.x = x.data_ptr();\n p.s = (readSigns || writeSigns) ? s.data_ptr() : 0;\n p.gain = gain;\n p.slope = slope;\n p.clamp = clamp;\n p.xShape = make_int4((int)x.size(3), (int)x.size(2), (int)x.size(1), (int)x.size(0));\n p.xStride = make_longlong4(x.stride(3), x.stride(2), x.stride(1), x.stride(0));\n p.sShape = (readSigns || writeSigns) ? make_int2((int)s.size(3) << 2, (int)s.size(2)) : make_int2(0, 0); // Width is in elements. Contiguous.\n p.sOfs = make_int2(sx, sy);\n\n // Choose CUDA kernel.\n void* func = 0;\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), \"filtered_lrelu_act_cuda\", [&]\n {\n if (writeSigns)\n func = choose_filtered_lrelu_act_kernel();\n else if (readSigns)\n func = choose_filtered_lrelu_act_kernel();\n else\n func = choose_filtered_lrelu_act_kernel();\n });\n TORCH_CHECK(func, \"internal error - CUDA kernel not found\");\n\n // Launch CUDA kernel.\n void* args[] = {&p};\n int bx = 128; // 4 warps per block.\n\n // Logical size of launch = writeSigns ? p.s : p.x\n uint32_t gx = writeSigns ? p.sShape.x : p.xShape.x;\n uint32_t gy = writeSigns ? p.sShape.y : p.xShape.y;\n uint32_t gz = p.xShape.z * p.xShape.w; // Same as in p.sShape if signs are in use.\n gx = (gx - 1) / bx + 1;\n\n // Make sure grid y and z dimensions are within CUDA launch limits. Kernel loops internally to do the rest.\n const uint32_t gmax = 65535;\n gy = std::min(gy, gmax);\n gz = std::min(gz, gmax);\n\n // Launch.\n AT_CUDA_CHECK(cudaLaunchKernel(func, dim3(gx, gy, gz), bx, args, 0, at::cuda::getCurrentCUDAStream()));\n return so;\n}\n\n//------------------------------------------------------------------------\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m)\n{\n m.def(\"filtered_lrelu\", &filtered_lrelu); // The whole thing.\n m.def(\"filtered_lrelu_act_\", &filtered_lrelu_act); // Activation and sign tensor handling only. Modifies data tensor in-place.\n}\n\n//------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/filtered_lrelu.cu",
"content": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors retain all intellectual property\n// and proprietary rights in and to this software, related documentation\n// and any modifications thereto. Any use, reproduction, disclosure or\n// distribution of this software and related documentation without an express\n// license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n#include \n#include \"filtered_lrelu.h\"\n#include \n\n//------------------------------------------------------------------------\n// Helpers.\n\nenum // Filter modes.\n{\n MODE_SUSD = 0, // Separable upsampling, separable downsampling.\n MODE_FUSD = 1, // Full upsampling, separable downsampling.\n MODE_SUFD = 2, // Separable upsampling, full downsampling.\n MODE_FUFD = 3, // Full upsampling, full downsampling.\n};\n\ntemplate struct InternalType;\ntemplate <> struct InternalType\n{\n typedef double scalar_t; typedef double2 vec2_t; typedef double4 vec4_t;\n __device__ __forceinline__ static vec2_t zero_vec2(void) { return make_double2(0, 0); }\n __device__ __forceinline__ static vec4_t zero_vec4(void) { return make_double4(0, 0, 0, 0); }\n __device__ __forceinline__ static double clamp(double x, double c) { return fmin(fmax(x, -c), c); }\n};\ntemplate <> struct InternalType\n{\n typedef float scalar_t; typedef float2 vec2_t; typedef float4 vec4_t;\n __device__ __forceinline__ static vec2_t zero_vec2(void) { return make_float2(0, 0); }\n __device__ __forceinline__ static vec4_t zero_vec4(void) { return make_float4(0, 0, 0, 0); }\n __device__ __forceinline__ static float clamp(float x, float c) { return fminf(fmaxf(x, -c), c); }\n};\ntemplate <> struct InternalType\n{\n typedef float scalar_t; typedef float2 vec2_t; typedef float4 vec4_t;\n __device__ __forceinline__ static vec2_t zero_vec2(void) { return make_float2(0, 0); }\n __device__ __forceinline__ static vec4_t zero_vec4(void) { return make_float4(0, 0, 0, 0); }\n __device__ __forceinline__ static float clamp(float x, float c) { return fminf(fmaxf(x, -c), c); }\n};\n\n#define MIN(A, B) ((A) < (B) ? (A) : (B))\n#define MAX(A, B) ((A) > (B) ? (A) : (B))\n#define CEIL_DIV(A, B) (((B)==1) ? (A) : \\\n ((B)==2) ? ((int)((A)+1) >> 1) : \\\n ((B)==4) ? ((int)((A)+3) >> 2) : \\\n (((A) + ((A) > 0 ? (B) - 1 : 0)) / (B)))\n\n// This works only up to blocks of size 256 x 256 and for all N that are powers of two.\ntemplate __device__ __forceinline__ void fast_div_mod(int& x, int& y, unsigned int i)\n{\n if ((N & (N-1)) && N <= 256)\n y = (i * ((1<<24)/N + 1)) >> 24; // Assumes N <= 256, i < N*256.\n else\n y = i/N;\n\n x = i - y*N;\n}\n\n// Type cast stride before reading it.\ntemplate __device__ __forceinline__ T get_stride(const int64_t& x)\n{\n return *reinterpret_cast(&x);\n}\n\n//------------------------------------------------------------------------\n// Filters, setup kernel, copying function.\n\n#define MAX_FILTER_SIZE 32\n\n// Combined up/down filter buffers so that transfer can be done with one copy.\n__device__ float g_fbuf[2 * MAX_FILTER_SIZE * MAX_FILTER_SIZE]; // Filters in global memory, written by setup kernel.\n__device__ __constant__ float c_fbuf[2 * MAX_FILTER_SIZE * MAX_FILTER_SIZE]; // Filters in constant memory, read by main kernel.\n\n// Accessors to combined buffers to index up/down filters individually.\n#define c_fu (c_fbuf)\n#define c_fd (c_fbuf + MAX_FILTER_SIZE * MAX_FILTER_SIZE)\n#define g_fu (g_fbuf)\n#define g_fd (g_fbuf + MAX_FILTER_SIZE * MAX_FILTER_SIZE)\n\n// Set up filters into global memory buffer.\nstatic __global__ void setup_filters_kernel(filtered_lrelu_kernel_params p)\n{\n for (int idx = threadIdx.x; idx < MAX_FILTER_SIZE * MAX_FILTER_SIZE; idx += blockDim.x)\n {\n int x, y;\n fast_div_mod(x, y, idx);\n\n int fu_x = p.flip ? x : (p.fuShape.x - 1 - x);\n int fu_y = p.flip ? y : (p.fuShape.y - 1 - y);\n if (p.fuShape.y > 0)\n g_fu[idx] = (x >= p.fuShape.x || y >= p.fuShape.y) ? 0.0f : p.fu[fu_x * p.fuStride.x + fu_y * p.fuStride.y];\n else\n g_fu[idx] = (x >= p.fuShape.x || y > 0) ? 0.0f : p.fu[fu_x * p.fuStride.x];\n\n int fd_x = p.flip ? x : (p.fdShape.x - 1 - x);\n int fd_y = p.flip ? y : (p.fdShape.y - 1 - y);\n if (p.fdShape.y > 0)\n g_fd[idx] = (x >= p.fdShape.x || y >= p.fdShape.y) ? 0.0f : p.fd[fd_x * p.fdStride.x + fd_y * p.fdStride.y];\n else\n g_fd[idx] = (x >= p.fdShape.x || y > 0) ? 0.0f : p.fd[fd_x * p.fdStride.x];\n }\n}\n\n// Host function to copy filters written by setup kernel into constant buffer for main kernel.\ntemplate static cudaError_t copy_filters(cudaStream_t stream)\n{\n void* src = 0;\n cudaError_t err = cudaGetSymbolAddress(&src, g_fbuf);\n if (err) return err;\n return cudaMemcpyToSymbolAsync(c_fbuf, src, 2 * MAX_FILTER_SIZE * MAX_FILTER_SIZE * sizeof(float), 0, cudaMemcpyDeviceToDevice, stream);\n}\n\n//------------------------------------------------------------------------\n// Coordinate spaces:\n// - Relative to input tensor: inX, inY, tileInX, tileInY\n// - Relative to input tile: relInX, relInY, tileInW, tileInH\n// - Relative to upsampled tile: relUpX, relUpY, tileUpW, tileUpH\n// - Relative to output tile: relOutX, relOutY, tileOutW, tileOutH\n// - Relative to output tensor: outX, outY, tileOutX, tileOutY\n//\n// Relationships between coordinate spaces:\n// - inX = tileInX + relInX\n// - inY = tileInY + relInY\n// - relUpX = relInX * up + phaseInX\n// - relUpY = relInY * up + phaseInY\n// - relUpX = relOutX * down\n// - relUpY = relOutY * down\n// - outX = tileOutX + relOutX\n// - outY = tileOutY + relOutY\n\nextern __shared__ char s_buf_raw[]; // When sharedKB <= 48, allocate shared memory statically inside the kernel, otherwise use the externally allocated shared memory buffer.\n\ntemplate \nstatic __global__ void filtered_lrelu_kernel(filtered_lrelu_kernel_params p)\n{\n // Check that we don't try to support non-existing filter modes.\n static_assert(up == 1 || up == 2 || up == 4, \"only up=1, up=2, up=4 scales supported\");\n static_assert(down == 1 || down == 2 || down == 4, \"only down=1, down=2, down=4 scales supported\");\n static_assert(fuSize >= up, \"upsampling filter size must be at least upsampling factor\");\n static_assert(fdSize >= down, \"downsampling filter size must be at least downsampling factor\");\n static_assert(fuSize % up == 0, \"upsampling filter size must be divisible with upsampling factor\");\n static_assert(fdSize % down == 0, \"downsampling filter size must be divisible with downsampling factor\");\n static_assert(fuSize <= MAX_FILTER_SIZE && fdSize <= MAX_FILTER_SIZE, \"filter size greater than MAX_FILTER_SIZE\");\n static_assert(up != 1 || (fuSize == 1 && (filterMode == MODE_FUFD || filterMode == MODE_FUSD)), \"up=1 supported only for 1x1 full filters\");\n static_assert(down != 1 || (fdSize == 1 && (filterMode == MODE_FUFD || filterMode == MODE_SUFD)), \"down=1 supported only for 1x1 full filters\");\n static_assert(!(up == 4 && (filterMode == MODE_FUFD || filterMode == MODE_FUSD)), \"full filters not supported for up=4\");\n static_assert(!(down == 4 && (filterMode == MODE_FUFD || filterMode == MODE_SUFD)), \"full filters not supported for down=4\");\n\n // Static definitions.\n typedef typename InternalType::scalar_t scalar_t;\n typedef typename InternalType::vec2_t vec2_t;\n typedef typename InternalType::vec4_t vec4_t;\n const int tileUpW = (tileOutW * down + (fdSize - 1) - (down - 1) + 3) & ~3; // Upsampled tile width, rounded up to multiple of 4.\n const int tileUpH = tileOutH * down + (fdSize - 1) - (down - 1); // Upsampled tile height.\n const int tileInW = CEIL_DIV(tileUpW + (fuSize - 1), up); // Input tile width.\n const int tileInH = CEIL_DIV(tileUpH + (fuSize - 1), up); // Input tile height.\n const int tileUpH_up = CEIL_DIV(tileUpH, up) * up; // Upsampled tile height rounded up to a multiple of up.\n const int tileInH_up = CEIL_DIV(tileUpH_up + (fuSize - 1), up); // For allocations only, to avoid shared memory read overruns with up=2 and up=4.\n\n // Merge 1x1 downsampling into last upsampling step for upf1 and ups2.\n const bool downInline = (down == 1) && ((up == 1 && filterMode == MODE_FUFD) || (up == 2 && filterMode == MODE_SUFD));\n\n // Sizes of logical buffers.\n const int szIn = tileInH_up * tileInW;\n const int szUpX = tileInH_up * tileUpW;\n const int szUpXY = downInline ? 0 : (tileUpH * tileUpW);\n const int szDownX = tileUpH * tileOutW;\n\n // Sizes for shared memory arrays.\n const int s_buf0_size_base =\n (filterMode == MODE_SUSD) ? MAX(szIn, szUpXY) :\n (filterMode == MODE_FUSD) ? MAX(szIn, szDownX) :\n (filterMode == MODE_SUFD) ? MAX(szIn, szUpXY) :\n (filterMode == MODE_FUFD) ? szIn :\n -1;\n const int s_buf1_size_base =\n (filterMode == MODE_SUSD) ? MAX(szUpX, szDownX) :\n (filterMode == MODE_FUSD) ? szUpXY :\n (filterMode == MODE_SUFD) ? szUpX :\n (filterMode == MODE_FUFD) ? szUpXY :\n -1;\n\n // Ensure U128 alignment.\n const int s_buf0_size = (s_buf0_size_base + 3) & ~3;\n const int s_buf1_size = (s_buf1_size_base + 3) & ~3;\n\n // Check at compile time that we don't use too much shared memory.\n static_assert((s_buf0_size + s_buf1_size) * sizeof(scalar_t) <= (sharedKB << 10), \"shared memory overflow\");\n\n // Declare shared memory arrays.\n scalar_t* s_buf0;\n scalar_t* s_buf1;\n if (sharedKB <= 48)\n {\n // Allocate shared memory arrays here.\n __shared__ scalar_t s_buf0_st[(sharedKB > 48) ? (1<<24) : (s_buf0_size + s_buf1_size)]; // Prevent launching if this isn't optimized away when unused.\n s_buf0 = s_buf0_st;\n s_buf1 = s_buf0 + s_buf0_size;\n }\n else\n {\n // Use the dynamically allocated shared memory array.\n s_buf0 = (scalar_t*)s_buf_raw;\n s_buf1 = s_buf0 + s_buf0_size;\n }\n\n // Pointers to the buffers.\n scalar_t* s_tileIn; // Input tile: [relInX * tileInH + relInY]\n scalar_t* s_tileUpX; // After horizontal upsampling: [relInY * tileUpW + relUpX]\n scalar_t* s_tileUpXY; // After upsampling: [relUpY * tileUpW + relUpX]\n scalar_t* s_tileDownX; // After horizontal downsampling: [relUpY * tileOutW + relOutX]\n if (filterMode == MODE_SUSD)\n {\n s_tileIn = s_buf0;\n s_tileUpX = s_buf1;\n s_tileUpXY = s_buf0;\n s_tileDownX = s_buf1;\n }\n else if (filterMode == MODE_FUSD)\n {\n s_tileIn = s_buf0;\n s_tileUpXY = s_buf1;\n s_tileDownX = s_buf0;\n }\n else if (filterMode == MODE_SUFD)\n {\n s_tileIn = s_buf0;\n s_tileUpX = s_buf1;\n s_tileUpXY = s_buf0;\n }\n else if (filterMode == MODE_FUFD)\n {\n s_tileIn = s_buf0;\n s_tileUpXY = s_buf1;\n }\n\n // Allow large grids in z direction via per-launch offset.\n int channelIdx = blockIdx.z + p.blockZofs;\n int batchIdx = channelIdx / p.yShape.z;\n channelIdx -= batchIdx * p.yShape.z;\n\n // Offset to output feature map. In bytes.\n index_t mapOfsOut = channelIdx * get_stride(p.yStride.z) + batchIdx * get_stride(p.yStride.w);\n\n // Sign shift amount.\n uint32_t signXo = ((threadIdx.x + p.sOfs.x) << 1) & 6;\n\n // Inner tile loop.\n #pragma unroll 1\n for (int tileIdx = 0; !enableXrep || (tileIdx < MIN(p.tilesXrep, p.tilesXdim - p.tilesXrep * blockIdx.y)); tileIdx++)\n {\n // Locate output tile.\n int tileX = enableXrep ? blockIdx.y * p.tilesXrep + tileIdx : blockIdx.x;\n int tileOutX = tileX * tileOutW;\n int tileOutY = (enableXrep ? blockIdx.x : blockIdx.y) * tileOutH;\n\n // Locate input tile.\n int tmpX = tileOutX * down - p.pad0.x;\n int tmpY = tileOutY * down - p.pad0.y;\n int tileInX = CEIL_DIV(tmpX, up);\n int tileInY = CEIL_DIV(tmpY, up);\n const int phaseInX = tileInX * up - tmpX;\n const int phaseInY = tileInY * up - tmpY;\n\n // Extra sync if input and output buffers are the same and we are not on first tile.\n if (enableXrep && tileIdx > 0 && (filterMode == MODE_FUSD || (filterMode == MODE_SUFD && !downInline) || (filterMode == MODE_FUFD && downInline)))\n __syncthreads();\n\n // Load input tile & apply bias. Unrolled.\n scalar_t b = (scalar_t)*(const T*)((const char*)p.b + (channelIdx * get_stride(p.bStride)));\n index_t mapOfsIn = channelIdx * get_stride(p.xStride.z) + batchIdx * get_stride(p.xStride.w);\n int idx = threadIdx.x;\n const int loopCountIN = CEIL_DIV(tileInW * tileInH, threadsPerBlock);\n #pragma unroll\n for (int loop = 0; loop < loopCountIN; loop++)\n {\n int relInX, relInY;\n fast_div_mod(relInX, relInY, idx);\n int inX = tileInX + relInX;\n int inY = tileInY + relInY;\n scalar_t v = 0;\n\n if ((uint32_t)inX < p.xShape.x && (uint32_t)inY < p.xShape.y)\n v = (scalar_t)*((const T*)((const char*)p.x + (inX * get_stride(p.xStride.x) + inY * get_stride(p.xStride.y) + mapOfsIn))) + b;\n\n bool skip = (loop == loopCountIN-1) && (idx >= tileInW * tileInH);\n if (!skip)\n s_tileIn[idx] = v;\n\n idx += threadsPerBlock;\n }\n\n if (filterMode == MODE_SUSD || filterMode == MODE_SUFD) // Separable upsampling filter.\n {\n // Horizontal upsampling.\n __syncthreads();\n if (up == 4)\n {\n for (int idx = threadIdx.x*up; idx < tileUpW * tileInH; idx += blockDim.x*up)\n {\n int relUpX0, relInY;\n fast_div_mod(relUpX0, relInY, idx);\n int relInX0 = relUpX0 / up;\n int src0 = relInX0 + tileInW * relInY;\n int dst = relInY * tileUpW + relUpX0;\n vec4_t v = InternalType::zero_vec4();\n scalar_t a = s_tileIn[src0];\n if (phaseInX == 0)\n {\n #pragma unroll\n for (int step = 0; step < fuSize / up; step++)\n {\n v.x += a * (scalar_t)c_fu[step * up + 0];\n a = s_tileIn[src0 + step + 1];\n v.y += a * (scalar_t)c_fu[step * up + 3];\n v.z += a * (scalar_t)c_fu[step * up + 2];\n v.w += a * (scalar_t)c_fu[step * up + 1];\n }\n }\n else if (phaseInX == 1)\n {\n #pragma unroll\n for (int step = 0; step < fuSize / up; step++)\n {\n v.x += a * (scalar_t)c_fu[step * up + 1];\n v.y += a * (scalar_t)c_fu[step * up + 0];\n a = s_tileIn[src0 + step + 1];\n v.z += a * (scalar_t)c_fu[step * up + 3];\n v.w += a * (scalar_t)c_fu[step * up + 2];\n }\n }\n else if (phaseInX == 2)\n {\n #pragma unroll\n for (int step = 0; step < fuSize / up; step++)\n {\n v.x += a * (scalar_t)c_fu[step * up + 2];\n v.y += a * (scalar_t)c_fu[step * up + 1];\n v.z += a * (scalar_t)c_fu[step * up + 0];\n a = s_tileIn[src0 + step + 1];\n v.w += a * (scalar_t)c_fu[step * up + 3];\n }\n }\n else // (phaseInX == 3)\n {\n #pragma unroll\n for (int step = 0; step < fuSize / up; step++)\n {\n v.x += a * (scalar_t)c_fu[step * up + 3];\n v.y += a * (scalar_t)c_fu[step * up + 2];\n v.z += a * (scalar_t)c_fu[step * up + 1];\n v.w += a * (scalar_t)c_fu[step * up + 0];\n a = s_tileIn[src0 + step + 1];\n }\n }\n s_tileUpX[dst+0] = v.x;\n s_tileUpX[dst+1] = v.y;\n s_tileUpX[dst+2] = v.z;\n s_tileUpX[dst+3] = v.w;\n }\n }\n else if (up == 2)\n {\n bool p0 = (phaseInX == 0);\n for (int idx = threadIdx.x*up; idx < tileUpW * tileInH; idx += blockDim.x*up)\n {\n int relUpX0, relInY;\n fast_div_mod(relUpX0, relInY, idx);\n int relInX0 = relUpX0 / up;\n int src0 = relInX0 + tileInW * relInY;\n int dst = relInY * tileUpW + relUpX0;\n vec2_t v = InternalType::zero_vec2();\n scalar_t a = s_tileIn[src0];\n if (p0) // (phaseInX == 0)\n {\n #pragma unroll\n for (int step = 0; step < fuSize / up; step++)\n {\n v.x += a * (scalar_t)c_fu[step * up + 0];\n a = s_tileIn[src0 + step + 1];\n v.y += a * (scalar_t)c_fu[step * up + 1];\n }\n }\n else // (phaseInX == 1)\n {\n #pragma unroll\n for (int step = 0; step < fuSize / up; step++)\n {\n v.x += a * (scalar_t)c_fu[step * up + 1];\n v.y += a * (scalar_t)c_fu[step * up + 0];\n a = s_tileIn[src0 + step + 1];\n }\n }\n s_tileUpX[dst+0] = v.x;\n s_tileUpX[dst+1] = v.y;\n }\n }\n\n // Vertical upsampling & nonlinearity.\n\n __syncthreads();\n int groupMask = 15 << ((threadIdx.x & 31) & ~3);\n int minY = tileOutY ? (tileOutY - tileOutH) * down + tileUpH : 0; // Skip already written signs.\n int sShapeMaxY = MIN(p.sShape.y, tileOutY * down + tileUpH); // Avoid out-of-tile sign writes.\n if (up == 4)\n {\n minY -= 3; // Adjust according to block height.\n for (int idx = threadIdx.x; idx < tileUpW * tileUpH_up / up; idx += blockDim.x)\n {\n int relUpX, relInY0;\n fast_div_mod(relUpX, relInY0, idx);\n int relUpY0 = relInY0 * up;\n int src0 = relInY0 * tileUpW + relUpX;\n int dst = relUpY0 * tileUpW + relUpX;\n vec4_t v = InternalType::zero_vec4();\n\n scalar_t a = s_tileUpX[src0];\n if (phaseInY == 0)\n {\n #pragma unroll\n for (int step = 0; step < fuSize / up; step++)\n {\n v.x += a * (scalar_t)c_fu[step * up + 0];\n a = s_tileUpX[src0 + (step + 1) * tileUpW];\n v.y += a * (scalar_t)c_fu[step * up + 3];\n v.z += a * (scalar_t)c_fu[step * up + 2];\n v.w += a * (scalar_t)c_fu[step * up + 1];\n }\n }\n else if (phaseInY == 1)\n {\n #pragma unroll\n for (int step = 0; step < fuSize / up; step++)\n {\n v.x += a * (scalar_t)c_fu[step * up + 1];\n v.y += a * (scalar_t)c_fu[step * up + 0];\n a = s_tileUpX[src0 + (step + 1) * tileUpW];\n v.z += a * (scalar_t)c_fu[step * up + 3];\n v.w += a * (scalar_t)c_fu[step * up + 2];\n }\n }\n else if (phaseInY == 2)\n {\n #pragma unroll\n for (int step = 0; step < fuSize / up; step++)\n {\n v.x += a * (scalar_t)c_fu[step * up + 2];\n v.y += a * (scalar_t)c_fu[step * up + 1];\n v.z += a * (scalar_t)c_fu[step * up + 0];\n a = s_tileUpX[src0 + (step + 1) * tileUpW];\n v.w += a * (scalar_t)c_fu[step * up + 3];\n }\n }\n else // (phaseInY == 3)\n {\n #pragma unroll\n for (int step = 0; step < fuSize / up; step++)\n {\n v.x += a * (scalar_t)c_fu[step * up + 3];\n v.y += a * (scalar_t)c_fu[step * up + 2];\n v.z += a * (scalar_t)c_fu[step * up + 1];\n v.w += a * (scalar_t)c_fu[step * up + 0];\n a = s_tileUpX[src0 + (step + 1) * tileUpW];\n }\n }\n\n int x = tileOutX * down + relUpX;\n int y = tileOutY * down + relUpY0;\n int signX = x + p.sOfs.x;\n int signY = y + p.sOfs.y;\n int signZ = blockIdx.z + p.blockZofs;\n int signXb = signX >> 2;\n index_t si0 = signXb + p.sShape.x * (signY + (index_t)p.sShape.y * signZ);\n index_t si1 = si0 + p.sShape.x;\n index_t si2 = si0 + p.sShape.x * 2;\n index_t si3 = si0 + p.sShape.x * 3;\n\n v.x *= (scalar_t)((float)up * (float)up * p.gain);\n v.y *= (scalar_t)((float)up * (float)up * p.gain);\n v.z *= (scalar_t)((float)up * (float)up * p.gain);\n v.w *= (scalar_t)((float)up * (float)up * p.gain);\n\n if (signWrite)\n {\n if (!enableWriteSkip)\n {\n // Determine and write signs.\n int sx = __float_as_uint(v.x) >> 31 << 0;\n int sy = __float_as_uint(v.y) >> 31 << 8;\n int sz = __float_as_uint(v.z) >> 31 << 16;\n int sw = __float_as_uint(v.w) >> 31 << 24;\n if (sx) v.x *= p.slope;\n if (sy) v.y *= p.slope;\n if (sz) v.z *= p.slope;\n if (sw) v.w *= p.slope;\n if (fabsf(v.x) > p.clamp) { sx = 2 << 0; v.x = InternalType::clamp(v.x, p.clamp); }\n if (fabsf(v.y) > p.clamp) { sy = 2 << 8; v.y = InternalType::clamp(v.y, p.clamp); }\n if (fabsf(v.z) > p.clamp) { sz = 2 << 16; v.z = InternalType::clamp(v.z, p.clamp); }\n if (fabsf(v.w) > p.clamp) { sw = 2 << 24; v.w = InternalType::clamp(v.w, p.clamp); }\n\n if ((uint32_t)signXb < p.swLimit && signY >= minY)\n {\n // Combine signs.\n uint32_t s = sx + sy + sw + sz;\n s <<= (signX & 3) << 1;\n s |= __shfl_xor_sync(groupMask, s, 1);\n s |= __shfl_xor_sync(groupMask, s, 2);\n\n // Write signs.\n if ((uint32_t)(signY + 0) < sShapeMaxY) { p.s[si0] = (unsigned char)(s >> 0); }\n if ((uint32_t)(signY + 1) < sShapeMaxY) { p.s[si1] = (unsigned char)(s >> 8); }\n if ((uint32_t)(signY + 2) < sShapeMaxY) { p.s[si2] = (unsigned char)(s >> 16); }\n if ((uint32_t)(signY + 3) < sShapeMaxY) { p.s[si3] = (unsigned char)(s >> 24); }\n }\n }\n else\n {\n // Determine and write signs.\n if ((uint32_t)signXb < p.swLimit && signY >= minY)\n {\n int sx = __float_as_uint(v.x) >> 31 << 0;\n int sy = __float_as_uint(v.y) >> 31 << 8;\n int sz = __float_as_uint(v.z) >> 31 << 16;\n int sw = __float_as_uint(v.w) >> 31 << 24;\n if (sx) v.x *= p.slope;\n if (sy) v.y *= p.slope;\n if (sz) v.z *= p.slope;\n if (sw) v.w *= p.slope;\n if (fabsf(v.x) > p.clamp) { sx = 2 << 0; v.x = InternalType::clamp(v.x, p.clamp); }\n if (fabsf(v.y) > p.clamp) { sy = 2 << 8; v.y = InternalType::clamp(v.y, p.clamp); }\n if (fabsf(v.z) > p.clamp) { sz = 2 << 16; v.z = InternalType::clamp(v.z, p.clamp); }\n if (fabsf(v.w) > p.clamp) { sw = 2 << 24; v.w = InternalType::clamp(v.w, p.clamp); }\n\n // Combine signs.\n uint32_t s = sx + sy + sw + sz;\n s <<= (signX & 3) << 1;\n s |= __shfl_xor_sync(groupMask, s, 1);\n s |= __shfl_xor_sync(groupMask, s, 2);\n\n // Write signs.\n if ((uint32_t)(signY + 0) < sShapeMaxY) { p.s[si0] = (unsigned char)(s >> 0); }\n if ((uint32_t)(signY + 1) < sShapeMaxY) { p.s[si1] = (unsigned char)(s >> 8); }\n if ((uint32_t)(signY + 2) < sShapeMaxY) { p.s[si2] = (unsigned char)(s >> 16); }\n if ((uint32_t)(signY + 3) < sShapeMaxY) { p.s[si3] = (unsigned char)(s >> 24); }\n }\n else\n {\n // Just compute the values.\n if (v.x < 0.f) v.x *= p.slope; v.x = InternalType::clamp(v.x, p.clamp);\n if (v.y < 0.f) v.y *= p.slope; v.y = InternalType::clamp(v.y, p.clamp);\n if (v.z < 0.f) v.z *= p.slope; v.z = InternalType::clamp(v.z, p.clamp);\n if (v.w < 0.f) v.w *= p.slope; v.w = InternalType::clamp(v.w, p.clamp);\n }\n }\n }\n else if (signRead) // Read signs and apply.\n {\n if ((uint32_t)signXb < p.swLimit)\n {\n int ss = (signX & 3) << 1;\n if ((uint32_t)(signY + 0) < p.sShape.y) { int s = p.s[si0] >> ss; if (s & 1) v.x *= p.slope; if (s & 2) v.x = 0.f; }\n if ((uint32_t)(signY + 1) < p.sShape.y) { int s = p.s[si1] >> ss; if (s & 1) v.y *= p.slope; if (s & 2) v.y = 0.f; }\n if ((uint32_t)(signY + 2) < p.sShape.y) { int s = p.s[si2] >> ss; if (s & 1) v.z *= p.slope; if (s & 2) v.z = 0.f; }\n if ((uint32_t)(signY + 3) < p.sShape.y) { int s = p.s[si3] >> ss; if (s & 1) v.w *= p.slope; if (s & 2) v.w = 0.f; }\n }\n }\n else // Forward pass with no sign write.\n {\n if (v.x < 0.f) v.x *= p.slope; v.x = InternalType::clamp(v.x, p.clamp);\n if (v.y < 0.f) v.y *= p.slope; v.y = InternalType::clamp(v.y, p.clamp);\n if (v.z < 0.f) v.z *= p.slope; v.z = InternalType::clamp(v.z, p.clamp);\n if (v.w < 0.f) v.w *= p.slope; v.w = InternalType::clamp(v.w, p.clamp);\n }\n\n s_tileUpXY[dst + 0 * tileUpW] = v.x;\n if (relUpY0 + 1 < tileUpH) s_tileUpXY[dst + 1 * tileUpW] = v.y;\n if (relUpY0 + 2 < tileUpH) s_tileUpXY[dst + 2 * tileUpW] = v.z;\n if (relUpY0 + 3 < tileUpH) s_tileUpXY[dst + 3 * tileUpW] = v.w;\n }\n }\n else if (up == 2)\n {\n minY -= 1; // Adjust according to block height.\n for (int idx = threadIdx.x; idx < tileUpW * tileUpH_up / up; idx += blockDim.x)\n {\n int relUpX, relInY0;\n fast_div_mod(relUpX, relInY0, idx);\n int relUpY0 = relInY0 * up;\n int src0 = relInY0 * tileUpW + relUpX;\n int dst = relUpY0 * tileUpW + relUpX;\n vec2_t v = InternalType::zero_vec2();\n\n scalar_t a = s_tileUpX[src0];\n if (phaseInY == 0)\n {\n #pragma unroll\n for (int step = 0; step < fuSize / up; step++)\n {\n v.x += a * (scalar_t)c_fu[step * up + 0];\n a = s_tileUpX[src0 + (step + 1) * tileUpW];\n v.y += a * (scalar_t)c_fu[step * up + 1];\n }\n }\n else // (phaseInY == 1)\n {\n #pragma unroll\n for (int step = 0; step < fuSize / up; step++)\n {\n v.x += a * (scalar_t)c_fu[step * up + 1];\n v.y += a * (scalar_t)c_fu[step * up + 0];\n a = s_tileUpX[src0 + (step + 1) * tileUpW];\n }\n }\n\n int x = tileOutX * down + relUpX;\n int y = tileOutY * down + relUpY0;\n int signX = x + p.sOfs.x;\n int signY = y + p.sOfs.y;\n int signZ = blockIdx.z + p.blockZofs;\n int signXb = signX >> 2;\n index_t si0 = signXb + p.sShape.x * (signY + (index_t)p.sShape.y * signZ);\n index_t si1 = si0 + p.sShape.x;\n\n v.x *= (scalar_t)((float)up * (float)up * p.gain);\n v.y *= (scalar_t)((float)up * (float)up * p.gain);\n\n if (signWrite)\n {\n if (!enableWriteSkip)\n {\n // Determine and write signs.\n int sx = __float_as_uint(v.x) >> 31 << 0;\n int sy = __float_as_uint(v.y) >> 31 << 8;\n if (sx) v.x *= p.slope;\n if (sy) v.y *= p.slope;\n if (fabsf(v.x) > p.clamp) { sx = 2 << 0; v.x = InternalType::clamp(v.x, p.clamp); }\n if (fabsf(v.y) > p.clamp) { sy = 2 << 8; v.y = InternalType::clamp(v.y, p.clamp); }\n\n if ((uint32_t)signXb < p.swLimit && signY >= minY)\n {\n // Combine signs.\n int s = sx + sy;\n s <<= signXo;\n s |= __shfl_xor_sync(groupMask, s, 1);\n s |= __shfl_xor_sync(groupMask, s, 2);\n\n // Write signs.\n if ((uint32_t)(signY + 0) < sShapeMaxY) { p.s[si0] = (unsigned char)(s >> 0); }\n if ((uint32_t)(signY + 1) < sShapeMaxY) { p.s[si1] = (unsigned char)(s >> 8); }\n }\n }\n else\n {\n // Determine and write signs.\n if ((uint32_t)signXb < p.swLimit && signY >= minY)\n {\n int sx = __float_as_uint(v.x) >> 31 << 0;\n int sy = __float_as_uint(v.y) >> 31 << 8;\n if (sx) v.x *= p.slope;\n if (sy) v.y *= p.slope;\n if (fabsf(v.x) > p.clamp) { sx = 2 << 0; v.x = InternalType::clamp(v.x, p.clamp); }\n if (fabsf(v.y) > p.clamp) { sy = 2 << 8; v.y = InternalType::clamp(v.y, p.clamp); }\n\n // Combine signs.\n int s = sx + sy;\n s <<= signXo;\n s |= __shfl_xor_sync(groupMask, s, 1);\n s |= __shfl_xor_sync(groupMask, s, 2);\n\n // Write signs.\n if ((uint32_t)(signY + 0) < sShapeMaxY) { p.s[si0] = (unsigned char)(s >> 0); }\n if ((uint32_t)(signY + 1) < sShapeMaxY) { p.s[si1] = (unsigned char)(s >> 8); }\n }\n else\n {\n // Just compute the values.\n if (v.x < 0.f) v.x *= p.slope; v.x = InternalType::clamp(v.x, p.clamp);\n if (v.y < 0.f) v.y *= p.slope; v.y = InternalType::clamp(v.y, p.clamp);\n }\n }\n }\n else if (signRead) // Read signs and apply.\n {\n if ((uint32_t)signXb < p.swLimit)\n {\n if ((uint32_t)(signY + 0) < p.sShape.y) { int s = p.s[si0] >> signXo; if (s & 1) v.x *= p.slope; if (s & 2) v.x = 0.f; }\n if ((uint32_t)(signY + 1) < p.sShape.y) { int s = p.s[si1] >> signXo; if (s & 1) v.y *= p.slope; if (s & 2) v.y = 0.f; }\n }\n }\n else // Forward pass with no sign write.\n {\n if (v.x < 0.f) v.x *= p.slope; v.x = InternalType::clamp(v.x, p.clamp);\n if (v.y < 0.f) v.y *= p.slope; v.y = InternalType::clamp(v.y, p.clamp);\n }\n\n if (!downInline)\n {\n // Write into temporary buffer.\n s_tileUpXY[dst] = v.x;\n if (relUpY0 < tileUpH - 1)\n s_tileUpXY[dst + tileUpW] = v.y;\n }\n else\n {\n // Write directly into output buffer.\n if ((uint32_t)x < p.yShape.x)\n {\n int ymax = MIN(p.yShape.y, tileUpH + tileOutY * down);\n index_t ofs = x * get_stride(p.yStride.x) + y * get_stride(p.yStride.y) + mapOfsOut;\n if ((uint32_t)y + 0 < p.yShape.y) *((T*)((char*)p.y + ofs)) = (T)(v.x * (scalar_t)c_fd[0]);\n if ((uint32_t)y + 1 < ymax) *((T*)((char*)p.y + ofs + get_stride(p.yStride.y))) = (T)(v.y * (scalar_t)c_fd[0]);\n }\n }\n }\n }\n }\n else if (filterMode == MODE_FUSD || filterMode == MODE_FUFD)\n {\n // Full upsampling filter.\n\n if (up == 2)\n {\n // 2 x 2-wide.\n __syncthreads();\n int minY = tileOutY ? (tileOutY - tileOutH) * down + tileUpH + p.sOfs.y : 0; // Skip already written signs.\n for (int idx = threadIdx.x * 4; idx < tileUpW * tileUpH; idx += blockDim.x * 4)\n {\n int relUpX0, relUpY0;\n fast_div_mod(relUpX0, relUpY0, idx);\n int relInX0 = CEIL_DIV(relUpX0 - phaseInX, up);\n int relInY0 = CEIL_DIV(relUpY0 - phaseInY, up);\n int src0 = relInX0 + tileInW * relInY0;\n int tap0y = (relInY0 * up + phaseInY - relUpY0);\n\n #define X_LOOP(TAPY, PX) \\\n for (int sx = 0; sx < fuSize / up; sx++) \\\n { \\\n v.x += a * (scalar_t)c_fu[(sx * up + (((PX) - 0) & (up - 1))) + (sy * up + (TAPY)) * MAX_FILTER_SIZE]; \\\n v.z += b * (scalar_t)c_fu[(sx * up + (((PX) - 0) & (up - 1))) + (sy * up + (TAPY)) * MAX_FILTER_SIZE]; if ((PX) == 0) { a = b; b = s_tileIn[src0 + 2 + sx + sy * tileInW]; } \\\n v.y += a * (scalar_t)c_fu[(sx * up + (((PX) - 1) & (up - 1))) + (sy * up + (TAPY)) * MAX_FILTER_SIZE]; \\\n v.w += b * (scalar_t)c_fu[(sx * up + (((PX) - 1) & (up - 1))) + (sy * up + (TAPY)) * MAX_FILTER_SIZE]; if ((PX) == 1) { a = b; b = s_tileIn[src0 + 2 + sx + sy * tileInW]; } \\\n }\n\n vec4_t v = InternalType::zero_vec4();\n if (tap0y == 0 && phaseInX == 0)\n #pragma unroll\n for (int sy = 0; sy < fuSize / up; sy++) { scalar_t a = s_tileIn[src0 + sy * tileInW]; scalar_t b = s_tileIn[src0 + sy * tileInW + 1];\n #pragma unroll\n X_LOOP(0, 0) }\n if (tap0y == 0 && phaseInX == 1)\n #pragma unroll\n for (int sy = 0; sy < fuSize / up; sy++) { scalar_t a = s_tileIn[src0 + sy * tileInW]; scalar_t b = s_tileIn[src0 + sy * tileInW + 1];\n #pragma unroll\n X_LOOP(0, 1) }\n if (tap0y == 1 && phaseInX == 0)\n #pragma unroll\n for (int sy = 0; sy < fuSize / up; sy++) { scalar_t a = s_tileIn[src0 + sy * tileInW]; scalar_t b = s_tileIn[src0 + sy * tileInW + 1];\n #pragma unroll\n X_LOOP(1, 0) }\n if (tap0y == 1 && phaseInX == 1)\n #pragma unroll\n for (int sy = 0; sy < fuSize / up; sy++) { scalar_t a = s_tileIn[src0 + sy * tileInW]; scalar_t b = s_tileIn[src0 + sy * tileInW + 1];\n #pragma unroll\n X_LOOP(1, 1) }\n\n #undef X_LOOP\n\n int x = tileOutX * down + relUpX0;\n int y = tileOutY * down + relUpY0;\n int signX = x + p.sOfs.x;\n int signY = y + p.sOfs.y;\n int signZ = blockIdx.z + p.blockZofs;\n int signXb = signX >> 2;\n index_t si = signXb + p.sShape.x * (signY + (index_t)p.sShape.y * signZ);\n\n v.x *= (scalar_t)((float)up * (float)up * p.gain);\n v.y *= (scalar_t)((float)up * (float)up * p.gain);\n v.z *= (scalar_t)((float)up * (float)up * p.gain);\n v.w *= (scalar_t)((float)up * (float)up * p.gain);\n\n if (signWrite)\n {\n if (!enableWriteSkip)\n {\n // Determine and write signs.\n int sx = __float_as_uint(v.x) >> 31;\n int sy = __float_as_uint(v.y) >> 31;\n int sz = __float_as_uint(v.z) >> 31;\n int sw = __float_as_uint(v.w) >> 31;\n if (sx) v.x *= p.slope; if (fabsf(v.x) > p.clamp) { sx = 2; v.x = InternalType::clamp(v.x, p.clamp); }\n if (sy) v.y *= p.slope; if (fabsf(v.y) > p.clamp) { sy = 2; v.y = InternalType::clamp(v.y, p.clamp); }\n if (sz) v.z *= p.slope; if (fabsf(v.z) > p.clamp) { sz = 2; v.z = InternalType::clamp(v.z, p.clamp); }\n if (sw) v.w *= p.slope; if (fabsf(v.w) > p.clamp) { sw = 2; v.w = InternalType::clamp(v.w, p.clamp); }\n\n if ((uint32_t)signXb < p.swLimit && (uint32_t)signY < p.sShape.y && signY >= minY)\n {\n p.s[si] = sx + (sy << 2) + (sz << 4) + (sw << 6);\n }\n }\n else\n {\n // Determine and write signs.\n if ((uint32_t)signXb < p.swLimit && (uint32_t)signY < p.sShape.y && signY >= minY)\n {\n int sx = __float_as_uint(v.x) >> 31;\n int sy = __float_as_uint(v.y) >> 31;\n int sz = __float_as_uint(v.z) >> 31;\n int sw = __float_as_uint(v.w) >> 31;\n if (sx) v.x *= p.slope; if (fabsf(v.x) > p.clamp) { sx = 2; v.x = InternalType::clamp(v.x, p.clamp); }\n if (sy) v.y *= p.slope; if (fabsf(v.y) > p.clamp) { sy = 2; v.y = InternalType::clamp(v.y, p.clamp); }\n if (sz) v.z *= p.slope; if (fabsf(v.z) > p.clamp) { sz = 2; v.z = InternalType::clamp(v.z, p.clamp); }\n if (sw) v.w *= p.slope; if (fabsf(v.w) > p.clamp) { sw = 2; v.w = InternalType::clamp(v.w, p.clamp); }\n\n p.s[si] = sx + (sy << 2) + (sz << 4) + (sw << 6);\n }\n else\n {\n // Just compute the values.\n if (v.x < 0.f) v.x *= p.slope; v.x = InternalType::clamp(v.x, p.clamp);\n if (v.y < 0.f) v.y *= p.slope; v.y = InternalType::clamp(v.y, p.clamp);\n if (v.z < 0.f) v.z *= p.slope; v.z = InternalType::clamp(v.z, p.clamp);\n if (v.w < 0.f) v.w *= p.slope; v.w = InternalType::clamp(v.w, p.clamp);\n }\n }\n }\n else if (signRead) // Read sign and apply.\n {\n if ((uint32_t)signY < p.sShape.y)\n {\n int s = 0;\n if ((uint32_t)signXb < p.swLimit) s = p.s[si];\n if ((uint32_t)signXb + 1 < p.swLimit) s |= p.s[si + 1] << 8;\n s >>= (signX & 3) << 1;\n if (s & 0x01) v.x *= p.slope; if (s & 0x02) v.x = 0.f;\n if (s & 0x04) v.y *= p.slope; if (s & 0x08) v.y = 0.f;\n if (s & 0x10) v.z *= p.slope; if (s & 0x20) v.z = 0.f;\n if (s & 0x40) v.w *= p.slope; if (s & 0x80) v.w = 0.f;\n }\n }\n else // Forward pass with no sign write.\n {\n if (v.x < 0.f) v.x *= p.slope; v.x = InternalType::clamp(v.x, p.clamp);\n if (v.y < 0.f) v.y *= p.slope; v.y = InternalType::clamp(v.y, p.clamp);\n if (v.z < 0.f) v.z *= p.slope; v.z = InternalType::clamp(v.z, p.clamp);\n if (v.w < 0.f) v.w *= p.slope; v.w = InternalType::clamp(v.w, p.clamp);\n }\n\n s_tileUpXY[idx + 0] = v.x;\n s_tileUpXY[idx + 1] = v.y;\n s_tileUpXY[idx + 2] = v.z;\n s_tileUpXY[idx + 3] = v.w;\n }\n }\n else if (up == 1)\n {\n __syncthreads();\n uint32_t groupMask = 15 << ((threadIdx.x & 31) & ~3);\n int minY = tileOutY ? (tileOutY - tileOutH) * down + tileUpH : 0; // Skip already written signs.\n for (int idx = threadIdx.x; idx < tileUpW * tileUpH; idx += blockDim.x)\n {\n int relUpX0, relUpY0;\n fast_div_mod(relUpX0, relUpY0, idx);\n scalar_t v = s_tileIn[idx] * (scalar_t)c_fu[0]; // 1x1 filter.\n\n int x = tileOutX * down + relUpX0;\n int y = tileOutY * down + relUpY0;\n int signX = x + p.sOfs.x;\n int signY = y + p.sOfs.y;\n int signZ = blockIdx.z + p.blockZofs;\n int signXb = signX >> 2;\n index_t si = signXb + p.sShape.x * (signY + (index_t)p.sShape.y * signZ);\n v *= (scalar_t)((float)up * (float)up * p.gain);\n\n if (signWrite)\n {\n if (!enableWriteSkip)\n {\n // Determine and write sign.\n uint32_t s = 0;\n uint32_t signXbit = (1u << signXo);\n if (v < 0.f)\n {\n s = signXbit;\n v *= p.slope;\n }\n if (fabsf(v) > p.clamp)\n {\n s = signXbit * 2;\n v = InternalType::clamp(v, p.clamp);\n }\n if ((uint32_t)signXb < p.swLimit && (uint32_t)signY < p.sShape.y && signY >= minY)\n {\n s += __shfl_xor_sync(groupMask, s, 1); // Coalesce.\n s += __shfl_xor_sync(groupMask, s, 2); // Coalesce.\n p.s[si] = s; // Write.\n }\n }\n else\n {\n // Determine and write sign.\n if ((uint32_t)signXb < p.swLimit && (uint32_t)signY < p.sShape.y && signY >= minY)\n {\n uint32_t s = 0;\n uint32_t signXbit = (1u << signXo);\n if (v < 0.f)\n {\n s = signXbit;\n v *= p.slope;\n }\n if (fabsf(v) > p.clamp)\n {\n s = signXbit * 2;\n v = InternalType::clamp(v, p.clamp);\n }\n s += __shfl_xor_sync(groupMask, s, 1); // Coalesce.\n s += __shfl_xor_sync(groupMask, s, 2); // Coalesce.\n p.s[si] = s; // Write.\n }\n else\n {\n // Just compute the value.\n if (v < 0.f) v *= p.slope;\n v = InternalType::clamp(v, p.clamp);\n }\n }\n }\n else if (signRead)\n {\n // Read sign and apply if within sign tensor bounds.\n if ((uint32_t)signXb < p.swLimit && (uint32_t)signY < p.sShape.y)\n {\n int s = p.s[si];\n s >>= signXo;\n if (s & 1) v *= p.slope;\n if (s & 2) v = 0.f;\n }\n }\n else // Forward pass with no sign write.\n {\n if (v < 0.f) v *= p.slope;\n v = InternalType::clamp(v, p.clamp);\n }\n\n if (!downInline) // Write into temporary buffer.\n s_tileUpXY[idx] = v;\n else if ((uint32_t)x < p.yShape.x && (uint32_t)y < p.yShape.y) // Write directly into output buffer\n *((T*)((char*)p.y + (x * get_stride(p.yStride.x) + y * get_stride(p.yStride.y) + mapOfsOut))) = (T)(v * (scalar_t)c_fd[0]);\n }\n }\n }\n\n // Downsampling.\n if (filterMode == MODE_SUSD || filterMode == MODE_FUSD)\n {\n // Horizontal downsampling.\n __syncthreads();\n if (down == 4 && tileOutW % 4 == 0)\n {\n // Calculate 4 pixels at a time.\n for (int idx = threadIdx.x * 4; idx < tileOutW * tileUpH; idx += blockDim.x * 4)\n {\n int relOutX0, relUpY;\n fast_div_mod(relOutX0, relUpY, idx);\n int relUpX0 = relOutX0 * down;\n int src0 = relUpY * tileUpW + relUpX0;\n vec4_t v = InternalType::zero_vec4();\n #pragma unroll\n for (int step = 0; step < fdSize; step++)\n {\n v.x += s_tileUpXY[src0 + 0 + step] * (scalar_t)c_fd[step];\n v.y += s_tileUpXY[src0 + 4 + step] * (scalar_t)c_fd[step];\n v.z += s_tileUpXY[src0 + 8 + step] * (scalar_t)c_fd[step];\n v.w += s_tileUpXY[src0 + 12 + step] * (scalar_t)c_fd[step];\n }\n s_tileDownX[idx+0] = v.x;\n s_tileDownX[idx+1] = v.y;\n s_tileDownX[idx+2] = v.z;\n s_tileDownX[idx+3] = v.w;\n }\n }\n else if ((down == 2 || down == 4) && (tileOutW % 2 == 0))\n {\n // Calculate 2 pixels at a time.\n for (int idx = threadIdx.x * 2; idx < tileOutW * tileUpH; idx += blockDim.x * 2)\n {\n int relOutX0, relUpY;\n fast_div_mod(relOutX0, relUpY, idx);\n int relUpX0 = relOutX0 * down;\n int src0 = relUpY * tileUpW + relUpX0;\n vec2_t v = InternalType::zero_vec2();\n #pragma unroll\n for (int step = 0; step < fdSize; step++)\n {\n v.x += s_tileUpXY[src0 + 0 + step] * (scalar_t)c_fd[step];\n v.y += s_tileUpXY[src0 + down + step] * (scalar_t)c_fd[step];\n }\n s_tileDownX[idx+0] = v.x;\n s_tileDownX[idx+1] = v.y;\n }\n }\n else\n {\n // Calculate 1 pixel at a time.\n for (int idx = threadIdx.x; idx < tileOutW * tileUpH; idx += blockDim.x)\n {\n int relOutX0, relUpY;\n fast_div_mod(relOutX0, relUpY, idx);\n int relUpX0 = relOutX0 * down;\n int src = relUpY * tileUpW + relUpX0;\n scalar_t v = 0.f;\n #pragma unroll\n for (int step = 0; step < fdSize; step++)\n v += s_tileUpXY[src + step] * (scalar_t)c_fd[step];\n s_tileDownX[idx] = v;\n }\n }\n\n // Vertical downsampling & store output tile.\n __syncthreads();\n for (int idx = threadIdx.x; idx < tileOutW * tileOutH; idx += blockDim.x)\n {\n int relOutX, relOutY0;\n fast_div_mod(relOutX, relOutY0, idx);\n int relUpY0 = relOutY0 * down;\n int src0 = relUpY0 * tileOutW + relOutX;\n scalar_t v = 0;\n #pragma unroll\n for (int step = 0; step < fdSize; step++)\n v += s_tileDownX[src0 + step * tileOutW] * (scalar_t)c_fd[step];\n\n int outX = tileOutX + relOutX;\n int outY = tileOutY + relOutY0;\n\n if (outX < p.yShape.x & outY < p.yShape.y)\n *((T*)((char*)p.y + (outX * get_stride(p.yStride.x) + outY * get_stride(p.yStride.y) + mapOfsOut))) = (T)v;\n }\n }\n else if (filterMode == MODE_SUFD || filterMode == MODE_FUFD)\n {\n // Full downsampling filter.\n if (down == 2)\n {\n // 2-wide.\n __syncthreads();\n for (int idx = threadIdx.x * 2; idx < tileOutW * tileOutH; idx += blockDim.x * 2)\n {\n int relOutX0, relOutY0;\n fast_div_mod(relOutX0, relOutY0, idx);\n int relUpX0 = relOutX0 * down;\n int relUpY0 = relOutY0 * down;\n int src0 = relUpY0 * tileUpW + relUpX0;\n vec2_t v = InternalType::zero_vec2();\n #pragma unroll\n for (int sy = 0; sy < fdSize; sy++)\n #pragma unroll\n for (int sx = 0; sx < fdSize; sx++)\n {\n v.x += s_tileUpXY[src0 + 0 + sx + sy * tileUpW] * (scalar_t)c_fd[sx + sy * MAX_FILTER_SIZE];\n v.y += s_tileUpXY[src0 + 2 + sx + sy * tileUpW] * (scalar_t)c_fd[sx + sy * MAX_FILTER_SIZE];\n }\n\n int outX = tileOutX + relOutX0;\n int outY = tileOutY + relOutY0;\n if ((uint32_t)outY < p.yShape.y)\n {\n index_t ofs = outX * get_stride(p.yStride.x) + outY * get_stride(p.yStride.y) + mapOfsOut;\n if (outX + 0 < p.yShape.x) *((T*)((char*)p.y + ofs)) = (T)v.x;\n if (outX + 1 < p.yShape.x) *((T*)((char*)p.y + ofs + get_stride(p.yStride.x))) = (T)v.y;\n }\n }\n }\n else if (down == 1 && !downInline)\n {\n // Thread per pixel.\n __syncthreads();\n for (int idx = threadIdx.x; idx < tileOutW * tileOutH; idx += blockDim.x)\n {\n int relOutX0, relOutY0;\n fast_div_mod(relOutX0, relOutY0, idx);\n scalar_t v = s_tileUpXY[idx] * (scalar_t)c_fd[0]; // 1x1 filter.\n\n int outX = tileOutX + relOutX0;\n int outY = tileOutY + relOutY0;\n if ((uint32_t)outX < p.yShape.x && (uint32_t)outY < p.yShape.y)\n *((T*)((char*)p.y + (outX * get_stride(p.yStride.x) + outY * get_stride(p.yStride.y) + mapOfsOut))) = (T)v;\n }\n }\n }\n\n if (!enableXrep)\n break;\n }\n}\n\n//------------------------------------------------------------------------\n// Compute activation function and signs for upsampled data tensor, modifying data tensor in-place. Used for accelerating the generic variant.\n// Sign tensor is known to be contiguous, and p.x and p.s have the same z, w dimensions. 64-bit indexing is always used.\n\ntemplate \nstatic __global__ void filtered_lrelu_act_kernel(filtered_lrelu_act_kernel_params p)\n{\n typedef typename InternalType::scalar_t scalar_t;\n\n // Indexing.\n int32_t x = threadIdx.x + blockIdx.x * blockDim.x;\n int32_t ymax = signWrite ? p.sShape.y : p.xShape.y;\n int32_t qmax = p.xShape.z * p.xShape.w; // Combined minibatch*channel maximum index.\n\n // Loop to accommodate oversized tensors.\n for (int32_t q = blockIdx.z; q < qmax; q += gridDim.z)\n for (int32_t y = blockIdx.y; y < ymax; y += gridDim.y)\n {\n // Extract z and w (channel, minibatch index).\n int32_t w = q / p.xShape.z;\n int32_t z = q - w * p.xShape.z;\n\n // Choose behavior based on sign read/write mode.\n if (signWrite)\n {\n // Process value if in p.x.\n uint32_t s = 0;\n if (x < p.xShape.x && y < p.xShape.y)\n {\n int64_t ix = x * p.xStride.x + y * p.xStride.y + z * p.xStride.z + w * p.xStride.w;\n T* pv = ((T*)p.x) + ix;\n scalar_t v = (scalar_t)(*pv);\n\n // Gain, LReLU, clamp.\n v *= p.gain;\n if (v < 0.f)\n {\n v *= p.slope;\n s = 1; // Sign.\n }\n if (fabsf(v) > p.clamp)\n {\n v = InternalType::clamp(v, p.clamp);\n s = 2; // Clamp.\n }\n\n *pv = (T)v; // Write value.\n }\n\n // Coalesce into threads 0 and 16 of warp.\n uint32_t m = (threadIdx.x & 16) ? 0xffff0000u : 0x0000ffffu;\n s <<= ((threadIdx.x & 15) << 1); // Shift into place.\n s |= __shfl_xor_sync(m, s, 1); // Distribute.\n s |= __shfl_xor_sync(m, s, 2);\n s |= __shfl_xor_sync(m, s, 4);\n s |= __shfl_xor_sync(m, s, 8);\n\n // Write signs if leader and in p.s.\n if (!(threadIdx.x & 15) && x < p.sShape.x) // y is always in.\n {\n uint64_t is = x + p.sShape.x * (y + (int64_t)p.sShape.y * q); // Contiguous.\n ((uint32_t*)p.s)[is >> 4] = s;\n }\n }\n else if (signRead)\n {\n // Process value if in p.x.\n if (x < p.xShape.x) // y is always in.\n {\n int64_t ix = x * p.xStride.x + y * p.xStride.y + z * p.xStride.z + w * p.xStride.w;\n T* pv = ((T*)p.x) + ix;\n scalar_t v = (scalar_t)(*pv);\n v *= p.gain;\n\n // Apply sign buffer offset.\n uint32_t sx = x + p.sOfs.x;\n uint32_t sy = y + p.sOfs.y;\n\n // Read and apply signs if we land inside valid region of sign buffer.\n if (sx < p.sShape.x && sy < p.sShape.y)\n {\n uint64_t is = (sx >> 2) + (p.sShape.x >> 2) * (sy + (uint64_t)p.sShape.y * q); // Contiguous.\n unsigned char s = p.s[is];\n s >>= (sx & 3) << 1; // Shift into place.\n if (s & 1) // Sign?\n v *= p.slope;\n if (s & 2) // Clamp?\n v = 0.f;\n }\n\n *pv = (T)v; // Write value.\n }\n }\n else\n {\n // Forward pass with no sign write. Process value if in p.x.\n if (x < p.xShape.x) // y is always in.\n {\n int64_t ix = x * p.xStride.x + y * p.xStride.y + z * p.xStride.z + w * p.xStride.w;\n T* pv = ((T*)p.x) + ix;\n scalar_t v = (scalar_t)(*pv);\n v *= p.gain;\n if (v < 0.f)\n v *= p.slope;\n if (fabsf(v) > p.clamp)\n v = InternalType::clamp(v, p.clamp);\n *pv = (T)v; // Write value.\n }\n }\n }\n}\n\ntemplate void* choose_filtered_lrelu_act_kernel(void)\n{\n return (void*)filtered_lrelu_act_kernel;\n}\n\n//------------------------------------------------------------------------\n// CUDA kernel selection.\n\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB)\n{\n filtered_lrelu_kernel_spec s = { 0 };\n\n // Return the first matching kernel.\n#define CASE(SH, U, FU, D, FD, MODE, TW, TH, W, XR, WS) \\\n if (sharedKB >= SH) \\\n if ((p.fuShape.y == 0 && (MODE == MODE_SUSD || MODE == MODE_SUFD)) || (p.fuShape.y > 0 && (MODE == MODE_FUSD || MODE == MODE_FUFD))) \\\n if ((p.fdShape.y == 0 && (MODE == MODE_SUSD || MODE == MODE_FUSD)) || (p.fdShape.y > 0 && (MODE == MODE_SUFD || MODE == MODE_FUFD))) \\\n if (p.up == U && p.fuShape.x <= FU && p.fuShape.y <= FU && p.down == D && p.fdShape.x <= FD && p.fdShape.y <= FD) \\\n { \\\n static_assert((D*TW % 4) == 0, \"down * tileWidth must be divisible by 4\"); \\\n static_assert(FU % U == 0, \"upscaling filter size must be multiple of upscaling factor\"); \\\n static_assert(FD % D == 0, \"downscaling filter size must be multiple of downscaling factor\"); \\\n s.setup = (void*)setup_filters_kernel; \\\n s.exec = (void*)filtered_lrelu_kernel; \\\n s.tileOut = make_int2(TW, TH); \\\n s.numWarps = W; \\\n s.xrep = XR; \\\n s.dynamicSharedKB = (SH == 48) ? 0 : SH; \\\n return s; \\\n }\n\n // Launch parameters for various kernel specializations.\n // Small filters must be listed before large filters, otherwise the kernel for larger filter will always match first.\n // Kernels that use more shared memory must be listed before those that use less, for the same reason.\n\n CASE(/*sharedKB*/48, /*up,fu*/1,1, /*down,fd*/1,1, /*mode*/MODE_FUFD, /*tw,th,warps,xrep,wskip*/64, 178, 32, 0, 0) // 1t-upf1-downf1\n CASE(/*sharedKB*/48, /*up,fu*/2,8, /*down,fd*/1,1, /*mode*/MODE_SUFD, /*tw,th,warps,xrep,wskip*/152, 95, 16, 0, 0) // 4t-ups2-downf1\n CASE(/*sharedKB*/48, /*up,fu*/1,1, /*down,fd*/2,8, /*mode*/MODE_FUSD, /*tw,th,warps,xrep,wskip*/56, 22, 16, 0, 0) // 4t-upf1-downs2\n CASE(/*sharedKB*/48, /*up,fu*/2,8, /*down,fd*/2,8, /*mode*/MODE_SUSD, /*tw,th,warps,xrep,wskip*/56, 29, 16, 11, 0) // 4t-ups2-downs2\n CASE(/*sharedKB*/48, /*up,fu*/2,8, /*down,fd*/2,8, /*mode*/MODE_FUSD, /*tw,th,warps,xrep,wskip*/60, 28, 16, 0, 0) // 4t-upf2-downs2\n CASE(/*sharedKB*/48, /*up,fu*/2,8, /*down,fd*/2,8, /*mode*/MODE_SUFD, /*tw,th,warps,xrep,wskip*/56, 28, 16, 0, 0) // 4t-ups2-downf2\n CASE(/*sharedKB*/48, /*up,fu*/4,16, /*down,fd*/2,8, /*mode*/MODE_SUSD, /*tw,th,warps,xrep,wskip*/56, 31, 16, 11, 0) // 4t-ups4-downs2\n CASE(/*sharedKB*/48, /*up,fu*/4,16, /*down,fd*/2,8, /*mode*/MODE_SUFD, /*tw,th,warps,xrep,wskip*/56, 36, 16, 0, 0) // 4t-ups4-downf2\n CASE(/*sharedKB*/48, /*up,fu*/2,8, /*down,fd*/4,16, /*mode*/MODE_SUSD, /*tw,th,warps,xrep,wskip*/16, 22, 16, 12, 0) // 4t-ups2-downs4\n CASE(/*sharedKB*/48, /*up,fu*/2,8, /*down,fd*/4,16, /*mode*/MODE_FUSD, /*tw,th,warps,xrep,wskip*/29, 15, 16, 0, 0) // 4t-upf2-downs4\n CASE(/*sharedKB*/48, /*up,fu*/2,12, /*down,fd*/1,1, /*mode*/MODE_SUFD, /*tw,th,warps,xrep,wskip*/96, 150, 28, 0, 0) // 6t-ups2-downf1\n CASE(/*sharedKB*/48, /*up,fu*/1,1, /*down,fd*/2,12, /*mode*/MODE_FUSD, /*tw,th,warps,xrep,wskip*/32, 35, 24, 0, 0) // 6t-upf1-downs2\n CASE(/*sharedKB*/48, /*up,fu*/2,12, /*down,fd*/2,12, /*mode*/MODE_SUSD, /*tw,th,warps,xrep,wskip*/32, 46, 16, 10, 0) // 6t-ups2-downs2\n CASE(/*sharedKB*/48, /*up,fu*/2,12, /*down,fd*/2,12, /*mode*/MODE_FUSD, /*tw,th,warps,xrep,wskip*/58, 28, 24, 8, 0) // 6t-upf2-downs2\n CASE(/*sharedKB*/48, /*up,fu*/2,12, /*down,fd*/2,12, /*mode*/MODE_SUFD, /*tw,th,warps,xrep,wskip*/52, 28, 16, 0, 0) // 6t-ups2-downf2\n CASE(/*sharedKB*/48, /*up,fu*/4,24, /*down,fd*/2,12, /*mode*/MODE_SUSD, /*tw,th,warps,xrep,wskip*/32, 51, 16, 5, 0) // 6t-ups4-downs2\n CASE(/*sharedKB*/48, /*up,fu*/4,24, /*down,fd*/2,12, /*mode*/MODE_SUFD, /*tw,th,warps,xrep,wskip*/32, 56, 16, 6, 0) // 6t-ups4-downf2\n CASE(/*sharedKB*/48, /*up,fu*/2,12, /*down,fd*/4,24, /*mode*/MODE_SUSD, /*tw,th,warps,xrep,wskip*/16, 18, 16, 12, 0) // 6t-ups2-downs4\n CASE(/*sharedKB*/96, /*up,fu*/2,12, /*down,fd*/4,24, /*mode*/MODE_FUSD, /*tw,th,warps,xrep,wskip*/27, 31, 32, 6, 0) // 6t-upf2-downs4 96kB\n CASE(/*sharedKB*/48, /*up,fu*/2,12, /*down,fd*/4,24, /*mode*/MODE_FUSD, /*tw,th,warps,xrep,wskip*/27, 13, 24, 0, 0) // 6t-upf2-downs4\n CASE(/*sharedKB*/48, /*up,fu*/2,16, /*down,fd*/1,1, /*mode*/MODE_SUFD, /*tw,th,warps,xrep,wskip*/148, 89, 24, 0, 0) // 8t-ups2-downf1\n CASE(/*sharedKB*/48, /*up,fu*/1,1, /*down,fd*/2,16, /*mode*/MODE_FUSD, /*tw,th,warps,xrep,wskip*/32, 31, 16, 5, 0) // 8t-upf1-downs2\n CASE(/*sharedKB*/48, /*up,fu*/2,16, /*down,fd*/2,16, /*mode*/MODE_SUSD, /*tw,th,warps,xrep,wskip*/32, 41, 16, 9, 0) // 8t-ups2-downs2\n CASE(/*sharedKB*/48, /*up,fu*/2,16, /*down,fd*/2,16, /*mode*/MODE_FUSD, /*tw,th,warps,xrep,wskip*/56, 26, 24, 0, 0) // 8t-upf2-downs2\n CASE(/*sharedKB*/48, /*up,fu*/2,16, /*down,fd*/2,16, /*mode*/MODE_SUFD, /*tw,th,warps,xrep,wskip*/32, 40, 16, 0, 0) // 8t-ups2-downf2\n CASE(/*sharedKB*/48, /*up,fu*/4,32, /*down,fd*/2,16, /*mode*/MODE_SUSD, /*tw,th,warps,xrep,wskip*/32, 46, 24, 5, 0) // 8t-ups4-downs2\n CASE(/*sharedKB*/48, /*up,fu*/4,32, /*down,fd*/2,16, /*mode*/MODE_SUFD, /*tw,th,warps,xrep,wskip*/32, 50, 16, 0, 0) // 8t-ups4-downf2\n CASE(/*sharedKB*/96, /*up,fu*/2,16, /*down,fd*/4,32, /*mode*/MODE_SUSD, /*tw,th,warps,xrep,wskip*/24, 24, 32, 12, 1) // 8t-ups2-downs4 96kB\n CASE(/*sharedKB*/48, /*up,fu*/2,16, /*down,fd*/4,32, /*mode*/MODE_SUSD, /*tw,th,warps,xrep,wskip*/16, 13, 16, 10, 1) // 8t-ups2-downs4\n CASE(/*sharedKB*/96, /*up,fu*/2,16, /*down,fd*/4,32, /*mode*/MODE_FUSD, /*tw,th,warps,xrep,wskip*/25, 28, 28, 4, 0) // 8t-upf2-downs4 96kB\n CASE(/*sharedKB*/48, /*up,fu*/2,16, /*down,fd*/4,32, /*mode*/MODE_FUSD, /*tw,th,warps,xrep,wskip*/25, 10, 24, 0, 0) // 8t-upf2-downs4\n\n #undef CASE\n return s; // No kernel found.\n}\n\n//------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/filtered_lrelu.h",
"content": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors retain all intellectual property\n// and proprietary rights in and to this software, related documentation\n// and any modifications thereto. Any use, reproduction, disclosure or\n// distribution of this software and related documentation without an express\n// license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n#include \n\n//------------------------------------------------------------------------\n// CUDA kernel parameters.\n\nstruct filtered_lrelu_kernel_params\n{\n // These parameters decide which kernel to use.\n int up; // upsampling ratio (1, 2, 4)\n int down; // downsampling ratio (1, 2, 4)\n int2 fuShape; // [size, 1] | [size, size]\n int2 fdShape; // [size, 1] | [size, size]\n\n int _dummy; // Alignment.\n\n // Rest of the parameters.\n const void* x; // Input tensor.\n void* y; // Output tensor.\n const void* b; // Bias tensor.\n unsigned char* s; // Sign tensor in/out. NULL if unused.\n const float* fu; // Upsampling filter.\n const float* fd; // Downsampling filter.\n\n int2 pad0; // Left/top padding.\n float gain; // Additional gain factor.\n float slope; // Leaky ReLU slope on negative side.\n float clamp; // Clamp after nonlinearity.\n int flip; // Filter kernel flip for gradient computation.\n\n int tilesXdim; // Original number of horizontal output tiles.\n int tilesXrep; // Number of horizontal tiles per CTA.\n int blockZofs; // Block z offset to support large minibatch, channel dimensions.\n\n int4 xShape; // [width, height, channel, batch]\n int4 yShape; // [width, height, channel, batch]\n int2 sShape; // [width, height] - width is in bytes. Contiguous. Zeros if unused.\n int2 sOfs; // [ofs_x, ofs_y] - offset between upsampled data and sign tensor.\n int swLimit; // Active width of sign tensor in bytes.\n\n longlong4 xStride; // Strides of all tensors except signs, same component order as shapes.\n longlong4 yStride; //\n int64_t bStride; //\n longlong3 fuStride; //\n longlong3 fdStride; //\n};\n\nstruct filtered_lrelu_act_kernel_params\n{\n void* x; // Input/output, modified in-place.\n unsigned char* s; // Sign tensor in/out. NULL if unused.\n\n float gain; // Additional gain factor.\n float slope; // Leaky ReLU slope on negative side.\n float clamp; // Clamp after nonlinearity.\n\n int4 xShape; // [width, height, channel, batch]\n longlong4 xStride; // Input/output tensor strides, same order as in shape.\n int2 sShape; // [width, height] - width is in elements. Contiguous. Zeros if unused.\n int2 sOfs; // [ofs_x, ofs_y] - offset between upsampled data and sign tensor.\n};\n\n//------------------------------------------------------------------------\n// CUDA kernel specialization.\n\nstruct filtered_lrelu_kernel_spec\n{\n void* setup; // Function for filter kernel setup.\n void* exec; // Function for main operation.\n int2 tileOut; // Width/height of launch tile.\n int numWarps; // Number of warps per thread block, determines launch block size.\n int xrep; // For processing multiple horizontal tiles per thread block.\n int dynamicSharedKB; // How much dynamic shared memory the exec kernel wants.\n};\n\n//------------------------------------------------------------------------\n// CUDA kernel selection.\n\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\ntemplate void* choose_filtered_lrelu_act_kernel(void);\ntemplate cudaError_t copy_filters(cudaStream_t stream);\n\n//------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/filtered_lrelu.py",
"content": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\nimport os\nimport numpy as np\nimport torch\nimport warnings\n\nfrom .. import custom_ops\nfrom .. import misc\nfrom . import upfirdn2d\nfrom . import bias_act\n\n#----------------------------------------------------------------------------\n\n_plugin = None\n\ndef _init():\n global _plugin\n if _plugin is None:\n _plugin = custom_ops.get_plugin(\n module_name='filtered_lrelu_plugin',\n sources=['filtered_lrelu.cpp', 'filtered_lrelu_wr.cu', 'filtered_lrelu_rd.cu', 'filtered_lrelu_ns.cu'],\n headers=['filtered_lrelu.h', 'filtered_lrelu.cu'],\n source_dir=os.path.dirname(__file__),\n extra_cuda_cflags=['--use_fast_math', '--allow-unsupported-compiler'],\n )\n return True\n\ndef _get_filter_size(f):\n if f is None:\n return 1, 1\n assert isinstance(f, torch.Tensor)\n assert 1 <= f.ndim <= 2\n return f.shape[-1], f.shape[0] # width, height\n\ndef _parse_padding(padding):\n if isinstance(padding, int):\n padding = [padding, padding]\n assert isinstance(padding, (list, tuple))\n assert all(isinstance(x, (int, np.integer)) for x in padding)\n padding = [int(x) for x in padding]\n if len(padding) == 2:\n px, py = padding\n padding = [px, px, py, py]\n px0, px1, py0, py1 = padding\n return px0, px1, py0, py1\n\n#----------------------------------------------------------------------------\n\ndef filtered_lrelu(x, fu=None, fd=None, b=None, up=1, down=1, padding=0, gain=np.sqrt(2), slope=0.2, clamp=None, flip_filter=False, impl='cuda'):\n r\"\"\"Filtered leaky ReLU for a batch of 2D images.\n\n Performs the following sequence of operations for each channel:\n\n 1. Add channel-specific bias if provided (`b`).\n\n 2. Upsample the image by inserting N-1 zeros after each pixel (`up`).\n\n 3. Pad the image with the specified number of zeros on each side (`padding`).\n Negative padding corresponds to cropping the image.\n\n 4. Convolve the image with the specified upsampling FIR filter (`fu`), shrinking it\n so that the footprint of all output pixels lies within the input image.\n\n 5. Multiply each value by the provided gain factor (`gain`).\n\n 6. Apply leaky ReLU activation function to each value.\n\n 7. Clamp each value between -clamp and +clamp, if `clamp` parameter is provided.\n\n 8. Convolve the image with the specified downsampling FIR filter (`fd`), shrinking\n it so that the footprint of all output pixels lies within the input image.\n\n 9. Downsample the image by keeping every Nth pixel (`down`).\n\n The fused op is considerably more efficient than performing the same calculation\n using standard PyTorch ops. It supports gradients of arbitrary order.\n\n Args:\n x: Float32/float16/float64 input tensor of the shape\n `[batch_size, num_channels, in_height, in_width]`.\n fu: Float32 upsampling FIR filter of the shape\n `[filter_height, filter_width]` (non-separable),\n `[filter_taps]` (separable), or\n `None` (identity).\n fd: Float32 downsampling FIR filter of the shape\n `[filter_height, filter_width]` (non-separable),\n `[filter_taps]` (separable), or\n `None` (identity).\n b: Bias vector, or `None` to disable. Must be a 1D tensor of the same type\n as `x`. The length of vector must must match the channel dimension of `x`.\n up: Integer upsampling factor (default: 1).\n down: Integer downsampling factor. (default: 1).\n padding: Padding with respect to the upsampled image. Can be a single number\n or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`\n (default: 0).\n gain: Overall scaling factor for signal magnitude (default: sqrt(2)).\n slope: Slope on the negative side of leaky ReLU (default: 0.2).\n clamp: Maximum magnitude for leaky ReLU output (default: None).\n flip_filter: False = convolution, True = correlation (default: False).\n impl: Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).\n\n Returns:\n Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.\n \"\"\"\n assert isinstance(x, torch.Tensor)\n assert impl in ['ref', 'cuda']\n if impl == 'cuda' and x.device.type == 'cuda' and _init():\n return _filtered_lrelu_cuda(up=up, down=down, padding=padding, gain=gain, slope=slope, clamp=clamp, flip_filter=flip_filter).apply(x, fu, fd, b, None, 0, 0)\n return _filtered_lrelu_ref(x, fu=fu, fd=fd, b=b, up=up, down=down, padding=padding, gain=gain, slope=slope, clamp=clamp, flip_filter=flip_filter)\n\n#----------------------------------------------------------------------------\n\n@misc.profiled_function\ndef _filtered_lrelu_ref(x, fu=None, fd=None, b=None, up=1, down=1, padding=0, gain=np.sqrt(2), slope=0.2, clamp=None, flip_filter=False):\n \"\"\"Slow and memory-inefficient reference implementation of `filtered_lrelu()` using\n existing `upfirdn2n()` and `bias_act()` ops.\n \"\"\"\n assert isinstance(x, torch.Tensor) and x.ndim == 4\n fu_w, fu_h = _get_filter_size(fu)\n fd_w, fd_h = _get_filter_size(fd)\n if b is not None:\n assert isinstance(b, torch.Tensor) and b.dtype == x.dtype\n misc.assert_shape(b, [x.shape[1]])\n assert isinstance(up, int) and up >= 1\n assert isinstance(down, int) and down >= 1\n px0, px1, py0, py1 = _parse_padding(padding)\n assert gain == float(gain) and gain > 0\n assert slope == float(slope) and slope >= 0\n assert clamp is None or (clamp == float(clamp) and clamp >= 0)\n\n # Calculate output size.\n batch_size, channels, in_h, in_w = x.shape\n in_dtype = x.dtype\n out_w = (in_w * up + (px0 + px1) - (fu_w - 1) - (fd_w - 1) + (down - 1)) // down\n out_h = (in_h * up + (py0 + py1) - (fu_h - 1) - (fd_h - 1) + (down - 1)) // down\n\n # Compute using existing ops.\n x = bias_act.bias_act(x=x, b=b) # Apply bias.\n x = upfirdn2d.upfirdn2d(x=x, f=fu, up=up, padding=[px0, px1, py0, py1], gain=up**2, flip_filter=flip_filter) # Upsample.\n x = bias_act.bias_act(x=x, act='lrelu', alpha=slope, gain=gain, clamp=clamp) # Bias, leaky ReLU, clamp.\n x = upfirdn2d.upfirdn2d(x=x, f=fd, down=down, flip_filter=flip_filter) # Downsample.\n\n # Check output shape & dtype.\n misc.assert_shape(x, [batch_size, channels, out_h, out_w])\n assert x.dtype == in_dtype\n return x\n\n#----------------------------------------------------------------------------\n\n_filtered_lrelu_cuda_cache = dict()\n\ndef _filtered_lrelu_cuda(up=1, down=1, padding=0, gain=np.sqrt(2), slope=0.2, clamp=None, flip_filter=False):\n \"\"\"Fast CUDA implementation of `filtered_lrelu()` using custom ops.\n \"\"\"\n assert isinstance(up, int) and up >= 1\n assert isinstance(down, int) and down >= 1\n px0, px1, py0, py1 = _parse_padding(padding)\n assert gain == float(gain) and gain > 0\n gain = float(gain)\n assert slope == float(slope) and slope >= 0\n slope = float(slope)\n assert clamp is None or (clamp == float(clamp) and clamp >= 0)\n clamp = float(clamp if clamp is not None else 'inf')\n\n # Lookup from cache.\n key = (up, down, px0, px1, py0, py1, gain, slope, clamp, flip_filter)\n if key in _filtered_lrelu_cuda_cache:\n return _filtered_lrelu_cuda_cache[key]\n\n # Forward op.\n class FilteredLReluCuda(torch.autograd.Function):\n @staticmethod\n def forward(ctx, x, fu, fd, b, si, sx, sy): # pylint: disable=arguments-differ\n assert isinstance(x, torch.Tensor) and x.ndim == 4\n\n # Replace empty up/downsample kernels with full 1x1 kernels (faster than separable).\n if fu is None:\n fu = torch.ones([1, 1], dtype=torch.float32, device=x.device)\n if fd is None:\n fd = torch.ones([1, 1], dtype=torch.float32, device=x.device)\n assert 1 <= fu.ndim <= 2\n assert 1 <= fd.ndim <= 2\n\n # Replace separable 1x1 kernels with full 1x1 kernels when scale factor is 1.\n if up == 1 and fu.ndim == 1 and fu.shape[0] == 1:\n fu = fu.square()[None]\n if down == 1 and fd.ndim == 1 and fd.shape[0] == 1:\n fd = fd.square()[None]\n\n # Missing sign input tensor.\n if si is None:\n si = torch.empty([0])\n\n # Missing bias tensor.\n if b is None:\n b = torch.zeros([x.shape[1]], dtype=x.dtype, device=x.device)\n\n # Construct internal sign tensor only if gradients are needed.\n write_signs = (si.numel() == 0) and (x.requires_grad or b.requires_grad)\n\n # Warn if input storage strides are not in decreasing order due to e.g. channels-last layout.\n strides = [x.stride(i) for i in range(x.ndim) if x.size(i) > 1]\n if any(a < b for a, b in zip(strides[:-1], strides[1:])):\n warnings.warn(\"low-performance memory layout detected in filtered_lrelu input\", RuntimeWarning)\n\n # Call C++/Cuda plugin if datatype is supported.\n if x.dtype in [torch.float16, torch.float32]:\n if torch.cuda.current_stream(x.device) != torch.cuda.default_stream(x.device):\n warnings.warn(\"filtered_lrelu called with non-default cuda stream but concurrent execution is not supported\", RuntimeWarning)\n y, so, return_code = _plugin.filtered_lrelu(x, fu, fd, b, si, up, down, px0, px1, py0, py1, sx, sy, gain, slope, clamp, flip_filter, write_signs)\n else:\n return_code = -1\n\n # No Cuda kernel found? Fall back to generic implementation. Still more memory efficient than the reference implementation because\n # only the bit-packed sign tensor is retained for gradient computation.\n if return_code < 0:\n warnings.warn(\"filtered_lrelu called with parameters that have no optimized CUDA kernel, using generic fallback\", RuntimeWarning)\n\n y = x.add(b.unsqueeze(-1).unsqueeze(-1)) # Add bias.\n y = upfirdn2d.upfirdn2d(x=y, f=fu, up=up, padding=[px0, px1, py0, py1], gain=up**2, flip_filter=flip_filter) # Upsample.\n so = _plugin.filtered_lrelu_act_(y, si, sx, sy, gain, slope, clamp, write_signs) # Activation function and sign handling. Modifies y in-place.\n y = upfirdn2d.upfirdn2d(x=y, f=fd, down=down, flip_filter=flip_filter) # Downsample.\n\n # Prepare for gradient computation.\n ctx.save_for_backward(fu, fd, (si if si.numel() else so))\n ctx.x_shape = x.shape\n ctx.y_shape = y.shape\n ctx.s_ofs = sx, sy\n return y\n\n @staticmethod\n def backward(ctx, dy): # pylint: disable=arguments-differ\n fu, fd, si = ctx.saved_tensors\n _, _, xh, xw = ctx.x_shape\n _, _, yh, yw = ctx.y_shape\n sx, sy = ctx.s_ofs\n dx = None # 0\n dfu = None; assert not ctx.needs_input_grad[1]\n dfd = None; assert not ctx.needs_input_grad[2]\n db = None # 3\n dsi = None; assert not ctx.needs_input_grad[4]\n dsx = None; assert not ctx.needs_input_grad[5]\n dsy = None; assert not ctx.needs_input_grad[6]\n\n if ctx.needs_input_grad[0] or ctx.needs_input_grad[3]:\n pp = [\n (fu.shape[-1] - 1) + (fd.shape[-1] - 1) - px0,\n xw * up - yw * down + px0 - (up - 1),\n (fu.shape[0] - 1) + (fd.shape[0] - 1) - py0,\n xh * up - yh * down + py0 - (up - 1),\n ]\n gg = gain * (up ** 2) / (down ** 2)\n ff = (not flip_filter)\n sx = sx - (fu.shape[-1] - 1) + px0\n sy = sy - (fu.shape[0] - 1) + py0\n dx = _filtered_lrelu_cuda(up=down, down=up, padding=pp, gain=gg, slope=slope, clamp=None, flip_filter=ff).apply(dy, fd, fu, None, si, sx, sy)\n\n if ctx.needs_input_grad[3]:\n db = dx.sum([0, 2, 3])\n\n return dx, dfu, dfd, db, dsi, dsx, dsy\n\n # Add to cache.\n _filtered_lrelu_cuda_cache[key] = FilteredLReluCuda\n return FilteredLReluCuda\n\n#----------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/filtered_lrelu_ns.cu",
"content": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors retain all intellectual property\n// and proprietary rights in and to this software, related documentation\n// and any modifications thereto. Any use, reproduction, disclosure or\n// distribution of this software and related documentation without an express\n// license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n#include \"filtered_lrelu.cu\"\n\n// Template/kernel specializations for no signs mode (no gradients required).\n\n// Full op, 32-bit indexing.\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\n\n// Full op, 64-bit indexing.\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\n\n// Activation/signs only for generic variant. 64-bit indexing.\ntemplate void* choose_filtered_lrelu_act_kernel(void);\ntemplate void* choose_filtered_lrelu_act_kernel(void);\ntemplate void* choose_filtered_lrelu_act_kernel(void);\n\n// Copy filters to constant memory.\ntemplate cudaError_t copy_filters(cudaStream_t stream);\n"
},
{
"path": "ADD/th_utils/ops/filtered_lrelu_rd.cu",
"content": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors retain all intellectual property\n// and proprietary rights in and to this software, related documentation\n// and any modifications thereto. Any use, reproduction, disclosure or\n// distribution of this software and related documentation without an express\n// license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n#include \"filtered_lrelu.cu\"\n\n// Template/kernel specializations for sign read mode.\n\n// Full op, 32-bit indexing.\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\n\n// Full op, 64-bit indexing.\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\n\n// Activation/signs only for generic variant. 64-bit indexing.\ntemplate void* choose_filtered_lrelu_act_kernel(void);\ntemplate void* choose_filtered_lrelu_act_kernel(void);\ntemplate void* choose_filtered_lrelu_act_kernel(void);\n\n// Copy filters to constant memory.\ntemplate cudaError_t copy_filters(cudaStream_t stream);\n"
},
{
"path": "ADD/th_utils/ops/filtered_lrelu_wr.cu",
"content": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors retain all intellectual property\n// and proprietary rights in and to this software, related documentation\n// and any modifications thereto. Any use, reproduction, disclosure or\n// distribution of this software and related documentation without an express\n// license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n#include \"filtered_lrelu.cu\"\n\n// Template/kernel specializations for sign write mode.\n\n// Full op, 32-bit indexing.\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\n\n// Full op, 64-bit indexing.\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\ntemplate filtered_lrelu_kernel_spec choose_filtered_lrelu_kernel(const filtered_lrelu_kernel_params& p, int sharedKB);\n\n// Activation/signs only for generic variant. 64-bit indexing.\ntemplate void* choose_filtered_lrelu_act_kernel(void);\ntemplate void* choose_filtered_lrelu_act_kernel(void);\ntemplate void* choose_filtered_lrelu_act_kernel(void);\n\n// Copy filters to constant memory.\ntemplate cudaError_t copy_filters(cudaStream_t stream);\n"
},
{
"path": "ADD/th_utils/ops/fma.py",
"content": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n\"\"\"Fused multiply-add, with slightly faster gradients than `torch.addcmul()`.\"\"\"\n\nimport torch\n\n#----------------------------------------------------------------------------\n\ndef fma(a, b, c): # => a * b + c\n return _FusedMultiplyAdd.apply(a, b, c)\n\n#----------------------------------------------------------------------------\n\nclass _FusedMultiplyAdd(torch.autograd.Function): # a * b + c\n @staticmethod\n def forward(ctx, a, b, c): # pylint: disable=arguments-differ\n out = torch.addcmul(c, a, b)\n ctx.save_for_backward(a, b)\n ctx.c_shape = c.shape\n return out\n\n @staticmethod\n def backward(ctx, dout): # pylint: disable=arguments-differ\n a, b = ctx.saved_tensors\n c_shape = ctx.c_shape\n da = None\n db = None\n dc = None\n\n if ctx.needs_input_grad[0]:\n da = _unbroadcast(dout * b, a.shape)\n\n if ctx.needs_input_grad[1]:\n db = _unbroadcast(dout * a, b.shape)\n\n if ctx.needs_input_grad[2]:\n dc = _unbroadcast(dout, c_shape)\n\n return da, db, dc\n\n#----------------------------------------------------------------------------\n\ndef _unbroadcast(x, shape):\n extra_dims = x.ndim - len(shape)\n assert extra_dims >= 0\n dim = [i for i in range(x.ndim) if x.shape[i] > 1 and (i < extra_dims or shape[i - extra_dims] == 1)]\n if len(dim):\n x = x.sum(dim=dim, keepdim=True)\n if extra_dims:\n x = x.reshape(-1, *x.shape[extra_dims+1:])\n assert x.shape == shape\n return x\n\n#----------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/grid_sample_gradfix.py",
"content": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all intellectual property\n# and proprietary rights in and to this software, related documentation\n# and any modifications thereto. Any use, reproduction, disclosure or\n# distribution of this software and related documentation without an express\n# license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n\"\"\"Custom replacement for `torch.nn.functional.grid_sample` that\nsupports arbitrarily high order gradients between the input and output.\nOnly works on 2D images and assumes\n`mode='bilinear'`, `padding_mode='zeros'`, `align_corners=False`.\"\"\"\n\nimport torch\nfrom pkg_resources import parse_version\n\n# pylint: disable=redefined-builtin\n# pylint: disable=arguments-differ\n# pylint: disable=protected-access\n\n#----------------------------------------------------------------------------\n\nenabled = False # Enable the custom op by setting this to true.\n_use_pytorch_1_11_api = parse_version(torch.__version__) >= parse_version('1.11.0a') # Allow prerelease builds of 1.11\n\n#----------------------------------------------------------------------------\n\ndef grid_sample(input, grid):\n if _should_use_custom_op():\n return _GridSample2dForward.apply(input, grid)\n return torch.nn.functional.grid_sample(input=input, grid=grid, mode='bilinear', padding_mode='zeros', align_corners=False)\n\n#----------------------------------------------------------------------------\n\ndef _should_use_custom_op():\n return enabled\n\n#----------------------------------------------------------------------------\n\nclass _GridSample2dForward(torch.autograd.Function):\n @staticmethod\n def forward(ctx, input, grid):\n assert input.ndim == 4\n assert grid.ndim == 4\n output = torch.nn.functional.grid_sample(input=input, grid=grid, mode='bilinear', padding_mode='zeros', align_corners=False)\n ctx.save_for_backward(input, grid)\n return output\n\n @staticmethod\n def backward(ctx, grad_output):\n input, grid = ctx.saved_tensors\n grad_input, grad_grid = _GridSample2dBackward.apply(grad_output, input, grid)\n return grad_input, grad_grid\n\n#----------------------------------------------------------------------------\n\nclass _GridSample2dBackward(torch.autograd.Function):\n @staticmethod\n def forward(ctx, grad_output, input, grid):\n op = torch._C._jit_get_operation('aten::grid_sampler_2d_backward')\n if _use_pytorch_1_11_api:\n output_mask = (ctx.needs_input_grad[1], ctx.needs_input_grad[2])\n grad_input, grad_grid = op(grad_output, input, grid, 0, 0, False, output_mask)\n else:\n grad_input, grad_grid = op(grad_output, input, grid, 0, 0, False)\n ctx.save_for_backward(grid)\n return grad_input, grad_grid\n\n @staticmethod\n def backward(ctx, grad2_grad_input, grad2_grad_grid):\n _ = grad2_grad_grid # unused\n grid, = ctx.saved_tensors\n grad2_grad_output = None\n grad2_input = None\n grad2_grid = None\n\n if ctx.needs_input_grad[0]:\n grad2_grad_output = _GridSample2dForward.apply(grad2_grad_input, grid)\n\n assert not ctx.needs_input_grad[2]\n return grad2_grad_output, grad2_input, grad2_grid\n\n#----------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/upfirdn2d.cpp",
"content": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors retain all intellectual property\n// and proprietary rights in and to this software, related documentation\n// and any modifications thereto. Any use, reproduction, disclosure or\n// distribution of this software and related documentation without an express\n// license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n#include \n#include \n#include \n#include \"upfirdn2d.h\"\n\n//------------------------------------------------------------------------\n\nstatic torch::Tensor upfirdn2d(torch::Tensor x, torch::Tensor f, int upx, int upy, int downx, int downy, int padx0, int padx1, int pady0, int pady1, bool flip, float gain)\n{\n // Validate arguments.\n TORCH_CHECK(x.is_cuda(), \"x must reside on CUDA device\");\n TORCH_CHECK(f.device() == x.device(), \"f must reside on the same device as x\");\n TORCH_CHECK(f.dtype() == torch::kFloat, \"f must be float32\");\n TORCH_CHECK(x.numel() <= INT_MAX, \"x is too large\");\n TORCH_CHECK(f.numel() <= INT_MAX, \"f is too large\");\n TORCH_CHECK(x.numel() > 0, \"x has zero size\");\n TORCH_CHECK(f.numel() > 0, \"f has zero size\");\n TORCH_CHECK(x.dim() == 4, \"x must be rank 4\");\n TORCH_CHECK(f.dim() == 2, \"f must be rank 2\");\n TORCH_CHECK((x.size(0)-1)*x.stride(0) + (x.size(1)-1)*x.stride(1) + (x.size(2)-1)*x.stride(2) + (x.size(3)-1)*x.stride(3) <= INT_MAX, \"x memory footprint is too large\");\n TORCH_CHECK(f.size(0) >= 1 && f.size(1) >= 1, \"f must be at least 1x1\");\n TORCH_CHECK(upx >= 1 && upy >= 1, \"upsampling factor must be at least 1\");\n TORCH_CHECK(downx >= 1 && downy >= 1, \"downsampling factor must be at least 1\");\n\n // Create output tensor.\n const at::cuda::OptionalCUDAGuard device_guard(device_of(x));\n int outW = ((int)x.size(3) * upx + padx0 + padx1 - (int)f.size(1) + downx) / downx;\n int outH = ((int)x.size(2) * upy + pady0 + pady1 - (int)f.size(0) + downy) / downy;\n TORCH_CHECK(outW >= 1 && outH >= 1, \"output must be at least 1x1\");\n torch::Tensor y = torch::empty({x.size(0), x.size(1), outH, outW}, x.options(), x.suggest_memory_format());\n TORCH_CHECK(y.numel() <= INT_MAX, \"output is too large\");\n TORCH_CHECK((y.size(0)-1)*y.stride(0) + (y.size(1)-1)*y.stride(1) + (y.size(2)-1)*y.stride(2) + (y.size(3)-1)*y.stride(3) <= INT_MAX, \"output memory footprint is too large\");\n\n // Initialize CUDA kernel parameters.\n upfirdn2d_kernel_params p;\n p.x = x.data_ptr();\n p.f = f.data_ptr();\n p.y = y.data_ptr();\n p.up = make_int2(upx, upy);\n p.down = make_int2(downx, downy);\n p.pad0 = make_int2(padx0, pady0);\n p.flip = (flip) ? 1 : 0;\n p.gain = gain;\n p.inSize = make_int4((int)x.size(3), (int)x.size(2), (int)x.size(1), (int)x.size(0));\n p.inStride = make_int4((int)x.stride(3), (int)x.stride(2), (int)x.stride(1), (int)x.stride(0));\n p.filterSize = make_int2((int)f.size(1), (int)f.size(0));\n p.filterStride = make_int2((int)f.stride(1), (int)f.stride(0));\n p.outSize = make_int4((int)y.size(3), (int)y.size(2), (int)y.size(1), (int)y.size(0));\n p.outStride = make_int4((int)y.stride(3), (int)y.stride(2), (int)y.stride(1), (int)y.stride(0));\n p.sizeMajor = (p.inStride.z == 1) ? p.inSize.w : p.inSize.w * p.inSize.z;\n p.sizeMinor = (p.inStride.z == 1) ? p.inSize.z : 1;\n\n // Choose CUDA kernel.\n upfirdn2d_kernel_spec spec;\n AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), \"upfirdn2d_cuda\", [&]\n {\n spec = choose_upfirdn2d_kernel(p);\n });\n\n // Set looping options.\n p.loopMajor = (p.sizeMajor - 1) / 16384 + 1;\n p.loopMinor = spec.loopMinor;\n p.loopX = spec.loopX;\n p.launchMinor = (p.sizeMinor - 1) / p.loopMinor + 1;\n p.launchMajor = (p.sizeMajor - 1) / p.loopMajor + 1;\n\n // Compute grid size.\n dim3 blockSize, gridSize;\n if (spec.tileOutW < 0) // large\n {\n blockSize = dim3(4, 32, 1);\n gridSize = dim3(\n ((p.outSize.y - 1) / blockSize.x + 1) * p.launchMinor,\n (p.outSize.x - 1) / (blockSize.y * p.loopX) + 1,\n p.launchMajor);\n }\n else // small\n {\n blockSize = dim3(256, 1, 1);\n gridSize = dim3(\n ((p.outSize.y - 1) / spec.tileOutH + 1) * p.launchMinor,\n (p.outSize.x - 1) / (spec.tileOutW * p.loopX) + 1,\n p.launchMajor);\n }\n\n // Launch CUDA kernel.\n void* args[] = {&p};\n AT_CUDA_CHECK(cudaLaunchKernel(spec.kernel, gridSize, blockSize, args, 0, at::cuda::getCurrentCUDAStream()));\n return y;\n}\n\n//------------------------------------------------------------------------\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m)\n{\n m.def(\"upfirdn2d\", &upfirdn2d);\n}\n\n//------------------------------------------------------------------------\n"
},
{
"path": "ADD/th_utils/ops/upfirdn2d.cu",
"content": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors retain all intellectual property\n// and proprietary rights in and to this software, related documentation\n// and any modifications thereto. Any use, reproduction, disclosure or\n// distribution of this software and related documentation without an express\n// license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n#include \n#include \"upfirdn2d.h\"\n\n//------------------------------------------------------------------------\n// Helpers.\n\ntemplate struct InternalType;\ntemplate <> struct InternalType { typedef double scalar_t; };\ntemplate <> struct InternalType { typedef float scalar_t; };\ntemplate <> struct InternalType { typedef float scalar_t; };\n\nstatic __device__ __forceinline__ int floor_div(int a, int b)\n{\n int t = 1 - a / b;\n return (a + t * b) / b - t;\n}\n\n//------------------------------------------------------------------------\n// Generic CUDA implementation for large filters.\n\ntemplate static __global__ void upfirdn2d_kernel_large(upfirdn2d_kernel_params p)\n{\n typedef typename InternalType::scalar_t scalar_t;\n\n // Calculate thread index.\n int minorBase = blockIdx.x * blockDim.x + threadIdx.x;\n int outY = minorBase / p.launchMinor;\n minorBase -= outY * p.launchMinor;\n int outXBase = blockIdx.y * p.loopX * blockDim.y + threadIdx.y;\n int majorBase = blockIdx.z * p.loopMajor;\n if (outXBase >= p.outSize.x | outY >= p.outSize.y | majorBase >= p.sizeMajor)\n return;\n\n // Setup Y receptive field.\n int midY = outY * p.down.y + p.up.y - 1 - p.pad0.y;\n int inY = min(max(floor_div(midY, p.up.y), 0), p.inSize.y);\n int h = min(max(floor_div(midY + p.filterSize.y, p.up.y), 0), p.inSize.y) - inY;\n int filterY = midY + p.filterSize.y - (inY + 1) * p.up.y;\n if (p.flip)\n filterY = p.filterSize.y - 1 - filterY;\n\n // Loop over major, minor, and X.\n for (int majorIdx = 0, major = majorBase; majorIdx < p.loopMajor & major < p.sizeMajor; majorIdx++, major++)\n for (int minorIdx = 0, minor = minorBase; minorIdx < p.loopMinor & minor < p.sizeMinor; minorIdx++, minor += p.launchMinor)\n {\n int nc = major * p.sizeMinor + minor;\n int n = nc / p.inSize.z;\n int c = nc - n * p.inSize.z;\n for (int loopX = 0, outX = outXBase; loopX < p.loopX & outX < p.outSize.x; loopX++, outX += blockDim.y)\n {\n // Setup X receptive field.\n int midX = outX * p.down.x + p.up.x - 1 - p.pad0.x;\n int inX = min(max(floor_div(midX, p.up.x), 0), p.inSize.x);\n int w = min(max(floor_div(midX + p.filterSize.x, p.up.x), 0), p.inSize.x) - inX;\n int filterX = midX + p.filterSize.x - (inX + 1) * p.up.x;\n if (p.flip)\n filterX = p.filterSize.x - 1 - filterX;\n\n // Initialize pointers.\n const T* xp = &((const T*)p.x)[inX * p.inStride.x + inY * p.inStride.y + c * p.inStride.z + n * p.inStride.w];\n const float* fp = &p.f[filterX * p.filterStride.x + filterY * p.filterStride.y];\n int filterStepX = ((p.flip) ? p.up.x : -p.up.x) * p.filterStride.x;\n int filterStepY = ((p.flip) ? p.up.y : -p.up.y) * p.filterStride.y;\n\n // Inner loop.\n scalar_t v = 0;\n for (int y = 0; y < h; y++)\n {\n for (int x = 0; x < w; x++)\n {\n v += (scalar_t)(*xp) * (scalar_t)(*fp);\n xp += p.inStride.x;\n fp += filterStepX;\n }\n xp += p.inStride.y - w * p.inStride.x;\n fp += filterStepY - w * filterStepX;\n }\n\n // Store result.\n v *= p.gain;\n ((T*)p.y)[outX * p.outStride.x + outY * p.outStride.y + c * p.outStride.z + n * p.outStride.w] = (T)v;\n }\n }\n}\n\n//------------------------------------------------------------------------\n// Specialized CUDA implementation for small filters.\n\ntemplate \nstatic __global__ void upfirdn2d_kernel_small(upfirdn2d_kernel_params p)\n{\n typedef typename InternalType::scalar_t scalar_t;\n const int tileInW = ((tileOutW - 1) * downx + filterW - 1) / upx + 1;\n const int tileInH = ((tileOutH - 1) * downy + filterH - 1) / upy + 1;\n __shared__ volatile scalar_t sf[filterH][filterW];\n __shared__ volatile scalar_t sx[tileInH][tileInW][loopMinor];\n\n // Calculate tile index.\n int minorBase = blockIdx.x;\n int tileOutY = minorBase / p.launchMinor;\n minorBase -= tileOutY * p.launchMinor;\n minorBase *= loopMinor;\n tileOutY *= tileOutH;\n int tileOutXBase = blockIdx.y * p.loopX * tileOutW;\n int majorBase = blockIdx.z * p.loopMajor;\n if (tileOutXBase >= p.outSize.x | tileOutY >= p.outSize.y | majorBase >= p.sizeMajor)\n return;\n\n // Load filter (flipped).\n for (int tapIdx = threadIdx.x; tapIdx < filterH * filterW; tapIdx += blockDim.x)\n {\n int fy = tapIdx / filterW;\n int fx = tapIdx - fy * filterW;\n scalar_t v = 0;\n if (fx < p.filterSize.x & fy < p.filterSize.y)\n {\n int ffx = (p.flip) ? fx : p.filterSize.x - 1 - fx;\n int ffy = (p.flip) ? fy : p.filterSize.y - 1 - fy;\n v = (scalar_t)p.f[ffx * p.filterStride.x + ffy * p.filterStride.y];\n }\n sf[fy][fx] = v;\n }\n\n // Loop over major and X.\n for (int majorIdx = 0, major = majorBase; majorIdx < p.loopMajor & major < p.sizeMajor; majorIdx++, major++)\n {\n int baseNC = major * p.sizeMinor + minorBase;\n int n = baseNC / p.inSize.z;\n int baseC = baseNC - n * p.inSize.z;\n for (int loopX = 0, tileOutX = tileOutXBase; loopX < p.loopX & tileOutX < p.outSize.x; loopX++, tileOutX += tileOutW)\n {\n // Load input pixels.\n int tileMidX = tileOutX * downx + upx - 1 - p.pad0.x;\n int tileMidY = tileOutY * downy + upy - 1 - p.pad0.y;\n int tileInX = floor_div(tileMidX, upx);\n int tileInY = floor_div(tileMidY, upy);\n __syncthreads();\n for (int inIdx = threadIdx.x; inIdx < tileInH * tileInW * loopMinor; inIdx += blockDim.x)\n {\n int relC = inIdx;\n int relInX = relC / loopMinor;\n int relInY = relInX / tileInW;\n relC -= relInX * loopMinor;\n relInX -= relInY * tileInW;\n int c = baseC + relC;\n int inX = tileInX + relInX;\n int inY = tileInY + relInY;\n scalar_t v = 0;\n if (inX >= 0 & inY >= 0 & inX < p.inSize.x & inY < p.inSize.y & c < p.inSize.z)\n v = (scalar_t)((const T*)p.x)[inX * p.inStride.x + inY * p.inStride.y + c * p.inStride.z + n * p.inStride.w];\n sx[relInY][relInX][relC] = v;\n }\n\n // Loop over output pixels.\n __syncthreads();\n for (int outIdx = threadIdx.x; outIdx < tileOutH * tileOutW * loopMinor; outIdx += blockDim.x)\n {\n int relC = outIdx;\n int relOutX = relC / loopMinor;\n int relOutY = relOutX / tileOutW;\n relC -= relOutX * loopMinor;\n relOutX -= relOutY * tileOutW;\n int c = baseC + relC;\n int outX = tileOutX + relOutX;\n int outY = tileOutY + relOutY;\n\n // Setup receptive field.\n int midX = tileMidX + relOutX * downx;\n int midY = tileMidY + relOutY * downy;\n int inX = floor_div(midX, upx);\n int inY = floor_div(midY, upy);\n int relInX = inX - tileInX;\n int relInY = inY - tileInY;\n int filterX = (inX + 1) * upx - midX - 1; // flipped\n int filterY = (inY + 1) * upy - midY - 1; // flipped\n\n // Inner loop.\n if (outX < p.outSize.x & outY < p.outSize.y & c < p.outSize.z)\n {\n scalar_t v = 0;\n #pragma unroll\n for (int y = 0; y < filterH / upy; y++)\n #pragma unroll\n for (int x = 0; x < filterW / upx; x++)\n v += sx[relInY + y][relInX + x][relC] * sf[filterY + y * upy][filterX + x * upx];\n v *= p.gain;\n ((T*)p.y)[outX * p.outStride.x + outY * p.outStride.y + c * p.outStride.z + n * p.outStride.w] = (T)v;\n }\n }\n }\n }\n}\n\n//------------------------------------------------------------------------\n// CUDA kernel selection.\n\ntemplate upfirdn2d_kernel_spec choose_upfirdn2d_kernel(const upfirdn2d_kernel_params& p)\n{\n int s = p.inStride.z, fx = p.filterSize.x, fy = p.filterSize.y;\n upfirdn2d_kernel_spec spec = {(void*)upfirdn2d_kernel_large, -1,-1,1, 4}; // contiguous\n if (s == 1) spec = {(void*)upfirdn2d_kernel_large, -1,-1,4, 1}; // channels_last\n\n // No up/downsampling.\n if (p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1)\n {\n // contiguous\n if (s != 1 && fx <= 24 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 64,32,1, 1};\n if (s != 1 && fx <= 16 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 64,32,1, 1};\n if (s != 1 && fx <= 7 && fy <= 7 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1};\n if (s != 1 && fx <= 6 && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1};\n if (s != 1 && fx <= 5 && fy <= 5 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1};\n if (s != 1 && fx <= 4 && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1};\n if (s != 1 && fx <= 3 && fy <= 3 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1};\n if (s != 1 && fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1};\n if (s != 1 && fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1};\n if (s != 1 && fx <= 8 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1};\n if (s != 1 && fx <= 1 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1};\n if (s != 1 && fx <= 1 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1};\n if (s != 1 && fx <= 1 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1};\n // channels_last\n if (s == 1 && fx <= 24 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1};\n if (s == 1 && fx <= 16 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1};\n if (s == 1 && fx <= 7 && fy <= 7 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1};\n if (s == 1 && fx <= 6 && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1};\n if (s == 1 && fx <= 5 && fy <= 5 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1};\n if (s == 1 && fx <= 4 && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1};\n if (s == 1 && fx <= 3 && fy <= 3 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1};\n if (s == 1 && fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1};\n if (s == 1 && fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1};\n if (s == 1 && fx <= 8 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1};\n if (s == 1 && fx <= 1 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1};\n if (s == 1 && fx <= 1 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1};\n if (s == 1 && fx <= 1 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1};\n }\n\n // 2x upsampling.\n if (p.up.x == 2 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1)\n {\n // contiguous\n if (s != 1 && fx <= 24 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 64,32,1, 1};\n if (s != 1 && fx <= 16 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 64,32,1, 1};\n if (s != 1 && fx <= 8 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1};\n if (s != 1 && fx <= 6 && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1};\n if (s != 1 && fx <= 4 && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1};\n if (s != 1 && fx <= 2 && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small, 64,16,1, 1};\n // channels_last\n if (s == 1 && fx <= 24 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1};\n if (s == 1 && fx <= 16 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1};\n if (s == 1 && fx <= 8 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1};\n if (s == 1 && fx <= 6 && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1};\n if (s == 1 && fx <= 4 && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1};\n if (s == 1 && fx <= 2 && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small, 16,16,8, 1};\n }\n if (p.up.x == 2 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1)\n {\n // contiguous\n if (s != 1 && fx <= 24 && fy <= 1) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1};\n if (s != 1 && fx <= 16 && fy <= 1) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1};\n if (s != 1 && fx <= 8 && fy <= 1) spec = {(void*)upfirdn2d_kernel_small, 128,8,1, 1};\n // channels_last\n if (s == 1 && fx <= 24 && fy <= 1) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1};\n if (s == 1 && fx <= 16 && fy <= 1) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1};\n if (s == 1 && fx <= 8 && fy <= 1) spec = {(void*)upfirdn2d_kernel_small, 128,1,16, 1};\n }\n if (p.up.x == 1 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1)\n {\n // contiguous\n if (s != 1 && fx <= 1 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1};\n if (s != 1 && fx <= 1 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1};\n if (s != 1 && fx <= 1 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 32,32,1, 1};\n // channels_last\n if (s == 1 && fx <= 1 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1};\n if (s == 1 && fx <= 1 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1};\n if (s == 1 && fx <= 1 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 1,128,16, 1};\n }\n\n // 2x downsampling.\n if (p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 2)\n {\n // contiguous\n if (s != 1 && fx <= 24 && fy <= 24) spec = {(void*)upfirdn2d_kernel_small, 32,16,1, 1};\n if (s != 1 && fx <= 16 && fy <= 16) spec = {(void*)upfirdn2d_kernel_small, 32,16,1, 1};\n if (s != 1 && fx <= 8 && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small, 32,8,1, 1};\n if (s != 1 && fx <= 6 && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small, 32,8,1, 1};\n if (s != 1 && fx <= 4 && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small