Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / symmetric_memory / aclshmem_ms / gen_ops_prim.py: 0%

1# Copyright 2023 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15 

16"""Operators definition generated by gen_ops.py, includes primitive classes.""" 

17 

18from mindspore.ops.primitive import Primitive, prim_arg_register 

19from mindspore.ops import signature as sig 

20from mindspore.common import dtype as mstype 

21from mindspore.common._decorator import deprecated 

22from mindspore.ops._primitive_cache import _get_cache_prim 

23from mindspore.ops._utils.arg_dtype_cast import type_it 

24from mindspore.ops._utils.arg_handler import * 

25from mindspore._c_expression import OpDtype 

26from mindspore.graph.jit_context import jit_context 

27from mindspore._checkparam import is_stub_tensor 

28import aclshmem_ms 

29 

30class Custom_put_mem(Primitive): 

31 r""" 

32 .. code-block:: 

33 

34 prim = ops.PutMem() 

35 out = prim(target, target_offset, src, src_offset, size, target_pe) 

36 

37 is equivalent to 

38 

39 .. code-block:: 

40 

41 ops.put_mem(target, target_offset, src, src_offset, size, target_pe) 

42 

43 Refer to :func:`mindspore.ops.put_mem` for more details. 

44 """ 

45 __mindspore_signature__ = ( 

46 sig.make_sig('target', sig.sig_rw.RW_WRITE), 

47 sig.make_sig('target_offset'), 

48 sig.make_sig('src'), 

49 sig.make_sig('src_offset'), 

50 sig.make_sig('size'), 

51 sig.make_sig('target_pe'), 

52 ) 

53 

54 @prim_arg_register 

55 def __init__(self, op_func): 

56 

57 self.custom_op_func = op_func 

58 self.add_prim_attr("side_effect_mem", True) 

59 

60 def __call__(self, target, target_offset, src, src_offset, size, target_pe): 

61 return self.custom_op_func(target, target_offset, src, src_offset, size, _scalar_tensor_to_int('PutMem', 'target_pe', target_pe)) 

62 

63put_mem_op=Custom_put_mem(aclshmem_ms.put_mem) 

64 

65 

66class Custom_get_mem(Primitive): 

67 r""" 

68 .. code-block:: 

69 

70 prim = ops.GetMem() 

71 out = prim(target, target_offset, src, src_offset, size, target_pe, non_blocking) 

72 

73 is equivalent to 

74 

75 .. code-block:: 

76 

77 ops.get_mem(target, target_offset, src, src_offset, size, target_pe, non_blocking) 

78 

79 Refer to :func:`mindspore.ops.get_mem` for more details. 

80 """ 

81 __mindspore_signature__ = ( 

82 sig.make_sig('target'), 

83 sig.make_sig('target_offset'), 

84 sig.make_sig('src'), 

85 sig.make_sig('src_offset'), 

86 sig.make_sig('size'), 

87 sig.make_sig('target_pe', default=0), 

88 sig.make_sig('non_blocking', default=False), 

89 ) 

90 

91 @prim_arg_register 

92 def __init__(self, op_func): 

93 

94 self.custom_op_func = op_func 

95 

96 def __call__(self, target, target_offset, src, src_offset, size, target_pe=0, non_blocking=False): 

97 return self.custom_op_func(target, target_offset, src, src_offset, size, _scalar_tensor_to_int('GetMem', 'target_pe', target_pe), non_blocking) 

98 

99get_mem_op=Custom_get_mem(aclshmem_ms.get_mem) 

100 

101 

102class Custom_signal_op(Primitive): 

103 r""" 

104 .. code-block:: 

105 

106 prim = ops.SignalOp() 

107 out = prim(signal, signal_offset, signal_value, signal_op, target_pe) 

108 

109 is equivalent to 

110 

111 .. code-block:: 

112 

113 ops.signal_op(signal, signal_offset, signal_value, signal_op, target_pe) 

114 

115 Refer to :func:`mindspore.ops.signal_op` for more details. 

116 """ 

