[Codegen][Tuner]: improve python binding to query target info (#21812)

On top of #21782, this PR enhances the IREE GPU Python bindings by
adding constructor support to the `TargetInfo` class. Also, add the
support of querying the mma intrinsics supported by arch.

Issue: nod-ai/shark-ai#2048

---------

Signed-off-by: Bangtian Liu <[email protected]>

Author: bangtianliu

compiler/bindings/c/iree/compiler/dialects/iree_gpu.h

@@ -14,6 +14,11 @@
 extern "C" {
 #endif
 
+// This typedef ensures consistency between the C API, C++ implementation, and
+// Python bindings. Update both this typedef and the static assertions if the
+// enum underlying types change.
+typedef uint32_t mma_intrinsic_enum_t;
+
 // The following C API is **NOT STABLE** and likely to change in the future.
 // It mirrors the IREE GPU Dialect which is not stable itself.
 
@@ -56,35 +61,37 @@ MLIR_CAPI_EXPORTED bool ireeAttributeIsAGPUMMAIntrinsicAttr(MlirAttribute attr);
 
 MLIR_CAPI_EXPORTED MlirTypeID ireeGPUMMAIntrinsicAttrGetTypeID(void);
 
-MLIR_CAPI_EXPORTED MlirAttribute ireeGPUMMAIntrinsicAttrGet(MlirContext mlirCtx,
-                                                            uint32_t value);
+MLIR_CAPI_EXPORTED MlirAttribute
+ireeGPUMMAIntrinsicAttrGet(MlirContext mlirCtx, mma_intrinsic_enum_t value);
 
-MLIR_CAPI_EXPORTED uint32_t ireeGPUMMAIntrinsicAttrGetValue(MlirAttribute attr);
+MLIR_CAPI_EXPORTED mma_intrinsic_enum_t
+ireeGPUMMAIntrinsicAttrGetValue(MlirAttribute attr);
 
 MLIR_CAPI_EXPORTED bool ireeAttributeIsAGPUMMAAttr(MlirAttribute attr);
 
 MLIR_CAPI_EXPORTED MlirTypeID ireeGPUMMAAttrGetTypeID(void);
 
 MLIR_CAPI_EXPORTED MlirAttribute ireeGPUMMAAttrGet(MlirContext mlirCtx,
-                                                   uint32_t value);
+                                                   mma_intrinsic_enum_t value);
 
 MLIR_CAPI_EXPORTED bool
 ireeAttributeIsAGPUVirtualMMAIntrinsicAttr(MlirAttribute attr);
 
 MLIR_CAPI_EXPORTED MlirTypeID ireeGPUVirtualMMAIntrinsicAttrGetTypeID(void);
 
-MLIR_CAPI_EXPORTED MlirAttribute
-ireeGPUVirtualMMAIntrinsicAttrGet(MlirContext mlirCtx, uint32_t value);
+MLIR_CAPI_EXPORTED MlirAttribute ireeGPUVirtualMMAIntrinsicAttrGet(
+    MlirContext mlirCtx, mma_intrinsic_enum_t value);
 
-MLIR_CAPI_EXPORTED uint32_t
+MLIR_CAPI_EXPORTED mma_intrinsic_enum_t
 ireeGPUVirtualMMAIntrinsicAttrGetValue(MlirAttribute attr);
 
 MLIR_CAPI_EXPORTED bool ireeAttributeIsAGPUVirtualMMAAttr(MlirAttribute attr);
 
 MLIR_CAPI_EXPORTED MlirTypeID ireeGPUVirtualMMAAttrGetTypeID(void);
 
-MLIR_CAPI_EXPORTED MlirAttribute ireeGPUVirtualMMAAttrGet(MlirContext mlirCtx,
-                                                          uint32_t value);
+MLIR_CAPI_EXPORTED MlirAttribute
+ireeGPUVirtualMMAAttrGet(MlirContext mlirCtx, mma_intrinsic_enum_t value);
+
 struct ireeGPUMMAInfo {
   MlirType aElementType;
   MlirType bElementType;
@@ -148,14 +155,32 @@ struct ireeGPUTargetInfo {
   MlirIdentifier arch;                // E.g., "gfx942".
   MlirAttribute subgroupSizeChoices;  // Subgroup size choices.
   MlirAttribute maxWorkgroupSizes;    // Max threads per X/Y/Z dimension.
-  int64_t maxThreadCountPerWorkgroup; // Max threads per workgroup.
-  int64_t maxWorkgroupMemoryBytes;    // Max workgroup memory.
+  int32_t maxThreadCountPerWorkgroup; // Max threads per workgroup.
+  int32_t maxWorkgroupMemoryBytes;    // Max workgroup memory.
+  MlirAttribute mmaIntrinsics;        // MMA Intrinsics.
 };
 
 // Queries GPU target info from the given `ExecutableTargetAttr` attribute.
 MLIR_CAPI_EXPORTED ireeGPUTargetInfo
 ireeHALExecutableTargetAttrGetGPUTargetInfo(MlirAttribute attr);
 
+MLIR_CAPI_EXPORTED ireeGPUTargetInfo ireeGPUTargetInfoGet(
+    MlirContext mlirCtx, const char *arch, const int32_t *subgroupChoices,
+    size_t numSubgroupChoices, const int32_t *workgroupSizes,
+    size_t numWorkgroupSizes, int32_t threadCount, int32_t memoryBytes,
+    const mma_intrinsic_enum_t *mmaIntrinsics, size_t numMmaIntrinsics);
+
+// Extracts MMA intrinsic values and their virtual status from an ArrayAttr.
+//
+// mmaIntrinsics: Array attribute containing MMA intrinsic attributes.
+// mmaIntrinsicVals: Output array for MMA intrinsic enum values.
+// virtualMmaIntrinsicTags: Output array - 1 if VirtualMMAIntrinsic, 0 if
+// MMAIntrinsic.
+MLIR_CAPI_EXPORTED void
+ireeGPUTargetInfoGetMMAIntrinsics(MlirAttribute mmaIntrinsics,
+                                  mma_intrinsic_enum_t *mmaIntrinsicVals,
+                                  uint8_t *virtualMmaIntrinsicTags);
+
 #ifdef __cplusplus
 }
 #endif

compiler/bindings/python/IREECompilerDialectsModule.cpp

@@ -10,6 +10,7 @@
 #include "iree/compiler/dialects/iree_codegen.h"
 #include "iree/compiler/dialects/iree_gpu.h"
 #include "mlir-c/BuiltinAttributes.h"
+#include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/IR.h"
 #include "mlir/Bindings/Python/Nanobind.h"
 #include "mlir/Bindings/Python/NanobindAdaptors.h"
@@ -20,6 +21,9 @@ static const char *kCodegenModuleImportPath =
 static const char *kGpuModuleImportPath =
     MAKE_MLIR_PYTHON_QUALNAME("dialects.iree_gpu");
 
+static const char *kMMAIntrinsicEnumName = "MMAIntrinsic";
+static const char *kVirtualMMAIntrinsicEnumName = "VirtualMMAIntrinsic";
+
 namespace py = nanobind;
 using namespace nanobind::literals;
 using namespace mlir::python::nanobind_adaptors;
@@ -42,7 +46,7 @@ ireeCodegenQueryMMAIntrinsicsBinding(MlirOperation op) {
   ireeCodegenQueryMMAIntrinsics(op, &numMMAs, mmaIntrinsics.data());
 
   py::object mmaIntrinsicEnum =
-      py::module_::import_(kGpuModuleImportPath).attr("MMAIntrinsic");
+      py::module_::import_(kGpuModuleImportPath).attr(kMMAIntrinsicEnumName);
   std::vector<py::object> mmaList(numMMAs);
   for (size_t i = 0; i < numMMAs; ++i) {
     mmaList[i] = mmaIntrinsicEnum(mmaIntrinsics[i]);
@@ -325,7 +329,7 @@ NB_MODULE(_ireeCompilerDialects, m) {
                                uint32_t rawValue =
                                    ireeGPUMMAIntrinsicAttrGetValue(self);
                                return py::module_::import_(kGpuModuleImportPath)
-                                   .attr("MMAIntrinsic")(rawValue);
+                                   .attr(kMMAIntrinsicEnumName)(rawValue);
                              })
       .def_property_readonly("mma", [](MlirAttribute self) -> MlirAttribute {
         uint32_t value = ireeGPUMMAIntrinsicAttrGetValue(self);
@@ -377,7 +381,7 @@ NB_MODULE(_ireeCompilerDialects, m) {
 
             static py::object virtualEnum =
                 py::module_::import_(kGpuModuleImportPath)
-                    .attr("VirtualMMAIntrinsic");
+                    .attr(kVirtualMMAIntrinsicEnumName);
 
             std::vector<py::object> result;
             for (int64_t val : getIntArrayAttrValues(rawArrayAttr)) {
@@ -404,13 +408,13 @@ NB_MODULE(_ireeCompilerDialects, m) {
           "Gets an #iree_gpu.virtual_mma_intrinsic from parameters.")
       .def_property_readonly("raw_value",
                              ireeGPUVirtualMMAIntrinsicAttrGetValue)
-      .def_property_readonly("value",
-                             [](MlirAttribute self) -> py::object {
-                               uint32_t rawValue =
-                                   ireeGPUVirtualMMAIntrinsicAttrGetValue(self);
-                               return py::module_::import_(kGpuModuleImportPath)
-                                   .attr("VirtualMMAIntrinsic")(rawValue);
-                             })
+      .def_property_readonly(
+          "value",
+          [](MlirAttribute self) -> py::object {
+            uint32_t rawValue = ireeGPUVirtualMMAIntrinsicAttrGetValue(self);
+            return py::module_::import_(kGpuModuleImportPath)
+                .attr(kVirtualMMAIntrinsicEnumName)(rawValue);
+          })
       .def_property_readonly("mma", [](MlirAttribute self) -> MlirAttribute {
         uint32_t value = ireeGPUVirtualMMAIntrinsicAttrGetValue(self);
         return ireeGPUVirtualMMAAttrGet(mlirAttributeGetContext(self), value);
@@ -509,6 +513,43 @@ NB_MODULE(_ireeCompilerDialects, m) {
   //===-------------------------------------------------------------------===//
 
   py::class_<ireeGPUTargetInfo>(iree_gpu_module, "TargetInfo")
+      .def(
+          "__init__",
+          [](ireeGPUTargetInfo *self, MlirContext context,
+             const std::string &arch,
+             const std::vector<int32_t> &subgroupChoices,
+             const std::vector<int32_t> &workgroupSizes, int32_t threadCount,
+             int32_t memoryBytes, const py::list &mmaIntrinsicObjs) {
+            std::vector<mma_intrinsic_enum_t> mmaIntrinsicVals;
+            py::module_ gpuModule = py::module_::import_(kGpuModuleImportPath);
+            py::object mmaIntrinsicClass =
+                gpuModule.attr(kMMAIntrinsicEnumName);
+            py::object virtualMmaIntrinsicClass =
+                gpuModule.attr(kVirtualMMAIntrinsicEnumName);
+
+            for (py::handle item : mmaIntrinsicObjs) {
+              if (!py::isinstance(item, mmaIntrinsicClass) &&
+                  !py::isinstance(item, virtualMmaIntrinsicClass)) {
+                throw py::type_error("All items must be MMA atributes");
+              }
+              mmaIntrinsicVals.push_back(
+                  py::cast<mma_intrinsic_enum_t>(item.attr("value")));
+            }
+
+            *self = ireeGPUTargetInfoGet(
+                context, arch.c_str(), subgroupChoices.data(),
+                subgroupChoices.size(), workgroupSizes.data(),
+                workgroupSizes.size(), threadCount, memoryBytes,
+                mmaIntrinsicVals.data(), mmaIntrinsicVals.size());
+          },
+          "context"_a, "arch"_a, "subgroup_size_choices"_a,
+          "max_workgroup_sizes"_a, "max_thread_count_per_workgroup"_a,
+          "max_workgroup_memory_bytes"_a, "mma_intrinsics"_a = py::list{},
+          "Create a GPUTargetInfo with the given parameters")
+      .def_static(
+          "get_gpu_target_info", &ireeHALExecutableTargetAttrGetGPUTargetInfo,
+          "executable_target_attr"_a,
+          "Get GPU target information from an executable target attribute")
       .def_prop_ro("arch",
                    [](const ireeGPUTargetInfo &self) -> std::string {
                      MlirStringRef strRef = mlirIdentifierStr(self.arch);
@@ -529,12 +570,44 @@ NB_MODULE(_ireeCompilerDialects, m) {
       .def_prop_ro("max_workgroup_memory_bytes",
                    [](const ireeGPUTargetInfo &self) -> int64_t {
                      return self.maxWorkgroupMemoryBytes;
-                   });
+                   })
+      .def_prop_ro(
+          "mma_intrinsics", [](const ireeGPUTargetInfo &self) -> py::list {
+            if (mlirAttributeIsNull(self.mmaIntrinsics) ||
+                !mlirAttributeIsAArray(self.mmaIntrinsics)) {
+              return py::list();
+            }
 
-  iree_gpu_module.def(
-      "get_gpu_target_info", &ireeHALExecutableTargetAttrGetGPUTargetInfo,
-      "Extracts GPU target information from an executable target attribute.",
-      py::arg("executable_target_attr"));
+            size_t numElements =
+                mlirArrayAttrGetNumElements(self.mmaIntrinsics);
+
+            std::vector<mma_intrinsic_enum_t> mmaIntrinsicVals(numElements);
+            // Use uint8_t instead of bool because std::vector<bool> is a
+            // specialized template that doesn't provide .data() method.
+            std::vector<uint8_t> virtualMmaIntrinsicTags(numElements);
+            ireeGPUTargetInfoGetMMAIntrinsics(self.mmaIntrinsics,
+                                              mmaIntrinsicVals.data(),
+                                              virtualMmaIntrinsicTags.data());
+
+            py::list intrinsics;
+            py::module_ gpuModule = py::module_::import_(kGpuModuleImportPath);
+            py::object mmaIntrinsicEnum = gpuModule.attr(kMMAIntrinsicEnumName);
+            py::object virtualMmaIntrinsicEnum =
+                gpuModule.attr(kVirtualMMAIntrinsicEnumName);
+
+            for (size_t i = 0; i < numElements; ++i) {
+              if (virtualMmaIntrinsicTags[i]) {
+                py::object virtualMmaIntrinsic =
+                    virtualMmaIntrinsicEnum(mmaIntrinsicVals[i]);
+                intrinsics.append(virtualMmaIntrinsic);
+                continue;
+              }
+              py::object mmaIntrinsic = mmaIntrinsicEnum(mmaIntrinsicVals[i]);
+              intrinsics.append(mmaIntrinsic);
+            }
+
+            return intrinsics;
+          });
 
   //===-------------------------------------------------------------------===//
   // Binding to utility function getSingleSubgroupLayout