cambricon: fix kernel (#1263)

Author: chenmiao1919

src/flag_gems/runtime/backend/_cambricon/__init__.py

@@ -24,6 +24,9 @@
     dispatch_key="PrivateUse1",
 )
 
-CUSTOMIZED_UNUSED_OPS = ("_upsample_bicubic2d_aa",)  # skip now
+CUSTOMIZED_UNUSED_OPS = (
+    "_upsample_bicubic2d_aa",  # skip now
+    "scatter_add_",
+)
 
 __all__ = ["*"]

src/flag_gems/runtime/backend/_cambricon/ops/log_softmax.py

@@ -408,7 +408,13 @@ def log_softmax_kernel_inner_k_partial_stats(
         ).to(tl.float32)
 
         tile_max = tl.max(tile, axis=1)
-        tile_sum = tl.sum(tl.exp(tile - tile_max[:, None]), axis=1)
+        all_neg_inf = tile_max == -float("inf")
+
+        tile_sum = tl.where(
+            all_neg_inf,
+            0.0,
+            tl.sum(tl.exp(tile - tile_max[:, None]), axis=1),
+        )
 
         tl.store(max_buf_ptr + offs_m * T + tile_id, tile_max, mask=(offs_m < M))
         tl.store(sum_buf_ptr + offs_m * T + tile_id, tile_sum, mask=(offs_m < M))
@@ -487,7 +493,13 @@ def log_softmax_kernel_inner_k_write_logsoftmax(
             other=-float("inf"),
         ).to(tl.float32)
 
-        o = tl.exp(tile - gmax[:, None]) / gsum[:, None]
+        valid = gsum[:, None] > 0
+
+        o = tl.where(
+            valid,
+            tl.exp(tile - gmax[:, None]) / gsum[:, None],
+            0.0,
+        )
         out = tl.log2(o) / log2e
 
         tl.store(y_ptr + offs_m[:, None] * N + offs_n[None, :], out, mask=mask)

src/flag_gems/runtime/backend/_cambricon/ops/softmax.py

@@ -380,7 +380,13 @@ def softmax_kernel_inner_k_partial_stats(
         ).to(tl.float32)
 
         tile_max = tl.max(tile, axis=1)
-        tile_sum = tl.sum(tl.exp(tile - tile_max[:, None]), axis=1)
+        all_neg_inf = tile_max == -float("inf")
+
+        tile_sum = tl.where(
+            all_neg_inf,
+            0.0,
+            tl.sum(tl.exp(tile - tile_max[:, None]), axis=1),
+        )
 
         tl.store(max_buf_ptr + offs_m * T + tile_id, tile_max, mask=(offs_m < M))
         tl.store(sum_buf_ptr + offs_m * T + tile_id, tile_sum, mask=(offs_m < M))
@@ -459,7 +465,13 @@ def softmax_kernel_inner_k_write_softmax(
             other=-float("inf"),
         ).to(tl.float32)
 
-        out = tl.exp(tile - gmax[:, None]) / gsum[:, None]
+        valid = gsum[:, None] > 0
+
+        out = tl.where(
+            valid,
+            tl.exp(tile - gmax[:, None]) / gsum[:, None],
+            0.0,
+        )
 
         tl.store(y_ptr + offs_m[:, None] * N + offs_n[None, :], out, mask=mask)