Correct cooperative TMATMUL shape and packed Shared operand views

asl/block/attributes/B.FPATR.asl

@@ -25,15 +25,26 @@
 // NDF-BEGIN: PTO-CUBE-SHARED-TRANSPOSE-001
 // ndf: kind=contract level=L1 layer=block status=accepted
 // B.FPATR bits 7 and 8 MUST independently select logical transpose for the
-// corresponding Shared A or B primary data operand only. Shared ScaleA is
-// always physically [M,G_A] and Shared ScaleB is always physically [N,G_B],
-// both K-group-major, regardless of the transpose controls. Local-side
-// transpose MUST remain illegal and neither source is mutated. Shared A [M,K]
-// is physically [M,K] at TransA=0 and [K,M] at TransA=1; Shared B [K,N] is
-// physically [N,K] at TransB=0 and [K,N] at TransB=1. Each Shared source MUST
-// expose the exact physical valid shape with a legal padded major pitch, and
-// every participating PE MUST pass metadata preflight before any payload
-// snapshot, destination allocation, or output effect.
+// corresponding Shared A or B primary data operand only. Shared ScaleA MUST
+// have stored-valid shape [group_M,G_A] and ScaleB [N,G_B], with each G equal
+// to ceil(K/group_size) for its own primary type (group-32 E8M0 or group-64
+// U32), independent of transpose. Local-side transpose MUST remain illegal.
+// Shared A [group_M,K] is physically [group_M,K] at TransA=0 and [K,group_M]
+// at TransA=1; Shared B [K,N] is physically [N,K] at TransB=0 and [K,N] at
+// TransB=1. Cooperative Shared forms MUST use group_M in 1..128 and N/K in
+// 1..65535 without power-of-two restrictions. If a supported packed-four-bit
+// Shared primary is present (E2M1X2, E1M2X2, HiF4X2 for Matrix-MX only, S4X2,
+// or U4X2), K and every such original parent's physical `columns` pitch MUST
+// be even under every transpose combination; valid M/group_M and N gain no
+// parity rule. This MUST not affect Local-only TMATMUL, TGEMV, generic Shared
+// allocation, or Local packed A when Shared B is not packed. Each Shared
+// source MUST expose its exact physical valid shape with a legal padded pitch;
+// Local CUBE roles retain existing descriptor/capacity rules, while Shared
+// roles are checked against original legal descriptors. Internal read views
+// MUST preserve original pitch and typed logical element/carrier/lane mapping
+// without a transient Local descriptor or 64 KiB capacity limit. Every
+// participating PE MUST pass metadata preflight before payload snapshots,
+// destination allocation, or output effects.
 // NDF-END: PTO-CUBE-SHARED-TRANSPOSE-001
 
 // NDF-BEGIN: PTO-CUBE-CSCALE-001

asl/block/model/dispatch/cube-tmatmul.asl

@@ -421,34 +421,37 @@ begin
             shared_ordinal = (shared_ordinal + 1) as integer {0..4};
         end;
     end;
-
     if TileMatrixFunctionUsesBias(function) then
         bias = BundleMatrixSourceAt(
             local_ordinal as integer {0..8});
         local_ordinal = (local_ordinal + 1) as integer {0..6};
     end;
-
     if _BundleFixedPointAttributes.c_scale_en then
         c_scale = BundleMatrixSourceAt(
             local_ordinal as integer {0..8});
     end;
-
     let right_group = TileMatrixRightGroupSourceCount(
         function, right_type);
+    let left_shared = shared_count != 0 && shared_count != right_group;
     let shape_legal = if shared_count == 0 then
         TileMatrixCubeInfosMatchDimensions(left, right, pe_m, n, k)
-    else if shared_count == right_group then
-        TileMatrixMixedInfosMatchDimensions(left, right, pe_m, n, k)
     else
-        TileMatrixInfosMatchDimensions(left, right, pe_m, n, k);
+        BundleMatrixCooperativeInfosMatchDimensions(
+            left, right, pe_m, n, k, left_shared,
+            left_type, right_type);
     let operand_types_legal = left.data_type == left_type &&
         right.data_type == right_type;
     let scales_legal = !TileMatrixFunctionUsesMX(function) ||
-        TileMatrixInfoOptionalScalesLegal(
-            left, left_scale, left_scale_present,
-            right, right_scale, right_scale_present);
+        if shared_count == 0 then
+            TileMatrixInfoOptionalScalesLegal(
+                left, left_scale, left_scale_present,
+                right, right_scale, right_scale_present)
+        else
+            BundleMatrixCooperativeScalesLegal(
+                left, left_scale, left_scale_present, left_shared,
+                right, right_scale, right_scale_present,
+                left_type, right_type);
     assert shape_legal && operand_types_legal && scales_legal;
