fix

Signed-off-by: Yaoyao Ding <[email protected]>

Author: yaoyaoding

python/tilus/ir/layout/inference/inference_rules/transform_shared.py

@@ -38,7 +38,10 @@ def inference(ctx: LayoutInferenceContext, inst: SliceSharedInst) -> dict[Shared
             outer_shape = []
             for i in range(len(a.shape)):
                 outer_shape.append(a.shape[i] // b_layout.shape[i])
-            return {a: shared_compose(shared_row_major(*outer_shape), b_layout).apply_swizzle(b.layout.swizzle)}
+            layout = shared_compose(shared_row_major(*outer_shape), b_layout)
+            if b.layout.optional_swizzle is not None:
+                layout = layout.apply_swizzle(b.layout.swizzle)
+            return {a: layout}
         else:
             return {}
 

python/tilus/ir/layout/ops/shared_ops.py

@@ -61,7 +61,7 @@ def shared_row_major(*shape: int) -> SharedLayout:
     """
     mode_shape = shape
     mode_strides = strides_from_ranks(shape=mode_shape, ranks=list(range(len(mode_shape))))
-    return shared_layout(shape=shape, mode_shape=mode_shape, mode_strides=mode_strides, swizzle=None)
+    return shared_layout(shape=shape, mode_shape=mode_shape, mode_strides=mode_strides, optional_swizzle=None)
 
 
 def shared_column_major(*shape: int) -> SharedLayout:
@@ -79,7 +79,7 @@ def shared_column_major(*shape: int) -> SharedLayout:
     """
     mode_shape = shape
     mode_strides = strides_from_ranks(shape=mode_shape, ranks=list(reversed(range(len(mode_shape)))))
-    return shared_layout(shape=shape, mode_shape=mode_shape, mode_strides=mode_strides, swizzle=None)
+    return shared_layout(shape=shape, mode_shape=mode_shape, mode_strides=mode_strides, optional_swizzle=None)
 
 
 def shared_compose(lhs: SharedLayout, rhs: SharedLayout) -> SharedLayout:
@@ -118,7 +118,7 @@ def shared_compose(lhs: SharedLayout, rhs: SharedLayout) -> SharedLayout:
         mode_strides.extend([stride * rhs_size for stride in (lhs.mode_strides[i] for i in lhs_group)])
         mode_strides.extend([rhs.mode_strides[i] for i in rhs_group])
 
-    return shared_layout(shape=shape, mode_shape=mode_shape, mode_strides=mode_strides, swizzle=None)
+    return shared_layout(shape=shape, mode_shape=mode_shape, mode_strides=mode_strides, optional_swizzle=None)
 
 
 def shared_permute(layout: SharedLayout, dims: Sequence[int]) -> SharedLayout:
@@ -151,7 +151,9 @@ def shared_permute(layout: SharedLayout, dims: Sequence[int]) -> SharedLayout:
         mode_shape.extend([layout.mode_shape[i] for i in layout_mode_groups[d]])
         mode_strides.extend([layout.mode_strides[i] for i in layout_mode_groups[d]])
 
-    return shared_layout(shape=shape, mode_shape=mode_shape, mode_strides=mode_strides, swizzle=layout.swizzle)
+    return shared_layout(
+        shape=shape, mode_shape=mode_shape, mode_strides=mode_strides, optional_swizzle=layout.optional_swizzle
+    )
 
 
 def shared_slice(layout: SharedLayout, retain_dims: Sequence[int]) -> SharedLayout:
@@ -184,7 +186,7 @@ def shared_slice(layout: SharedLayout, retain_dims: Sequence[int]) -> SharedLayo
         shape=shape,
         mode_shape=mode_shape,
         mode_strides=mode_strides,
-        swizzle=layout.swizzle,
+        optional_swizzle=layout.optional_swizzle,
     )
 
 
@@ -211,7 +213,7 @@ def shared_unsqueeze(layout: SharedLayout, dims: Sequence[int]) -> SharedLayout:
         shape=shape,
         mode_shape=layout.mode_shape,
         mode_strides=layout.mode_strides,
-        swizzle=layout.swizzle,
+        optional_swizzle=layout.optional_swizzle,
     )
 
 
@@ -373,7 +375,7 @@ def shared_row_major_swizzle(shape: Sequence[int], dtype_nbytes: int) -> SharedL
         shape=shape,
         mode_shape=mode_shape,
         mode_strides=mode_strides,
-        swizzle=swizzle,
+        optional_swizzle=swizzle,
     )
 
 

python/tilus/ir/layout/shared_layout.py

@@ -126,7 +126,10 @@ def swizzle(self) -> Swizzle:
 
     @staticmethod
     def create(
-        shape: Sequence[int], mode_shape: Sequence[int], mode_strides: Sequence[int], swizzle: Optional[Swizzle]
+        shape: Sequence[int],
+        mode_shape: Sequence[int],
+        mode_strides: Sequence[int],
+        optional_swizzle: Optional[Swizzle],
     ) -> SharedLayout:
         """
         Create a SharedLayout from shape, mode_shape, and mode_strides.
@@ -154,7 +157,10 @@ def create(
         if prod(mode_shape) != prod(shape):
             raise ValueError("The product of mode_shape must equal to the product of shape.")
         return SharedLayout(
-            shape=tuple(shape), mode_shape=tuple(mode_shape), mode_strides=tuple(mode_strides), optional_swizzle=swizzle
+            shape=tuple(shape),
+            mode_shape=tuple(mode_shape),
+            mode_strides=tuple(mode_strides),
+            optional_swizzle=optional_swizzle,
         )
 
     def as_numpy_grid(self) -> np.ndarray:
@@ -195,7 +201,7 @@ def apply_swizzle(self, swizzle: Swizzle) -> SharedLayout:
             shape=self.shape,
             mode_shape=self.mode_shape,
             mode_strides=self.mode_strides,
-            swizzle=swizzle,
+            optional_swizzle=swizzle,
         )
 
     def prepend_dim(self, extent: int) -> SharedLayout:
@@ -211,7 +217,7 @@ def prepend_dim(self, extent: int) -> SharedLayout:
             shape=shape,
             mode_shape=mode_shape,
             mode_strides=mode_strides,
-            swizzle=self.optional_swizzle,
+            optional_swizzle=self.optional_swizzle,
         )
 
     def transpose(self) -> SharedLayout:
@@ -238,7 +244,7 @@ def shared_layout(
     shape: Sequence[int],
     mode_shape: Sequence[int],
     mode_strides: Sequence[int],
-    swizzle: Optional[Swizzle] = None,
+    optional_swizzle: Optional[Swizzle] = None,
 ) -> SharedLayout:
     """Create a SharedLayout from shape, mode_shape, and mode_strides.
 
@@ -270,7 +276,9 @@ def shared_layout(
         mode_strides = updated_mode_strides
 
     # canonicalize swizzle: if swizzle has 0 bits, set it to None (both mean no swizzle)
-    if swizzle is not None and swizzle.bits == 0:
-        swizzle = None
+    if optional_swizzle is not None and optional_swizzle.bits == 0:
+        optional_swizzle = None
 
-    return SharedLayout.create(shape=shape, mode_shape=mode_shape, mode_strides=mode_strides, swizzle=swizzle)
+    return SharedLayout.create(
+        shape=shape, mode_shape=mode_shape, mode_strides=mode_strides, optional_swizzle=optional_swizzle
+    )

python/tilus/ir/layout/utils/cute.py

@@ -156,7 +156,10 @@ def flatten_int_tuple(t: IntTuple) -> list[Int]:
         mode_strides = [int(s) for s in flat_strides]
 
         return shared_layout(
-            shape=tensor_shape, mode_shape=mode_shape, mode_strides=mode_strides, swizzle=self.swizzle.as_swizzle()
+            shape=tensor_shape,
+            mode_shape=mode_shape,
+            mode_strides=mode_strides,
+            optional_swizzle=self.swizzle.as_swizzle(),
         )
 
 

python/tilus/ir/tools/printer.py

@@ -452,7 +452,7 @@ def visit_SharedLayout(self, node: SharedLayout) -> Doc:
             "shape=[" + self(node.shape) + "]",
             "mode_shape=[" + self(node.mode_shape) + "]",
             "mode_strides=[" + self(node.mode_strides) + "]",
-            "swizzle=" + (str(node.swizzle) if node.swizzle is not None else "None"),
+            "swizzle=" + (str(node.swizzle) if node.optional_swizzle is not None else "None"),
         ]
         doc = Text("SharedLayout(") + doc_join(items, ", ") + ")"
         return self.add_key_comment("shared_layout", doc)