[feat] support converter for torch.log2

Author: bowang007

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -2670,3 +2670,20 @@ def aten_ops_flip(
         args[0],
         args[1],
     )
+
+
+@dynamo_tensorrt_converter(torch.log2)
+def log2(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.unary.log2(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+    )

py/torch_tensorrt/dynamo/conversion/impl/unary/ops.py

@@ -9,6 +9,7 @@
     cast_trt_tensor,
     get_trt_tensor,
 )
+from torch_tensorrt.dynamo.conversion.impl.elementwise import div
 from torch_tensorrt.dynamo.conversion.impl.unary.base import convert_unary
 from torch_tensorrt.fx.types import TRTTensor
 
@@ -58,6 +59,20 @@ def log(
     )
 
 
+def log2(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input_val: TRTTensor,
+) -> TRTTensor:
+    log_layer_output = log(ctx, target, source_ir, f"{name}_log", input_val)
+
+    ln2 = 0.693147180559945309
+
+    return div(ctx, target, source_ir, f"{name}_div", log_layer_output, ln2)
+
+
 def sqrt(
     ctx: ConversionContext,
     target: Target,

tests/py/dynamo/conversion/test_log2.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+from .harness import DispatchTestCase
+
+
+class TestLogConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ((10,), torch.float),
+            ((1, 20), torch.float),
+            ((2, 3, 4), torch.float),
+            ((2, 3, 4, 5), torch.float),
+        ]
+    )
+    def test_log_float(self, input_shape, dtype):
+        class log(nn.Module):
+            def forward(self, input):
+                return torch.log2(input)
+
+        inputs = [torch.randn(input_shape, dtype=dtype)]
+        self.run_test(
+            log(),
+            inputs,
+        )
+
+    @parameterized.expand(
+        [
+            ((10,), torch.int, 0, 5),
+            ((1, 20), torch.int32, -10, 10),
+            ((2, 3, 4), torch.int, -5, 5),
+        ]
+    )
+    def test_log_int(self, input_shape, dtype, low, high):
+        class log(nn.Module):
+            def forward(self, input):
+                return torch.log2(input)
+
+        inputs = [torch.randint(low, high, input_shape, dtype=dtype)]
+        self.run_test(
+            log(),
+            inputs,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()