Removes TorchBind Custom Classes (#183)

* Refactoring PyTorch tensor product.

* Tensor product refactored successfully.

* Eliminated opaque bindings, batch tests working.

* Ready to run convolution tests.

* Ready to fix fakes.

* DeviceBuffer is gone.

* Fixed the fakes.

* Fixing compile tests.

* Fixed JAX error.

* Fixed compile issues.

* Ruff.

* Eliminated raw pointer functions and generic classes.

* Ruff.

* Fixed benchmark.py.

* Updated JAX extension version due to backend change.

Author: vbharadwaj-bk

openequivariance/openequivariance/_torch/TensorProduct.py

@@ -51,7 +51,7 @@
 
         extra_cflags = ["-O3"]
         generic_sources = ["generic_module.cpp"]
-        torch_sources = ["libtorch_tp_jit.cpp"]
+        torch_sources = ["libtorch_tp_jit.cpp", "json11/json11.cpp"]
 
         include_dirs, extra_link_args = (["util"], ["-Wl,--no-as-needed"])
 
@@ -149,22 +149,13 @@ def torch_ext_so_path():
 
 if BUILT_EXTENSION:
     from generic_module import (
-        JITTPImpl,
-        JITConvImpl,
         GroupMM_F32,
         GroupMM_F64,
         DeviceProp,
-        DeviceBuffer,
         GPUTimer,
     )
 else:
 
-    def JITTPImpl(*args, **kwargs):
-        _raise_import_error_helper("JITTPImpl")
-
-    def JITConvImpl(*args, **kwargs):
-        _raise_import_error_helper("JITConvImpl")
-
     def GroupMM_F32(*args, **kwargs):
         _raise_import_error_helper("GroupMM_F32")
 
@@ -174,8 +165,5 @@ def GroupMM_F64(*args, **kwargs):
     def DeviceProp(*args, **kwargs):
         _raise_import_error_helper("DeviceProp")
 
-    def DeviceBuffer(*args, **kwargs):
-        _raise_import_error_helper("DeviceBuffer")
-
     def GPUTimer(*args, **kwargs):
         _raise_import_error_helper("GPUTimer")

openequivariance/openequivariance/_torch/TensorProductConv.py

@@ -1,4 +1,5 @@
 import torch
+import numpy as np
 from types import MappingProxyType
 from openequivariance.core.utils import DTypeEnum
 
@@ -66,3 +67,11 @@ def reorder_torch(schedule, weights_in, direction, has_batch_dim):
         DTypeEnum.UINT8: torch.uint8,
     }
 )
+
+
+def string_to_tensor(text: str) -> torch.Tensor:
+    bytes_data = text.encode("utf-8")
+    np_bytes = np.frombuffer(bytes_data, dtype=np.uint8)
+    result = torch.tensor(np_bytes, device="cpu")
+    result.requires_grad = False
+    return result

openequivariance/openequivariance/_torch/extlib/__init__.py

@@ -8,6 +8,7 @@
 import openequivariance as oeq
 from openequivariance.benchmark.logging_utils import getLogger
 from openequivariance.core.ConvolutionBase import CoordGraph
+from openequivariance.benchmark.benchmark_utils import NpEncoder
 
 logger = getLogger()
 
@@ -145,7 +146,7 @@ def run(
                     f"{output_folder}/{self.exp_count}_{impl.name()}_{graph.name}.json"
                 )
                 with open(fname, "w") as f:
-                    json.dump(result, f, indent=2)
+                    json.dump(result, f, indent=2, cls=NpEncoder)
                 self.exp_count += 1
 
                 logger.info(f"Finished {tc_name}, graph {graph.name}")

openequivariance/openequivariance/_torch/utils.py

@@ -21,6 +21,7 @@
     benchmark_forward,
     benchmark_backward,
     benchmark_double_backward,
+    NpEncoder,
 )
 
 logger = getLogger()
