diff --git a/src/backend/common/pto_ops_memory.cpp b/src/backend/common/pto_ops_memory.cpp index 0c0bc3828b..57f36b8c19 100644 --- a/src/backend/common/pto_ops_memory.cpp +++ b/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); diff --git a/src/ir/op/tile_ops/memory.cpp b/src/ir/op/tile_ops/memory.cpp index c008e22ae3..0ff907e702 100644 --- a/src/ir/op/tile_ops/memory.cpp +++ b/src/ir/op/tile_ops/memory.cpp @@ -215,12 +215,21 @@ TypePtr DeduceTileLoadType(const std::vector& 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(expr); @@ -230,12 +239,16 @@ TypePtr DeduceTileLoadType(const std::vector& 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. diff --git a/tests/st/runtime/ops/test_matmul_mx.py b/tests/st/runtime/ops/test_matmul_mx.py index 1f24eaa470..33d38de704 100644 --- a/tests/st/runtime/ops/test_matmul_mx.py +++ b/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) diff --git a/tests/ut/codegen/test_mx_ops_codegen.py b/tests/ut/codegen/test_mx_ops_codegen.py index b37b27a8d9..30b3c54363 100644 --- a/tests/ut/codegen/test_mx_ops_codegen.py +++ b/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: