[FxImporter] Add aten._scaled_dot_product_flash_attention_for_cpu to …

…default decomposition table (#3456)
llvm · Jun 14, 2024 · a02e14e · a02e14e
1 parent 919b599
commit a02e14e
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 4 deletions.
diff --git a/projects/pt1/e2e_testing/xfail_sets.py b/projects/pt1/e2e_testing/xfail_sets.py
@@ -248,8 +248,6 @@
     # Lowering Torch Backend IR -> Linalg-on-Tensors Backend IR failed
     # 'linalg.depthwise_conv_2d_nchw_chw' op inferred input/output operand #1 has shape's dimension #0 to be 4, but found 8
     "Conv2dWithPaddingDilationStrideStaticModule_depthwise_multiplier",
-    # AssertionError: Unregistered operation: torch.aten._scaled_dot_product_flash_attention_for_cpu
-    "ScaledDotProductAttentionDifferentModule_basic",
     # AssertionError: Unregistered operation: torch.aten._embedding_bag_forward_only
     "AtenEmbeddingBagStaticModule_basic",
     # Lowering not present for this case
@@ -731,7 +729,6 @@
     "RsubInt0d_NumToTensor_Module_basic",
     "ScalarConstantTupleModule_basic",
     "ScalarImplicitFloatModule_basic",
-    "ScaledDotProductAttentionDifferentModule_basic",
     "ScatterReduceFloatMaxModule",
     "ScatterReduceFloatMaxModuleIncludeSelf",
     "ScatterReduceFloatMeanModule",
@@ -1978,6 +1975,7 @@
         "ViewSizeDimLedAndFollowedByCollapsedOnesModule_basic",
         "ViewSizeDimLedByCollapsedOnesModule_basic",
         "ViewSizeFromOtherTensor_basic",
+        "ScaledDotProductAttentionDifferentModule_basic",
     }
 ) - {
     ### Test failing in make_fx_tosa but not in tosa
@@ -3349,7 +3347,6 @@
     "ScalarConstantTupleModule_basic",
     "ScalarImplicitFloatModule_basic",
     "ScalarImplicitIntModule_basic",
-    "ScaledDotProductAttentionDifferentModule_basic",
     "ScatterReduceFloatMaxModule",
     "ScatterReduceFloatMaxModuleIncludeSelf",
     "ScatterReduceFloatMeanModule",

diff --git a/projects/pt1/python/torch_mlir/dynamo.py b/projects/pt1/python/torch_mlir/dynamo.py
@@ -65,6 +65,7 @@ def _get_decomposition_table():
         aten.sigmoid_backward,
         aten._native_batch_norm_legit,
         aten.squeeze,
+        aten._scaled_dot_product_flash_attention_for_cpu,
     ]
     # TODO: enable test once 2.1.0 is stable
     if torch_version_for_comparison() >= version.parse("2.1.0.dev"):

diff --git a/python/torch_mlir/extras/fx_decomp_util.py b/python/torch_mlir/extras/fx_decomp_util.py
@@ -48,6 +48,7 @@
     torch.ops.aten.triu.default,
     torch.ops.aten.nan_to_num.default,
     torch.ops.aten.unbind,
+    torch.ops.aten._scaled_dot_product_flash_attention_for_cpu,
 ]