Refactor layout implementation (#491)

Summary:
TODO

Test Plan:
TODO

Reviewers:

Subscribers:

Tasks:

Tags:

Author: jerryzh168

torchao/dtypes/__init__.py

@@ -1,12 +1,21 @@
 from .nf4tensor import NF4Tensor, to_nf4
 # from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor
 from .uint4 import UInt4Tensor
-from .affine_quantized_tensor import AffineQuantizedTensor, to_affine_quantized
+from .affine_quantized_tensor import (
+    AffineQuantizedTensor,
+    to_affine_quantized,
+    LayoutType,
+    PlainLayoutType,
+    TensorCoreTiledLayoutType,
+)
 
 __all__ = [
     "NF4Tensor",
     "to_nf4",
     "UInt4Tensor"
     "AffineQuantizedTensor",
     "to_affine_quantized",
+    "LayoutType",
+    "PlainLayoutType",
+    "TensorCoreTiledLayoutType",
 ]

torchao/dtypes/affine_quantized_tensor.py

@@ -20,10 +20,35 @@
     _ATEN_OP_OR_TORCH_FN_TABLE,
     _register_layout_cls,
     _get_layout_tensor_constructor,
+    LayoutType,
 )
+from typing import ClassVar
+from dataclasses import dataclass
 
 aten = torch.ops.aten
 
