Add ctc_loss operator (flagos-ai#2723)

Author: HyosungSink

benchmark/core_shapes.yaml

@@ -10,6 +10,22 @@ single_dim_shapes: &single_dim_shapes
 vdot:
   <<: *single_dim_shapes
 
+ctc_loss:
+  shapes:
+    - [64, 4, 32, 16]
+    - [256, 16, 64, 48]
+    - [512, 32, 64, 48]
+    - [1024, 32, 128, 96]
+  shape_desc: "T, N, C, S"
+
+ctc_loss_backward:
+  shapes:
+    - [64, 4, 32, 16]
+    - [256, 16, 64, 48]
+    - [512, 32, 64, 48]
+    - [1024, 32, 128, 96]
+  shape_desc: "T, N, C, S"
+
 randperm:
   <<: *single_dim_shapes
 

benchmark/test_ctc_loss.py

@@ -0,0 +1,131 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+import flag_gems
+
+from . import base, consts
+
+CTC_DTYPES = [torch.float32, torch.float16]
+
+
+def _ctc_loss_reference(
+    log_probs,
+    targets,
+    input_lengths,
+    target_lengths,
+    blank=0,
+    reduction="mean",
+    zero_infinity=False,
+):
+    original_dtype = log_probs.dtype
+    work_log_probs = log_probs
+    if work_log_probs.dtype in (torch.float16, torch.bfloat16):
+        work_log_probs = work_log_probs.float()
+    out = F.ctc_loss(
+        work_log_probs,
+        targets,
+        input_lengths,
+        target_lengths,
+        blank=blank,
+        reduction=reduction,
+        zero_infinity=zero_infinity,
+    )
+    return out.to(original_dtype) if out.dtype != original_dtype else out
+
+
+def _make_targets(batch, max_target, classes, device, target_layout):
+    target_lengths = torch.empty(batch, device=device, dtype=torch.long)
+    padded = torch.zeros(batch, max_target, device=device, dtype=torch.long)
+    pieces = []
+    for row in range(batch):
+        length = max(1, max_target - (row % 5))
+        target_lengths[row] = length
+        values = (torch.arange(length, device=device, dtype=torch.long) + row) % (
+            classes - 1
+        )
+        values = values + 1
+        padded[row, :length] = values
+        pieces.append(values)
+
+    if target_layout == "padded":
+        targets = padded
+    else:
+        targets = torch.cat(pieces)
+    return targets, target_lengths
+
+
+def ctc_loss_input_fn(shape, dtype, device):
+    t_steps, batch, classes, max_target = shape
+    raw = torch.randn(t_steps, batch, classes, dtype=torch.float32, device=device)
+    log_probs = raw.log_softmax(-1).to(dtype)
+    input_lengths = torch.full((batch,), t_steps, dtype=torch.long, device=device)
+
+    targets, target_lengths = _make_targets(
+        batch, max_target, classes, device, "padded"
+    )
+    yield (
+        log_probs,
+        targets,
+        input_lengths,
+        target_lengths,
+        {"blank": 0, "reduction": "mean", "zero_infinity": False},
+    )
+
+    targets, target_lengths = _make_targets(
+        batch, max_target, classes, device, "concatenated"
+    )
+    yield (
+        log_probs,
+        targets,
+        input_lengths,
+        target_lengths,
+        {"blank": 0, "reduction": "mean", "zero_infinity": False},
+    )
+
+    if base.Config.bench_level.value == consts.BenchLevel.COMPREHENSIVE.value:
+        yield (
+            log_probs,
+            targets,
+            input_lengths,
+            target_lengths,
+            {"blank": 0, "reduction": "sum", "zero_infinity": False},
+        )
+
+
+class CtcLossBenchmark(base.GenericBenchmark):
+    DEFAULT_SHAPES = [
+        (64, 4, 32, 16),
+        (256, 16, 64, 48),
+        (512, 32, 64, 48),
+        (1024, 32, 128, 96),
+    ]
+    DEFAULT_SHAPE_DESC = "T, N, C, S"
+
+    def set_more_shapes(self):
+        return []
+
+
+@pytest.mark.ctc_loss
+def test_ctc_loss():
+    bench = CtcLossBenchmark(
+        op_name="ctc_loss",
+        input_fn=ctc_loss_input_fn,
+        torch_op=_ctc_loss_reference,
+        dtypes=CTC_DTYPES,
+    )
+    bench.set_gems(flag_gems.ctc_loss)
+    bench.run()
+
+
+@pytest.mark.ctc_loss
+def test_ctc_loss_backward():
+    bench = CtcLossBenchmark(
+        op_name="ctc_loss",
+        input_fn=ctc_loss_input_fn,
+        torch_op=_ctc_loss_reference,
+        dtypes=CTC_DTYPES,
+        is_backward=True,
+    )
+    bench.set_gems(flag_gems.ctc_loss)
+    bench.run()

benchmark/test_trunc_divide.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from . import base, consts, utils
+from . import base, utils
 
 
 def _binary_input_fn(shape, dtype, device):

src/flag_gems/__init__.py

@@ -22,6 +22,7 @@
 registrar = Register
 current_work_registrar = None
 runtime.replace_customized_ops(globals())
+AUTOGRAD_DISPATCH_KEY = torch._C.DispatchKey.Autograd.name
 
 
 def torch_ge(v):
@@ -184,6 +185,8 @@ def torch_ge(v):
     ("copysign", copysign),
     ("copysign.out", copysign_out),
     ("count_nonzero", count_nonzero),
+    ("ctc_loss.IntList", ctc_loss, None, (AUTOGRAD_DISPATCH_KEY,)),
+    ("ctc_loss.Tensor", ctc_loss, None, (AUTOGRAD_DISPATCH_KEY,)),
     ("cudnn_convolution", cudnn_convolution),
     ("cummax", cummax),
     ("cummin", cummin),

src/flag_gems/ops/__init__.py

@@ -90,6 +90,7 @@
 from flag_gems.ops.cos import cos, cos_
 from flag_gems.ops.cosh import cosh, cosh_, cosh_out
 from flag_gems.ops.count_nonzero import count_nonzero
+from flag_gems.ops.ctc_loss import ctc_loss
 from flag_gems.ops.cudnn_convolution import cudnn_convolution
 from flag_gems.ops.cummax import cummax
 from flag_gems.ops.cummin import cummin
@@ -495,6 +496,7 @@
     "cosh_",
     "cosh_out",
     "count_nonzero",
+    "ctc_loss",
     "cudnn_convolution",
     "cummax",
     "cummin",