update core.ops with pyboost (#1884)

Author: lvyufeng

.github/workflows/ci_pipeline.yaml

@@ -80,7 +80,7 @@ jobs:
         pip install -r download.txt
     - name: Test with pytest
       run: |
-        pytest -c pytest.ini -m 'not download and not gpu_only' --ignore=tests/transformers tests/ut
+        pytest -c pytest.ini -m 'not download and not gpu_only' --ignore=tests/transformers tests
 
   release-test:
     needs: pylint-check
@@ -104,8 +104,8 @@ jobs:
         pip install https://ms-release.obs.cn-north-4.myhuaweicloud.com/${{matrix.ms_version}}/MindSpore/unified/x86_64/mindspore-${{matrix.ms_version}}-cp39-cp39-linux_x86_64.whl --trusted-host ms-release.obs.cn-north-4.myhuaweicloud.com
     - name: Test with pytest
       run: |
-        pytest -c pytest.ini -m 'not download and not gpu_only' --ignore=tests/transformers tests/ut 
-        # pytest -c pytest.ini -m 'not download and not gpu_only' tests/ut
+        pytest -c pytest.ini -m 'not download and not gpu_only' --ignore=tests/transformers tests
+        # pytest -c pytest.ini -m 'not download and not gpu_only' tests
 
   transformers-model-test:
     needs: pylint-check

mindnlp/core/ops/array.py

@@ -19,12 +19,14 @@ def argwhere(input):
     return ops.argwhere(input)
 
 # cat
+has_cat = hasattr(mindspore.mint, 'cat')
 def cat(tensors, dim=0):
-    if use_pyboost():
+    if use_pyboost() and has_cat:
         return mindspore.mint.cat(tensors, dim)
     return ops.cat(tensors, dim)
 
 # concat
+has_concat = hasattr(mindspore.mint, 'concat')
 def concat(tensors, dim=0):
     return cat(tensors, dim)
 
@@ -37,7 +39,10 @@ def conj(input):
     return ops.conj(input)
 
 # chunk
+has_chunk = hasattr(mindspore.mint, 'chunk')
 def chunk(input, chunks, dim=0):
+    if use_pyboost() and has_chunk:
+        return mindspore.mint.chunk(input, chunks, dim)
     return ops.chunk(input, chunks, dim)
 
 # dsplit
@@ -50,8 +55,9 @@ def chunk(input, chunks, dim=0):
 
 
 # gather
+has_gather = hasattr(mindspore.mint, 'gather')
 def gather(input, dim, index):
-    if use_pyboost():
+    if use_pyboost() and has_gather:
         return mindspore.mint.gather(input, dim, index)
     index = ops.where(index < input.shape[dim], index, index - input.shape[dim])
     return ops.gather_elements(input, dim, index)
@@ -91,13 +97,17 @@ def inplace_index_add(input, dim, index, source):
 
 
 # index_select
+has_index_select = hasattr(mindspore.mint, 'index_select')
 def index_select(input, dim, index):
-    if use_pyboost():
+    if use_pyboost() and has_index_select:
         return mindspore.mint.index_select(input, dim, index)
     return ops.index_select(input, dim, index)
 
 # masked_select
+has_masked_select = hasattr(mindspore.mint, 'masked_select')
 def masked_select(input, mask):
+    if use_pyboost() and has_masked_select:
+        return mindspore.mint.masked_select(input, mask)
     return ops.masked_select(input, mask)
 
 # movedim
@@ -107,17 +117,19 @@ def masked_select(input, mask):
 
 
 # narrow
+has_narrow = hasattr(mindspore.mint, 'narrow')
 def narrow(input, dim, start, length):
-    if use_pyboost():
+    if use_pyboost() and has_narrow:
         return mindspore.mint.narrow(input, dim, start, length)
     return ops.narrow(input, dim, start, length)
 
 # narrow_copy
 
 
 # nonzero
+has_nonzero = hasattr(mindspore.mint, 'nonzero')
 def nonzero(input, *, as_tuple=False):
-    if use_pyboost():
+    if use_pyboost() and has_nonzero:
         return mindspore.mint.nonzero(input, as_tuple)
     _nonzero = _get_cache_prim(ops.NonZero)()
     out = _nonzero(input)
@@ -128,37 +140,40 @@ def nonzero(input, *, as_tuple=False):
     return out
 
 # permute
+has_permute = hasattr(mindspore.mint, 'permute')
 def permute(input, dims):
-    if use_pyboost():
+    if use_pyboost() and has_permute:
         return mindspore.mint.permute(input, dims)
     return ops.permute(input, dims)
 
 # reshape
+has_reshape = hasattr(mindspore.mint, 'reshape')
 def reshape(input, shape):
+    if use_pyboost() and has_reshape:
+        return mindspore.mint.reshape(input, shape)
     return ops.reshape(input, shape)
 
 def view(input, *shape):
-    # if use_pyboost():
-    #     return mindspore.ops.auto_generate.gen_ops_prim.view_op(input, shape)
-    return ops.reshape(input, shape)
+    return reshape(input, shape)
 
 # row_stack
 
 # select
+has_select = hasattr(mindspore.mint, 'select')
 def select(input, dim, index):
+    if use_pyboost() and has_select:
+        return mindspore.mint.select(input, dim, index)
     slices = ()
     for _ in range(dim):
         slices += (slice(None, None, None),)
     slices += (index,)
     return input[slices]
 
 # scatter
+has_scatter = hasattr(mindspore.mint, 'scatter')
 def scatter(input, dim, index, src):
-    if use_pyboost():
-        try:
-            return mindspore.mint.scatter(input, dim, index, src)
-        except:
-            return mindspore.ops.auto_generate.gen_ops_prim.scatter_op(input, dim, index, src, 0)
+    if use_pyboost() and has_scatter:
+        return mindspore.mint.scatter(input, dim, index, src)
     if not isinstance(src, mindspore.Tensor):
         src = ops.full(index.shape, src, dtype=input.dtype)
     return ops.tensor_scatter_elements(input, index, src, dim)
@@ -179,8 +194,9 @@ def tf_scatter_nd(indices, updates, shape):
 
 
 # scatter_add
+has_scatter_add = hasattr(mindspore.mint, 'scatter_add')
 def scatter_add(input, dim, index, src):
-    if use_pyboost():
+    if use_pyboost() and has_scatter_add:
         return mindspore.mint.scatter_add(input, dim, index, src)
     return ops.tensor_scatter_elements(input, index, src, dim, 'add')
 
@@ -196,8 +212,9 @@ def scatter_update(input, indices, updates):
     return ops.scatter_update(input, indices, updates)
 
 # split
+has_split = hasattr(mindspore.mint, 'split')
 def split(tensor, split_size_or_sections, dim=0):
-    if use_pyboost():
+    if use_pyboost() and has_split:
         return mindspore.mint.split(tensor, split_size_or_sections, dim)
     return ops.split(tensor, split_size_or_sections, dim)
 
@@ -206,8 +223,9 @@ def squeeze(input, dim=None):
     return ops.squeeze(input, dim)
 
 # stack
+has_stack = hasattr(mindspore.mint, 'stack')
 def stack(tensors, dim=0):
-    if use_pyboost():
+    if use_pyboost() and has_stack:
         return mindspore.mint.stack(tensors, dim)
     return ops.stack(tensors, dim)
 
@@ -235,20 +253,22 @@ def take(input, index):
 
 
 # tile
+has_tile = hasattr(mindspore.mint, 'tile')
 def tile(input, dims):
-    if use_pyboost():
+    if use_pyboost() and has_tile:
         return mindspore.mint.tile(input, dims)
     return ops.tile(input, dims)
 
 # transpose
+has_transpose = hasattr(mindspore.mint, 'transpose')
 def transpose(input, dim0, dim1):
+    if use_pyboost() and has_transpose:
+        return mindspore.mint.transpose(input, dim0, dim1)
     ranks = list(range(input.ndim))
     rank0 = ranks[dim0]
     rank1 = ranks[dim1]
     ranks[dim0] = rank1
     ranks[dim1] = rank0
-    if use_pyboost():
-        return mindspore.ops.auto_generate.gen_ops_prim.transpose_op(input, tuple(ranks))
     return permute(input, tuple(ranks))
 
 # unbind
@@ -258,7 +278,10 @@ def unbind(input, dim=0):
 # unravel_index
 
 # unsqueeze
+has_unsqueeze = hasattr(mindspore.mint, 'unsqueeze')
 def unsqueeze(input, dim):
+    if use_pyboost() and has_unsqueeze:
+        return mindspore.mint.unsqueeze(input, dim)
     return ops.expand_dims(input, dim)
 
 # vsplit
@@ -273,10 +296,11 @@ def vstack(input):
 )
 
 # where
