Addings c-api and py-apis

Author: sdasgup3

mlir/lib/Bindings/Python/DialectQuant.cpp

@@ -9,10 +9,11 @@
 #include <cstdint>
 #include <vector>
 
+#include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/Dialect/Quant.h"
 #include "mlir-c/IR.h"
-#include "mlir/Bindings/Python/NanobindAdaptors.h"
 #include "mlir/Bindings/Python/Nanobind.h"
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
 
 namespace nb = nanobind;
 using namespace llvm;
@@ -292,7 +293,7 @@ static void populateDialectQuantSubmodule(const nb::module_ &m) {
       mlirTypeIsAUniformQuantizedSubChannelType, quantizedType.get_class());
   uniformQuantizedSubChannelType.def_classmethod(
       "get",
-      [](py::object cls, unsigned flags, MlirType storageType,
+      [](nb::object cls, unsigned flags, MlirType storageType,
          MlirType expressedType, MlirAttribute scales, MlirAttribute zeroPoints,
          std::vector<int32_t> quantizedDimensions,
          std::vector<int64_t> blockSizes, int64_t storageTypeMin,
@@ -305,10 +306,10 @@ static void populateDialectQuantSubmodule(const nb::module_ &m) {
       },
       "Gets an instance of UniformQuantizedSubChannel in the same context as "
       "the provided storage type.",
-      py::arg("cls"), py::arg("flags"), py::arg("storage_type"),
-      py::arg("expressed_type"), py::arg("scales"), py::arg("zero_points"),
-      py::arg("quantized_dimensions"), py::arg("block_sizes"),
-      py::arg("storage_type_min"), py::arg("storage_type_max"));
+      nb::arg("cls"), nb::arg("flags"), nb::arg("storage_type"),
+      nb::arg("expressed_type"), nb::arg("scales"), nb::arg("zero_points"),
+      nb::arg("quantized_dimensions"), nb::arg("block_sizes"),
+      nb::arg("storage_type_min"), nb::arg("storage_type_max"));
   uniformQuantizedSubChannelType.def_property_readonly(
       "quantized_dimensions",
       [](MlirType type) {

mlir/lib/CAPI/Dialect/Quant.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir-c/Dialect/Quant.h"
+#include "mlir-c/BuiltinAttributes.h"
 #include "mlir/CAPI/Registration.h"
 #include "mlir/Dialect/Quant/IR/Quant.h"
 #include "mlir/Dialect/Quant/IR/QuantTypes.h"

mlir/python/mlir/_mlir_libs/_mlir/dialects/quant.pyi

@@ -3,7 +3,7 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 
-from mlir.ir import Type
+from mlir.ir import DenseElementsAttr, Type
 
 __all__ = [
   "QuantizedType",
@@ -109,6 +109,26 @@ class UniformQuantizedPerAxisType(QuantizedType):
   @property
   def is_fixed_point(self) -> bool: ...
 
+class UniformQuantizedSubChannelType(QuantizedType):
+
+  @classmethod
+  def get(cls, flags: int, storage_type: Type, expressed_type: Type,
+          scales: DenseElementsAttr, zero_points: DenseElementsAttr,
+          quantized_dimensions: list[int], block_sizes: list[int],
+          storage_type_min: int, storage_type_max: int):
+    ...
+
+  @property
+  def quantized_dimensions(self) -> list[int]: ...
+
+  @property
+  def block_sizes(self) -> list[int]: ...
+
+  @property
+  def scales(self) -> DenseElementsAttr: ...
+
+  @property
+  def zero_points(self) -> DenseElementsAttr: ...
 
 def CalibratedQuantizedType(QuantizedType):
 

mlir/test/CAPI/quant.c

@@ -10,6 +10,7 @@
 // RUN: mlir-capi-quant-test 2>&1 | FileCheck %s
 
 #include "mlir-c/Dialect/Quant.h"
+#include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 #include "mlir-c/IR.h"
 
@@ -357,6 +358,7 @@ int main(void) {
   testAnyQuantizedType(ctx);
   testUniformType(ctx);
   testUniformPerAxisType(ctx);
+  testUniformSubChannelType(ctx);
   testCalibratedType(ctx);
   mlirContextDestroy(ctx);
   return EXIT_SUCCESS;

mlir/test/python/dialects/quant.py

@@ -1,5 +1,6 @@
 # RUN: %PYTHON %s | FileCheck %s
 
+import numpy as np
 from mlir.ir import *
 from mlir.dialects import quant
 
@@ -18,21 +19,28 @@ def test_type_hierarchy():
         any = Type.parse("!quant.any<i8<-8:7>:f32>")
         uniform = Type.parse("!quant.uniform<i8<-8:7>:f32, 0.99872:127>")
         per_axis = Type.parse("!quant.uniform<i8:f32:1, {2.0e+2,0.99872:120}>")
+        sub_channel = Type.parse(
+            "!quant.uniform<i8:f32:{0:1,1:2}, {{2.0:10, 3.0:20}, {4.0:30, 5.0:40}}>"
+        )
         calibrated = Type.parse("!quant.calibrated<f32<-0.998:1.2321>>")
 
         assert not quant.QuantizedType.isinstance(i8)
         assert quant.QuantizedType.isinstance(any)
         assert quant.QuantizedType.isinstance(uniform)
         assert quant.QuantizedType.isinstance(per_axis)
+        assert quant.QuantizedType.isinstance(sub_channel)
         assert quant.QuantizedType.isinstance(calibrated)
 
         assert quant.AnyQuantizedType.isinstance(any)
         assert quant.UniformQuantizedType.isinstance(uniform)
         assert quant.UniformQuantizedPerAxisType.isinstance(per_axis)
+        assert quant.UniformQuantizedSubChannelType.isinstance(sub_channel)
         assert quant.CalibratedQuantizedType.isinstance(calibrated)
 
         assert not quant.AnyQuantizedType.isinstance(uniform)
         assert not quant.UniformQuantizedType.isinstance(per_axis)
+        assert not quant.UniformQuantizedType.isinstance(sub_channel)
+        assert not quant.UniformQuantizedPerAxisType.isinstance(sub_channel)
 
 
 # CHECK-LABEL: TEST: test_any_quantized_type
@@ -121,6 +129,47 @@ def test_uniform_per_axis_type():
         assert per_axis == Type.parse("!quant.uniform<i8:f32:1, {2.0e+2,0.99872:120}>")
 
 
+# CHECK-LABEL: TEST: test_uniform_sub_channel_type
+@run
+def test_uniform_sub_channel_type():
+    with Context():
+        i8 = IntegerType.get_signless(8)
+        f32 = F32Type.get()
+        sub_channel = quant.UniformQuantizedSubChannelType.get(
+            quant.QuantizedType.FLAG_SIGNED,
+            i8,
+            f32,
+            DenseElementsAttr.get(
+                np.asarray([2.0, 3.0, 4.0, 5.0], np.float32).reshape(2, 2)
+            ),
+            DenseElementsAttr.get(np.asarray([10, 20, 30, 40], np.int8).reshape(2, 2)),
+            [0, 1],
+            [1, 2],
+            storage_type_min=quant.QuantizedType.default_minimum_for_integer(
+                is_signed=True, integral_width=8
+            ),
+            storage_type_max=quant.QuantizedType.default_maximum_for_integer(
+                is_signed=True, integral_width=8
+            ),
+        )
+
+        # CHECK: quantized dimensions: [0, 1]
+        print(f"quantized dimensions: {sub_channel.quantized_dimensions}")
+        # CHECK: block sizes: [1, 2]
+        print(f"block sizes: {sub_channel.block_sizes}")
+        # CHECK: scales: {{\[}}[2. 3.]
+        # CHECK:               [4. 5.]]
+        print(f"scales: {np.asarray(sub_channel.scales)}")
+        # CHECK: zero-points: {{\[}}[10 20]
+        # CHECK:                    [30 40]]
+        print(f"zero-points: {np.asarray(sub_channel.zero_points)}")
+        # CHECK: !quant.uniform<i8:f32:{0:1,1:2}, {{\{}}{2.000000e+00:10, 3.000000e+00:20}, {4.000000e+00:30, 5.000000e+00:40}}>
+        print(sub_channel)
+        assert sub_channel == Type.parse(
+            "!quant.uniform<i8:f32:{0:1,1:2},{{2.0:10, 3.0:20}, {4.0:30, 5.0:40}}>"
+        )
+
+
 # CHECK-LABEL: TEST: test_calibrated_type
 @run
 def test_calibrated_type():