add inplace pointwise operators (#224)

* add inplace pointwise operators
* add inplace autograd function(pay attention to save_for_backward & mark_dirty); add tests for inplace unary pointwise operations
* add test cases for inplace binary pointwise operators
* check out size && is_tensor list
* change broadcast
* check out_tensor overlapping
* modify tag for overlapping
* test internal_overlapping for only pointwise operation

---------

Co-authored-by: Jiang Bin <[email protected]>
Co-authored-by: FatJhon <[email protected]>

Author: 3 people

src/flag_gems/__init__.py

@@ -21,58 +21,93 @@ def enable(lib=aten_lib, unused=None, registrar=registrar):
     current_work_registrar = registrar(
         (
             ("abs", abs, Autograd.disable),
+            ("abs_", abs_, Autograd.disable),
             ("add.Tensor", add, Autograd.disable),
+            ("add_.Tensor", add_, Autograd.disable),
             ("addmm", addmm, Autograd.disable),
             ("arange.start_step", arange_start, Autograd.disable),
             ("arange.start", arange_start, Autograd.disable),
             ("arange", arange, Autograd.disable),
             ("batch_norm", batch_norm, Autograd.enable),
             ("bitwise_and.Tensor", bitwise_and_tensor, Autograd.disable),
+            ("bitwise_and_.Tensor_", bitwise_and_tensor_, Autograd.disable),
             ("bitwise_and.Scalar", bitwise_and_scalar, Autograd.disable),
+            ("bitwise_and_.Scalar", bitwise_and_scalar_, Autograd.disable),
             ("bitwise_and.Scalar_Tensor", bitwise_and_scalar_tensor, Autograd.disable),
             ("bitwise_not", bitwise_not, Autograd.disable),
+            ("bitwise_not_", bitwise_not_, Autograd.disable),
             ("bitwise_or.Tensor", bitwise_or_tensor, Autograd.disable),
+            ("bitwise_or_.Tensor", bitwise_or_tensor_, Autograd.disable),
             ("bitwise_or.Scalar", bitwise_or_scalar, Autograd.disable),
+            ("bitwise_or_.Scalar", bitwise_or_scalar_, Autograd.disable),
             ("bitwise_or.Scalar_Tensor", bitwise_or_scalar_tensor, Autograd.disable),
             ("bmm", bmm, Autograd.disable),
             ("clamp", clamp, Autograd.disable),
+            ("clamp_", clamp_, Autograd.disable),
             ("clamp.Tensor", clamp_tensor, Autograd.disable),
+            ("clamp_.Tensor", clamp_tensor_, Autograd.disable),
             ("cos", cos, Autograd.disable),
+            ("cos_", cos_, Autograd.disable),
             ("pad", pad, Autograd.disable),
             ("constant_pad_nd", constant_pad_nd, Autograd.disable),
             ("cumsum", cumsum, Autograd.disable),
             ("cummin", cummin, Autograd.disable),
             ("div.Tensor", true_divide, Autograd.disable),
+            ("div_.Tensor", true_divide_, Autograd.disable),
             ("div.Scalar", true_divide, Autograd.disable),
+            ("div_.Scalar", true_divide_, Autograd.disable),
             ("div.Tensor_mode", div_mode, Autograd.disable),
+            ("div_.Tensor_mode", div_mode_, Autograd.disable),
             ("div.Scalar_mode", div_mode, Autograd.disable),
+            ("div_.Scalar_mode", div_mode_, Autograd.disable),
             (
                 "divide.Tensor",
                 true_divide,
                 Autograd.disable,
             ),  # divide, an alias for div
+            (
+                "divide_.Tensor",
+                true_divide_,
+                Autograd.disable,
+            ),  # divide, an alias for div
             ("divide.Scalar", true_divide, Autograd.disable),
+            ("divide_.Scalar", true_divide_, Autograd.disable),
             ("divide.Tensor_mode", div_mode, Autograd.disable),
+            ("divide_.Tensor_mode", div_mode_, Autograd.disable),
             ("divide.Scalar_mode", div_mode, Autograd.disable),
+            ("divide_.Scalar_mode", div_mode_, Autograd.disable),
             (
                 "true_divide.Tensor",
                 true_divide,
                 Autograd.disable,
             ),  # true_divide, an alias for div
+            (
+                "true_divide_.Tensor",
+                true_divide_,
+                Autograd.disable,
+            ),  # true_divide, an alias for div
             ("true_divide.Scalar", true_divide, Autograd.disable),
+            ("true_divide_.Scalar", true_divide_, Autograd.disable),
             ("floor_divide", floor_divide, Autograd.disable),
             ("floor_divide.Scalar", floor_divide, Autograd.disable),
             ("remainder.Tensor", remainder, Autograd.disable),
+            ("remainder_.Tensor", remainder_, Autograd.disable),
+            ("remainder.Scalar", remainder, Autograd.disable),
+            ("remainder_.Scalar", remainder_, Autograd.disable),
+            ("remainder.Scalar_Tensor", remainder, Autograd.disable),
             ("native_dropout", native_dropout, Autograd.enable),
             ("erf", erf, Autograd.disable),
+            ("erf_", erf_, Autograd.disable),
             ("embedding", embedding, Autograd.enable),
             ("eq.Tensor", eq, Autograd.disable),
             ("eq.Scalar", eq_scalar, Autograd.disable),
             ("exp", exp, Autograd.disable),
+            ("exp_", exp_, Autograd.disable),
             ("exponential_", exponential_, Autograd.disable),
             ("ge.Tensor", ge, Autograd.disable),
             ("ge.Scalar", ge_scalar, Autograd.disable),
             ("gelu", gelu, Autograd.enable),
+            ("gelu_", gelu_, Autograd.enable),
             ("native_group_norm", group_norm, Autograd.enable),
             ("_weight_norm_interface", weight_norm_interface, Autograd.enable),
             ("_weight_norm", weight_norm, Autograd.enable),
@@ -118,19 +153,30 @@ def enable(lib=aten_lib, unused=None, registrar=registrar):
             ("ne.Tensor", ne, Autograd.disable),
             ("ne.Scalar", ne_scalar, Autograd.disable),
             ("neg", neg, Autograd.disable),
+            ("neg_", neg_, Autograd.disable),
             ("pow.Scalar", pow_scalar, Autograd.disable),
             ("pow.Tensor_Scalar", pow_tensor_scalar, Autograd.disable),
+            ("pow_.Scalar", pow_tensor_scalar_, Autograd.disable),
             ("pow.Tensor_Tensor", pow_tensor_tensor, Autograd.disable),
+            ("pow_.Tensor", pow_tensor_tensor_, Autograd.disable),
             ("reciprocal", reciprocal, Autograd.disable),
+            ("reciprocal_", reciprocal_, Autograd.disable),
             ("relu", relu, Autograd.enable),
+            ("relu_", relu_, Autograd.enable),
             ("rsqrt", rsqrt, Autograd.disable),
+            ("rsqrt_", rsqrt_, Autograd.disable),
             ("sigmoid", sigmoid, Autograd.enable),
+            ("sigmoid_", sigmoid_, Autograd.enable),
             ("silu", silu, Autograd.enable),
+            ("silu_", silu_, Autograd.enable),
             ("sin", sin, Autograd.disable),
+            ("sin_", sin_, Autograd.disable),
             ("softmax.int", softmax, Autograd.enable),
             ("sort", sort, Autograd.disable),
             ("sub.Tensor", sub, Autograd.disable),
+            ("sub_.Tensor", sub_, Autograd.disable),
             ("tanh", tanh, Autograd.enable),
+            ("tanh_", tanh_, Autograd.enable),
             ("triu", triu, Autograd.disable),
             # ("topk", topk, Autograd.disable),
             ("var_mean.correction", var_mean, Autograd.disable),

src/flag_gems/ops/__init__.py

@@ -1,5 +1,5 @@
-from .abs import abs
-from .add import add
+from .abs import abs, abs_
+from .add import add, add_
 from .addmm import addmm
 from .all import all, all_dim, all_dims
 from .amax import amax
@@ -11,40 +11,57 @@
 from .batch_norm import batch_norm
 from .bitwise_and import (
     bitwise_and_scalar,
+    bitwise_and_scalar_,
     bitwise_and_scalar_tensor,
     bitwise_and_tensor,
+    bitwise_and_tensor_,
+)
+from .bitwise_not import bitwise_not, bitwise_not_
+from .bitwise_or import (
+    bitwise_or_scalar,
+    bitwise_or_scalar_,
+    bitwise_or_scalar_tensor,
+    bitwise_or_tensor,
+    bitwise_or_tensor_,
 )
-from .bitwise_not import bitwise_not
-from .bitwise_or import bitwise_or_scalar, bitwise_or_scalar_tensor, bitwise_or_tensor
 from .bmm import bmm
 from .cat import cat
-from .clamp import clamp, clamp_tensor
+from .clamp import clamp, clamp_, clamp_tensor, clamp_tensor_
 from .conv1d import conv1d
 from .conv2d import conv2d
 from .conv_depthwise2d import _conv_depthwise2d
-from .cos import cos
+from .cos import cos, cos_
 from .count_nonzero import count_nonzero
 from .cross_entropy_loss import cross_entropy_loss
 from .cummin import cummin
 from .cumsum import cumsum, normed_cumsum
 from .diag import diag
 from .diag_embed import diag_embed
 from .diagonal import diagonal_backward
-from .div import div_mode, floor_divide, remainder, true_divide
+from .div import (
+    div_mode,
+    div_mode_,
+    floor_divide,
+    floor_divide_,
+    remainder,
+    remainder_,
+    true_divide,
+    true_divide_,
+)
 from .dropout import native_dropout
 from .elu import elu
 from .embedding import embedding
 from .eq import eq, eq_scalar
-from .erf import erf
-from .exp import exp
+from .erf import erf, erf_
+from .exp import exp, exp_
 from .exponential_ import exponential_
 from .fill import fill_scalar, fill_tensor
 from .flip import flip
 from .full import full
 from .full_like import full_like
 from .gather import gather, gather_backward
 from .ge import ge, ge_scalar
-from .gelu import gelu
+from .gelu import gelu, gelu_
 from .groupnorm import group_norm
 from .gt import gt, gt_scalar
 from .hstack import hstack
@@ -76,11 +93,11 @@
 from .minimum import minimum
 from .mm import mm
 from .mse_loss import mse_loss
-from .mul import mul
+from .mul import mul, mul_
 from .multinomial import multinomial
 from .mv import mv
 from .ne import ne, ne_scalar
-from .neg import neg
+from .neg import neg, neg_
 from .nllloss import (
     nll_loss2d_backward,
     nll_loss2d_forward,
@@ -93,16 +110,22 @@
 from .ones_like import ones_like
 from .outer import outer
 from .pad import constant_pad_nd, pad
-from .pow import pow_scalar, pow_tensor_scalar, pow_tensor_tensor
+from .pow import (
+    pow_scalar,
+    pow_tensor_scalar,
+    pow_tensor_scalar_,
+    pow_tensor_tensor,
+    pow_tensor_tensor_,
+)
 from .prod import prod, prod_dim
 from .quantile import quantile
 from .rand import rand
 from .rand_like import rand_like
 from .randn import randn
 from .randn_like import randn_like
 from .randperm import randperm
-from .reciprocal import reciprocal
-from .relu import relu
+from .reciprocal import reciprocal, reciprocal_
+from .relu import relu, relu_
 from .repeat import repeat
 from .repeat_interleave import (
     repeat_interleave_self_int,
@@ -112,19 +135,19 @@
 from .resolve_conj import resolve_conj
 from .resolve_neg import resolve_neg
 from .rms_norm import rms_norm
-from .rsqrt import rsqrt
+from .rsqrt import rsqrt, rsqrt_
 from .scatter import scatter
 from .select_scatter import select_scatter
-from .sigmoid import sigmoid
-from .silu import silu
-from .sin import sin
+from .sigmoid import sigmoid, sigmoid_
+from .silu import silu, silu_
+from .sin import sin, sin_
 from .slice_scatter import slice_scatter
 from .softmax import softmax
 from .sort import sort
 from .stack import stack
-from .sub import sub
+from .sub import sub, sub_
 from .sum import sum, sum_dim
-from .tanh import tanh
+from .tanh import tanh, tanh_
 from .tile import tile
 from .topk import topk
 from .triu import triu
@@ -151,22 +174,32 @@
     "any_dim",
     "any_dims",
     "add",
+    "add_",
     "abs",
+    "abs_",
     "addmm",
     "arange",
     "arange_start",
     "batch_norm",
     "bitwise_and_tensor",
+    "bitwise_and_tensor_",
     "bitwise_and_scalar",
+    "bitwise_and_scalar_",
     "bitwise_and_scalar_tensor",
     "bitwise_not",
+    "bitwise_not_",
     "bitwise_or_tensor",
+    "bitwise_or_tensor_",
     "bitwise_or_scalar",
+    "bitwise_or_scalar_",
     "bitwise_or_scalar_tensor",
     "bmm",
     "clamp",
+    "clamp_",
     "clamp_tensor",
+    "clamp_tensor_",
     "cos",
+    "cos_",
     "count_nonzero",
     "diag",
     "diag_embed",
@@ -178,18 +211,24 @@
     "cumsum",
     "normed_cumsum",
     "true_divide",
+    "true_divide_",
     "div_mode",
+    "div_mode_",
     "floor_divide",
+    "floor_divide_",
     "remainder",
+    "remainder_",
     "zeros",
     "ones",
     "full",
     "native_dropout",
     "erf",
+    "erf_",
     "embedding",
     "eq",
     "eq_scalar",
     "exp",
+    "exp_",
     "fill_scalar",
     "fill_tensor",
     "exponential_",
@@ -202,6 +241,7 @@
     "ge",
     "ge_scalar",
     "gelu",
+    "gelu_",
     "group_norm",
     "gt",
     "gt_scalar",
@@ -224,6 +264,7 @@
     "mean_dim",
     "mm",
     "mul",
+    "mul_",
     "multinomial",
     "maximum",
     "minimum",
@@ -242,19 +283,30 @@
     "ne",
     "ne_scalar",
     "neg",
+    "neg_",
     "pow_scalar",
     "pow_tensor_scalar",
     "pow_tensor_tensor",
+    "pow_tensor_scalar_",
+    "pow_tensor_tensor_",
     "reciprocal",
+    "reciprocal_",
     "relu",
+    "relu_",
     "rsqrt",
+    "rsqrt_",
     "scatter",
     "sigmoid",
+    "sigmoid_",
     "silu",
+    "silu_",
     "sin",
+    "sin_",
     "softmax",
     "sub",
+    "sub_",
     "tanh",
+    "tanh_",
     "tile",
     "triu",
     "topk",

src/flag_gems/ops/abs.py

@@ -15,3 +15,9 @@ def abs_func(x):
 def abs(A):
     logging.debug("GEMS ABS")
     return abs_func(A)
+
+
+def abs_(A):
+    logging.debug("GEMS ABS_")
+    abs_func(A, out0=A)
+    return A

src/flag_gems/ops/add.py

@@ -38,3 +38,15 @@ def add(A, B, *, alpha=1):
         return add_func_scalar_tensor(A, B, alpha)
     else:
         return torch.tensor(A + B * alpha)
+
+
+def add_(A, B, *, alpha=1):
+    logging.debug("GEMS ADD_")
+    if isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor):
+        return add_func(A, B, alpha, out0=A)
+    elif isinstance(A, torch.Tensor):
+        return add_func_tensor_scalar(A, B, alpha, out0=A)
+    # elif isinstance(B, torch.Tensor):
+    #     return add_func_scalar_tensor(A, B, alpha, out0=A)
+    else:
+        raise ValueError("Unreachable.")