fix tensor pack -> linalg.pack (#123)

Author: makslevental

mlir/extras/dialects/ext/linalg.py

@@ -2,6 +2,16 @@
 from ...util import get_user_code_loc
 from ....dialects import linalg
 
+from ....dialects._ods_common import (
+    _dispatch_mixed_values,
+    _cext,
+    get_op_results_or_values,
+    get_default_loc_context,
+    get_op_result_or_op_results,
+    get_default_loc_context,
+    segmented_accessor,
+)
+
 # noinspection PyUnresolvedReferences
 from ....dialects.linalg import *
 from ....extras import types as T
@@ -524,3 +534,185 @@ def vecmat(y, A, x, *, loc=None, ip=None):
     if loc is None:
         loc = get_user_code_loc()
     return linalg.vecmat(y, A, loc=loc, ip=ip, outs=[x])
+
+
+@_cext.register_operation(linalg.Dialect)
+class PackOp(ir.OpView):
+    OPERATION_NAME = "linalg.pack"
+
+    _ODS_OPERAND_SEGMENTS = [
+        1,
+        1,
+        0,
+        -1,
+    ]
+
+    _ODS_REGIONS = (0, True)
+
+    def __init__(
+        self,
+        source,
+        dest,
+        inner_dims_pos,
+        inner_tiles,
+        static_inner_tiles,
+        *,
+        padding_value=None,
+        outer_dims_perm=None,
+        loc=None,
+        ip=None,
+    ):
+        operands = []
+        results = []
+        attributes = {}
+        regions = None
+        operands.append(source)
+        operands.append(dest)
+        operands.append(padding_value)
+        operands.append(get_op_results_or_values(inner_tiles))
+        _ods_context = get_default_loc_context(loc)
+        if outer_dims_perm is not None:
+            attributes["outer_dims_perm"] = (
+                outer_dims_perm
+                if (
+                    isinstance(outer_dims_perm, ir.Attribute)
+                    or not ir.AttrBuilder.contains("DenseI64ArrayAttr")
+                )
+                else ir.AttrBuilder.get("DenseI64ArrayAttr")(
+                    outer_dims_perm, context=_ods_context
+                )
+            )
+        attributes["inner_dims_pos"] = (
+            inner_dims_pos
+            if (
+                isinstance(inner_dims_pos, ir.Attribute)
+                or not ir.AttrBuilder.contains("DenseI64ArrayAttr")
+            )
+            else ir.AttrBuilder.get("DenseI64ArrayAttr")(
+                inner_dims_pos, context=_ods_context
+            )
+        )
+        attributes["static_inner_tiles"] = (
+            static_inner_tiles
+            if (
+                isinstance(static_inner_tiles, ir.Attribute)
+                or not ir.AttrBuilder.contains("DenseI64ArrayAttr")
+            )
+            else ir.AttrBuilder.get("DenseI64ArrayAttr")(
+                static_inner_tiles, context=_ods_context
+            )
+        )
+        _ods_successors = None
+        super().__init__(
+            self.OPERATION_NAME,
+            self._ODS_REGIONS,
+            self._ODS_OPERAND_SEGMENTS,
+            self._ODS_RESULT_SEGMENTS,
+            attributes=attributes,
+            operands=operands,
+            successors=_ods_successors,
+            regions=regions,
+            loc=loc,
+            ip=ip,
+        )
+
+    @property
+    def source(self):
+        operand_range = segmented_accessor(
+            self.operation.operands, self.operation.attributes["operandSegmentSizes"], 0
+        )
+        return operand_range[0]
+
+    @property
+    def dest(self):
+        operand_range = segmented_accessor(
+            self.operation.operands, self.operation.attributes["operandSegmentSizes"], 1
+        )
+        return operand_range[0]
+
+    @property
+    def padding_value(self):
+        operand_range = segmented_accessor(
+            self.operation.operands, self.operation.attributes["operandSegmentSizes"], 2
+        )
+        return operand_range[0] if len(operand_range) > 0 else None
+
+    @property
+    def inner_tiles(self):
+        operand_range = segmented_accessor(
+            self.operation.operands, self.operation.attributes["operandSegmentSizes"], 3
+        )
+        return operand_range
+
+    @property
+    def outer_dims_perm(self):
+        if "outer_dims_perm" not in self.operation.attributes:
+            return None
+        return self.operation.attributes["outer_dims_perm"]
+
+    @outer_dims_perm.setter
+    def outer_dims_perm(self, value):
+        if value is not None:
+            self.operation.attributes["outer_dims_perm"] = value
+        elif "outer_dims_perm" in self.operation.attributes:
+            del self.operation.attributes["outer_dims_perm"]
+
+    @outer_dims_perm.deleter
+    def outer_dims_perm(self):
+        del self.operation.attributes["outer_dims_perm"]
+
+    @property
+    def inner_dims_pos(self):
+        return self.operation.attributes["inner_dims_pos"]
+
+    @inner_dims_pos.setter
+    def inner_dims_pos(self, value):
+        if value is None:
+            raise ValueError("'None' not allowed as value for mandatory attributes")
+        self.operation.attributes["inner_dims_pos"] = value
+
+    @property
+    def static_inner_tiles(self):
+        return self.operation.attributes["static_inner_tiles"]
+
+    @static_inner_tiles.setter
+    def static_inner_tiles(self, value):
+        if value is None:
+            raise ValueError("'None' not allowed as value for mandatory attributes")
+        self.operation.attributes["static_inner_tiles"] = value
+
+    @property
+    def result(self):
+        return self.operation.results[0]
+
+
+def pack(
+    source,
+    dest,
+    inner_dims_pos,
+    inner_tiles,
+    *,
+    padding_value=None,
+    outer_dims_perm=None,
+    loc=None,
+    ip=None,
+) -> ir.Value:
+
+    (
+        dynamic_inner_tiles,
+        # packed here means %1:2 packing (results packing)
+        _inner_tiles,
+        static_inner_tiles,
+    ) = _dispatch_mixed_values(inner_tiles)
+
+    return PackOp(
+        source=source,
+        dest=dest,
+        inner_dims_pos=inner_dims_pos,
+        inner_tiles=dynamic_inner_tiles,
+        static_inner_tiles=static_inner_tiles,
+        padding_value=padding_value,
+        outer_dims_perm=outer_dims_perm,
+        loc=loc,
+        ip=ip,
+    ).result

mlir/extras/dialects/ext/tensor.py

@@ -689,36 +689,3 @@ def pad_(
 generate = region_op(
     lambda result, dynamic_extents: tensor.GenerateOp(result, dynamic_extents)
 )
-
-_pack = pack
-
-
-def pack(
-    source,
-    dest,
-    inner_dims_pos,
-    inner_tiles,
-    *,
-    padding_value=None,
-    outer_dims_perm=None,
-    loc=None,
-    ip=None,
-):
-    (
-        dynamic_inner_tiles,
-        # packed here means %1:2 packing (results packing)
-        _inner_tiles,
-        static_inner_tiles,
-    ) = _dispatch_mixed_values(inner_tiles)
-
-    return _pack(
-        source=source,
-        dest=dest,
-        inner_dims_pos=inner_dims_pos,
-        inner_tiles=dynamic_inner_tiles,
-        static_inner_tiles=static_inner_tiles,
-        padding_value=padding_value,
-        outer_dims_perm=outer_dims_perm,
-        loc=loc,
-        ip=ip,
-    )

tests/test_transform.py

@@ -990,22 +990,22 @@ def main(variant_op: any_op_t()):
               %c0_i32 = arith.constant 0 : i32
               %0 = tensor.empty() : tensor<16x256xi8>
               %1 = tensor.empty() : tensor<1x4x16x64xi8>
-              %pack = tensor.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [16, 64] into %1 : tensor<16x256xi8> -> tensor<1x4x16x64xi8>
+              %pack = linalg.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [16, 64] into %1 : tensor<16x256xi8> -> tensor<1x4x16x64xi8>
               %2 = tensor.empty() : tensor<4x1x64x64xi8>
               %3 = tensor.empty() : tensor<1x1x16x64xi8>
               %4 = linalg.fill ins(%c0_i32 : i32) outs(%3 : tensor<1x1x16x64xi8>) -> tensor<1x1x16x64xi8>
               %5 = scf.forall (%arg2, %arg3) in (1, 4) shared_outs(%arg4 = %0) -> (tensor<16x256xi8>) {
                 %6 = affine.apply #map(%arg3)
                 %extracted_slice = tensor.extract_slice %arg1[0, %6] [256, 64] [1, 1] : tensor<256x256xi8> to tensor<256x64xi8>
                 %extracted_slice_0 = tensor.extract_slice %arg4[0, %6] [16, 64] [1, 1] : tensor<16x256xi8> to tensor<16x64xi8>
-                %pack_1 = tensor.pack %extracted_slice outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [64, 64] into %2 : tensor<256x64xi8> -> tensor<4x1x64x64xi8>
+                %pack_1 = linalg.pack %extracted_slice outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [64, 64] into %2 : tensor<256x64xi8> -> tensor<4x1x64x64xi8>
                 %7 = linalg.generic {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]} ins(%pack, %pack_1 : tensor<1x4x16x64xi8>, tensor<4x1x64x64xi8>) outs(%4 : tensor<1x1x16x64xi8>) {
                 ^bb0(%in: i8, %in_2: i8, %out: i8):
                   %8 = arith.muli %in, %in_2 : i8
                   %9 = arith.addi %out, %8 : i8
                   linalg.yield %9 : i8
                 } -> tensor<1x1x16x64xi8>
-                %unpack = tensor.unpack %7 inner_dims_pos = [0, 1] inner_tiles = [16, 64] into %extracted_slice_0 : tensor<1x1x16x64xi8> -> tensor<16x64xi8>
+                %unpack = linalg.unpack %7 inner_dims_pos = [0, 1] inner_tiles = [16, 64] into %extracted_slice_0 : tensor<1x1x16x64xi8> -> tensor<16x64xi8>
                 scf.forall.in_parallel {
                   tensor.parallel_insert_slice %unpack into %arg4[0, %6] [16, 64] [1, 1] : tensor<16x64xi8> into tensor<16x256xi8>
                 }
@@ -1050,7 +1050,7 @@ def tensor_pack(
         src: T.tensor(129, 47, 16, 16, T.f32()),
         dst: T.tensor(17, 2, 16, 16, 32, 8, T.f32()),
     ):
-        return tensor.pack(
+        return linalg.pack(
             src,
             dst,
             inner_dims_pos=[1, 0],
@@ -1068,8 +1068,8 @@ def mod_transform():
         def main(variant_op: any_op_t()):
             packed = match(
                 variant_op,
-                ops=["tensor.pack"],
-                matched_op=transform_op_t("tensor.pack"),
+                ops=["linalg.pack"],
+                matched_op=transform_op_t("linalg.pack"),
             )
             lowered_pack = transform.structured.lower_pack(packed)
 
@@ -1097,8 +1097,8 @@ def main(variant_op: any_op_t()):
           }
           module attributes {transform.with_named_sequence} {
             transform.named_sequence @main(%arg0: !transform.any_op) {
-              %0 = transform.structured.match ops{["tensor.pack"]} in %arg0 : (!transform.any_op) -> !transform.op<"tensor.pack">
-              %pad_op, %expand_shape_op, %transpose_op = transform.structured.lower_pack %0 : (!transform.op<"tensor.pack">) -> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
+              %0 = transform.structured.match ops{["linalg.pack"]} in %arg0 : (!transform.any_op) -> !transform.op<"linalg.pack">
+              %pad_op, %expand_shape_op, %transpose_op = transform.structured.lower_pack %0 : (!transform.op<"linalg.pack">) -> (!transform.op<"tensor.pad">, !transform.op<"tensor.expand_shape">, !transform.op<"linalg.transpose">)
               transform.yield 
             }
           }