Update fused kernels and call _safe_softmax from SDPA

Differential Revision: D61418679

Pull Request resolved: #4772

Author: drisspg

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -68,15 +68,23 @@ def test_vit_skip_conv(self):
             )
         )
 
+        conv_block = ["aten.convolution.default", "executorch_call_delegate"]
+        safe_softmax_block = [
+            "getitem",
+            "getitem",
+            "getitem",
+            "getitem",
+            "aten.any.dim",
+            "executorch_call_delegate",
+        ]
+        final_block = ["getitem"]
+        total = conv_block + 12 * safe_softmax_block + final_block
+
         assert [
             node.target.__name__
             for node in delegated_program_manager.exported_program().graph.nodes
             if node.op == "call_function"
-        ] == [
-            "aten.convolution.default",
-            "executorch_call_delegate",
-            "getitem",
-        ]
+        ] == total
 
 
 if __name__ == "__main__":