Use shapes for MX physical TLoad box; re-enable host-prequant matmul_mx ST

src/backend/common/pto_ops_memory.cpp

@@ -119,14 +119,17 @@ int64_t CheckedMxMultiply(int64_t lhs, int64_t rhs, std::string_view name, const
 // explicit rank-5 view (and skip the generic parameter view — see
 // EmitMakeTensorViews).  The physical view is row-major; its strides match
 // pto-isa BaseShape2D for the corresponding MX layout.
+//
+// Physical partition sizes come from load *shapes* (fractal-aligned). A narrower
+// *valid_shape* is tile metadata only and must not shrink the TLoad box.
 MxPhysicalView EmitMxPhysicalView(const CallPtr& op, const ir::VarPtr& tensor,
                                   const ir::TensorTypePtr& tensor_type, const ir::MakeTuplePtr& offsets,
-                                  const ir::MakeTuplePtr& valid_shape, std::string_view pto_layout,
+                                  const ir::MakeTuplePtr& shapes, std::string_view pto_layout,
                                   codegen::PTOCodegen& codegen) {
   INTERNAL_CHECK_SPAN(tensor_type->shape_.size() == 2, op->span_)
       << "MX rank-5 tensor view requires a logical rank-2 tensor";
-  INTERNAL_CHECK_SPAN(offsets->elements_.size() == 2 && valid_shape->elements_.size() == 2, op->span_)
-      << "MX rank-5 tensor view requires rank-2 offsets and valid_shape";
+  INTERNAL_CHECK_SPAN(offsets->elements_.size() == 2 && shapes->elements_.size() == 2, op->span_)
+      << "MX rank-5 tensor view requires rank-2 offsets and shapes";
 
   const bool is_a = pto_layout == "mx_a_zz";
   INTERNAL_CHECK_SPAN(is_a || pto_layout == "mx_b_nn", op->span_)
@@ -140,10 +143,10 @@ MxPhysicalView EmitMxPhysicalView(const CallPtr& op, const ir::VarPtr& tensor,
                                                        "tensor block dimension", op->span_);
   const int64_t group_extent = GetStaticAlignedMxValue(tensor_type->shape_[group_axis], kSCols, kSCols,
                                                        "tensor group dimension", op->span_);
-  const int64_t block_size = GetStaticAlignedMxValue(valid_shape->elements_[block_axis], kSRows, kSRows,
-                                                     "load block size", op->span_);
-  const int64_t group_size = GetStaticAlignedMxValue(valid_shape->elements_[group_axis], kSCols, kSCols,
-                                                     "load group size", op->span_);
+  const int64_t block_size =
+      GetStaticAlignedMxValue(shapes->elements_[block_axis], kSRows, kSRows, "load block shape", op->span_);
+  const int64_t group_size =
+      GetStaticAlignedMxValue(shapes->elements_[group_axis], kSCols, kSCols, "load group shape", op->span_);
   const int64_t block_offset =
       GetStaticAlignedMxValue(offsets->elements_[block_axis], kSRows, 0, "load block offset", op->span_);
   const int64_t group_offset =
@@ -283,7 +286,7 @@ static std::string MakeTileLoadCodegenPTO(const CallPtr& op, codegen::CodegenBas
   std::string tensor_view_type;
   if (is_mx_load) {
     auto physical =
-        EmitMxPhysicalView(op, tensor, tensor_type, offsets_tuple, valid_shape_tuple, pto_layout, codegen);
+        EmitMxPhysicalView(op, tensor, tensor_type, offsets_tuple, shapes_tuple, pto_layout, codegen);
     tensor_view = std::move(physical.tensor_view);
     tensor_view_type = std::move(physical.tensor_view_type);
     partition_dims = std::move(physical.partition_dims);

src/ir/op/tile_ops/memory.cpp

@@ -215,12 +215,21 @@ TypePtr DeduceTileLoadType(const std::vector<ExprPtr>& args,
     const size_t block_axis = is_mx_a ? 0 : 1;
     const size_t group_axis = is_mx_a ? 1 : 0;
     const std::string layout_name = TensorLayoutToString(source_view.layout);
+    constexpr int64_t kRows = tile_view_semantics::kMXSFractalRows;
+    constexpr int64_t kCols = tile_view_semantics::kMXSFractalCols;
     // PTOAS / pto-isa special requirement (feeds EmitMxPhysicalView):
     //   Physical MX GlobalTensor needs SFractal axes [16, 2], so every logical
-    //   block/group extent, load size, and offset must be static and divisible
-    //   by 16 (block) / 2 (group).  Example: logical MX_A_ZZ [64, 4] load at
-    //   [0,0] size [64,4] -> physical [1,4,2,16,2]; a dynamic or misaligned
-    //   group size cannot form that box and A5 TLoad static_asserts.
+    //   block/group *extent*, physical *load shapes*, and *offset* must be
+    //   static and divisible by 16 (block) / 2 (group).  Example: logical
+    //   MX_A_ZZ [64, 4] load shapes [64,4] -> physical [1,4,2,16,2]; a dynamic
+    //   or misaligned shapes extent cannot form that box and A5 TLoad
+    //   static_asserts.
+    //
+    //   *valid_shape* may narrow the tile relative to shapes (partial M/N). The
+    //   hardware transfer still uses fractal-aligned shapes; tile metadata
+    //   carries the narrower valid region for consumers such as matmul_mx.
+    //   Bounds 0 < valid_shape[i] <= shapes[i] are enforced by
+    //   InferWindowReadValidShape.
     auto check_static_aligned = [&](const ExprPtr& expr, int64_t alignment, int64_t minimum, const char* name,
                                     const Span& span) {
       auto value = As<ConstInt>(expr);
@@ -230,12 +239,16 @@ TypePtr DeduceTileLoadType(const std::vector<ExprPtr>& args,
           << "The operator " << op_name << " of an " << layout_name << " tensor requires " << name
           << " >= " << minimum << " and divisible by " << alignment << ", but got " << value->value_;
     };
-    check_static_aligned(tensor_type->shape_[block_axis], 16, 16, "tensor block dimension", args[0]->span_);
-    check_static_aligned(tensor_type->shape_[group_axis], 2, 2, "tensor group dimension", args[0]->span_);
-    check_static_aligned(valid_shape_tuple->elements_[block_axis], 16, 16, "load block size", args[3]->span_);
-    check_static_aligned(valid_shape_tuple->elements_[group_axis], 2, 2, "load group size", args[3]->span_);
-    check_static_aligned(offsets_tuple->elements_[block_axis], 16, 0, "load block offset", args[1]->span_);
-    check_static_aligned(offsets_tuple->elements_[group_axis], 2, 0, "load group offset", args[1]->span_);
+    check_static_aligned(tensor_type->shape_[block_axis], kRows, kRows, "tensor block dimension",
+                         args[0]->span_);
+    check_static_aligned(tensor_type->shape_[group_axis], kCols, kCols, "tensor group dimension",
+                         args[0]->span_);
+    check_static_aligned(shapes_tuple->elements_[block_axis], kRows, kRows, "load block shape",
+                         args[2]->span_);
+    check_static_aligned(shapes_tuple->elements_[group_axis], kCols, kCols, "load group shape",
+                         args[2]->span_);
+    check_static_aligned(offsets_tuple->elements_[block_axis], kRows, 0, "load block offset", args[1]->span_);
+    check_static_aligned(offsets_tuple->elements_[group_axis], kCols, 0, "load group offset", args[1]->span_);
     // MX cube scale loads are Mat-only (TLoadMxCube*) and require the caller to
     // spell the target explicitly. The public load interface keeps its ordinary
     // Vec default, so an omitted target fails instead of being silently changed.

tests/st/runtime/ops/test_matmul_mx.py

@@ -23,13 +23,6 @@
 from harness.core.harness import DataType, PTOTestCase, TensorSpec
 from pypto.runtime.runner import RunConfig
 
-pytestmark = pytest.mark.skip(
-    reason=(
-        "PTOAS temporarily pinned to v0.57 (revert #2523 level3 TMP for pypto#2558); "
-        "MX pack/layout requires PTOAS v0.60 — re-enable after tile.ci root-cause fix"
-    ),
-)
-
 _REQUIRED_TORCH_DTYPES = ("float4_e2m1fn_x2", "float8_e4m3fn", "float8_e8m0fnu")
 if not all(hasattr(torch, name) for name in _REQUIRED_TORCH_DTYPES):
     pytest.skip("torch MXFP4/MXFP8/E8M0 dtypes required", allow_module_level=True)

tests/ut/codegen/test_mx_ops_codegen.py

@@ -89,6 +89,60 @@ def main(
         with pytest.raises(ValueError, match=r"matmul_mx.*only supported.*Ascend950.*a5.*a2a3"):
             _run_default_pipeline(Program, BackendType.Ascend910B)
 
+    def test_mx_scale_load_accepts_narrowed_valid_shape(self):
+        """Physical shapes stay fractal-aligned; valid_shape may narrow M."""
+
+        @pl.program
+        class Program:
+            @pl.function(type=pl.FunctionType.InCore)
+            def main(
+                self,
+                a_s: pl.Tensor[[128, 8], pl.FP8E8M0, pl.MX_A_ZZ],
+            ):
+                _ = pl.load(
+                    a_s,
+                    [0, 0],
+                    [16, 2],
+                    valid_shape=[8, 2],
+                    target_memory=pl.Mem.Mat,
+                )
+
+        mlir = _emit_incore_mlir(Program)
+        assert "mx5d_view" in mlir
+        assert "pto.tload" in mlir
+        partitions = [line for line in mlir.splitlines() if "partition_view" in line]
+        # Physical TLoad box remains shapes=[16,2] -> one SFractal block row.
+        assert any(
+            "sizes = [%c1_index, %c1_index, %c1_index, %c16_index, %c2_index]" in line for line in partitions
+        ), mlir
+
+    def test_mx_scale_load_accepts_dynamic_narrowed_valid_shape(self):
+        """Dynamic valid_shape must not shrink the physical partition box."""
+
+        @pl.program
+        class Program:
+            @pl.function(type=pl.FunctionType.InCore)
+            def main(
+                self,
+                a_s: pl.Tensor[[128, 8], pl.FP8E8M0, pl.MX_A_ZZ],
+                valid_rows: pl.Scalar[pl.INDEX],
+            ):
+                _ = pl.load(
+                    a_s,
+                    [0, 0],
+                    [16, 2],
+                    valid_shape=[valid_rows, 2],
+                    target_memory=pl.Mem.Mat,
+                )
+
+        mlir = _emit_incore_mlir(Program)
+        assert "mx5d_view" in mlir
+        assert "pto.tload" in mlir
+        partitions = [line for line in mlir.splitlines() if "partition_view" in line]
+        assert any(
+            "sizes = [%c1_index, %c1_index, %c1_index, %c16_index, %c2_index]" in line for line in partitions
+        ), mlir
+
     def test_emits_tmatmul_mx_and_tget(self):
         @pl.program
         class Program: