Refactor code generation for pointwise operation & PointwiseDynamicFunction

.github/workflows/python-test.yaml

@@ -37,6 +37,7 @@ jobs:
             "CUDA_VISIBLE_DEVICES=3 pytest -s tests/test_reduction_ops.py &"
             "CUDA_VISIBLE_DEVICES=4 pytest -s tests/test_special_ops.py &"
             "CUDA_VISIBLE_DEVICES=5 pytest -s tests/test_libentry.py &"
+            "CUDA_VISIBLE_DEVICES=5 pytest -s tests/test_pointwise_dynamic.py &"
           )
 
           declare -a exit_statuses

.gitignore

@@ -21,3 +21,4 @@ build/
 # Project files, i.e. `.project`, `.actionScriptProperties` and `.flexProperties`
 # should NOT be excluded as they contain compiler settings and other important
 # information for Eclipse / Flash Builder.
+playground/

src/flag_gems/ops/flip.py

@@ -4,11 +4,12 @@
 import triton
 
 from ..utils import pointwise_dynamic
+from ..utils.tensor_wrapper import StridedBuffer
 
 
 @pointwise_dynamic(is_tensor=[True], promotion_methods=[(0, "DEFAULT")])
 @triton.jit
-def flip_func(x, **kwargs):
+def copy_func(x):
     return x
 
 
@@ -29,10 +30,17 @@ def flip(A: torch.Tensor, dims) -> torch.Tensor:
     n = 0
     offset = 0
     for i in range(len(flip_dims_b)):
-        if flip_dims_b[i] and A.size()[i] > 1 and A.stride()[i] != 0:
+        if flip_dims_b[i] and A.size(i) > 1 and A.stride(i) != 0:
             offset += strides[i] * (A.shape[i] - 1)
             strides[i] = -strides[i]
             n += 1
     if n == 0 or A.numel() <= 1:
         return A.clone()
-    return flip_func(A, out0_offset=offset, out0_strides=strides)
+    out = torch.empty_like(A)
+    # a flipped view of A
+    flipped_A = StridedBuffer(A, strides=strides, offset=offset)
+
+    # TODO: flip op can have a custom task simplification method, but we skip it now and just use A's rank.
+    overload = copy_func.instantiate(A.ndim)
+    overload(flipped_A, out0=out)
+    return out

src/flag_gems/utils/code_utils.py

@@ -102,6 +102,9 @@ def writelines(self, lines):
         for line in lines:
             self.writeline(line)
 
+    def writemultiline(self, s):
+        self.writelines(s.splitlines())
+
     def indent(self, offset=1):
         @contextlib.contextmanager
         def ctx():