Support for hipifying CuPy and add extra_extensions argument in hipify (ROCm#52)

* added optional extensions for hipification

* add cupy special map with cupy related replacements

* update cupy mappings

* add support for custom mappings

* update based on review comments and update readme

* Update README.md

---------

Co-authored-by: Jithun Nair <[email protected]>

Author: lcskrishna

README.md

@@ -11,6 +11,7 @@ It can also "hipify" the header include statements in your source code to ensure
   - [Through python](#through-python)
 - [Utilities](#utilities)
   - [CMake utility function](#cmake-utility-function)
+- [Custom hipify mapping](#custom-hipify-mapping)
 - [Intended users](#intended-users)
 
 <!-- tocstop -->
@@ -100,6 +101,23 @@ get_hipified_list("${TP_CUDA_SRCS}" TP_CUDA_SRCS)
 Here the `TP_CUDA_SRCS` in the input list containing cuda source files and doing a inplace update with  output list `TP_CUDA_SRCS`
 For the file suffix unique string, list variable name itself is passed as a string.
 
+# Custom hipify mapping
+
+Users can define their own custom mapping by adding a custom_hipify_mapping.json from project_directory from where the hipify() function is being called.
+To use a JSON file from a different directory, users can pass in the JSON file path via ```custom_map``` argument in the hipify method.
+The custom hipify mappings will be applied *before* any other default hipify mappings. 
+The below is the sample JSON file:
+
+```
+{
+    "custom_map" : {
+        "src mapping 1" : "dst mapping 1",
+        "src mapping 2" : "dst mapping 2",
+        ...
+    }
+}
+```
+
 # Intended users
 
 This module can be used to

hipify_torch/hipify_python.py

@@ -29,6 +29,7 @@
 import shutil
 import sys
 import os
+import json
 
 from . import constants
 from .cuda_to_hip_mappings import CUDA_TO_HIP_MAPPINGS
@@ -46,6 +47,8 @@
 to their actual types."""
 PYTORCH_TEMPLATE_MAP = {"Dtype": "scalar_t", "T": "scalar_t"}
 
+# Custom mapping json file (default), if the file is not available hipify call doesn't process it.
+custom_mapping_file = "custom_hipify_mappings.json"
 
 class InputError(Exception):
     # Exception raised for errors in the input.
@@ -682,7 +685,8 @@ def pattern(self):
 # When a file contains "sparse" in the filename, a mapping marked with API_SPARSE is preferred over other choices.
 # Similarly, "linalg" files require rocBLAS -> hipSOLVER so they also need special handling.
 PYTORCH_SPECIAL_MAP = {}
-
+CUSTOM_TRIE = Trie()
+CUSTOM_SPECIAL_MAP = {}
 
 for mapping in CUDA_TO_HIP_MAPPINGS:
     assert isinstance(mapping, Mapping)
@@ -708,6 +712,17 @@ def pattern(self):
 RE_THC_GENERIC_FILE = re.compile(r'#define THC_GENERIC_FILE "([^"]+)"')
 RE_CU_SUFFIX = re.compile(r'\.cu\b')  # be careful not to pick up .cuh
 
+#  This function takes in the custom json file and looks for custom_mappings available.
+def update_custom_mappings():
+    if os.path.exists(custom_mapping_file):
+        with open(custom_mapping_file, 'r') as f:
+            json_data = json.load(f)
+        custom_map_dict = json_data['custom_map']
+        for src in custom_map_dict.keys():
+            CUSTOM_TRIE.add(src)
+            dst = custom_map_dict[src]
+            CUSTOM_SPECIAL_MAP[src] = dst
+
 """
 Returns a dict with the following keys:
     "hipified_path" : absolute path of hipified source file
@@ -754,6 +769,12 @@ def pt_special_repl(m):
         # checks SPARSE map first, and if a miss occurs, falls back to pytorch mappings
         return PYTORCH_SPECIAL_MAP.get(m.group(0), pt_repl(m))
 
+    # replace all custom mappings.
+    if len(CUSTOM_SPECIAL_MAP) > 0:
+        RE_CUSTOM_PREPROCESSOR = re.compile(CUSTOM_TRIE.pattern())
+        def custom_repl(m):
+            return CUSTOM_SPECIAL_MAP[m.group(0)]
+        output_source = RE_CUSTOM_PREPROCESSOR.sub(custom_repl, output_source)
 
     if is_pytorch_extension:
         output_source = RE_PYTORCH_PREPROCESSOR.sub(pt_repl, output_source)
@@ -962,9 +983,11 @@ def hipify(
     show_detailed: bool = False,
     extensions: Iterable = (".cu", ".cuh", ".c", ".cc", ".cpp", ".h", ".in", ".hpp"),
     header_extensions: Iterable = (".cuh", ".h", ".hpp"),
+    extra_extensions: Iterable = (),
     output_directory: str = "",
     header_include_dirs: Iterable = (),
     includes: Iterable = ('*',),
+    custom_map_list: str = "",
     extra_files: Iterable = (),
     out_of_place_only: bool = False,
     ignores: Iterable = (),
@@ -982,6 +1005,12 @@ def hipify(
         print("The project folder specified does not exist.")
         sys.exit(1)
 
+    # custom mapping json file that is provided by user.
+    if custom_map_list:
+        global custom_mapping_file
+        custom_mapping_file = os.path.abspath(custom_map_list)
+    update_custom_mappings()
+
     # If no output directory, provide a default one.
     if not output_directory:
         project_directory.rstrip("/")
@@ -991,6 +1020,9 @@ def hipify(
         includes = [include.replace(project_directory, output_directory) for include in includes]
         ignores = [ignore.replace(project_directory, output_directory) for ignore in ignores]
 
+    # update extensions with optional extensions.
+    extensions = extensions + extra_extensions
+
     # Copy from project directory to output directory if not done already.
     if not os.path.exists(output_directory):
         shutil.copytree(project_directory, output_directory)