feat(ir): Added expand_clone ops and substitute_tiles pass

Author: wuzhf9

CMakeLists.txt

@@ -148,6 +148,7 @@ set(PYPTO_SOURCES
     src/ir/transforms/expand_mixed_kernel_pass.cpp
     src/ir/transforms/split_vector_kernel_pass.cpp
     src/ir/transforms/flatten_tile_nd_to_2d_pass.cpp
+    src/ir/transforms/substitute_tiles_pass.cpp
     src/ir/transforms/pass_context.cpp
     src/ir/transforms/passes.cpp
     src/ir/transforms/resolve_backend_op_layouts_pass.cpp

docs/en/dev/passes/00-pass_manager.md

@@ -68,6 +68,7 @@ struct PassProperties {
 | OutlineIncoreScopes | TypeChecked, SSAForm | SplitIncoreOrch | — |
 | OutlineClusterScopes | TypeChecked, SSAForm | ClusterOutlined | — |
 | ConvertTensorToTileOps | SplitIncoreOrch | IncoreTileOps | — |
+| SubstituteTiles | SSAForm, IncoreTileOps | SSAForm, IncoreTileOps | — |
 | FlattenTileNdTo2D | SSAForm, IncoreTileOps | SSAForm, TileOps2D | — |
 | ResolveBackendOpLayouts | SSAForm, IncoreTileOps, SplitIncoreOrch, TileOps2D | SSAForm, IncoreTileOps, SplitIncoreOrch, TileOps2D | NormalizedStmtStructure |
 | ExpandMixedKernel | SSAForm, IncoreTileOps, SplitIncoreOrch, TileOps2D | SSAForm, MixedKernelExpanded | — |
@@ -359,15 +360,16 @@ with passes.PassContext([passes.VerificationInstrument(passes.VerificationMode.A
 
 The PTO-oriented tile stage shared by `Default` and `DebugTileOptimization` is:
 
-1. `FlattenTileNdTo2D`
-2. `InferTileMemorySpace`
-3. `ResolveTransposeLayout`
-4. `ResolveBackendOpLayouts`
-5. `ExpandMixedKernel`
-6. `InitMemRef`
-7. `MemoryReuse`
-8. `LegalizePTOBufferReuse`
-9. `AllocateMemoryAddr`
+1. `SubstituteTiles`
+2. `FlattenTileNdTo2D`
+3. `InferTileMemorySpace`
+4. `ResolveTransposeLayout`
+5. `ResolveBackendOpLayouts`
+6. `ExpandMixedKernel`
+7. `InitMemRef`
+8. `MemoryReuse`
+9. `LegalizePTOBufferReuse`
+10. `AllocateMemoryAddr`
 
 `DebugTileOptimization` is a debug-only strategy for inspecting this tile stage
 without the tensor-only prefix passes. Use `Default` for normal compilation and

docs/zh-cn/dev/passes/00-pass_manager.md

@@ -68,6 +68,7 @@ struct PassProperties {
 | OutlineIncoreScopes | TypeChecked, SSAForm | SplitIncoreOrch | — |
 | OutlineClusterScopes | TypeChecked, SSAForm | ClusterOutlined | — |
 | ConvertTensorToTileOps | SplitIncoreOrch | IncoreTileOps | — |
+| SubstituteTiles | SSAForm, IncoreTileOps | SSAForm, IncoreTileOps | — |
 | FlattenTileNdTo2D | SSAForm, IncoreTileOps | SSAForm, TileOps2D | — |
 | ResolveBackendOpLayouts | SSAForm, IncoreTileOps, SplitIncoreOrch, TileOps2D | SSAForm, IncoreTileOps, SplitIncoreOrch, TileOps2D | NormalizedStmtStructure |
 | ExpandMixedKernel | SSAForm, IncoreTileOps, SplitIncoreOrch, TileOps2D | SSAForm, MixedKernelExpanded | — |
@@ -359,15 +360,16 @@ with passes.PassContext([passes.VerificationInstrument(passes.VerificationMode.A
 
 `Default` 和 `DebugTileOptimization` 共享的 PTO tile 阶段顺序为：
 
-1. `FlattenTileNdTo2D`
-2. `InferTileMemorySpace`
-3. `ResolveTransposeLayout`
-4. `ResolveBackendOpLayouts`
-5. `ExpandMixedKernel`
-6. `InitMemRef`
-7. `MemoryReuse`
-8. `LegalizePTOBufferReuse`
-9. `AllocateMemoryAddr`
+1. `SubstituteTiles`
+2. `FlattenTileNdTo2D`
+3. `InferTileMemorySpace`
+4. `ResolveTransposeLayout`
+5. `ResolveBackendOpLayouts`
+6. `ExpandMixedKernel`
+7. `InitMemRef`
+8. `MemoryReuse`
+9. `LegalizePTOBufferReuse`
+10. `AllocateMemoryAddr`
 
 `DebugTileOptimization` 只是用于排查 PTO tile 阶段的调试策略，会跳过
 tensor-only 前缀 pass。正常编译和非 strategy 专项测试都应优先使用

include/pypto/ir/transforms/pass_properties.h

@@ -83,6 +83,12 @@ inline const PassProperties kConvertTensorToTileOpsProperties{
     .required = {IRProperty::SSAForm, IRProperty::SplitIncoreOrch, IRProperty::NormalizedStmtStructure},
     .produced = {IRProperty::SSAForm, IRProperty::IncoreTileOps, IRProperty::NormalizedStmtStructure}};
 
+// -- Tile op substitution pass ----------------------------------------------
+
+inline const PassProperties kSubstituteTilesProperties{
+    .required = {IRProperty::SSAForm, IRProperty::IncoreTileOps, IRProperty::NormalizedStmtStructure},
+    .produced = {IRProperty::SSAForm, IRProperty::IncoreTileOps, IRProperty::NormalizedStmtStructure}};
+
 // -- Tile ND-to-2D flattening pass --------------------------------------------
 
 inline const PassProperties kFlattenTileNdTo2DProperties{

include/pypto/ir/transforms/passes.h

@@ -265,6 +265,17 @@ Pass OutlineClusterScopes();
  */
 Pass ConvertTensorToTileOps();
 
+/**
+ * @brief Substitute unsupported tile ops with PTO-supported tile ops
+ *
+ * Rewrites tile ops that lack direct PTO instruction support into equivalent
+ * combinations of supported tile ops.
+ *
+ * Requirements:
+ * - Input IR must have tile ops (run ConvertTensorToTileOps first)
+ */
+Pass SubstituteTiles();
+
 /**
  * @brief Flatten ND tile ops to 2D in InCore functions
  *

include/pypto/ir/type_inference.h

@@ -21,6 +21,7 @@
 #ifndef PYPTO_IR_TYPE_INFERENCE_H_
 #define PYPTO_IR_TYPE_INFERENCE_H_
 
+#include <cstddef>
 #include <cstdint>
 #include <memory>
 #include <optional>
@@ -187,6 +188,17 @@ bool IsBroadcastable(const ExprPtr& source_dim, const ExprPtr& target_dim);
  */
 std::string FormatShape(const std::vector<ExprPtr>& shape);
 
+int NormalizeAxis(int axis, size_t ndim);
+
+int64_t ComputeShapeProduct(const std::vector<ExprPtr>& shape);
+
+bool IsIndexLikeDtype(DataType dtype);
+
+TileLayout InferTileLayoutFromShape(const std::vector<ExprPtr>& shape);
+
+void ValidateIndexTupleElements(const TupleTypePtr& tuple_type, const std::string& op_name,
+                                const std::string& arg_name);
+
 /**
  * @brief Propagate blayout and pad from a source TileType's tile_view into a new TileView
  *

python/bindings/modules/passes.cpp

@@ -312,6 +312,8 @@ void BindPass(nb::module_& m) {
              "Create a pass that outlines Hierarchy scopes into separate level/role functions");
   passes.def("convert_tensor_to_tile_ops", &pass::ConvertTensorToTileOps,
              "Create a pass that converts tensor ops to tile ops in InCore functions");
+  passes.def("substitute_tiles", &pass::SubstituteTiles,
+             "Create a pass that substitutes unsupported tile ops with PTO-supported tile ops");
   passes.def("flatten_tile_nd_to_2d", &pass::FlattenTileNdTo2D,
              "Create a pass that flattens ND tile ops to 2D in InCore functions\n\n"
              "Merges all dimensions except the last into a single dimension.\n"

python/pypto/ir/op/tensor_ops.py

@@ -681,6 +681,33 @@ def expands(target: Expr, scalar: int | float | Expr, span: Span | None = None)
     return _ir_core.create_op_call("tensor.expands", [target, scalar_expr], {}, actual_span)
 
 
+def expand_clone(
+    tensor: Expr,
+    shape: list[int | Expr] | _ir_core.MakeTuple,
+    valid_shape: list[int | Expr] | _ir_core.MakeTuple | None = None,
+    span: Span | None = None,
+) -> Call:
+    """Expand tensor to new shape.
+
+    Args:
+        tensor: Input tensor expression
+        shape: New shape dimensions, or a MakeTuple
+        valid_shape: Valid shape dimensions (optional, defaults to empty)
+        span: Optional source span for debugging (auto-captured if not provided)
+
+    Returns:
+        Call expression for tensor expand_clone
+    """
+    actual_span = _get_span_or_capture(span)
+
+    shape_tuple = _to_make_tuple(shape, actual_span)
+
+    args = [tensor, shape_tuple]
+    if valid_shape is not None:
+        args.append(_to_make_tuple(valid_shape, actual_span))
+    return _ir_core.create_op_call("tensor.expand_clone", args, {}, actual_span)
+
+
 def exp(input: Expr, span: Span | None = None) -> Call:
     """Element-wise exponential operation.
 

python/pypto/ir/op/tile_ops.py

@@ -1506,6 +1506,29 @@ def expands(target: Expr, scalar: int | float | Expr, span: Span | None = None)
     return _ir_core.create_op_call("tile.expands", [target, scalar_expr], {}, actual_span)
 
 
+def expand_clone(
+    tile: Expr,
+    shape: Sequence[int | Expr] | _ir_core.MakeTuple,
+    span: Span | None = None,
+) -> Call:
+    """Expand tile to new shape by cloning elements.
+
+    Args:
+        tile: Input tile expression
+        shape: New shape dimensions, or a MakeTuple
+        span: Optional source span for debugging (auto-captured if not provided)
+
+    Returns:
+        Call expression for tile expand_clone
+    """
+    actual_span = _get_span_or_capture(span)
+
+    shape_tuple = _to_make_tuple(shape, actual_span)
+
+    args = [tile, shape_tuple]
+    return _ir_core.create_op_call("tile.expand_clone", args, {}, actual_span)
+
+
 def maximum(lhs: Expr, rhs: Expr, span: Span | None = None) -> Call:
     """Element-wise maximum of two tiles.
 

python/pypto/ir/pass_manager.py

@@ -131,6 +131,7 @@ def _register_passes(cls):
             ("ConvertTensorToTileOps", lambda: passes.convert_tensor_to_tile_ops()),
         ]
         tile_pto_passes: list[PassSpec] = [
+            ("SubstituteTiles", lambda: passes.substitute_tiles()),
             ("FlattenTileNdTo2D", lambda: passes.flatten_tile_nd_to_2d()),
             ("InferTileMemorySpace", lambda: passes.infer_tile_memory_space()),
             ("ResolveTransposeLayout", lambda: passes.resolve_transpose_layout()),