Add TargetModel column and row shift to python bindings (#1560)

Xilinx · Jun 14, 2024 · c5fc35d · c5fc35d
1 parent 413a021
commit c5fc35d
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 2 deletions.
diff --git a/include/aie-c/TargetModel.h b/include/aie-c/TargetModel.h
@@ -140,6 +140,12 @@ aieTargetModelGetMemTileSize(AieTargetModel targetModel);
 /// Returns true if this is an NPU target model.
 MLIR_CAPI_EXPORTED bool aieTargetModelIsNPU(AieTargetModel targetModel);
 
+MLIR_CAPI_EXPORTED uint32_t
+aieTargetModelGetColumnShift(AieTargetModel targetModel);
+
+MLIR_CAPI_EXPORTED uint32_t
+aieTargetModelGetRowShift(AieTargetModel targetModel);
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/lib/CAPI/TargetModel.cpp b/lib/CAPI/TargetModel.cpp
@@ -151,4 +151,12 @@ uint32_t aieTargetModelGetMemTileSize(AieTargetModel targetModel) {
 
 bool aieTargetModelIsNPU(AieTargetModel targetModel) {
   return unwrap(targetModel).isNPU();
+}
+
+uint32_t aieTargetModelGetColumnShift(AieTargetModel targetModel) {
+  return unwrap(targetModel).getColumnShift();
+}
+
+uint32_t aieTargetModelGetRowShift(AieTargetModel targetModel) {
+  return unwrap(targetModel).getRowShift();
 }
diff --git a/python/AIEMLIRModule.cpp b/python/AIEMLIRModule.cpp
@@ -295,7 +295,15 @@ PYBIND11_MODULE(_aie, m) {
       // .def("get_num_dest_shim_mux_connections", int col, int row)
       // .def("get_num_source_shim_mux_connections", int col, int row)
       // .def("is_legal_memtile_connection")
-      .def("is_npu", [](PyAieTargetModel &self) {
-        return aieTargetModelIsNPU(self.get());
+      .def("is_npu",
+           [](PyAieTargetModel &self) {
+             return aieTargetModelIsNPU(self.get());
+           })
+      .def("get_column_shift",
+           [](PyAieTargetModel &self) {
+             return aieTargetModelGetColumnShift(self.get());
+           })
+      .def("get_row_shift", [](PyAieTargetModel &self) {
+        return aieTargetModelGetRowShift(self.get());
       });
 }