(idtype)
| 1199 | ) |
| 1200 | @parametrize_idtype |
| 1201 | def test_pin_memory_(idtype): |
| 1202 | # TODO: rewrite this test case to accept different graphs so we |
| 1203 | # can test reverse graph and batched graph |
| 1204 | g = create_test_heterograph(idtype) |
| 1205 | g.nodes["user"].data["h"] = F.ones((3, 5)) |
| 1206 | g.nodes["game"].data["i"] = F.ones((2, 5)) |
| 1207 | g.edges["plays"].data["e"] = F.ones((4, 4)) |
| 1208 | g = g.to(F.cpu()) |
| 1209 | assert not g.is_pinned() |
| 1210 | |
| 1211 | # unpin an unpinned CPU graph, directly return |
| 1212 | g.unpin_memory_() |
| 1213 | assert not g.is_pinned() |
| 1214 | assert g.device == F.cpu() |
| 1215 | |
| 1216 | # pin a CPU graph |
| 1217 | g.pin_memory_() |
| 1218 | assert g.is_pinned() |
| 1219 | assert g.device == F.cpu() |
| 1220 | assert g.nodes["user"].data["h"].is_pinned() |
| 1221 | assert g.nodes["game"].data["i"].is_pinned() |
| 1222 | assert g.edges["plays"].data["e"].is_pinned() |
| 1223 | assert F.context(g.nodes["user"].data["h"]) == F.cpu() |
| 1224 | assert F.context(g.nodes["game"].data["i"]) == F.cpu() |
| 1225 | assert F.context(g.edges["plays"].data["e"]) == F.cpu() |
| 1226 | for ntype in g.ntypes: |
| 1227 | assert F.context(g.batch_num_nodes(ntype)) == F.cpu() |
| 1228 | for etype in g.canonical_etypes: |
| 1229 | assert F.context(g.batch_num_edges(etype)) == F.cpu() |
| 1230 | |
| 1231 | # it's fine to clone with new formats, but new graphs are not pinned |
| 1232 | # >>> g.formats() |
| 1233 | # {'created': ['coo'], 'not created': ['csr', 'csc']} |
| 1234 | assert not g.formats("csc").is_pinned() |
| 1235 | assert not g.formats("csr").is_pinned() |
| 1236 | # 'coo' formats is already created and thus not cloned |
| 1237 | assert g.formats("coo").is_pinned() |
| 1238 | |
| 1239 | # pin a pinned graph, directly return |
| 1240 | g.pin_memory_() |
| 1241 | assert g.is_pinned() |
| 1242 | assert g.device == F.cpu() |
| 1243 | |
| 1244 | # unpin a pinned graph |
| 1245 | g.unpin_memory_() |
| 1246 | assert not g.is_pinned() |
| 1247 | assert g.device == F.cpu() |
| 1248 | |
| 1249 | g1 = g.to(F.cuda()) |
| 1250 | |
| 1251 | # unpin an unpinned GPU graph, directly return |
| 1252 | g1.unpin_memory_() |
| 1253 | assert not g1.is_pinned() |
| 1254 | assert g1.device == F.cuda() |
| 1255 | |
| 1256 | # error pinning a GPU graph |
| 1257 | with pytest.raises(DGLError): |
| 1258 | g1.pin_memory_() |
nothing calls this directly
no test coverage detected