+has_where = hasattr(mindspore.mint, 'where')
 def where(condition, input, other):
     if ON_ORANGE_PI:
         return condition * input + (~condition) * other
-    if use_pyboost():
+    if use_pyboost() and has_where:
         return mindspore.mint.where(condition, input, other)
     return ops.where(condition, input, other)
 

mindnlp/core/ops/blas.py

@@ -19,15 +19,19 @@ def addmm(input, mat1, mat2, *, beta=1, alpha=1):
 
 
 # baddbmm
+has_baddbmm = hasattr(mindspore.mint, 'baddbmm')
 def baddbmm(input, batch1, batch2, *, beta=1, alpha=1):
+    if use_pyboost() and has_baddbmm:
+        return mindspore.mint.baddbmm(input, batch1, batch2, beta=beta, alpha=alpha)
     return ops.baddbmm(input, batch1, batch2, beta=beta, alpha=alpha)
 
 # bmm
+has_bmm = hasattr(mindspore.mint, 'bmm')
 def bmm(input, other):
     if ON_ORANGE_PI:
         input = input.to(mindspore.float16)
         other = input.to(mindspore.float16)
-    if use_pyboost():
+    if use_pyboost() and has_bmm:
         return mindspore.mint.bmm(input, other)
     return ops.bmm(input, other)
 