-
     let accumulator_legal = !TileMatrixFunctionUsesAccumulator(function) ||
         TileMatrixLocalCubeAccumulatorSchemaLegal(
             accumulator, pe_m, n, result_type, primary_layout,

asl/block/model/dispatch/matrix-scale.asl

@@ -1,4 +1,4 @@
-// PTO-UNIT: {"id":"PTO-BLOCK-MODEL-DISPATCH-MATRIX-SCALE","surface":"block","classification":["model","dispatch","matrix-scale"],"depends_on":["PTO-TILE-MODEL-LEGALITY-MATRIX-POSTPROCESS"]}
+// PTO-UNIT: {"id":"PTO-BLOCK-MODEL-DISPATCH-MATRIX-SCALE","surface":"block","classification":["model","dispatch","matrix-scale"],"depends_on":["PTO-BLOCK-MODEL-DISPATCH-SHARED-CUBE-MATRIX","PTO-TILE-MODEL-LEGALITY-MATRIX-FUNCTIONS","PTO-TILE-MODEL-LEGALITY-MATRIX-POSTPROCESS"]}
 
 readonly func BundleMatrixCScaleDestinationIndicesDistinct(
     c_scale_ordinal: integer {0..8}) => boolean
@@ -13,3 +13,71 @@ begin
     end;
     return TRUE;
 end;
+
+readonly func BundleMatrixCooperativeInfosMatchDimensions(
+    left: TileInfo,
+    right: TileInfo,
+    m: integer {1..65535},
+    n: integer {1..65535},
+    k: integer {1..65535},
+    left_shared: boolean,
+    left_type: TileDataType,
+    right_type: TileDataType) => boolean
+begin
+    let left_legal = if left_shared then
+        BundleMatrixPrivateSharedViewLegal(
+            left, m, k, left_type, FALSE)
+    else
+        TileCubeDescriptorLegal(left) && left.contents_defined &&
+        left.valid_rows == m && left.valid_columns == k;
+    return left_legal && BundleMatrixPrivateSharedViewLegal(
+        right, k, n, right_type, TRUE);
+end;
+
+readonly func BundleMatrixCooperativeScalesLegal(
+    left: TileInfo,
+    left_scale: TileInfo,
+    left_scale_present: boolean,
+    left_scale_shared: boolean,
+    right: TileInfo,
+    right_scale: TileInfo,
+    right_scale_present: boolean,
+    left_type: TileDataType,
+    right_type: TileDataType) => boolean
+begin
+    if !TileMXOperandPairLegal(left_type, right_type) ||
+       left_scale_present != TileMXInputTypeNeedsScale(left_type) ||
+       right_scale_present != TileMXInputTypeNeedsScale(right_type) then
+        return FALSE;
+    end;
+    let left_groups = if left_scale_present then
+        TileMXScaleGroupCount(
+            left.valid_columns as integer {1..65535}, left_type)
+    else 1;
+    let right_groups = if right_scale_present then
+        TileMXScaleGroupCount(
+            left.valid_columns as integer {1..65535}, right_type)
+    else 1;
+    let left_scale_type = if left_scale_present then
+        TileMXScaleCarrierType(left_type) else left_type;
+    let right_scale_type = if right_scale_present then
+        TileMXScaleCarrierType(right_type) else right_type;
+    let left_scale_legal = !left_scale_present ||
+        (left_scale_shared && BundleMatrixPrivateSharedViewLegal(
+            left_scale, left.valid_rows as integer {1..65535},
+            left_groups, left_scale_type, FALSE)) ||
+        (!left_scale_shared &&
+         ((left_scale.layout == TileLayout_CUBE_M32 &&
+           TileCubeDescriptorLegal(left_scale)) ||
+          (left_scale.layout == TileLayout_RowMajor &&
+           TileInfoDescriptorLegal(left_scale))) &&
+         left_scale.data_type == left_scale_type &&
+         left_scale.valid_rows == left.valid_rows &&
+         left_scale.valid_columns == left_groups);
+    let right_scale_legal = !right_scale_present ||
+        BundleMatrixPrivateSharedViewLegal(
+            right_scale, right_groups,
+            right.valid_columns as integer {1..65535},
+            right_scale_type, FALSE);
+    return left_scale_legal && right_scale_legal;
+end;

asl/block/model/dispatch/shared-cube-matrix.asl

@@ -1,5 +1,4 @@
 // PTO-UNIT: {"id":"PTO-BLOCK-MODEL-DISPATCH-SHARED-CUBE-MATRIX","surface":"block","classification":["model","dispatch","shared-cube-matrix"],"depends_on":["PTO-BLOCK-MODEL-OPERANDS-SHARED-GENERATION","PTO-TILE-MODEL-LEGALITY-MATRIX-FUNCTIONS","PTO-TILE-MODEL-STATE-SHARED-REGISTERS"]}
-
 readonly func BundleMatrixSharedSourceSchemaLegal(
     ordinal: integer {0..3},
     valid_rows: integer {1..65535},
@@ -30,10 +29,11 @@ begin
            shared.tile.valid_rows == valid_rows &&
            shared.tile.valid_columns == valid_columns &&
            shared.tile.columns >= columns &&
+           (!TileDataTypeIsFourBit(data_type) ||
+            shared.tile.columns MOD 2 == 0) &&
            shared.tile.data_type == data_type &&
            shared.tile.layout == TileLayout_RowMajor;
 end;
-
 // Matrix schema preflight derives Shared CELL views without reading payload.
 readonly func BundleSharedSubviewMatrixMetadataLegalForPE(
     binding: BundleSharedBindingIndex, pe_identity: MemoryAgentId,
@@ -55,6 +55,9 @@ begin
         modifier.size_code as integer {1..12});
     let element_bits = TileElementBits(parent.data_type);
     let bounded_columns = parent.columns as integer {1..65535};
+    if TileDataTypeIsFourBit(data_type) && bounded_columns MOD 2 != 0 then
+        return FALSE;
+    end;
     let raw_offset = UInt(BundleSharedSubviewOffsetRawForPE(
         binding, pe_identity));
     if raw_offset > 2047 then return FALSE; end;
@@ -95,7 +98,6 @@ begin
     return clipped_rows == valid_rows && clipped_columns == valid_columns &&
            selected_columns >= minimum_columns;
 end;
-
 readonly func BundleMatrixSharedSourcesReady(
     shared_count: integer {0..4}) => boolean
 begin
@@ -110,7 +112,6 @@ begin
     end;
     return TRUE;
 end;
-
 readonly func BundleMatrixSharedPrimarySchemaLegal(
     ordinal: integer {0..3},
     logical_rows: integer {1..65535},
@@ -124,7 +125,6 @@ begin
         ordinal, stored_rows, stored_columns,
         stored_columns, data_type);
 end;
-
 readonly func BundleMatrixSharedBPrimarySchemaLegal(
     ordinal: integer {0..3},
     logical_rows: integer {1..65535},
@@ -148,7 +148,6 @@ begin
     else return 0;
     end;
 end;
-
 pure func BundleMatrixCooperativeValidM(
     group_m: integer {1..65535},
     pe_identity: MemoryAgentId) => integer {0..32}
@@ -161,7 +160,6 @@ begin
     if remaining < m_per_pe then return remaining as integer {1..31}; end;
     return m_per_pe as integer {16,32};
 end;
-
 pure func BundleMatrixCooperativeCurrentPEMask(
     group_m: integer {1..65535},
     pe_identity: MemoryAgentId) => bits(4)
@@ -172,7 +170,6 @@ begin
     end;
     return mask;
 end;
-
 readonly func BundleMatrixSharedLeftPrimarySchemaLegal(
     ordinal: integer {0..3},
     group_m: integer {1..65535},
@@ -184,7 +181,6 @@ begin
     return BundleMatrixSharedPrimarySchemaLegal(
         ordinal, group_m, k, data_type, transpose);
 end;
-
 readonly func BundleMatrixSharedSchemasLegal(
     function: integer {0..31},
     left_type: TileDataType,
@@ -209,11 +205,15 @@ begin
     let right_group = TileMatrixRightGroupSourceCount(
         function, right_type);
     let left_shared = shared_count != right_group;
+    if !TileMatrixFunctionIsGEMV(function) &&
+       ((left_shared && TileDataTypeIsFourBit(left_type)) ||
+        TileDataTypeIsFourBit(right_type)) && k MOD 2 != 0 then
+        return FALSE;
+    end;
     if _BundleFixedPointAttributes.trans_a && !left_shared then
         return FALSE;
     end;
     var ordinal: integer {0..4} = 0;
-
     if left_shared then
         if !BundleMatrixSharedLeftPrimarySchemaLegal(
                ordinal as integer {0..3},
@@ -232,7 +232,6 @@ begin
             ordinal = (ordinal + 1) as integer {0..4};
         end;
     end;
-
     if !BundleMatrixSharedBPrimarySchemaLegal(
            ordinal as integer {0..3},
            k, n, right_type,
@@ -251,6 +250,25 @@ begin
     end;
     return ordinal == shared_count;
 end;
+readonly func BundleMatrixPrivateSharedViewLegal(
+    tile: TileInfo,
+    rows: integer {1..65535},
+    columns: integer {1..65535},
+    data_type: TileDataType,
+    allow_column_major_packed: boolean) => boolean
+begin
+    return tile.allocated && tile.contents_defined &&
+        SharedTileCapacityIsLegal(tile.capacity_bytes) &&
+        (tile.layout == TileLayout_RowMajor ||
+         (allow_column_major_packed &&
+          tile.layout == TileLayout_ColumnMajor &&
+          TileDataTypeIsFourBit(data_type))) &&
+        tile.rows == rows && tile.columns == columns &&
+        tile.valid_rows == rows && tile.valid_columns == columns &&
+        tile.data_type == data_type &&
+        rows * columns <= TileLogicalElementCapacity(
+            tile.capacity_bytes, data_type);
+end;
 
 readonly func MaterializeBundleSharedMatrixLeftPrimary(
     ordinal: integer {0..3},
@@ -270,14 +288,14 @@ begin
     let shared_tile_id = BundleSharedBindingId(ordinal);
     let source = if
         _BundleSharedBindings[[ordinal]].source0_subview.valid then
-        MaterializeBundleSharedSubview(ordinal)
+        MaterializeBundleSharedSubviewForPE(ordinal, pe_identity)
     else SharedTileRecord(shared_tile_id).tile;
     var tile = source;
     tile.contents_defined = FALSE;
     tile.defined_elements = Zeros{PTO_MODEL_TILE_ELEMENTS};
     tile.packed_defined_elements = zero_packed_tile_elements;
     tile.defined_valid_elements = 0;
-    tile.rows = DerivedTileRows(tile.capacity_bytes, k, data_type);
+    tile.rows = pe_m;
     tile.columns = k;
     tile.valid_rows = pe_m;
     tile.valid_columns = k;
@@ -310,7 +328,6 @@ begin
     tile.defined_valid_elements = (pe_m * k) as integer {0..524288};
     return tile;
 end;
-
 readonly func MaterializeBundleSharedMatrixLeftScale(
     ordinal: integer {0..3},
     group_m: integer {1..65535},
@@ -328,7 +345,7 @@ begin
     let shared_tile_id = BundleSharedBindingId(ordinal);
     let source = if
         _BundleSharedBindings[[ordinal]].source0_subview.valid then
-        MaterializeBundleSharedSubview(ordinal)
+        MaterializeBundleSharedSubviewForPE(ordinal, pe_identity)
     else SharedTileRecord(shared_tile_id).tile;
     assert BundleMatrixSharedPrimarySchemaLegal(
         ordinal, group_m, scale_groups, scale_type, FALSE);
@@ -337,8 +354,7 @@ begin
     tile.defined_elements = Zeros{PTO_MODEL_TILE_ELEMENTS};
     tile.packed_defined_elements = zero_packed_tile_elements;
     tile.defined_valid_elements = 0;
-    tile.rows = DerivedTileRows(
-        tile.capacity_bytes, scale_groups, scale_type);
+    tile.rows = pe_m;
     tile.columns = scale_groups;
     tile.valid_rows = pe_m;
     tile.valid_columns = scale_groups;
@@ -371,7 +387,6 @@ begin
         (pe_m * scale_groups) as integer {0..524288};
     return tile;
 end;
-
 readonly func MaterializeBundleSharedMatrixSource(
     ordinal: integer {0..3},
     valid_rows: integer {1..65535},
@@ -397,7 +412,6 @@ begin
         (valid_rows * valid_columns) as integer {0..524288};
     return tile;
 end;
-
 readonly func MaterializeBundleSharedMatrixPrimary(
     ordinal: integer {0..3},
     logical_rows: integer {1..65535},
@@ -419,14 +433,15 @@ begin
     tile.defined_elements = Zeros{PTO_MODEL_TILE_ELEMENTS};
     tile.packed_defined_elements = zero_packed_tile_elements;
     tile.defined_valid_elements = 0;
-    tile.rows = DerivedTileRows(
-        tile.capacity_bytes, logical_columns, data_type);
+    tile.rows = logical_rows;
     tile.columns = logical_columns;
     tile.valid_rows = logical_rows;
     tile.valid_columns = logical_columns;
     tile.data_type = data_type;
     tile.predicate_basis_type = data_type;
-    tile.layout = TileLayout_RowMajor;
+    // Four-bit right primaries use dense ColumnMajor private views; source pitches stay architectural.
+    tile.layout = if TileDataTypeIsFourBit(data_type) then
+        TileLayout_ColumnMajor else TileLayout_RowMajor;
     tile.cube_k_repeat = 0;
     tile.cube_n_repeat = 0;
     tile.cube_cell_count = 0;
@@ -452,7 +467,6 @@ begin
         (logical_rows * logical_columns) as integer {0..524288};
     return tile;
 end;
-
 readonly func BundleMatrixCooperativeMLayout(
     function: integer {0..31},
     right_type: TileDataType,