117 __mindspore_signature__ = ( 

118 sig.make_sig('signal'), 

119 sig.make_sig('signal_offset'), 

120 sig.make_sig('signal_value'), 

121 sig.make_sig('signal_op', default=0), 

122 sig.make_sig('target_pe', default=0), 

123 ) 

124 

125 @prim_arg_register 

126 def __init__(self, op_func): 

127 

128 self.custom_op_func = op_func 

129 

130 def __call__(self, signal, signal_offset, signal_value, signal_op=0, target_pe=0): 

131 return self.custom_op_func(signal, signal_offset, signal_value, _scalar_tensor_to_int('SignalOp', 'signal_op', signal_op), _scalar_tensor_to_int('SignalOp', 'target_pe', target_pe)) 

132 

133signal_op_op=Custom_signal_op(aclshmem_ms.signal_op) 

134 

135 

136class Custom_put_mem_signal(Primitive): 

137 r""" 

138 .. code-block:: 

139 

140 prim = ops.PutMemSignal() 

141 out = prim(target, target_offset, src, src_offset, size, signal, signal_offset, signal_value, signal_op, target_pe) 

142 

143 is equivalent to 

144 

145 .. code-block:: 

146 

147 ops.put_mem_signal(target, target_offset, src, src_offset, size, signal, signal_offset, signal_value, signal_op, target_pe) 

148 

149 Refer to :func:`mindspore.ops.put_mem_signal` for more details. 

150 """ 

151 __mindspore_signature__ = ( 

152 sig.make_sig('target'), 

153 sig.make_sig('target_offset'), 

154 sig.make_sig('src'), 

155 sig.make_sig('src_offset'), 

156 sig.make_sig('size'), 

157 sig.make_sig('signal'), 

158 sig.make_sig('signal_offset'), 

159 sig.make_sig('signal_value'), 

160 sig.make_sig('signal_op', default=0), 

161 sig.make_sig('target_pe', default=0), 

162 ) 

163 

164 @prim_arg_register 

165 def __init__(self, op_func): 

166 

167 self.custom_op_func = op_func 

168 

169 def __call__(self, target, target_offset, src, src_offset, size, signal, signal_offset, signal_value, signal_op=0, target_pe=0): 

170 return self.custom_op_func(target, target_offset, src, src_offset, size, signal, signal_offset, signal_value, _scalar_tensor_to_int('PutMemSignal', 'signal_op', signal_op), _scalar_tensor_to_int('PutMemSignal', 'target_pe', target_pe)) 

171 

172put_mem_signal_op=Custom_put_mem_signal(aclshmem_ms.put_mem_signal) 

173 

174 

175class Custom_signal_wait_until(Primitive): 

176 r""" 

177 .. code-block:: 

178 

179 prim = ops.SignalWaitUntil() 

180 out = prim(depend_target, signal, signal_offset, compare_value, compare_op) 

181 

182 is equivalent to 

183 

184 .. code-block:: 

185 

186 ops.signal_wait_until(depend_target, signal, signal_offset, compare_value, compare_op) 

187 

188 Refer to :func:`mindspore.ops.signal_wait_until` for more details. 

189 """ 

190 __mindspore_signature__ = ( 

191 sig.make_sig('depend_target'), 

192 sig.make_sig('signal'), 

193 sig.make_sig('signal_offset'), 

194 sig.make_sig('compare_value'), 

195 sig.make_sig('compare_op', default=0), 

196 ) 

197 

198 @prim_arg_register 

199 def __init__(self, op_func): 

200 

201 self.custom_op_func = op_func 

202 

203 def __call__(self, depend_target, signal, signal_offset, compare_value, compare_op=0): 

204 return self.custom_op_func(depend_target, signal, signal_offset, compare_value, _scalar_tensor_to_int('SignalWaitUntil', 'compare_op', compare_op)) 

205 

206signal_wait_until_op=Custom_signal_wait_until(aclshmem_ms.signal_wait_until)