@@ -235,10 +236,12 @@ def run(
 
             fname = pathlib.Path(f"{output_folder}/{test_ID}_{impl.name()}.json")
 
-            pretty_result = json.dumps(obj=result, indent=2).replace("\\n", "\n")
+            pretty_result = json.dumps(obj=result, indent=2, cls=NpEncoder).replace(
+                "\\n", "\n"
+            )
             logger.debug(pretty_result)
             with open(fname, "w") as f:
-                json.dump(result, f, indent=2)
+                json.dump(result, f, indent=2, cls=NpEncoder)
 
             self.results.append(result)
             logger.info(f"Finished Test ID: {test_ID}")

openequivariance/openequivariance/benchmark/ConvBenchmarkSuite.py

@@ -1,3 +1,4 @@
+import json
 import numpy as np
 
 from openequivariance.benchmark.random_buffer_utils import (
@@ -290,3 +291,14 @@ def benchmark_double_backward(
     )
 
     return result
+
+
+class NpEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.integer):
+            return int(obj)
+        if isinstance(obj, np.floating):
+            return float(obj)
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return super(NpEncoder, self).default(obj)

openequivariance/openequivariance/benchmark/TestBenchmarkSuite.py

@@ -1,14 +1,18 @@
 import numpy as np
+import json
 
 from openequivariance.core.ConvolutionBase import ConvolutionBase
 from openequivariance.core.ComputationSchedule import (
     ComputationSchedule,
     SMEMCapacityException,
 )
 
-from openequivariance.core.utils import dtype_to_enum
 from openequivariance.templates.jinja_utils import get_jinja_environment
-from openequivariance.core.utils import filter_and_analyze_problem
+from openequivariance.core.utils import (
+    filter_and_analyze_problem,
+    dtype_to_enum,
+    hash_str_64,
+)
 
 
 class LoopUnrollConv(ConvolutionBase):
@@ -203,5 +207,15 @@ def generate_double_backward_schedule(warps_per_block):
         )
         self.jit_kernel = postprocess_kernel(self.jit_kernel)
 
-        # with open("scratch.txt", "w") as f:
-        #    f.write(self.jit_kernel)
+        self.kernel_string = json.dumps(
+            {
+                "kernel": self.jit_kernel,
+                "forward_config": vars(self.forward_schedule.launch_config),
+                "backward_config": vars(self.backward_schedule.launch_config),
+                "double_backward_config": vars(
+                    self.double_backward_schedule.launch_config
+                ),
+                "kernel_prop": self.kernel_prop,
+            }
+        )
+        self.hash = hash_str_64(self.kernel_string)

openequivariance/openequivariance/benchmark/benchmark_utils.py

@@ -1,15 +1,19 @@
 import numpy as np
+import json
 
 from openequivariance.templates.jinja_utils import get_jinja_environment
 from openequivariance.core.ComputationSchedule import ComputationSchedule
 from openequivariance.core.TensorProductBase import TensorProductBase
-from openequivariance.core.utils import dtype_to_enum
+from openequivariance.benchmark.logging_utils import getLogger
+from openequivariance.core.utils import dtype_to_enum, hash_str_64
 
 from openequivariance.core.utils import (
     filter_and_analyze_problem,
     count_cg_non_zero,
 )
 
+logger = getLogger()
+
 
 class LoopUnrollTP(TensorProductBase):
     def __init__(self, config, dp, postprocess_kernel, torch_op):
@@ -91,7 +95,7 @@ def generate_double_backward_schedule(warps_per_block):
             )
         )
 
-        self.kernelProp = {
+        self.kernel_prop = {
             "L1_dim": self.L1.dim,
             "L2_dim": self.L2.dim,
             "L3_dim": self.L3.dim,
@@ -106,6 +110,20 @@ def generate_double_backward_schedule(warps_per_block):
             "idx_dtype": 0,
         }
 
+        self.kernel_string = json.dumps(
+            {
+                "kernel": self.jit_kernel,
+                "forward_config": vars(self.forward_schedule.launch_config),
+                "backward_config": vars(self.backward_schedule.launch_config),
+                "double_backward_config": vars(
+                    self.double_backward_schedule.launch_config
+                ),
+                "kernel_prop": self.kernel_prop,
+            }
+        )
+        self.hash = hash_str_64(self.kernel_string)
+        logger.info(f"Kernel File Size: {len(self.jit_kernel) // 1024} KB")
+
     def calculate_flops_forward(self, batch_size: int) -> dict:
         if self.is_uvw:
             return super().calculate_flops_forward(batch_size)

openequivariance/openequivariance/core/LoopUnrollConv.py

@@ -7,9 +7,9 @@
 
 import json
 import tempfile
+import hashlib
 
 from enum import IntEnum
-import hashlib
 
 
 class DTypeEnum(IntEnum):