Added maybe_downcast & hardened TT Attrs and Types to include better …

…support (#1253)

* Added maybe_downcast & hardened TT Attrs and Types to include better support

* Removed manual maybe_downcast, added tt_class

* Removed redundant imports

* Lint Fixes

include/ttmlir-c/TTAttrs.h

@@ -75,6 +75,14 @@ MLIR_CAPI_EXPORTED MlirAttribute ttmlirTTOperandConstraintArrayAttrGet(
     MlirContext ctx, uint32_t *OperandConstraints,
     size_t OperandConstraintsSize);
 
+MLIR_CAPI_EXPORTED MlirAttribute ttmlirTTTileSizeAttrGet(MlirContext ctx,
+                                                         int64_t y, int64_t x);
+
+MLIR_CAPI_EXPORTED MlirAttribute ttmlirTTChipPhysicalCoresAttrGet(
+    MlirContext ctx, MlirAttribute *worker, size_t workerSize,
+    MlirAttribute *dram, size_t dramSize, MlirAttribute *eth, size_t ethSize,
+    MlirAttribute *eth_inactive, size_t eth_inactiveSize);
+
 #ifdef __cplusplus
 }
 #endif

include/ttmlir/Bindings/Python/TTMLIRModule.h

@@ -21,9 +21,40 @@
 #include "ttmlir/RegisterAll.h"
 #include "llvm/Support/CommandLine.h"
 
+#include <variant>
+
 namespace py = pybind11;
 
 namespace mlir::ttmlir::python {
+
+template <typename T>
+py::class_<T> tt_attribute_class(py::module &m, const char *class_name) {
+  py::class_<T> cls(m, class_name);
+  cls.def_static("maybe_downcast",
+                 [](MlirAttribute attr) -> std::variant<T, py::object> {
+                   auto res = mlir::dyn_cast<T>(unwrap(attr));
+                   if (res) {
+                     return res;
+                   }
+                   return py::none();
+                 });
+  return cls;
+}
+
+template <typename T>
+py::class_<T> tt_type_class(py::module &m, const char *class_name) {
+  py::class_<T> cls(m, class_name);
+  cls.def_static("maybe_downcast",
+                 [](MlirType type) -> std::variant<T, py::object> {
+                   auto res = mlir::dyn_cast<T>(unwrap(type));
+                   if (res) {
+                     return res;
+                   }
+                   return py::none();
+                 });
+  return cls;
+}
+
 void populateTTModule(py::module &m);
 void populateTTIRModule(py::module &m);
 void populateTTKernelModule(py::module &m);

lib/CAPI/TTAttrs.cpp

@@ -182,4 +182,34 @@ ttmlirTTOperandConstraintArrayAttrGet(MlirContext ctx,
   return wrap(ArrayAttr::get(unwrap(ctx), operandConstraintsArray));
 }
 
+MlirAttribute ttmlirTTTileSizeAttrGet(MlirContext ctx, int64_t y, int64_t x) {
+  return wrap(TileSizeAttr::get(unwrap(ctx), y, x));
+}
+
+MlirAttribute ttmlirTTChipPhysicalCoresAttrGet(
+    MlirContext ctx, MlirAttribute *worker, size_t workerSize,
+    MlirAttribute *dram, size_t dramSize, MlirAttribute *eth, size_t ethSize,
+    MlirAttribute *eth_inactive, size_t eth_inactiveSize) {
+  std::vector<CoreCoordAttr> workerVec, dramVec, ethVec, ethInactiveVec;
+  for (size_t i = 0; i < workerSize; i++) {
+    workerVec.push_back(mlir::cast<CoreCoordAttr>(unwrap(worker[i])));
+  }
+
+  for (size_t i = 0; i < dramSize; i++) {
+    dramVec.push_back(mlir::cast<CoreCoordAttr>(unwrap(dram[i])));
+  }
+
+  for (size_t i = 0; i < ethSize; i++) {
+    ethVec.push_back(mlir::cast<CoreCoordAttr>(unwrap(eth[i])));
+  }
+
+  for (size_t i = 0; i < eth_inactiveSize; i++) {
+    ethInactiveVec.push_back(
+        mlir::cast<CoreCoordAttr>(unwrap(eth_inactive[i])));
+  }
+
+  return wrap(ChipPhysicalCoresAttr::get(unwrap(ctx), workerVec, dramVec,
+                                         ethVec, ethInactiveVec));
+}
+
 } // namespace mlir::tt

python/CMakeLists.txt

@@ -102,6 +102,7 @@ declare_mlir_python_extension(TTMLIRPythonExtensions.Main
 
 set(TTMLIR_PYTHON_SOURCES
   MLIRPythonSources.Core
+  MLIRPythonSources.Dialects.affine
   MLIRPythonSources.Dialects.arith
   MLIRPythonSources.Dialects.func
   MLIRPythonSources.Dialects.tensor