Specializing slice_scatter. (FlagOpen#270)

* specializing slice_scatter. WIP. * polish and refine 2d_inner cases. * fix slice_scatter error on 1d inputs. * test slice_scatter fallback
2niuhe · Nov 5, 2024 · 189f645 · 189f645
1 parent af0a3a6
commit 189f645
Show file tree

Hide file tree

Showing 5 changed files with 599 additions and 16 deletions.
diff --git a/src/flag_gems/__init__.py b/src/flag_gems/__init__.py
@@ -142,7 +142,7 @@ def enable(lib=aten_lib):
     lib.impl("fill.Scalar", fill_scalar, "CUDA")
     lib.impl("fill.Tensor", fill_tensor, "CUDA")
     lib.impl("flip", flip, "CUDA")
-    lib.impl("slice_scatter", slice_scatter, "CUDA")
+    lib.impl("slice_scatter", slice_scatter_v2, "CUDA")
     lib.impl("select_scatter", select_scatter, "CUDA")
     lib.impl("index_select", index_select, "CUDA")
     lib.impl("tile", tile, "CUDA")

diff --git a/src/flag_gems/ops/__init__.py b/src/flag_gems/ops/__init__.py
@@ -88,7 +88,7 @@
 from .sigmoid import sigmoid
 from .silu import silu
 from .sin import sin
-from .slice_scatter import slice_scatter
+from .slice_scatter import slice_scatter, slice_scatter_v2
 from .softmax import softmax
 from .stack import stack
 from .sub import sub
@@ -234,6 +234,7 @@
     "where_scalar_other",
     "select_scatter",
     "slice_scatter",
+    "slice_scatter_v2",
     "masked_fill",
     "_unique2",
     "_upsample_bicubic2d_aa",