| 206 | get_device_count("gpu") < 1, reason="matmul scale operator is only supported on GPU" |
| 207 | ) |
| 208 | def test_custom_op(): |
| 209 | org_op_list = custom._get_custom_op_list() |
| 210 | assert len(custom._get_custom_op_lib_info()) == 0 |
| 211 | |
| 212 | assert "ElemAddSmoothForward" not in custom._get_custom_op_list() |
| 213 | assert not hasattr(custom, "ElemAddSmoothForward") |
| 214 | assert "MatMulScaleForward" not in custom._get_custom_op_list() |
| 215 | assert not hasattr(custom, "MatMulScaleForward") |
| 216 | |
| 217 | srcs1 = [os.path.join(cur_dir_path, "custom_opsrc", "elem_add.cpp")] |
| 218 | lib_path1 = custom_op_tools.build_and_load( |
| 219 | "elem", |
| 220 | srcs1, |
| 221 | extra_include_paths=extra_include_paths, |
| 222 | build_dir=build_path, |
| 223 | extra_ldflags=extra_ld_flags, |
| 224 | verbose=True, |
| 225 | ) |
| 226 | assert "ElemAddSmoothForward" in custom._get_custom_op_list() |
| 227 | assert hasattr(custom, "ElemAddSmoothForward") |
| 228 | assert lib_path1 in custom._get_custom_op_lib_info() |
| 229 | assert "ElemAddSmoothForward" in custom._get_custom_op_lib_info()[lib_path1] |
| 230 | |
| 231 | srcs2 = [ |
| 232 | os.path.join(cur_dir_path, "custom_opsrc", src) |
| 233 | for src in ["matmul_scale.cpp", "matmul_scale.cu"] |
| 234 | ] |
| 235 | lib_path2 = custom_op_tools.build_and_load( |
| 236 | "matmul", |
| 237 | srcs2, |
| 238 | extra_include_paths=extra_include_paths, |
| 239 | build_dir=build_path, |
| 240 | extra_ldflags=extra_ld_flags, |
| 241 | verbose=True, |
| 242 | ) |
| 243 | |
| 244 | assert "MatMulScaleForward" in custom._get_custom_op_list() |
| 245 | assert hasattr(custom, "MatMulScaleForward") |
| 246 | assert lib_path2 in custom._get_custom_op_lib_info() |
| 247 | assert "MatMulScaleForward" in custom._get_custom_op_lib_info()[lib_path2] |
| 248 | |
| 249 | assert len(custom._get_custom_op_list()) == len(org_op_list) + 4 |
| 250 | |
| 251 | custom.unload(lib_path1) |
| 252 | assert "ElemAddSmoothForward" not in custom._get_custom_op_list() |
| 253 | assert not hasattr(custom, "ElemAddSmoothForward") |
| 254 | assert lib_path1 not in custom._get_custom_op_lib_info() |
| 255 | |
| 256 | custom.unload(lib_path2) |
| 257 | assert "MatMulScaleForward" not in custom._get_custom_op_list() |
| 258 | assert not hasattr(custom, "MatMulScaleForward") |
| 259 | assert lib_path1 not in custom._get_custom_op_lib_info() |
| 260 | |
| 261 | assert len(custom._get_custom_op_lib_info()) == 0 |
| 262 | assert set(custom._get_custom_op_list()) == set(org_op_list) |
| 263 | |
| 264 | custom.load(lib_path2) |
| 265 | assert "MatMulScaleForward" in custom._get_custom_op_list() |