@@ -66,11 +70,12 @@ def dot(input, other):
 # lu_unpack
 
 # matmul
+has_matmul = hasattr(mindspore.mint, 'matmul')
 def matmul(input, other):
     if ON_ORANGE_PI:
         input = input.to(mindspore.float16)
         other = other.to(mindspore.float16)
-    if use_pyboost():
+    if use_pyboost() and has_matmul:
         return mindspore.mint.matmul(input, other)
     return ops.matmul(input, other)
 
@@ -90,7 +95,10 @@ def mm(input, other):
 # ormqr
 
 # outer
+has_outer = hasattr(mindspore.mint, 'outer')
 def outer(input, vec2):
+    if use_pyboost() and has_outer:
+        return mindspore.mint.outer(input, vec2)
     return ops.outer(input, vec2)
 
 # pinverse

mindnlp/core/ops/comparison.py

@@ -13,8 +13,9 @@ def argsort(input, dim=-1, descending=False, stable=False):
     return sort(input, dim=dim, descending=descending, stable=stable)[1]
 
 # eq
+has_eq = hasattr(mindspore.mint, 'eq')
 def eq(input, other):
-    if use_pyboost():
+    if use_pyboost() and has_eq:
         return mindspore.mint.eq(input, other)
     return ops.eq(input, other)
 
@@ -27,24 +28,28 @@ def ge(input, other):
     return ops.ge(input, other)
 
 # gt
+has_gt = hasattr(mindspore.mint, 'gt')
 def gt(input, other):
-    if use_pyboost():
+    if use_pyboost() and has_gt:
         return mindspore.mint.gt(input, other)
     return ops.gt(input, other)
 
 # greater
+has_greater = hasattr(mindspore.mint, 'greater')
 def greater(input, other):
     return gt(input, other)
 
 # isclose
+has_isclose = hasattr(mindspore.mint, 'isclose')
 def isclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False):
-    if use_pyboost():
+    if use_pyboost() and has_isclose:
         return mindspore.mint.isclose(input, other, rtol, atol, equal_nan)
     return mindspore.tensor(np.isclose(input.numpy(), other.numpy(), rtol, atol, equal_nan))
 
 # isfinite
+has_isfinite = hasattr(mindspore.mint, 'isfinite')
 def isfinite(input):
-    if use_pyboost():
+    if use_pyboost() and has_isfinite:
         return mindspore.mint.isfinite(input)
     return ops.isfinite(input)
 
@@ -79,8 +84,9 @@ def isnan(input):
 # kthvalue
 
 # le
+has_le = hasattr(mindspore.mint, 'le')
 def le(input, other):
-    if use_pyboost():
+    if use_pyboost() and has_le:
         return mindspore.mint.le(input, other)
     return ops.le(input, other)
 
@@ -89,8 +95,9 @@ def less_equal(input, other):
     return le(input, other)
 
 # lt
+has_lt = hasattr(mindspore.mint, 'lt')
 def lt(input, other):
-    if use_pyboost():
+    if use_pyboost() and has_lt:
         return mindspore.mint.lt(input, other)
     return ops.lt(input, other)
 
@@ -99,14 +106,16 @@ def less(input, other):
     return lt(input, other)
 
 # maximum
+has_maximum = hasattr(mindspore.mint, 'maximum')
 def maximum(input, other):
-    if use_pyboost():
+    if use_pyboost() and has_maximum:
         return mindspore.mint.maximum(input, other)
     return ops.maximum(input, other)
 
 # minimum
+has_minimum = hasattr(mindspore.mint, 'minimum')
 def minimum(input, other):
-    if use_pyboost():
+    if use_pyboost() and has_minimum:
         return mindspore.mint.minimum(input, other)
     return ops.minimum(input, other)
 
@@ -120,8 +129,9 @@ def fmin(input, other):
     return ops.fmin(input, other)
 
 # ne
+has_ne = hasattr(mindspore.mint, 'ne')
 def ne(input, other):
-    if use_pyboost():
+    if use_pyboost() and has_ne:
         return mindspore.mint.ne(input, other)
     return ops.ne(input, other)
 
@@ -130,14 +140,16 @@ def not_equal(input, other):
     return ne(input, other)
 
 # sort
+has_sort = hasattr(mindspore.mint, 'sort')
 def sort(input, *, dim=-1, descending=False, stable=False):
-    if use_pyboost():
+    if use_pyboost() and has_sort:
         return mindspore.mint.sort(input, dim=dim, descending=descending, stable=stable)
     return ops.sort(input, dim, descending)
 
 # topk
+has_topk = hasattr(mindspore.mint, 'topk')
 def topk(input, k, dim=-1, largest=True, sorted=True):
-    if use_pyboost():
+    if use_pyboost() and has_topk:
         return mindspore.mint.topk(input, k, dim, largest, sorted)
     return ops.topk(input, k, dim, largest, sorted)