[usmanmateen]

Summary

The optimized _log_softmax.out kernel only handled Float, Half and BFloat16, with a TODO: support Double as well. A float64 log_softmax fails at runtime wherever the optimized kernels are linked, including the Python runtime, while the same program runs with the portable kernels:

import torch
from executorch.exir import to_edge
from executorch.runtime import Runtime


class LogSoftmax(torch.nn.Module):
    def forward(self, x):
        return torch.log_softmax(x, dim=1)


x = torch.randn(3, 4, dtype=torch.float64)
prog = to_edge(torch.export.export(LogSoftmax(), (x,))).to_executorch()
Runtime.get().load_program(prog.buffer).load_method("forward").execute([x])
# [op_log_softmax.cpp:171] Check failed (false):
# RuntimeError: method->execute() failed with error 0x12

The ATen vectorized helpers the kernel calls (serial_vec_log_softmax_lastdim_range, serial_vec_logsoftmax_range) support double, so this adds a Double case that goes through the same same-dtype path as Half and BFloat16. It also removes the op_log_softmax: dtype_double: false entry from kernels/optimized/test/supported_features_def.yaml, so the existing double tests run for the optimized kernels too.

Test plan

I built optimized_kernels_test with CMake (EXECUTORCH_BUILD_TESTS=ON, EXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON) on Linux x86_64 (WSL2).

With the yaml entry removed but without the kernel change, AllDtypesSupported, TestWithLargeNumber and DoubleCase fail. With this change, all 13 OpLogSoftmaxOutTest tests pass (4 of them were skipped before), and the full optimized_kernels_test run passes (1657 tests). clang-format 18.1.3 reports no changes.

The repro above was run against executorch 1.5.1; op_log_softmax.cpp is the same on main.

This PR was authored with AI assistance (Claude Code).

cc @larryliu0820 @manuelcandales @JakeStevens @nil-is-all

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23686

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit f23eb3b with merge base 6cc8edf ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[usmanmateen]

@pytorchbot label "release notes: ops & kernels"

[JakeStevens]

a few nits:

(1) can you add test_dtype_noncontiguous_dim<double, ScalarType::Double>() to NonContiguous in kernels/test/op_log_softmax_test.cpp so we test that pathway too?
(2) The non-last-dim tiling hardcodes vec_logsoftmax_chunk_size_and_num_chunks<float, 64 * 1024> . Pre-existing, but for double the byte budget doubles; templating on IN_T would match ATen's own caller. The new test above would cover either choice.
(3) Please update fbcode/executorch/kernels/test/supported_features_summary.md to reflect the yaml changes

[usmanmateen]

Thanks! Done in f23eb3b, on top of a rebase onto main (6cc8edf):

1. NonContiguous now also runs test_dtype_noncontiguous_dim<double, ScalarType::Double>(), guarded by op_log_softmax_dtype_double like AllDtypesSupported.
2. The non-last-dim path now calls vec_logsoftmax_chunk_size_and_num_chunks<IN_T, 64 * 1024>, so the chunk budget follows the element size as in ATen. Half and BFloat16 follow IN_T too now; their tests still pass.
3. Updated kernels/test/supported_features_summary.md: optimized is now True for op_log_softmax dtype_double.

Locally: optimized_kernels_test 1657 passed, portable_kernels_test 1648 passed, clang-format 18 clean.