+@dataclass(frozen=True)
+class PlainLayoutType(LayoutType):
+    pass
+
+@dataclass(frozen=True)
+class TensorCoreTiledLayoutType(LayoutType):
+    inner_k_tiles: int = 8
+
+    def pre_process(self, input: torch.Tensor) -> torch.Tensor:
+        orig_out_features, orig_in_features = input.shape
+        in_features = find_multiple(orig_in_features, 1024)
+        out_features = find_multiple(orig_out_features, 8)
+        input = torch.nn.functional.pad(
+            input,
+            (0, in_features - orig_in_features, 0, out_features - orig_out_features),
+        )
+        return input
+
+    def extra_repr(self):
+        return f"inner_k_tiles={self.inner_k_tiles}"
+
+
 def _aqt_is_int8(aqt):
     """Check if an AffineQuantizedTensor is int8 quantized Tensor"""
     return (
@@ -52,10 +77,10 @@ class AQTLayout(torch.Tensor):
     """
     Base class for the layout tensor for `AffineQuantizedTensor`
     """
-    # this should be set for each layout class during registration
-    extended_layout: Optional[str] = None
+    def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        pass
 
-    def get_plain() -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def get_layout_type(self) -> LayoutType:
         pass
 
     @classmethod
@@ -64,9 +89,15 @@ def from_plain(
         int_data: torch.Tensor,
         scale: torch.Tensor,
         zero_point: torch.Tensor,
+        layout_type: LayoutType,
     ):
         pass
 
+    def __repr__(self):
+        int_data, scale, zero_point = self.get_plain()
+        layout_type = self.get_layout_type()
+        return f"{self.__class__.__name__}(int_data={int_data}, scale={scale}, zero_point={zero_point}, layout_type={layout_type})"
+
     def _get_to_kwargs(self, *args, **kwargs):
         device, dtype, _, memory_format = torch._C._nn._parse_to(*args, **kwargs)
         device = self.device if device is None else device
@@ -194,30 +225,17 @@ def from_float(
         zero_point_dtype: Optional[torch.dtype] = None,
         preserve_zero: bool = True,
         zero_point_domain: ZeroPointDomain = ZeroPointDomain.INT,
-        extended_layout: str = "plain",
-        # TODO: this is only for "tensor_core_tiled", need to figure out
-        # the proper API for this arg
-        inner_k_tiles: Optional[int] = None,
+        layout_type: LayoutType = PlainLayoutType(),
     ):
         original_shape = input_float.shape
-        if extended_layout == "tensor_core_tiled":
-            orig_out_features, orig_in_features = input_float.shape
-            in_features = find_multiple(orig_in_features, 1024)
-            out_features = find_multiple(orig_out_features, 8)
-            input_float = torch.nn.functional.pad(
-                input_float,
-                (0, in_features - orig_in_features, 0, out_features - orig_out_features),
-            )
+        input_float = layout_type.pre_process(input_float)
 
         scale, zero_point = choose_qparams_affine(input_float, mapping_type, block_size, target_dtype, quant_min, quant_max, eps, scale_dtype, zero_point_dtype, preserve_zero, zero_point_domain)
         int_data = quantize_affine(input_float, block_size, scale, zero_point, target_dtype, quant_min, quant_max, zero_point_domain)
+        int_data = layout_type.post_process(int_data)
 
-        layout_cls_ctr = get_layout_tensor_constructor(extended_layout)
-        # TODO: this is temporary, need to come up with the proper UX
-        if extended_layout == "tensor_core_tiled":
-            layout_tensor = layout_cls_ctr(int_data, scale, zero_point, inner_k_tiles)
-        else:
-            layout_tensor = layout_cls_ctr(int_data, scale, zero_point)
+        layout_tensor_ctr = get_layout_tensor_constructor(type(layout_type))
+        layout_tensor = layout_tensor_ctr(int_data, scale, zero_point, layout_type)
         return cls(
             layout_tensor,
             block_size,
@@ -229,8 +247,8 @@ def from_float(
         )
 
     @property
-    def extended_layout(self) -> str:
-        return self.layout_tensor.extended_layout
+    def layout_type(self) -> str:
+        return self.layout_tensor.layout_type
 
     @classmethod
     def __torch_function__(cls, func, types, args=(), kwargs=None):
@@ -308,13 +326,13 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
 def implements(aten_ops_or_torch_fn):
     return _implements(AffineQuantizedTensor, aten_ops_or_torch_fn)
 
-def register_layout_cls(extended_layout: str):
-    return _register_layout_cls(AffineQuantizedTensor, extended_layout)
+def register_layout_cls(layout_type_class: type(LayoutType)):
+    return _register_layout_cls(AffineQuantizedTensor, layout_type_class)
 
-def get_layout_tensor_constructor(extended_layout: str):
-    return _get_layout_tensor_constructor(AffineQuantizedTensor, extended_layout)
+def get_layout_tensor_constructor(layout_type_class: type(LayoutType)):
+    return _get_layout_tensor_constructor(AffineQuantizedTensor, layout_type_class)
 
-@register_layout_cls("plain")
+@register_layout_cls(PlainLayoutType)
 class PlainAQTLayout(AQTLayout):
     """
     Layout storage class for plain layout for affine quantized tensor, it stores int_data, scale, zero_point
@@ -330,6 +348,7 @@ def __new__(
         int_data: torch.Tensor,
         scale: torch.Tensor,
         zero_point: torch.Tensor,
+        layout_type: LayoutType,
     ):
         kwargs = {}
         kwargs["device"] = int_data.device
@@ -346,34 +365,39 @@ def __init__(
         int_data: torch.Tensor,
         scale: torch.Tensor,
         zero_point: torch.Tensor,
+        layout_type: LayoutType,
     ):
         self.int_data = int_data
         self.scale = scale
         self.zero_point = zero_point
+        self.layout_type = layout_type
 
     def __tensor_flatten__(self):
-        return ["int_data", "scale", "zero_point"], []
+        return ["int_data", "scale", "zero_point"], [self.layout_type]
 
     @classmethod
     def __tensor_unflatten__(
         cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
     ):
         int_data, scale, zero_point = tensor_data_dict["int_data"], tensor_data_dict["scale"], tensor_data_dict["zero_point"]
-        return cls(int_data, scale, zero_point)
+        layout_type, = tensor_attributes
+        return cls(int_data, scale, zero_point, layout_type)
 
     def to(self, *args, **kwargs):
         kwargs = self._get_to_kwargs(*args, **kwargs)
         return self.__class__(
             self.int_data.to(kwargs["device"]),
             self.scale.to(kwargs["device"]),
             self.zero_point.to(kwargs["device"]),
+            self.layout_type,
         )
 
     def _apply_fn_to_data(self, fn):
         return self.__class__(
             fn(self.int_data),
             fn(self.scale),
             fn(self.zero_point),
+            self.layout_type,
         )
 
     @classmethod
@@ -398,19 +422,24 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
 
     __torch_function__ = torch._C._disabled_torch_function_impl
 
-    def get_plain(self):
+    def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         return self.int_data, self.scale, self.zero_point
 
+    def get_layout_type(self) -> LayoutType:
+        return self.layout_type
+
     @classmethod
     def from_plain(
         cls,
         int_data: torch.Tensor,
         scale: torch.Tensor,
         zero_point: torch.Tensor,
+        layout_type: LayoutType,
     ):
-        return cls(int_data, scale, zero_point)
+        assert isinstance(layout_type, PlainLayoutType)
+        return cls(int_data, scale, zero_point, layout_type)
 
-@register_layout_cls("tensor_core_tiled")
+@register_layout_cls(TensorCoreTiledLayoutType)
 class TensorCoreTiledAQTLayout(AQTLayout):
     """
     Layout storage class for tensor_core_tiled layout for affine quantized tensor, this is for int4 only,
@@ -427,6 +456,7 @@ def __new__(
         packed_weight: torch.Tensor,
         scale_and_zero: torch.Tensor,
         transposed: bool,
+        layout_type: LayoutType,
     ):
         kwargs = {}
         kwargs["device"] = packed_weight.device
@@ -443,31 +473,40 @@ def __init__(
         packed_weight: torch.Tensor,
         scale_and_zero: torch.Tensor,
         transposed: bool,
+        layout_type: LayoutType,
     ):
         self.packed_weight = packed_weight
         self.scale_and_zero = scale_and_zero
         self.transposed = False
+        self.layout_type = layout_type
 
     def __tensor_flatten__(self):
-        return ["packed_weight", "scale_and_zero"], [self.transposed]
+        return ["packed_weight", "scale_and_zero"], [self.transposed, self.layout_type]
 
     @classmethod
     def __tensor_unflatten__(
         cls, tensor_data_dict, tensor_attributes, outer_size, outer_stride
     ):
         packed_weight, scale_and_zero = tensor_data_dict["packed_weight"], tensor_data_dict["scale_and_zero"]
-        transposed, = tensor_attributes
-        return cls(packed_weight, scale_and_zero, transposed)
+        transposed, layout_type, = tensor_attributes
+        return cls(packed_weight, scale_and_zero, transposed, layout_type)
 
     @classmethod
-    def from_plain(cls, int_data, scale, zero_point, inner_k_tiles=8):
+    def from_plain(
+        cls,
+        int_data: torch.Tensor,
+        scale: torch.Tensor,
+        zero_point: torch.Tensor,
+        layout_type: LayoutType
+    ):
+        assert isinstance(layout_type, TensorCoreTiledLayoutType)
         # assert int_data.dtype == torch.uint8, "torch.ops.aten._convert_weight_to_int4pack expects `uint8` dtype"
         # packed_weight = torch.ops.aten._convert_weight_to_int4pack(int_data, inner_k_tiles)
-        packed_weight = torch.ops.aten._convert_weight_to_int4pack(int_data.to(torch.int32), inner_k_tiles)
+        packed_weight = torch.ops.aten._convert_weight_to_int4pack(int_data.to(torch.int32), layout_type.inner_k_tiles)
         scale = scale.reshape(int_data.shape[0], -1)
         zero_point = zero_point.reshape(int_data.shape[0], -1)
         scale_and_zero = pack_tinygemm_scales_and_zeros(scale, zero_point)
-        return cls(packed_weight, scale_and_zero, False)
+        return cls(packed_weight, scale_and_zero, False, layout_type)
 
     def to(self, *args, **kwargs):
         kwargs = self._get_to_kwargs(*args, **kwargs)
@@ -477,18 +516,15 @@ def to(self, *args, **kwargs):
         return self.__class__(
             self.packed_weight.to(device),
             self.scale_and_zero.to(device),
-            self.transposed
+            self.transposed,
+            self.layout_type,
         )
 
     def _apply_fn_to_data(self, fn):
         self.packed_weight = fn(self.packed_weight)
         self.scale_and_zero = fn(self.scale_and_zero)
         return self
 
-    def __repr__(self):
-        int_data, scale, zero_point = self.get_plain()
-        return f"TensorCoreTiledAQTLayout(int_data={int_data}, scale={scale}, zero_point={zero_point})"
-
     @classmethod
     def __torch_dispatch__(cls, func, types, args, kwargs):
         kwargs = {} if kwargs is None else kwargs
@@ -511,7 +547,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
 
     __torch_function__ = torch._C._disabled_torch_function_impl
 
-    def get_plain(self):
+    def get_plain(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         from torchao.quantization.quant_primitives import (
             ZeroPointDomain,
             quantize_affine,
@@ -542,6 +578,9 @@ def get_plain(self):
         int_data = quantize_affine(dequantized, block_size, scale, zero, target_dtype, quant_min, quant_max, zero_point_domain)
         return int_data, scale, zero
 
+    def get_layout_type(self) -> LayoutType:
+        return self.layout_type
+
 def _quantized_linear_op(input_tensor, weight_qtensor, bias):
     """
     Quantized version of F.linear operator
@@ -565,8 +604,8 @@ def _quantized_linear_op(input_tensor, weight_qtensor, bias):
                 is_cuda and
                 input_is_int8 and
                 input_tensor.dtype == weight_qtensor.dtype and
-                input_tensor.extended_layout == "plain" and
-                weight_qtensor.extended_layout == "plain"
+                isinstance(input_tensor.layout_type, PlainLayoutType) and
+                isinstance(weight_qtensor.layout_type, PlainLayoutType)
             ):
                 #
                 # 1. do the matrix form of dot(X_i, W_j)
@@ -608,7 +647,7 @@ def _quantized_linear_op(input_tensor, weight_qtensor, bias):
             weight_qtensor.dtype == torch.bfloat16 and
             len(weight_qtensor.shape) == 2 and
             weight_qtensor.zero_point_domain == ZeroPointDomain.FLOAT and
-            weight_qtensor.extended_layout == "tensor_core_tiled"
+            isinstance(weight_qtensor.layout_type, TensorCoreTiledLayoutType)
         ):
             assert weight_qtensor.block_size[0] == 1, f"Requires groupwise quantization, got block_size: {block_size}"
             assert input_tensor.shape[-1] == weight_qtensor.shape[1], (
@@ -651,7 +690,7 @@ def _quantized_linear_op(input_tensor, weight_qtensor, bias):
             weight_qtensor.block_size[0] == 1 and
             weight_qtensor.block_size[1] == weight_qtensor.shape[1] and
             weight_qtensor.zero_point_domain == ZeroPointDomain.INT and
-            weight_qtensor.extended_layout == "plain"
+            isinstance(weight_qtensor.layout_type, PlainLayoutType)
         ):
             # TODO: enable cpu and mps efficient path
             # per channel int8 weight only quantizated mm