Skip to content
This repository was archived by the owner on Apr 28, 2023. It is now read-only.

Commit 5faf98a

Browse files
author
Jules Pondard
committed
Add getDict function to PyBinds
Given a CudaMappingOptions instance, return a Python dictionary that represents all the given options and their values.
1 parent 4807746 commit 5faf98a

File tree

1 file changed

+32
-0
lines changed

1 file changed

+32
-0
lines changed

tensor_comprehensions/pybinds/tclib.cc

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -654,6 +654,38 @@ PYBIND11_MODULE(tclib, m) {
654654
return str;
655655
},
656656
"Returns the CudaMappingOptions as a human-readable string")
657+
.def(
658+
"getDict",
659+
[](tc::CudaMappingOptions& instance) {
660+
py::dict rv;
661+
rv["outerScheduleFusionStrategy"] = FusionStrategy_Name(
662+
instance.generic.outerScheduleOptions.proto.fusion_strategy());
663+
if (instance.generic.proto.has_intra_tile_schedule_options())
664+
rv["intraTileScheduleFusionStrategy"] =
665+
FusionStrategy_Name(instance.generic.intraTileScheduleOptions
666+
.proto.fusion_strategy());
667+
rv["fixParametersBeforeScheduling"] =
668+
instance.generic.proto.fix_parameters_before_scheduling();
669+
if (instance.generic.proto.has_tiling())
670+
rv["tile"] = instance.generic.tiling.extractVector();
671+
if (instance.generic.proto.has_unroll())
672+
rv["unroll"] = instance.generic.proto.unroll();
673+
rv["tileImperfectlyNested"] =
674+
instance.generic.proto.tile_imperfectly_nested();
675+
rv["matchLibraryCalls"] =
676+
instance.generic.proto.match_library_calls();
677+
rv["mapToThreads"] = instance.block.extractVector();
678+
rv["mapToBlocks"] = instance.grid.extractVector();
679+
rv["useSharedMemory"] = instance.proto().use_shared_memory();
680+
rv["usePrivateMemory"] = instance.proto().use_private_memory();
681+
rv["unrollCopyShared"] = instance.proto().unroll_copy_shared();
682+
rv["useReadOnlyCache"] = instance.proto().use_readonly_cache();
683+
if (instance.proto().has_max_shared_memory())
684+
rv["maxSharedMemory"] = instance.proto().max_shared_memory();
685+
rv["privateDepth"] = instance.proto().private_depth();
686+
return rv;
687+
},
688+
"Returns a dictionary with the CudaMappingOptions")
657689
.def(
658690
"serialize",
659691
[](tc::CudaMappingOptions& instance) {

0 commit comments

Comments
 (0)