| // Copyright 2024 The IREE Authors | |
| // | |
| // Licensed under the Apache License v2.0 with LLVM Exceptions. | |
| // See https://llvm.org/LICENSE.txt for license information. | |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |
| #include <cassert> | |
| #include <cstdint> | |
| #include <optional> | |
| #include <string> | |
| #include <vector> | |
| #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-c/Target/ExportSMTLIB.h" | |
| #include "mlir-c/Target/LLVMIR.h" | |
| #include "mlir/Bindings/Python/Nanobind.h" | |
| #include "mlir/Bindings/Python/NanobindAdaptors.h" | |
| #include "mlir/CAPI/Support.h" | |
| 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; | |
| static std::vector<MlirOperation> | |
| ireeCodegenGetExecutableVariantOpsBinding(MlirModule module) { | |
| size_t numOps = 0; | |
| ireeCodegenGetExecutableVariantOps(module, &numOps, nullptr); | |
| std::vector<MlirOperation> ops(numOps); | |
| ireeCodegenGetExecutableVariantOps(module, &numOps, ops.data()); | |
| return ops; | |
| } | |
| static std::vector<MlirOperation> | |
| ireeCodegenGetTunerRootOpsBinding(MlirModule module) { | |
| size_t numOps = 0; | |
| ireeCodegenGetTunerRootOps(module, &numOps, nullptr); | |
| std::vector<MlirOperation> ops(numOps); | |
| ireeCodegenGetTunerRootOps(module, &numOps, ops.data()); | |
| return ops; | |
| } | |
| static std::vector<int64_t> getIntArrayAttrValues(MlirAttribute attr) { | |
| if (mlirAttributeIsNull(attr) || !mlirAttributeIsAArray(attr)) { | |
| return {}; | |
| } | |
| std::vector<int64_t> result; | |
| size_t n = mlirArrayAttrGetNumElements(attr); | |
| result.reserve(n); | |
| for (size_t i = 0; i < n; ++i) { | |
| MlirAttribute elem = mlirArrayAttrGetElement(attr, i); | |
| int64_t val = mlirIntegerAttrGetValueInt(elem); | |
| result.push_back(val); | |
| } | |
| return result; | |
| } | |
| NB_MODULE(_ireeCompilerDialects, m) { | |
| m.doc() = "iree-compiler dialects python extension"; | |
| auto iree_codegen_module = | |
| m.def_submodule("iree_codegen", "iree_codegen dialect bindings"); | |
| //===-------------------------------------------------------------------===// | |
| // Codegen pipeline attrs | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_codegen_module, "VMVXPipelineAttr", | |
| ireeAttributeIsACodegenVMVXPipelineAttr, | |
| ireeCodegenVMVXPipelineAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, MlirContext ctx) { | |
| return ireeCodegenVMVXPipelineAttrGet(ctx); | |
| }, | |
| "cls"_a, "ctx"_a = py::none(), | |
| "Gets an #iree_codegen.vmvx_pipeline attribute."); | |
| mlir_attribute_subclass( | |
| iree_codegen_module, "TransformDialectCodegenPipelineAttr", | |
| ireeAttributeIsACodegenTransformDialectCodegenPipelineAttr, | |
| ireeCodegenTransformDialectCodegenPipelineAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, MlirContext ctx) { | |
| return ireeCodegenTransformDialectCodegenPipelineAttrGet(ctx); | |
| }, | |
| "cls"_a, "ctx"_a = py::none(), | |
| "Gets an #iree_codegen.transform_dialect_codegen attribute."); | |
| mlir_attribute_subclass(iree_codegen_module, "NoPipelineAttr", | |
| ireeAttributeIsACodegenNoPipelineAttr, | |
| ireeCodegenNoPipelineAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, MlirContext ctx) { | |
| return ireeCodegenNoPipelineAttrGet(ctx); | |
| }, | |
| "cls"_a, "ctx"_a = py::none(), | |
| "Gets an #iree_codegen.no_pipeline attribute."); | |
| mlir_attribute_subclass(iree_codegen_module, "PassPipelineAttr", | |
| ireeAttributeIsACodegenPassPipelineAttr, | |
| ireeCodegenPassPipelineAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, std::string pipeline, MlirContext ctx) { | |
| MlirStringRef pipelineRef = | |
| mlirStringRefCreate(pipeline.data(), pipeline.size()); | |
| return ireeCodegenPassPipelineAttrGet(ctx, pipelineRef); | |
| }, | |
| "cls"_a, "pipeline"_a, "ctx"_a = py::none(), | |
| "Gets an #iree_codegen.pass_pipeline attribute.") | |
| .def_property_readonly("pipeline", | |
| ireeCodegenPassPipelineAttrGetPipeline); | |
| //===-------------------------------------------------------------------===// | |
| // CodegenTranslationInfoAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_codegen_module, "TranslationInfoAttr", | |
| ireeAttributeIsACodegenTranslationInfoAttr, | |
| ireeCodegenTranslationInfoAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, MlirAttribute passPipeline, | |
| std::optional<MlirAttribute> codegenSpec, | |
| std::optional<std::vector<int64_t>> workgroupSize, | |
| std::optional<int64_t> subgroupSize, | |
| std::optional<MlirAttribute> configuration, MlirContext ctx) { | |
| ireeCodegenTranslationInfoParameters parameters = {}; | |
| parameters.passPipeline = passPipeline; | |
| parameters.codegenSpec = | |
| codegenSpec.value_or(mlirAttributeGetNull()); | |
| if (workgroupSize.has_value()) { | |
| parameters.workgroupSize = workgroupSize->data(); | |
| parameters.numWorkgroupSizeElements = workgroupSize->size(); | |
| } | |
| parameters.subgroupSize = subgroupSize.value_or(0); | |
| parameters.configuration = | |
| configuration.value_or(mlirAttributeGetNull()); | |
| return ireeCodegenTranslationInfoAttrGet(ctx, parameters); | |
| }, | |
| "cls"_a, "pass_pipeline"_a, "codegen_spec"_a = py::none(), | |
| "workgroup_size"_a = py::none(), "subgroup_size"_a = py::none(), | |
| "configuration"_a = py::none(), py::kw_only(), "ctx"_a = py::none(), | |
| "Gets an #iree_codegen.translation_info from parameters.") | |
| .def_property_readonly( | |
| "pass_pipeline", | |
| [](MlirAttribute self) -> MlirAttribute { | |
| auto parameters = ireeCodegenTranslationInfoAttrGetParameters(self); | |
| return parameters.passPipeline; | |
| }) | |
| .def_property_readonly( | |
| "codegen_spec", | |
| [](MlirAttribute self) -> std::optional<MlirAttribute> { | |
| auto parameters = ireeCodegenTranslationInfoAttrGetParameters(self); | |
| if (mlirAttributeIsNull(parameters.codegenSpec)) { | |
| return std::nullopt; | |
| } | |
| return parameters.codegenSpec; | |
| }) | |
| .def_property_readonly( | |
| "workgroup_size", | |
| [](MlirAttribute self) -> std::vector<int64_t> { | |
| auto parameters = ireeCodegenTranslationInfoAttrGetParameters(self); | |
| return {parameters.workgroupSize, | |
| parameters.workgroupSize + | |
| parameters.numWorkgroupSizeElements}; | |
| }) | |
| .def_property_readonly( | |
| "subgroup_size", | |
| [](MlirAttribute self) -> int64_t { | |
| auto parameters = ireeCodegenTranslationInfoAttrGetParameters(self); | |
| return parameters.subgroupSize; | |
| }) | |
| .def_property_readonly( | |
| "configuration", | |
| [](MlirAttribute self) -> std::optional<MlirAttribute> { | |
| auto parameters = ireeCodegenTranslationInfoAttrGetParameters(self); | |
| if (mlirAttributeIsNull(parameters.configuration)) { | |
| return std::nullopt; | |
| } | |
| return parameters.configuration; | |
| }); | |
| //===-------------------------------------------------------------------===// | |
| // CodegenCompilationInfoAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_codegen_module, "CompilationInfoAttr", | |
| ireeAttributeIsACodegenCompilationInfoAttr, | |
| ireeCodegenCompilationInfoAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, MlirAttribute loweringConfig, | |
| MlirAttribute translationInfo, MlirContext ctx) { | |
| ireeCodegenCompilationInfoParameters parameters = {}; | |
| parameters.loweringConfig = loweringConfig; | |
| parameters.translationInfo = translationInfo; | |
| return ireeCodegenCompilationInfoAttrGet(ctx, parameters); | |
| }, | |
| "cls"_a, "lowering_config"_a, "translation_info"_a, | |
| "ctx"_a = py::none(), | |
| "Gets an #iree_codegen.compilation_info from parameters.") | |
| .def_property_readonly( | |
| "lowering_config", | |
| [](MlirAttribute self) -> MlirAttribute { | |
| auto parameters = ireeCodegenCompilationInfoAttrGetParameters(self); | |
| return parameters.loweringConfig; | |
| }) | |
| .def_property_readonly( | |
| "translation_info", [](MlirAttribute self) -> MlirAttribute { | |
| auto parameters = ireeCodegenCompilationInfoAttrGetParameters(self); | |
| return parameters.translationInfo; | |
| }); | |
| //===-------------------------------------------------------------------===// | |
| // CodegenRootOpAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_codegen_module, "RootOpAttr", | |
| ireeAttributeIsACodegenRootOpAttr, | |
| ireeCodegenRootOpAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, int64_t set, MlirContext ctx) { | |
| return ireeCodegenRootOpAttrGet(ctx, set); | |
| }, | |
| "cls"_a, "set"_a = 0, py::kw_only(), "ctx"_a = py::none(), | |
| "Gets an #iree_codegen.root_op attribute.") | |
| .def_property_readonly("set", ireeCodegenRootOpAttrGetSet); | |
| //===-------------------------------------------------------------------===// | |
| // CodegenIntKnobAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_codegen_module, "IntKnobAttr", | |
| ireeAttributeIsACodegenIntKnobAttr, | |
| ireeCodegenIntKnobAttrGetTypeID) | |
| .def_property_readonly("name", [](MlirAttribute self) -> MlirStringRef { | |
| return mlirStringAttrGetValue(ireeCodegenIntKnobAttrGetName(self)); | |
| }); | |
| //===-------------------------------------------------------------------===// | |
| // CodegenOneOfKnobAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_codegen_module, "OneOfKnobAttr", | |
| ireeAttributeIsACodegenOneOfKnobAttr, | |
| ireeCodegenOneOfKnobAttrGetTypeID) | |
| .def_property_readonly("name", | |
| [](MlirAttribute self) -> MlirStringRef { | |
| return mlirStringAttrGetValue( | |
| ireeCodegenOneOfKnobAttrGetName(self)); | |
| }) | |
| .def_property_readonly("options", [](MlirAttribute self) { | |
| intptr_t n = 0; | |
| ireeCodegenOneOfKnobAttrGetOptions(self, &n, nullptr); | |
| std::vector<MlirAttribute> opts(n); | |
| ireeCodegenOneOfKnobAttrGetOptions(self, &n, opts.data()); | |
| return opts; | |
| }); | |
| //===--------------------------------------------------------------------===// | |
| auto iree_gpu_module = | |
| m.def_submodule("iree_gpu", "iree_gpu dialect bindings"); | |
| //===-------------------------------------------------------------------===// | |
| // GPUReorderWorkgroupsStrategyAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_gpu_module, "ReorderWorkgroupsStrategyAttr", | |
| ireeAttributeIsAGPUReorderWorkgroupsStrategyAttr, | |
| ireeGPUReorderWorkgroupsStrategyAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, uint32_t value, MlirContext ctx) { | |
| return ireeGPUReorderWorkgroupsStrategyAttrGet(ctx, value); | |
| }, | |
| "cls"_a, "value"_a, "ctx"_a = py::none(), | |
| "Gets an #iree_gpu.reorder_workgroups_strategy from parameters.") | |
| .def_property_readonly("raw_value", | |
| ireeGPUReorderWorkgroupsStrategyAttrGetValue) | |
| .def_property_readonly("value", [](MlirAttribute self) -> py::object { | |
| uint32_t rawValue = ireeGPUReorderWorkgroupsStrategyAttrGetValue(self); | |
| return py::module_::import_(kGpuModuleImportPath) | |
| .attr("ReorderWorkgroupsStrategy")(rawValue); | |
| }); | |
| //===-------------------------------------------------------------------===// | |
| // GPUPipelineAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_gpu_module, "PipelineAttr", | |
| ireeAttributeIsAGPUPipelineAttr, | |
| ireeGPUPipelineAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, uint32_t value, MlirContext ctx) { | |
| return ireeGPUPipelineAttrGet(ctx, value); | |
| }, | |
| "cls"_a, "value"_a, "ctx"_a = py::none(), | |
| "Gets an #iree_gpu.pipeline from parameters.") | |
| .def_property_readonly("raw_value", ireeGPUPipelineAttrGetValue) | |
| .def_property_readonly("value", [](MlirAttribute self) -> py::object { | |
| uint32_t rawValue = ireeGPUPipelineAttrGetValue(self); | |
| return py::module_::import_(kGpuModuleImportPath) | |
| .attr("LoweringPipeline")(rawValue); | |
| }); | |
| //===-------------------------------------------------------------------===// | |
| // GPUPipelineOptionsAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_gpu_module, "PipelineOptionsAttr", | |
| ireeAttributeIsAGPUPipelineOptionsAttr, | |
| ireeGPUPipelineOptionsAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, std::optional<int64_t> prefetchNumStages, | |
| std::optional<bool> noReduceSharedMemoryBankConflicts, | |
| std::optional<bool> useIgemmConvolution, | |
| std::optional<MlirAttribute> reorderWorkgroupsStrategy, | |
| MlirContext ctx) { | |
| return ireeGPUPipelineOptionsAttrGet( | |
| ctx, | |
| prefetchNumStages.has_value() ? &*prefetchNumStages : nullptr, | |
| noReduceSharedMemoryBankConflicts.has_value() | |
| ? &*noReduceSharedMemoryBankConflicts | |
| : nullptr, | |
| useIgemmConvolution.has_value() ? &*useIgemmConvolution | |
| : nullptr, | |
| reorderWorkgroupsStrategy.has_value() | |
| ? &*reorderWorkgroupsStrategy | |
| : nullptr); | |
| }, | |
| "cls"_a, "prefetch_num_stages"_a = py::none(), | |
| "no_reduce_shared_memory_bank_conflicts"_a = py::none(), | |
| "use_igemm_convolution"_a = py::none(), | |
| "reorder_workgroups_strategy"_a = py::none(), py::kw_only(), | |
| "ctx"_a = py::none(), | |
| "Gets an #iree_gpu.pipeline_options from parameters.") | |
| .def_property_readonly( | |
| "prefetch_num_stages", | |
| [](MlirAttribute self) -> std::optional<int64_t> { | |
| auto attr = ireeGPUPipelineOptionsAttrGetPrefetchNumStages(self); | |
| if (!mlirAttributeIsNull(attr)) { | |
| return mlirIntegerAttrGetValueInt(attr); | |
| } | |
| return std::nullopt; | |
| }) | |
| .def_property_readonly( | |
| "no_reduce_shared_memory_bank_conflicts", | |
| [](MlirAttribute self) -> std::optional<bool> { | |
| auto attr = | |
| ireeGPUPipelineOptionsAttrGetNoReduceSharedMemoryBankConflicts( | |
| self); | |
| if (!mlirAttributeIsNull(attr)) { | |
| return mlirBoolAttrGetValue(attr); | |
| } | |
| return std::nullopt; | |
| }) | |
| .def_property_readonly( | |
| "use_igemm_convolution", | |
| [](MlirAttribute self) -> std::optional<bool> { | |
| auto attr = ireeGPUPipelineOptionsAttrGetUseIgemmConvolution(self); | |
| if (!mlirAttributeIsNull(attr)) { | |
| return mlirBoolAttrGetValue(attr); | |
| } | |
| return std::nullopt; | |
| }) | |
| .def_property_readonly( | |
| "reorder_workgroups_strategy", | |
| [](MlirAttribute self) -> std::optional<MlirAttribute> { | |
| auto attr = | |
| ireeGPUPipelineOptionsAttrGetReorderWorkgroupsStrategy(self); | |
| if (!mlirAttributeIsNull(attr)) { | |
| return attr; | |
| } | |
| return std::nullopt; | |
| }); | |
| //===-------------------------------------------------------------------===// | |
| // GPUMMAIntrinsicAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_gpu_module, "MMAIntrinsicAttr", | |
| ireeAttributeIsAGPUMMAIntrinsicAttr, | |
| ireeGPUMMAIntrinsicAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, uint32_t value, MlirContext ctx) { | |
| return ireeGPUMMAIntrinsicAttrGet(ctx, value); | |
| }, | |
| "cls"_a, "value"_a, "ctx"_a = py::none(), | |
| "Gets an #iree_gpu.mma_intrinsic from parameters.") | |
| .def_property_readonly("raw_value", ireeGPUMMAIntrinsicAttrGetValue) | |
| .def_property_readonly("value", | |
| [](MlirAttribute self) -> py::object { | |
| uint32_t rawValue = | |
| ireeGPUMMAIntrinsicAttrGetValue(self); | |
| return py::module_::import_(kGpuModuleImportPath) | |
| .attr(kMMAIntrinsicEnumName)(rawValue); | |
| }) | |
| .def_property_readonly("mma", [](MlirAttribute self) -> MlirAttribute { | |
| uint32_t value = ireeGPUMMAIntrinsicAttrGetValue(self); | |
| return ireeGPUMMAAttrGet(mlirAttributeGetContext(self), value, | |
| /*colMajor=*/false); | |
| }); | |
| //===-------------------------------------------------------------------===// | |
| // GPUMMAAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_gpu_module, "MMAAttr", | |
| ireeAttributeIsAGPUMMAAttr, ireeGPUMMAAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, uint32_t value, bool colMajor, | |
| MlirContext ctx) { | |
| return ireeGPUMMAAttrGet(ctx, value, colMajor); | |
| }, | |
| // col_major defaults to false for backward compatibility. | |
| "cls"_a, "value"_a, "col_major"_a = false, py::kw_only(), | |
| "ctx"_a = py::none(), "Gets an #iree_gpu.mma from parameters.") | |
| .def_property_readonly("col_major", ireeGPUMMAAttrGetColMajor) | |
| .def_property_readonly( | |
| "abc_element_types", | |
| [](MlirAttribute self) -> py::tuple { | |
| ireeGPUMMAInfo info = ireeGPUMMAAttrGetInfo(self); | |
| return py::make_tuple(info.aElementType, info.bElementType, | |
| info.cElementType); | |
| }) | |
| .def_property_readonly( | |
| "abc_vector_types", | |
| [](MlirAttribute self) -> py::tuple { | |
| ireeGPUMMAInfo info = ireeGPUMMAAttrGetInfo(self); | |
| return py::make_tuple(info.aVectorType, info.bVectorType, | |
| info.cVectorType); | |
| }) | |
| .def_property_readonly( | |
| "mnk_shape", | |
| [](MlirAttribute self) -> py::tuple { | |
| ireeGPUMMAInfo info = ireeGPUMMAAttrGetInfo(self); | |
| return py::make_tuple(info.mElements, info.nElements, | |
| info.kElements); | |
| }) | |
| .def( | |
| "get_virtual_intrinsics", | |
| [](MlirAttribute self) { | |
| MlirAttribute rawArrayAttr = | |
| ireeGPUMMAAttrGetVirtualMMAIntrinsic(self); | |
| if (mlirAttributeIsNull(rawArrayAttr)) { | |
| return std::vector<py::object>{}; | |
| } | |
| static py::object virtualEnum = | |
| py::module_::import_(kGpuModuleImportPath) | |
| .attr(kVirtualMMAIntrinsicEnumName); | |
| std::vector<py::object> result; | |
| for (int64_t val : getIntArrayAttrValues(rawArrayAttr)) { | |
| result.push_back(virtualEnum(static_cast<uint32_t>(val))); | |
| } | |
| return result; | |
| }, | |
| "Returns a list of virtual intrinsics associated with this " | |
| "MMAAttr."); | |
| //===-------------------------------------------------------------------===// | |
| // GPUVirtualMMAIntrinsicAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_gpu_module, "VirtualMMAIntrinsicAttr", | |
| ireeAttributeIsAGPUVirtualMMAIntrinsicAttr, | |
| ireeGPUVirtualMMAIntrinsicAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, uint32_t value, MlirContext ctx) { | |
| return ireeGPUVirtualMMAIntrinsicAttrGet(ctx, value); | |
| }, | |
| "cls"_a, "value"_a, "ctx"_a = py::none(), | |
| "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(kVirtualMMAIntrinsicEnumName)(rawValue); | |
| }) | |
| .def_property_readonly("mma", [](MlirAttribute self) -> MlirAttribute { | |
| uint32_t value = ireeGPUVirtualMMAIntrinsicAttrGetValue(self); | |
| return ireeGPUVirtualMMAAttrGet(mlirAttributeGetContext(self), value, | |
| /*colMajor=*/false); | |
| }); | |
| //===-------------------------------------------------------------------===// | |
| // GPUVirtualMMAAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_gpu_module, "VirtualMMAAttr", | |
| ireeAttributeIsAGPUVirtualMMAAttr, | |
| ireeGPUVirtualMMAAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, uint32_t value, bool colMajor, | |
| MlirContext ctx) { | |
| return ireeGPUVirtualMMAAttrGet(ctx, value, colMajor); | |
| }, | |
| // col_major defaults to false for backward compatibility. | |
| "cls"_a, "value"_a, "col_major"_a = false, py::kw_only(), | |
| "ctx"_a = py::none(), "Gets an #iree_gpu.virtualmma from parameters.") | |
| .def_property_readonly("col_major", ireeGPUVirtualMMAAttrGetColMajor) | |
| .def_property_readonly( | |
| "abc_element_types", | |
| [](MlirAttribute self) -> py::tuple { | |
| ireeGPUMMAInfo info = ireeGPUMMAAttrGetInfo(self); | |
| return py::make_tuple(info.aElementType, info.bElementType, | |
| info.cElementType); | |
| }) | |
| .def_property_readonly( | |
| "abc_vector_types", | |
| [](MlirAttribute self) -> py::tuple { | |
| ireeGPUMMAInfo info = ireeGPUMMAAttrGetInfo(self); | |
| return py::make_tuple(info.aVectorType, info.bVectorType, | |
| info.cVectorType); | |
| }) | |
| .def_property_readonly("mnk_shape", [](MlirAttribute self) -> py::tuple { | |
| ireeGPUMMAInfo info = ireeGPUMMAAttrGetInfo(self); | |
| return py::make_tuple(info.mElements, info.nElements, info.kElements); | |
| }); | |
| //===-------------------------------------------------------------------===// | |
| // GPULoweringConfigAttr | |
| //===-------------------------------------------------------------------===// | |
| mlir_attribute_subclass(iree_gpu_module, "LoweringConfigAttr", | |
| ireeAttributeIsAGPULoweringConfigAttr, | |
| ireeGPULoweringConfigAttrGetTypeID) | |
| .def_classmethod( | |
| "get", | |
| [](const py::object &, MlirAttribute attributeDictionary, | |
| MlirContext ctx) { | |
| return ireeGPULoweringConfigAttrGet(ctx, attributeDictionary); | |
| }, | |
| "cls"_a, "value"_a, "ctx"_a = py::none(), | |
| "Gets an #iree_gpu.lowering_config from parameters.") | |
| .def_property_readonly("attributes", | |
| ireeGPULoweringConfigAttrGetAttributes) | |
| .def_property_readonly( | |
| "workgroup_tile_sizes", | |
| [](MlirAttribute self) -> std::vector<int64_t> { | |
| auto tilesizes = ireeGPULoweringConfigAttrGetTileSizes(self); | |
| return getIntArrayAttrValues(tilesizes.workgroupAttr); | |
| }) | |
| .def_property_readonly( | |
| "reduction_tile_sizes", | |
| [](MlirAttribute self) -> std::vector<int64_t> { | |
| auto tilesizes = ireeGPULoweringConfigAttrGetTileSizes(self); | |
| return getIntArrayAttrValues(tilesizes.reductionAttr); | |
| }) | |
| .def_property_readonly( | |
| "subgroup_basis", | |
| [](MlirAttribute self) | |
| -> std::tuple<std::vector<int64_t>, std::vector<int64_t>> { | |
| ireeGPUSubgroupBasisInfo basisInfo = | |
| ireeGPULoweringConfigAttrGetSubgroupBasis(self); | |
| std::vector<int64_t> counts; | |
| std::vector<int64_t> mapping; | |
| if (!mlirAttributeIsNull(basisInfo.countsAttr)) { | |
| counts = getIntArrayAttrValues(basisInfo.countsAttr); | |
| } | |
| if (!mlirAttributeIsNull(basisInfo.mappingAttr)) { | |
| mapping = getIntArrayAttrValues(basisInfo.mappingAttr); | |
| } | |
| return std::tuple(counts, mapping); | |
| }) | |
| .def_property_readonly( | |
| "mma_kind", [](MlirAttribute self) -> std::optional<MlirAttribute> { | |
| auto attr = ireeGPULoweringConfigAttrGetMmaKind(self); | |
| if (!mlirAttributeIsNull(attr)) { | |
| return attr; | |
| } | |
| return std::nullopt; | |
| }); | |
| //===-------------------------------------------------------------------===// | |
| // Binding to query target info | |
| //===-------------------------------------------------------------------===// | |
| 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, int32_t wgpCount, int32_t simdsPerWgp, | |
| 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 attributes"); | |
| } | |
| mmaIntrinsicVals.push_back( | |
| py::cast<mma_intrinsic_enum_t>(item.attr("value"))); | |
| } | |
| if (wgpCount < 0) { | |
| throw py::value_error("workgroup_count must be non-negative"); | |
| } | |
| *self = ireeGPUTargetInfoGet( | |
| context, arch.c_str(), subgroupChoices.data(), | |
| subgroupChoices.size(), workgroupSizes.data(), | |
| workgroupSizes.size(), threadCount, memoryBytes, wgpCount, | |
| simdsPerWgp, 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, "workgroup_count"_a, | |
| "simds_per_workgroup"_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) -> MlirStringRef { | |
| return mlirIdentifierStr(self.arch); | |
| }) | |
| .def_prop_ro("subgroup_size_choices", | |
| [](const ireeGPUTargetInfo &self) -> std::vector<int64_t> { | |
| return getIntArrayAttrValues(self.subgroupSizeChoices); | |
| }) | |
| .def_prop_ro("max_thread_count_per_workgroup", | |
| [](const ireeGPUTargetInfo &self) -> int64_t { | |
| return self.maxThreadCountPerWorkgroup; | |
| }) | |
| .def_prop_ro("max_workgroup_sizes", | |
| [](const ireeGPUTargetInfo &self) -> std::vector<int64_t> { | |
| return getIntArrayAttrValues(self.maxWorkgroupSizes); | |
| }) | |
| .def_prop_ro("max_workgroup_memory_bytes", | |
| [](const ireeGPUTargetInfo &self) -> int64_t { | |
| return self.maxWorkgroupMemoryBytes; | |
| }) | |
| .def_prop_ro("workgroup_count", | |
| [](const ireeGPUTargetInfo &self) -> int64_t { | |
| return self.wgpCount; | |
| }) | |
| .def_prop_ro("simds_per_workgroup", | |
| [](const ireeGPUTargetInfo &self) -> int64_t { | |
| return self.simdsPerWgp; | |
| }) | |
| .def_prop_ro( | |
| "mma_intrinsics", [](const ireeGPUTargetInfo &self) -> py::list { | |
| if (mlirAttributeIsNull(self.mmaIntrinsics) || | |
| !mlirAttributeIsAArray(self.mmaIntrinsics)) { | |
| return py::list(); | |
| } | |
| 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 | |
| //===-------------------------------------------------------------------===// | |
| py::class_<ireeGPUMMASingleSubgroupLayout>(iree_gpu_module, | |
| "GPUMMASingleSubgroupLayout") | |
| .def_prop_ro("outer", | |
| [](const ireeGPUMMASingleSubgroupLayout &self) { | |
| return getIntArrayAttrValues(self.outer); | |
| }) | |
| .def_prop_ro("thread", | |
| [](const ireeGPUMMASingleSubgroupLayout &self) { | |
| return getIntArrayAttrValues(self.thread); | |
| }) | |
| .def_prop_ro("tstrides", | |
| [](const ireeGPUMMASingleSubgroupLayout &self) { | |
| return getIntArrayAttrValues(self.tstrides); | |
| }) | |
| .def_prop_ro("element", [](const ireeGPUMMASingleSubgroupLayout &self) { | |
| return getIntArrayAttrValues(self.element); | |
| }); | |
| iree_gpu_module.def( | |
| "get_single_subgroup_layout", ireeGPUGetSingleSubgroupLayout, | |
| "Returns the single subgroup layout (element, thread, outer, " | |
| "tstrides) for a given MMA or VirtualMMA intrinsic and fragment. ", | |
| py::arg("attr"), py::arg("fragment")); | |
| //===-------------------------------------------------------------------===// | |
| // Binding to XOR shuffle utility functions | |
| //===-------------------------------------------------------------------===// | |
| iree_gpu_module.def( | |
| "get_xor_shuffle_bounds", | |
| [](MlirAttribute mmaIntrinsic, | |
| int operandIndex) -> std::optional<std::tuple<int64_t, int64_t>> { | |
| int64_t minAccessElems = 0, totalTileElems = 0; | |
| if (ireeGPUGetXorShuffleBounds(mmaIntrinsic, operandIndex, | |
| &minAccessElems, &totalTileElems)) { | |
| return std::tuple(minAccessElems, totalTileElems); | |
| } | |
| return std::nullopt; | |
| }, | |
| "Returns the bounds for valid XOR shuffle parameters (min_access_elems, " | |
| "total_tile_elems) for the given MMA intrinsic and operand index. See " | |
| "GPUUtils for sweep semantics. Returns (min_access_elems, " | |
| "total_tile_elems) or None on failure.", | |
| py::arg("mmaIntrinsic"), py::arg("operand_index")); | |
| iree_gpu_module.def( | |
| "is_xor_shuffle_valid", ireeGPUIsXORShuffleValid, | |
| "Returns true if the XOR shuffle is valid for the given parameters.", | |
| py::arg("num_row_elems"), py::arg("num_access_elems"), | |
| py::arg("total_tile_elems")); | |
| //===-------------------------------------------------------------------===// | |
| // Binding to utility function getExecutableVariantOps | |
| //===-------------------------------------------------------------------===// | |
| iree_codegen_module.def( | |
| "get_executable_variant_ops", &ireeCodegenGetExecutableVariantOpsBinding, | |
| "Gets the executable variant operations from a module.", | |
| py::arg("module")); | |
| //===-------------------------------------------------------------------===// | |
| // Binding to utility function ireeCodegenConvertConstraintsOpToSMTLIB | |
| //===-------------------------------------------------------------------===// | |
| iree_codegen_module.def( | |
| "convert_constraints_op_to_smtlib", | |
| [](MlirOperation op, bool emitReset) -> MlirStringRef { | |
| MlirAttribute strAttr = | |
| ireeCodegenConvertConstraintsOpToSMTLIB(op, emitReset); | |
| if (mlirAttributeIsNull(strAttr)) { | |
| throw std::runtime_error("SMT-LIB export failed"); | |
| } | |
| return mlirStringAttrGetValue(strAttr); | |
| }, | |
| "Convert an iree_codegen.smt.constraints op to an SMT-LIB string.", | |
| py::arg("constraints_op"), py::arg("emit_reset") = false); | |
| iree_codegen_module.def( | |
| "materialize_compilation_info", | |
| [](MlirOperation op, py::dict assignments) -> MlirAttribute { | |
| std::vector<MlirStringRef> nameRefs; | |
| std::vector<int64_t> values; | |
| nameRefs.reserve(assignments.size()); | |
| values.reserve(assignments.size()); | |
| for (auto [key, value] : assignments) { | |
| Py_ssize_t nameSize = 0; | |
| const char *nameData = PyUnicode_AsUTF8AndSize(key.ptr(), &nameSize); | |
| if (!nameData) { | |
| throw py::python_error(); | |
| } | |
| nameRefs.push_back( | |
| mlirStringRefCreate(nameData, static_cast<size_t>(nameSize))); | |
| values.push_back(py::cast<int64_t>(value)); | |
| } | |
| MlirAttribute diagnosticMessage = mlirAttributeGetNull(); | |
| MlirStringRef configurationAttrName = | |
| mlirStringRefCreateFromCString("compilation_info"); | |
| MlirAttribute attr = | |
| ireeCodegenMaterializeConfigurationAttrFromConstraintsOp( | |
| op, configurationAttrName, nameRefs.size(), nameRefs.data(), | |
| values.data(), &diagnosticMessage); | |
| if (mlirAttributeIsNull(attr)) { | |
| if (!mlirAttributeIsNull(diagnosticMessage)) { | |
| throw std::runtime_error( | |
| unwrap(mlirStringAttrGetValue(diagnosticMessage)).str()); | |
| } | |
| throw std::runtime_error( | |
| "compilation_info materialization from constraints failed"); | |
| } | |
| return attr; | |
| }, | |
| "Materialize a compilation_info attr from a constraints op and flat knob " | |
| "assignment dictionary.", | |
| py::arg("constraints_op"), py::arg("assignments")); | |
| iree_codegen_module.def( | |
| "materialize_configuration_attr", | |
| [](MlirOperation op, std::string configurationAttrName, | |
| py::dict assignments) -> MlirAttribute { | |
| std::vector<MlirStringRef> nameRefs; | |
| std::vector<int64_t> values; | |
| nameRefs.reserve(assignments.size()); | |
| values.reserve(assignments.size()); | |
| for (auto [key, value] : assignments) { | |
| Py_ssize_t nameSize = 0; | |
| const char *nameData = PyUnicode_AsUTF8AndSize(key.ptr(), &nameSize); | |
| if (!nameData) { | |
| throw py::python_error(); | |
| } | |
| nameRefs.push_back( | |
| mlirStringRefCreate(nameData, static_cast<size_t>(nameSize))); | |
| values.push_back(py::cast<int64_t>(value)); | |
| } | |
| MlirAttribute diagnosticMessage = mlirAttributeGetNull(); | |
| MlirStringRef configurationAttrNameRef = mlirStringRefCreate( | |
| configurationAttrName.data(), configurationAttrName.size()); | |
| MlirAttribute attr = | |
| ireeCodegenMaterializeConfigurationAttrFromConstraintsOp( | |
| op, configurationAttrNameRef, nameRefs.size(), nameRefs.data(), | |
| values.data(), &diagnosticMessage); | |
| if (mlirAttributeIsNull(attr)) { | |
| if (!mlirAttributeIsNull(diagnosticMessage)) { | |
| throw std::runtime_error( | |
| unwrap(mlirStringAttrGetValue(diagnosticMessage)).str()); | |
| } | |
| throw std::runtime_error( | |
| "configuration attr materialization from constraints failed"); | |
| } | |
| return attr; | |
| }, | |
| "Materialize a named configuration attr from a constraints op and flat " | |
| "knob " | |
| "assignment dictionary.", | |
| py::arg("constraints_op"), py::arg("configuration_attr_name"), | |
| py::arg("assignments")); | |
| //===-------------------------------------------------------------------===// | |
| // Binding to utility function ireeCodegenGetTunerRootOps | |
| //===-------------------------------------------------------------------===// | |
| iree_codegen_module.def("get_tuner_root_ops", | |
| &ireeCodegenGetTunerRootOpsBinding, | |
| "Get the operations marked with the tuner root op " | |
| "attribute from a module.", | |
| py::arg("module")); | |
| //===-------------------------------------------------------------------===// | |
| // Binding to utility function ireeCodegenGetAttentionOpDetail | |
| //===-------------------------------------------------------------------===// | |
| py::class_<ireeCodegenAttentionOpDetail>(iree_codegen_module, | |
| "AttentionOpDetail") | |
| .def_prop_ro("batch_dims", | |
| [](const ireeCodegenAttentionOpDetail &self) { | |
| return getIntArrayAttrValues(self.batch); | |
| }) | |
| .def_prop_ro("m_dims", | |
| [](const ireeCodegenAttentionOpDetail &self) { | |
| return getIntArrayAttrValues(self.m); | |
| }) | |
| .def_prop_ro("k1_dims", | |
| [](const ireeCodegenAttentionOpDetail &self) { | |
| return getIntArrayAttrValues(self.k1); | |
| }) | |
| .def_prop_ro("k2_dims", | |
| [](const ireeCodegenAttentionOpDetail &self) { | |
| return getIntArrayAttrValues(self.k2); | |
| }) | |
| .def_prop_ro("n_dims", | |
| [](const ireeCodegenAttentionOpDetail &self) { | |
| return getIntArrayAttrValues(self.n); | |
| }) | |
| .def_prop_ro("domain_rank", [](const ireeCodegenAttentionOpDetail &self) { | |
| return self.domainRank; | |
| }); | |
| iree_codegen_module.def( | |
| "get_attention_op_detail", &ireeCodegenGetAttentionOpDetail, | |
| "Infers the structure of an attention operation from affine indexing " | |
| "maps.", | |
| py::arg("q"), py::arg("k"), py::arg("v"), py::arg("o")); | |
| iree_codegen_module.def( | |
| "isa_attention_op", &ireeCodegenMlirOperationIsACodegenAttentionOp, | |
| "Checks if the given operation is an IREE LinalgExt attention op.", | |
| py::arg("op")); | |
| //===-------------------------------------------------------------------===// | |
| // Binding to utility function ireeCodegenGetIGEMMGenericConvDetails | |
| //===-------------------------------------------------------------------===// | |
| py::class_<ireeCodegenIGEMMGenericConvDetails>(iree_codegen_module, | |
| "IGEMMGenericConvDetails") | |
| .def_prop_ro("igemm_contraction_maps", | |
| [](const ireeCodegenIGEMMGenericConvDetails &self) { | |
| return self.igemmContractionMaps; | |
| }) | |
| .def_prop_ro("igemm_loop_bounds", | |
| [](const ireeCodegenIGEMMGenericConvDetails &self) { | |
| return getIntArrayAttrValues(self.igemmLoopBounds); | |
| }) | |
| .def_prop_ro("igemm_loop_iterators", | |
| [](const ireeCodegenIGEMMGenericConvDetails &self) { | |
| return self.igemmLoopIterators; | |
| }) | |
| .def_prop_ro("im2col_output_perm", | |
| [](const ireeCodegenIGEMMGenericConvDetails &self) { | |
| return getIntArrayAttrValues(self.im2colOutputPerm); | |
| }) | |
| .def_prop_ro( | |
| "filter_reassoc_indices", | |
| [](const ireeCodegenIGEMMGenericConvDetails &self) | |
| -> std::vector<std::vector<int64_t>> { | |
| MlirAttribute attr = self.filterReassocIndices; | |
| assert(!mlirAttributeIsNull(attr) && mlirAttributeIsAArray(attr) && | |
| "filterReassocIndices should be a valid ArrayAttr"); | |
| size_t n = mlirArrayAttrGetNumElements(attr); | |
| std::vector<std::vector<int64_t>> result; | |
| result.reserve(n); | |
| for (size_t i = 0; i < n; ++i) { | |
| MlirAttribute innerArrayAttr = mlirArrayAttrGetElement(attr, i); | |
| result.push_back(getIntArrayAttrValues(innerArrayAttr)); | |
| } | |
| return result; | |
| }) | |
| .def_prop_ro("is_output_channel_first", | |
| [](const ireeCodegenIGEMMGenericConvDetails &self) { | |
| return self.isOutputChannelFirst; | |
| }) | |
| .def_prop_ro( | |
| "conv_to_igemm_dim_map", | |
| [](const ireeCodegenIGEMMGenericConvDetails &self) -> py::dict { | |
| py::dict result; | |
| MlirAttribute attr = self.convToIgemmDimMap; | |
| assert(!mlirAttributeIsNull(attr) && mlirAttributeIsAArray(attr) && | |
| "convToIgemmDimMap should be a valid ArrayAttr"); | |
| size_t n = mlirArrayAttrGetNumElements(attr); | |
| for (size_t i = 0; i < n; ++i) { | |
| MlirAttribute pairAttr = mlirArrayAttrGetElement(attr, i); | |
| assert(mlirAttributeIsAArray(pairAttr) && | |
| mlirArrayAttrGetNumElements(pairAttr) == 2 && | |
| "Each pair should be [conv_dim, igemm_dim]"); | |
| MlirAttribute keyAttr = mlirArrayAttrGetElement(pairAttr, 0); | |
| MlirAttribute valueAttr = mlirArrayAttrGetElement(pairAttr, 1); | |
| int64_t key = mlirIntegerAttrGetValueInt(keyAttr); | |
| int64_t value = mlirIntegerAttrGetValueInt(valueAttr); | |
| result[py::int_(key)] = py::int_(value); | |
| } | |
| return result; | |
| }); | |
| iree_codegen_module.def( | |
| "get_igemm_generic_conv_details", | |
| [](MlirOperation op) | |
| -> std::optional<ireeCodegenIGEMMGenericConvDetails> { | |
| if (!ireeCodegenHasIGEMMGenericConvDetails(op)) { | |
| return std::nullopt; | |
| } | |
| return ireeCodegenGetIGEMMGenericConvDetails(op); | |
| }, | |
| "Gets IGEMM details for a linalg operation. " | |
| "Returns None if failed to infer IGEMM convolution details.", | |
| py::arg("linalg_op")); | |
| //===-------------------------------------------------------------------===// | |
| // Binding to utility function ireeCodegenGetScaledContractionDetails | |
| //===-------------------------------------------------------------------===// | |
| iree_codegen_module.def("isa_scaled_contraction_op", | |
| &ireeCodegenMlirOperationIsAScaledContractionOp, | |
| "Checks if the given operation is an IREE LinalgExt " | |
| "scaled contraction op.", | |
| py::arg("op")); | |
| //===-------------------------------------------------------------------===// | |
| // Binding to struct ireeCodegenScaledContractionDimensions | |
| //===-------------------------------------------------------------------===// | |
| py::class_<ireeCodegenScaledContractionDimensions>( | |
| iree_codegen_module, "ScaledContractionDimensions") | |
| .def_prop_ro("batch", | |
| [](const ireeCodegenScaledContractionDimensions &self) { | |
| return getIntArrayAttrValues(self.batch); | |
| }) | |
| .def_prop_ro("m", | |
| [](const ireeCodegenScaledContractionDimensions &self) { | |
| return getIntArrayAttrValues(self.m); | |
| }) | |
| .def_prop_ro("n", | |
| [](const ireeCodegenScaledContractionDimensions &self) { | |
| return getIntArrayAttrValues(self.n); | |
| }) | |
| .def_prop_ro("k", | |
| [](const ireeCodegenScaledContractionDimensions &self) { | |
| return getIntArrayAttrValues(self.k); | |
| }) | |
| .def_prop_ro("kB", | |
| [](const ireeCodegenScaledContractionDimensions &self) { | |
| return getIntArrayAttrValues(self.kB); | |
| }); | |
| iree_codegen_module.def( | |
| "infer_scaled_contraction_dimensions", | |
| &ireeCodegenInferScaledContractionDimensions, | |
| "Infers the scaled contraction dimensions for a given operation.", | |
| py::arg("op")); | |
| } |