diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 70ade82bbf..415017bdce 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -1465,10 +1465,16 @@ def TMovOp : PTO_TOp<"tmov", [ return ::mlir::pto::PIPE::PIPE_V; } - // MAT -> L0 (Left/Right/Bias/Scaling) are MTE1 moves. + // MAT -> Scaling moves L1 data into the fixpipe parameter buffer. + if (s == ::mlir::pto::AddressSpace::MAT && + d == ::mlir::pto::AddressSpace::SCALING) { + return ::mlir::pto::PIPE::PIPE_FIX; + } + + // MAT -> L0 (Left/Right/Bias) moves are executed by MTE1. if ((s == ::mlir::pto::AddressSpace::MAT && (d == ::mlir::pto::AddressSpace::LEFT || d == ::mlir::pto::AddressSpace::RIGHT || - d == ::mlir::pto::AddressSpace::BIAS || d == ::mlir::pto::AddressSpace::SCALING))) { + d == ::mlir::pto::AddressSpace::BIAS))) { return ::mlir::pto::PIPE::PIPE_MTE1; } @@ -6576,6 +6582,7 @@ def TSqrtOp: PTO_TOp<"tsqrt", [ def TStoreFPOp: PTO_TOp<"tstore_fp", [ PTO_DpsInitOpInterface, OpPipeInterface, + DeclareOpInterfaceMethods ]> { let summary = "TSTORE_FP: Store an accumulator tile into global memory using a scaling (fp) tile for vector quantization parameters."; diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index c63e216e24..aa9ff29624 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -13693,6 +13693,14 @@ void TStoreOp::getEffects(SmallVectorImpl> &effects) { + addEffect(effects, &getSrcMutable(), MemoryEffects::Read::get()); + addEffect(effects, &getFpMutable(), MemoryEffects::Read::get()); + addEffect(effects, &getDstMutable(), MemoryEffects::Write::get()); +} + // === TMovOp === // Read: src, Write: dst void TMovOp::getEffects(SmallVectorImpl> &effects) { diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index 908c604a16..c8715ec606 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -2830,9 +2830,11 @@ static LogicalResult verifyStructuredAccStoreLike( if (!isStructuredAccStoreScalingPayload(preQuant)) return op->emitOpError("vector pre_quant mode requires scaling pointer payload"); if (!isStructuredAccStoreFloatScalarPayloadType( - getStructuredAccStoreScalingElementType(preQuant))) + getStructuredAccStoreScalingElementType(preQuant)) && + getStructuredAccStoreScalingElementType(preQuant) != + IntegerType::get(op->getContext(), 64, IntegerType::Unsigned)) return op->emitOpError( - "vector pre_quant mode requires scaling pointer element type to be f16, bf16, or f32"); + "vector pre_quant mode requires scaling pointer element type to be f16, bf16, f32, or packed ui64 Quant_PRE entries"); } else if (!isStructuredAccStoreFloatScalarPayload(preQuant)) { return op->emitOpError( "scalar pre_quant mode requires f16/bf16/f32 payload"); diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 087d8d62c6..887a797562 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -1048,6 +1048,14 @@ func::FuncOp ExpandState::invokeTilelangDaemon(const SpecKey &key, func::FuncOp ExpandState::invokeTilelangDSL(const SpecKey &key, Operation *tileOp, ModuleOp mod, MLIRContext *ctx) { + std::string uniqueName = buildUniqueFunctionBaseName(key); + SymbolTable targetSymTable(mod); + if (auto existingFunc = targetSymTable.lookup(uniqueName)) { + llvm::errs() << "ExpandTileOp: reuse existing function @" << uniqueName + << "\n"; + return cast(existingFunc); + } + // Try daemon first if daemon socket path is provided. if (!daemonSocketPath.empty()) { func::FuncOp daemonResult = invokeTilelangDaemon(key, tileOp, mod, ctx); @@ -1189,16 +1197,6 @@ func::FuncOp ExpandState::invokeTilelangDSL(const SpecKey &key, SmallVector clonedFuncs; llvm::StringMap renamedSymbols; - std::string uniqueName = buildUniqueFunctionBaseName(key); - - // Check if function already exists in module (deduplication) - SymbolTable targetSymTable(mod); - if (auto existingFunc = targetSymTable.lookup(uniqueName)) { - // Function already exists, return it directly (avoid redefinition) - llvm::errs() << "ExpandTileOp: reuse existing function @" << uniqueName << "\n"; - return cast(existingFunc); - } - std::vector newNameStorage; for (auto [index, fn] : llvm::enumerate(parsedFuncs)) { IRMapping mapping; diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index 9235d85f56..5246e1281b 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -831,6 +831,8 @@ static std::string getCopyElementFragment(Type elementType) { return intType.isUnsigned() ? "u16" : "s16"; case 32: return intType.isUnsigned() ? "u32" : "s32"; + case 64: + return "s32"; default: return {}; } diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index cee5d4cd70..9f5f2f4937 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -799,6 +799,8 @@ static std::string getCopyElementFragment(Type elementType) { return intType.isUnsigned() ? "u16" : "s16"; case 32: return intType.isUnsigned() ? "u32" : "s32"; + case 64: + return "s32"; default: return {}; } diff --git a/lib/TileOps/tload_template.py b/lib/TileOps/tload_template.py index 9e40bedb38..846ec879b7 100644 --- a/lib/TileOps/tload_template.py +++ b/lib/TileOps/tload_template.py @@ -31,6 +31,21 @@ def _known_le(lhs, rhs) -> bool: return lhs_value <= rhs_value +def _memory_space_value(value): + memory_space = getattr(value, "memory_space", None) + if memory_space is None: + return None + return memory_space.value if hasattr(memory_space, "value") else memory_space + + +def _is_scaling_tile(value) -> bool: + return _memory_space_value(value) in {"scaling", "SCALING"} + + +def _is_ub_tile(value) -> bool: + return _memory_space_value(value) in {"ub", "UB"} + + def _match_tile_layout(dst, *, row_major: bool, s_layout) -> bool: b_layout_ok = ( dst.config.b_layout == pto.BLayout.ROW_MAJOR @@ -62,6 +77,8 @@ def _check_load_bounds(src, dst, *, logical_rows, logical_cols=None, stride_axis def _tload_preconditions_nd2nd(src, dst) -> bool: + if not _is_ub_tile(dst): + return False logical_rows = src.shape[0] * src.shape[1] * src.shape[2] * src.shape[3] logical_cols = src.shape[4] return _match_tile_layout( @@ -72,6 +89,8 @@ def _tload_preconditions_nd2nd(src, dst) -> bool: def _tload_preconditions_dn2dn(src, dst) -> bool: + if not _is_ub_tile(dst): + return False logical_rows = src.shape[3] logical_cols = src.shape[0] * src.shape[1] * src.shape[2] * src.shape[4] return _match_tile_layout( @@ -81,6 +100,8 @@ def _tload_preconditions_dn2dn(src, dst) -> bool: ) def _tload_preconditions_nz2nz(src, dst) -> bool: + if not _is_ub_tile(dst): + return False logical_rows = src.shape[2] return _match_tile_layout( dst, row_major=False, s_layout=pto.SLayout.ROW_MAJOR @@ -330,12 +351,41 @@ def _constraint_tload_mat_base(src, dst) -> bool: if dst_dtype is None: return False dtype_name = dst_dtype.name if hasattr(dst_dtype, "name") else str(dst_dtype) - supported_dtypes = {"f16", "bf16", "f32", "i8", "si8", "ui8", "i16", "si16", "ui16", "i32", "si32"} + supported_dtypes = { + "f16", + "bf16", + "f32", + "i8", + "si8", + "ui8", + "i16", + "si16", + "ui16", + "i32", + "si32", + "ui32", + "i64", + "si64", + "ui64", + } if dtype_name not in supported_dtypes: return False return True +def _layout_value(config): + if config is None or not hasattr(config, "layout"): + return None + layout = config.layout + return layout.value if hasattr(layout, "value") else layout + + +def _last_two_shape_dims(value): + if not hasattr(value, "shape") or value.shape is None or len(value.shape) < 2: + return None + return value.shape[-2], value.shape[-1] + + def _constraint_tload_mat_nd2nz(src, dst) -> bool: """TLOAD.MAT ND2NZ fractal load constraint""" if not _constraint_tload_mat_base(src, dst): @@ -356,13 +406,19 @@ def _constraint_tload_mat_nd2nz(src, dst) -> bool: # ND2NZ: source is in ND (row-major) format where the inner dimension (g4) # corresponds to the tile column count. Disambiguates from DN format where # g4 corresponds to the tile row count. - if hasattr(src, 'rank') and src.rank == 5: - dst_valid_cols = dst.valid_shape[1] if hasattr(dst, 'valid_shape') and dst.valid_shape is not None else None - if dst_valid_cols is not None and hasattr(src, 'shape') and src.shape is not None: - src_inner = src.shape[4] if len(src.shape) >= 5 else None - if src_inner is not None: - if not _known_eq(dst_valid_cols, src_inner): - return False + src_layout = _layout_value(src.config) + if src_layout is not None and src_layout in {"dn", "DN"}: + return False + src_tail = _last_two_shape_dims(src) + if src_tail is not None and hasattr(dst, 'valid_shape') and dst.valid_shape is not None: + src_rows, src_cols = src_tail + dst_valid_rows, dst_valid_cols = dst.valid_shape + is_nd_shape = _known_eq(dst_valid_rows, src_rows) and _known_eq(dst_valid_cols, src_cols) + is_dn_shape = _known_eq(dst_valid_rows, src_cols) and _known_eq(dst_valid_cols, src_rows) + if not is_nd_shape: + return False + if is_dn_shape and src_layout is not None and src_layout in {"dn", "DN"}: + return False return True @@ -384,16 +440,139 @@ def _constraint_tload_mat_dn2nz(src, dst) -> bool: # DN2NZ: source is in DN (col-major) format where the inner dimension (g4) # corresponds to the tile row count. Disambiguates from ND format where # g4 corresponds to the tile column count. - if hasattr(src, 'rank') and src.rank == 5: - dst_valid_rows = dst.valid_shape[0] if hasattr(dst, 'valid_shape') and dst.valid_shape is not None else None - if dst_valid_rows is not None and hasattr(src, 'shape') and src.shape is not None: - src_inner = src.shape[4] if len(src.shape) >= 5 else None - if src_inner is not None: - if not _known_eq(dst_valid_rows, src_inner): - return False + src_tail = _last_two_shape_dims(src) + if src_tail is not None and hasattr(dst, 'valid_shape') and dst.valid_shape is not None: + src_rows, src_cols = src_tail + dst_valid_rows, dst_valid_cols = dst.valid_shape + is_nd_shape = _known_eq(dst_valid_rows, src_rows) and _known_eq(dst_valid_cols, src_cols) + is_dn_shape = _known_eq(dst_valid_rows, src_cols) and _known_eq(dst_valid_cols, src_rows) + src_layout = _layout_value(src.config) + if src_layout is not None and src_layout in {"dn", "DN"} and is_nd_shape: + return is_dn_shape + if not is_dn_shape: + return False + if is_nd_shape and (src_layout is None or src_layout not in {"dn", "DN"}): + return False + return True + + +def _constraint_tload_mat_dn2nz_layout(src, dst) -> bool: + """TLOAD.MAT DN2NZ for source views explicitly marked as DN layout.""" + if not _constraint_tload_mat_base(src, dst): + return False + config = dst.config + if config is None: + return False + b_layout = config.b_layout + s_layout = config.s_layout + if b_layout is None or s_layout is None: + return False + b_layout_value = b_layout.value if hasattr(b_layout, "value") else b_layout + s_layout_value = s_layout.value if hasattr(s_layout, "value") else s_layout + if b_layout_value not in {"col_major", "COL_MAJOR"} or s_layout_value not in {"row_major", "ROW_MAJOR"}: + return False + if _layout_value(src.config) not in {"dn", "DN"}: + return False + src_tail = _last_two_shape_dims(src) + if src_tail is not None and hasattr(dst, 'valid_shape') and dst.valid_shape is not None: + src_rows, src_cols = src_tail + dst_valid_rows, dst_valid_cols = dst.valid_shape + return _known_eq(dst_valid_rows, src_rows) and _known_eq(dst_valid_cols, src_cols) + return True + + +def _constraint_tload_gm_to_mat_linear(src, dst) -> bool: + """TLOAD.MAT linear GM -> L1 load constraint.""" + if not _constraint_tload_mat_base(src, dst): + return False + config = dst.config + if config is None: + return False + b_layout = config.b_layout + s_layout = config.s_layout + if b_layout is None or s_layout is None: + return False + b_layout_value = b_layout.value if hasattr(b_layout, "value") else b_layout + s_layout_value = s_layout.value if hasattr(s_layout, "value") else s_layout + if b_layout_value not in {"row_major", "ROW_MAJOR"}: + return False + if s_layout_value not in {"none_box", "NONE_BOX"}: + return False + src_tail = _last_two_shape_dims(src) + if src_tail is not None and hasattr(dst, "valid_shape") and dst.valid_shape is not None: + src_rows, src_cols = src_tail + dst_valid_rows, dst_valid_cols = dst.valid_shape + if not (_known_eq(dst_valid_rows, src_rows) and _known_eq(dst_valid_cols, src_cols)): + return False + if hasattr(src, "strides") and src.strides is not None and len(src.strides) >= 1: + if not _known_eq(src.strides[-1], 1): + return False + if len(src.strides) >= 2: + src_tail = _last_two_shape_dims(src) + expected_row_stride = None + if src_tail is not None: + expected_row_stride = src_tail[1] + elif hasattr(dst, "valid_shape") and dst.valid_shape is not None: + expected_row_stride = dst.valid_shape[1] + if expected_row_stride is not None and not _known_eq(src.strides[-2], expected_row_stride): + return False return True +@pto.ckernel( + target="a5", + op="pto.tload", + priority=2, + dtypes=[ + (pto.f16, pto.f16), + (pto.bf16, pto.bf16), + (pto.f32, pto.f32), + (pto.ui64, pto.ui64), + ], + constraints=[_constraint_tload_gm_to_mat_linear], + name="tload_gm_to_mat_linear", +) +def template_tload_gm_to_mat_linear(src: pto.PartitionTensorView, dst: pto.Tile): + """GM -> MAT linear load template for row-major MAT tiles.""" + rows, cols = dst.valid_shape + elem_bytes = pto.bytewidth(dst.element_type) + len_burst = rows * cols * elem_bytes + pto.mte_gm_l1( + src.as_ptr(), + dst.as_ptr(), + len_burst, + nburst=(1, 0, 0), + ) + return + + +@pto.ckernel( + target="a5", + op="pto.tload", + priority=2, + dtypes=[ + (pto.f16, pto.f16), + (pto.bf16, pto.bf16), + (pto.f32, pto.f32), + (pto.i8, pto.i8), + ], + constraints=[_constraint_tload_mat_dn2nz_layout], + name="tload_gm_to_mat_dn2nz_layout", +) +def template_tload_gm_to_mat_dn2nz_layout(src: pto.PartitionTensorView, dst: pto.Tile): + """GM -> MAT DN2NZ load for DN-layout views with logical shape.""" + m, k = dst.valid_shape + + pto.mte_gm_l1_frac( + src.as_ptr(), dst.as_ptr(), pto.FractalMode.DN2NZ, + shape=(m, k), + src_layout=(m * pto.bytewidth(dst.element_type),), + dst_group=(1, 1, k, 0), + ctrl=(0, False) + ) + return + + @pto.ckernel( target="a5", op="pto.tload", @@ -402,6 +581,7 @@ def _constraint_tload_mat_dn2nz(src, dst) -> bool: (pto.f16, pto.f16), (pto.bf16, pto.bf16), (pto.f32, pto.f32), + (pto.i8, pto.i8), ], constraints=[_constraint_tload_mat_nd2nz], name="tload_gm_to_mat_nd2nz", @@ -428,8 +608,8 @@ def template_tload_gm_to_mat_nd2nz(src: pto.PartitionTensorView, dst: pto.Tile): n_value = m d_value = k - # src_layout: inner stride = K (number of elements per row) - src_inner_stride = k + # src_layout uses byte stride in GM. + src_inner_stride = k * pto.bytewidth(dst.element_type) # dst_group: (group_count, loop2_stride, loop3_stride, loop4_stride) # For simple single-block case: (1, 1, m, 0) @@ -455,6 +635,7 @@ def template_tload_gm_to_mat_nd2nz(src: pto.PartitionTensorView, dst: pto.Tile): (pto.f16, pto.f16), (pto.bf16, pto.bf16), (pto.f32, pto.f32), + (pto.i8, pto.i8), ], constraints=[_constraint_tload_mat_dn2nz], name="tload_gm_to_mat_dn2nz", @@ -478,15 +659,11 @@ def template_tload_gm_to_mat_dn2nz(src: pto.PartitionTensorView, dst: pto.Tile): gm_ptr = src.as_ptr() mat_ptr = dst.as_ptr() - # DN2NZ parameter calculation - # For DN format, the original shape is (K, M) -- no logical conversion - # needed. dn2nz writes the same logical N x D result into NZ layout. - # n_value = K, d_value = M n_value = k d_value = m - # src_layout: inner stride = M (number of elements per column) - src_inner_stride = m + # src_layout uses byte stride in GM. + src_inner_stride = m * pto.bytewidth(dst.element_type) # dst_group: (group_count, loop2_stride, loop3_stride, loop4_stride) # (1, 1, k, 0) diff --git a/lib/TileOps/tmov2scale_template.py b/lib/TileOps/tmov2scale_template.py index f997a176fd..ac13574bc3 100644 --- a/lib/TileOps/tmov2scale_template.py +++ b/lib/TileOps/tmov2scale_template.py @@ -22,10 +22,8 @@ Constraint: This template is selected when dst.memory_space == SCALING. Data format: - - Each f32 scale value is stored as ui64 (f32 bits in lower 32 bits, upper 32 bits = 0) - - Hardware interprets each ui64's lower 32 bits as one f32 scale value - - Total bytes = N * 8 (N ui64 elements) - - len_burst = N (number of ui64 elements to transfer) + - The source and destination tile element types match. + - len_burst is expressed in 64-byte copy_cbuf_to_fbuf burst units. """ import tilelang_dsl as pto @@ -65,7 +63,10 @@ def _tmov_m2s_constraint(src: pto.Tile, dst: pto.Tile) -> bool: op="pto.tmov", constraints=[_tmov_m2s_constraint], dtypes=[ + (pto.f16, pto.f16), + (pto.bf16, pto.bf16), (pto.f32, pto.f32), + (pto.ui64, pto.ui64), ], ) def template_tmov_m2s(src: pto.Tile, dst: pto.Tile): @@ -78,19 +79,16 @@ def template_tmov_m2s(src: pto.Tile, dst: pto.Tile): The scale parameters are loaded into FB for fixpipe quantization. mte_l1_fb semantics: - - len_burst = M (number of rows in the scaling tile) + - len_burst = rows * cols * sizeof(dtype) / 64 - Each row contains one set of scale parameters for all columns - - For 16x16 f32 scaling: len_burst = 16 (see textract_fp.pto:128) - - For 1x16 f32 scaling: len_burst = 1 (single row of parameters) nburst = (n_burst, src_gap, dst_gap) - single burst with no gaps """ - # Scale tile has shape MxN (typically 1xN for per-column scales) - m, _ = dst.valid_shape - # len_burst = M (number of rows/parameter sets) - len_burst = m + # copy_cbuf_to_fbuf uses 64-byte burst units. + m, n = src.valid_shape + len_burst = (m * n * pto.bytewidth(src.element_type)) // 64 n_burst = 1 src_gap = 0 dst_gap = 0 pto.mte_l1_fb(src.as_ptr(), dst.as_ptr(), len_burst, nburst=(n_burst, src_gap, dst_gap)) - return \ No newline at end of file + return diff --git a/lib/TileOps/tstore_template.py b/lib/TileOps/tstore_template.py index 59c2ee86cc..32922b4f3c 100644 --- a/lib/TileOps/tstore_template.py +++ b/lib/TileOps/tstore_template.py @@ -329,13 +329,20 @@ def _extract_s_layout_value(config) -> str | None: """Extract the effective s_layout string from a config object. For TileConfig this is the s_layout attribute. - For ViewConfig there is no slayout distinction; returns None. + For ViewConfig, NZ maps to the boxed row-major secondary layout used by + cube accumulator stores. """ if config is None: return None if hasattr(config, "s_layout") and config.s_layout is not None: sl = config.s_layout return sl.value if hasattr(sl, "value") else sl + if hasattr(config, "layout") and config.layout is not None: + layout = config.layout + layout_str = layout.value if hasattr(layout, "value") else str(layout) + normalized = layout_str.strip().lower().replace("-", "_") + if normalized == "nz": + return "row_major" return None @@ -414,17 +421,36 @@ def template_tstore_acc_to_gm_nz2nd(src: pto.Tile, dst: pto.PartitionTensorView) acc_ptr = src.as_ptr() gm_ptr = dst.as_ptr() - # src_stride: ACC buffer stride (N under NZ format) + # src_stride: ACC buffer stride (M under NZ format) # dst_stride: GM stride (N under ND format) - src_stride = n + src_stride = m dst_stride = n - - pto.mte_l0c_gm( - acc_ptr, gm_ptr, - m, n, src_stride, dst_stride, - 0, 0, - layout="nz2nd" - ) + src_elem = src.element_type + dst_elem = dst.element_type + + if pto.constexpr(src_elem == pto.f32 and dst_elem == pto.f16): + pto.mte_l0c_gm( + acc_ptr, gm_ptr, + m, n, src_stride, dst_stride, + 0, 0, + layout="nz2nd", + pre_quant=(pto.f16(1.0), "f32_f16"), + ) + elif pto.constexpr(src_elem == pto.f32 and dst_elem == pto.bf16): + pto.mte_l0c_gm( + acc_ptr, gm_ptr, + m, n, src_stride, dst_stride, + 0, 0, + layout="nz2nd", + pre_quant=(pto.bf16(1.0), "f32_bf16"), + ) + else: + pto.mte_l0c_gm( + acc_ptr, gm_ptr, + m, n, src_stride, dst_stride, + 0, 0, + layout="nz2nd" + ) return @@ -460,16 +486,35 @@ def template_tstore_acc_to_gm_nz2dn(src: pto.Tile, dst: pto.PartitionTensorView) gm_ptr = dst.as_ptr() # NZ2DN requires an additional loop0_src_stride parameter - src_stride = n + src_stride = m dst_stride = m # Under DN format, stride is M loop0_src_stride = 1 # loop0_src_stride for NZ2DN - - pto.mte_l0c_gm( - acc_ptr, gm_ptr, - m, n, src_stride, dst_stride, - 0, 0, - layout=("nz2dn", loop0_src_stride) - ) + src_elem = src.element_type + dst_elem = dst.element_type + + if pto.constexpr(src_elem == pto.f32 and dst_elem == pto.f16): + pto.mte_l0c_gm( + acc_ptr, gm_ptr, + m, n, src_stride, dst_stride, + 0, 0, + layout=("nz2dn", loop0_src_stride), + pre_quant=(pto.f16(1.0), "f32_f16"), + ) + elif pto.constexpr(src_elem == pto.f32 and dst_elem == pto.bf16): + pto.mte_l0c_gm( + acc_ptr, gm_ptr, + m, n, src_stride, dst_stride, + 0, 0, + layout=("nz2dn", loop0_src_stride), + pre_quant=(pto.bf16(1.0), "f32_bf16"), + ) + else: + pto.mte_l0c_gm( + acc_ptr, gm_ptr, + m, n, src_stride, dst_stride, + 0, 0, + layout=("nz2dn", loop0_src_stride) + ) return @@ -503,16 +548,35 @@ def template_tstore_acc_to_gm_nz2nz(src: pto.Tile, dst: pto.PartitionTensorView) acc_ptr = src.as_ptr() gm_ptr = dst.as_ptr() - src_stride = n + src_stride = m dst_stride = n - split = 1 # NZ2NZ requires a split parameter - - pto.mte_l0c_gm( - acc_ptr, gm_ptr, - m, n, src_stride, dst_stride, - 0, 0, - layout=("nz2nz", split) - ) + split = 0 # NZ2NZ requires a split parameter + src_elem = src.element_type + dst_elem = dst.element_type + + if pto.constexpr(src_elem == pto.f32 and dst_elem == pto.f16): + pto.mte_l0c_gm( + acc_ptr, gm_ptr, + m, n, src_stride, dst_stride, + 0, 0, + layout=("nz2nz", split), + pre_quant=(pto.f16(1.0), "f32_f16"), + ) + elif pto.constexpr(src_elem == pto.f32 and dst_elem == pto.bf16): + pto.mte_l0c_gm( + acc_ptr, gm_ptr, + m, n, src_stride, dst_stride, + 0, 0, + layout=("nz2nz", split), + pre_quant=(pto.bf16(1.0), "f32_bf16"), + ) + else: + pto.mte_l0c_gm( + acc_ptr, gm_ptr, + m, n, src_stride, dst_stride, + 0, 0, + layout=("nz2nz", split) + ) return @@ -575,6 +639,10 @@ def _constraint_tstore_fp(src, fp, dst) -> bool: dtypes=[ (pto.f32, pto.f16, pto.f16), (pto.f32, pto.bf16, pto.bf16), + (pto.f32, pto.f32, pto.f16), + (pto.f32, pto.f32, pto.bf16), + (pto.f32, pto.ui64, pto.f16), + (pto.f32, pto.ui64, pto.bf16), ], constraints=[_constraint_tstore_fp], name="tstore_fp_acc_to_gm", @@ -602,21 +670,16 @@ def template_tstore_fp_acc_to_gm(src: pto.Tile, fp: pto.Tile, dst: pto.Partition fp_ptr = fp.as_ptr() gm_ptr = dst.as_ptr() - # Determine pre_quant mode based on fp (scaling) buffer dtype: - # Use qf322*_pre_vec modes which support vector scale rows (scaling buffer). - # f32_f16/f32_bf16 modes only accept scalar payload and cannot be used here. - # Note: dst.dtype is not supported by the TileLang DSL v1 frontend for - # PartitionTensorView, so we use fp.element_type instead. - fp_dtype = fp.element_type + dst_dtype = dst.element_type - if pto.constexpr(fp_dtype == pto.bf16): + if pto.constexpr(dst_dtype == pto.bf16): quant_mode = "qf322bf16_pre_vec" else: quant_mode = "qf322f16_pre_vec" # TODO: Replace with tstore_fp DSL surface once available. # Currently using mte_l0c_gm + pre_quant as a temporary workaround. - src_stride = n + src_stride = m dst_stride = n pto.mte_l0c_gm( diff --git a/test/lit/pto/tmov_m2s_pipe_selection.pto b/test/lit/pto/tmov_m2s_pipe_selection.pto new file mode 100644 index 0000000000..4e73b76f9c --- /dev/null +++ b/test/lit/pto/tmov_m2s_pipe_selection.pto @@ -0,0 +1,38 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch a5 --enable-insert-sync %s | FileCheck %s + +module attributes {"pto.target_arch" = "a5"} { + func.func @tmov_m2s_waits_for_l1_load( + %scale_gm: memref<1x32xui64, #pto.address_space>) { + %scale_mat = pto.alloc_tile + : !pto.tile_buf + %scale_fb = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%scale_gm : memref<1x32xui64, #pto.address_space>) + outs(%scale_mat : !pto.tile_buf) + pto.tmov ins(%scale_mat : !pto.tile_buf) + outs(%scale_fb : !pto.tile_buf) + return + } +} + +// CHECK-LABEL: AICORE void tmov_m2s_waits_for_l1_load( +// CHECK: set_flag(PIPE_MTE2, PIPE_FIX, EVENT_ID{{[0-9]+}}); +// CHECK: wait_flag(PIPE_MTE2, PIPE_FIX, EVENT_ID{{[0-9]+}}); +// CHECK-NOT: wait_flag(PIPE_MTE2, PIPE_MTE1 diff --git a/test/lit/vpto/acc_store_verify_invalid_pre_quant_vec_payload_type.pto b/test/lit/vpto/acc_store_verify_invalid_pre_quant_vec_payload_type.pto index 6431bea657..a7ead80c62 100644 --- a/test/lit/vpto/acc_store_verify_invalid_pre_quant_vec_payload_type.pto +++ b/test/lit/vpto/acc_store_verify_invalid_pre_quant_vec_payload_type.pto @@ -6,16 +6,16 @@ module attributes {"pto.target_arch" = "a5"} { %c0_i64 = arith.constant 0 : i64 %c16_i64 = arith.constant 16 : i64 %c32_i64 = arith.constant 32 : i64 - %payload = pto.castptr %c0_i64 : i64 -> !pto.ptr + %payload = pto.castptr %c0_i64 : i64 -> !pto.ptr pto.mte_l0c_ub %src, %dst, %c16_i64, %c16_i64, %c16_i64, %c32_i64, dst_mode(%c0_i64), pre_quant(%payload, mode = qf322f16_pre_vec) : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, - !pto.ptr + !pto.ptr return } } -// CHECK: vector pre_quant mode requires scaling pointer element type to be f16, bf16, or f32 +// CHECK: vector pre_quant mode requires scaling pointer element type to be f16, bf16, f32, or packed ui64 Quant_PRE entries diff --git a/test/lit/vpto/tstore_fp_sync_effects.pto b/test/lit/vpto/tstore_fp_sync_effects.pto new file mode 100644 index 0000000000..c21ee944e5 --- /dev/null +++ b/test/lit/vpto/tstore_fp_sync_effects.pto @@ -0,0 +1,39 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-insert-sync --emit-pto-ir %s -o - | FileCheck %s + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @tstore_fp_after_tmatmul_sync(%dst_gm: !pto.ptr) { + %c0 = arith.constant 0 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c1 = arith.constant 1 : index + + %lhs = pto.alloc_tile : !pto.tile_buf + %rhs = pto.alloc_tile : !pto.tile_buf + %acc = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + + pto.tmatmul ins(%lhs, %rhs : !pto.tile_buf, + !pto.tile_buf) + outs(%acc : !pto.tile_buf) + + %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view + %sv_dst = pto.partition_view %tv_dst, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xf16> + pto.tstore_fp ins(%acc, %fp : !pto.tile_buf, !pto.tile_buf) + outs(%sv_dst : !pto.partition_tensor_view<16x32xf16>) + return + } +} + +// CHECK-LABEL: func.func @tstore_fp_after_tmatmul_sync +// CHECK: pto.tmatmul +// CHECK: pto.set_flag[, +// CHECK: pto.wait_flag[, +// CHECK: pto.tstore_fp diff --git a/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/CMakeLists.txt new file mode 100644 index 0000000000..70b2a80a5f --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/CMakeLists.txt @@ -0,0 +1,9 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +pto_tilelang_vec_st(fp8_load_test) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/compare.py b/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/compare.py new file mode 100644 index 0000000000..3c5bd54aab --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/compare.py @@ -0,0 +1,48 @@ +#!/usr/bin/python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import os +import sys +import numpy as np + + +CASES = [ + ("fp8_load_16x64", (16, 64)), + ("fp8_load_32x64", (32, 64)), + ("fp8_load_128x64", (128, 64)), +] + + +def compare_case(name, shape): + golden = np.fromfile(os.path.join(name, "golden.bin"), dtype=np.uint8).reshape(shape) + output = np.fromfile(os.path.join(name, "output.bin"), dtype=np.uint8).reshape(shape) + if np.array_equal(golden, output): + print(f"[INFO] {name}: compare passed") + return True + + mismatch = golden != output + count = int(np.count_nonzero(mismatch)) + print(f"[ERROR] {name}: {count}/{golden.size} byte mismatches") + for row, col in np.argwhere(mismatch)[:5]: + print(f" [{row},{col}]: expected=0x{golden[row, col]:02x}, actual=0x{output[row, col]:02x}") + return False + + +def main(): + all_passed = True + for name, shape in CASES: + if not compare_case(name, shape): + all_passed = False + if not all_passed: + sys.exit(2) + print("[INFO] all cases passed") + + +if __name__ == "__main__": + main() diff --git a/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/fp8_load_test.pto b/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/fp8_load_test.pto new file mode 100644 index 0000000000..13c94f9b21 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/fp8_load_test.pto @@ -0,0 +1,114 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @FP8_LOAD_16x64(%src_ptr: !pto.ptr, %dst_ptr: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 64 : index + %c1024 = arith.constant 1024 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x64xf8E4M3FN> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x64xf8E4M3FN> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x64xf8E4M3FN> -> !pto.partition_tensor_view<1x1x1x16x64xf8E4M3FN> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x64xf8E4M3FN> -> !pto.partition_tensor_view<1x1x1x16x64xf8E4M3FN> + + %tile = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x64xf8E4M3FN>) + outs(%tile : !pto.tile_buf) + pto.tstore ins(%tile : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x64xf8E4M3FN>) + return + } + + func.func @FP8_LOAD_32x64(%src_ptr: !pto.ptr, %dst_ptr: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 32 : index + %c64 = arith.constant 64 : index + %c2048 = arith.constant 2048 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x32x64xf8E4M3FN> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x32x64xf8E4M3FN> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x32x64xf8E4M3FN> -> !pto.partition_tensor_view<1x1x1x32x64xf8E4M3FN> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x32x64xf8E4M3FN> -> !pto.partition_tensor_view<1x1x1x32x64xf8E4M3FN> + + %tile = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x32x64xf8E4M3FN>) + outs(%tile : !pto.tile_buf) + pto.tstore ins(%tile : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x32x64xf8E4M3FN>) + return + } + + func.func @FP8_LOAD_128x64(%src_ptr: !pto.ptr, %dst_ptr: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c64 = arith.constant 64 : index + %c128 = arith.constant 128 : index + %c8192 = arith.constant 8192 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c128, %c64], + strides = [%c8192, %c8192, %c8192, %c64, %c1] + : !pto.tensor_view<1x1x1x128x64xf8E4M3FN> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c128, %c64], + strides = [%c8192, %c8192, %c8192, %c64, %c1] + : !pto.tensor_view<1x1x1x128x64xf8E4M3FN> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c128, %c64] + : !pto.tensor_view<1x1x1x128x64xf8E4M3FN> -> !pto.partition_tensor_view<1x1x1x128x64xf8E4M3FN> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c128, %c64] + : !pto.tensor_view<1x1x1x128x64xf8E4M3FN> -> !pto.partition_tensor_view<1x1x1x128x64xf8E4M3FN> + + %tile = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x128x64xf8E4M3FN>) + outs(%tile : !pto.tile_buf) + pto.tstore ins(%tile : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x128x64xf8E4M3FN>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/gen_data.py new file mode 100644 index 0000000000..2039d8d988 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/gen_data.py @@ -0,0 +1,33 @@ +#!/usr/bin/python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import os +import numpy as np + + +CASES = [ + ("fp8_load_16x64", (16, 64)), + ("fp8_load_32x64", (32, 64)), + ("fp8_load_128x64", (128, 64)), +] + + +def main(): + rng = np.random.default_rng(20260707) + for name, shape in CASES: + case_dir = name + os.makedirs(case_dir, exist_ok=True) + data = rng.integers(0, 256, size=shape, dtype=np.uint8) + data.tofile(os.path.join(case_dir, "input.bin")) + data.tofile(os.path.join(case_dir, "golden.bin")) + print(f"[INFO] Generated {name}: shape={shape}, dtype=uint8") + + +if __name__ == "__main__": + main() diff --git a/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/launch.cpp new file mode 100644 index 0000000000..92b084dcfa --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/launch.cpp @@ -0,0 +1,29 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include + +#ifndef AICORE +#define AICORE [aicore] +#endif + +extern "C" __global__ AICORE void FP8_LOAD_16x64(__gm__ float8_e4m3_t *src, __gm__ float8_e4m3_t *dst); +extern "C" __global__ AICORE void FP8_LOAD_32x64(__gm__ float8_e4m3_t *src, __gm__ float8_e4m3_t *dst); +extern "C" __global__ AICORE void FP8_LOAD_128x64(__gm__ float8_e4m3_t *src, __gm__ float8_e4m3_t *dst); + +void LaunchFP8_LOAD_16x64(uint8_t *src, uint8_t *dst, void *stream) { + FP8_LOAD_16x64<<<1, nullptr, stream>>>((__gm__ float8_e4m3_t *)src, (__gm__ float8_e4m3_t *)dst); +} + +void LaunchFP8_LOAD_32x64(uint8_t *src, uint8_t *dst, void *stream) { + FP8_LOAD_32x64<<<1, nullptr, stream>>>((__gm__ float8_e4m3_t *)src, (__gm__ float8_e4m3_t *)dst); +} + +void LaunchFP8_LOAD_128x64(uint8_t *src, uint8_t *dst, void *stream) { + FP8_LOAD_128x64<<<1, nullptr, stream>>>((__gm__ float8_e4m3_t *)src, (__gm__ float8_e4m3_t *)dst); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/main.cpp new file mode 100644 index 0000000000..ea73c6b7c4 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/fp8_load_test/main.cpp @@ -0,0 +1,120 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include +#include +#include + +using namespace PtoTestCommon; + +void LaunchFP8_LOAD_16x64(uint8_t *src, uint8_t *dst, void *stream); +void LaunchFP8_LOAD_32x64(uint8_t *src, uint8_t *dst, void *stream); +void LaunchFP8_LOAD_128x64(uint8_t *src, uint8_t *dst, void *stream); + +using LaunchFn = void (*)(uint8_t *, uint8_t *, void *); + +struct TestCase { + const char *name; + size_t rows; + size_t cols; + LaunchFn launch; +}; + +static const TestCase kCases[] = { + {"fp8_load_16x64", 16, 64, LaunchFP8_LOAD_16x64}, + {"fp8_load_32x64", 32, 64, LaunchFP8_LOAD_32x64}, + {"fp8_load_128x64", 128, 64, LaunchFP8_LOAD_128x64}, +}; +static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); + +static int RunCase(const TestCase &tc, aclrtStream stream) { + int rc = 0; + const size_t bytes = tc.rows * tc.cols * sizeof(uint8_t); + + std::printf("[INFO] === case: %s (%zux%zu) ===\n", tc.name, tc.rows, tc.cols); + + std::string caseDir = std::string("./") + tc.name; + size_t inputFileSize = bytes; + + void *inputHost = nullptr; + void *outputHost = nullptr; + void *inputDevice = nullptr; + void *outputDevice = nullptr; + + aclrtMallocHost(&inputHost, bytes); + aclrtMallocHost(&outputHost, bytes); + aclrtMalloc(&inputDevice, bytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&outputDevice, bytes, ACL_MEM_MALLOC_HUGE_FIRST); + + if (!ReadFile((caseDir + "/input.bin").c_str(), inputFileSize, inputHost, bytes)) { + std::fprintf(stderr, "[ERROR] read input failed\n"); + rc = 1; + } + + if (rc == 0) { + aclrtMemcpy(inputDevice, bytes, inputHost, bytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemset(outputDevice, bytes, 0, bytes); + + tc.launch(static_cast(inputDevice), static_cast(outputDevice), stream); + aclrtSynchronizeStream(stream); + + aclrtMemcpy(outputHost, bytes, outputDevice, bytes, ACL_MEMCPY_DEVICE_TO_HOST); + } + + if (rc == 0 && !WriteFile((caseDir + "/output.bin").c_str(), outputHost, bytes)) { + std::fprintf(stderr, "[ERROR] write output failed\n"); + rc = 1; + } + + if (inputDevice != nullptr) + aclrtFree(inputDevice); + if (outputDevice != nullptr) + aclrtFree(outputDevice); + if (inputHost != nullptr) + aclrtFreeHost(inputHost); + if (outputHost != nullptr) + aclrtFreeHost(outputHost); + + if (rc == 0) + std::printf("[INFO] case %s done\n", tc.name); + return rc; +} + +int main(int argc, char *argv[]) { + const char *caseFilter = (argc > 1) ? argv[1] : nullptr; + int rc = 0; + int deviceId = 0; + aclrtStream stream = nullptr; + + aclInit(nullptr); + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) + deviceId = std::atoi(envDevice); + aclrtSetDevice(deviceId); + aclrtCreateStream(&stream); + + for (size_t i = 0; i < kNumCases; ++i) { + if (caseFilter != nullptr && std::strcmp(kCases[i].name, caseFilter) != 0) + continue; + int ret = RunCase(kCases[i], stream); + if (ret != 0) { + std::fprintf(stderr, "[ERROR] case %s failed\n", kCases[i].name); + rc = 1; + break; + } + } + + if (stream != nullptr) + aclrtDestroyStream(stream); + aclrtResetDevice(deviceId); + aclFinalize(); + return rc; +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/st_common.py b/test/tilelang_st/npu/a5/src/st/testcase/st_common.py index 5d67504e5c..25be2cf264 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/st_common.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/st_common.py @@ -94,6 +94,70 @@ def save_case_data(case_name, data_dict): arr.tofile(os.path.join(case_name, f"{name}.bin")) +# --------------------------------------------------------------------------- +# FP4 helpers +# --------------------------------------------------------------------------- + +FP4_E1M2 = "f4e1m2x2" + + +def _fp4_e2m1_dtype(): + import ml_dtypes + + return ml_dtypes.float4_e2m1fn + + +FP4_E2M1 = _fp4_e2m1_dtype() + + +def is_fp4_e1m2(dtype): + return dtype == FP4_E1M2 + + +def is_fp4_e2m1(dtype): + return dtype == FP4_E2M1 + + +def make_fp4_random(dtype, shape): + if is_fp4_e2m1(dtype): + return np.random.randint(-6, 6, shape).astype(FP4_E2M1) + if is_fp4_e1m2(dtype): + return np.random.randint(-1, 2, shape).astype(np.float32) + raise TypeError(f"unsupported fp4 dtype: {dtype!r}") + + +def zeros_fp4(dtype, shape): + if is_fp4_e2m1(dtype): + return np.zeros(shape, dtype=FP4_E2M1) + if is_fp4_e1m2(dtype): + return np.zeros(shape, dtype=np.float32) + raise TypeError(f"unsupported fp4 dtype: {dtype!r}") + + +def _e1m2_nibbles(values): + values = np.asarray(values, dtype=np.float32) + nibbles = np.zeros(values.shape, dtype=np.uint8) + nibbles[values > 0] = 0x4 + nibbles[values < 0] = 0xC + return nibbles + + +def fp4_nibbles(values, dtype): + if is_fp4_e2m1(dtype): + return np.asarray(values).reshape(-1).view(np.uint8).reshape(np.asarray(values).shape) & 0x0F + if is_fp4_e1m2(dtype): + return _e1m2_nibbles(values) + raise TypeError(f"unsupported fp4 dtype: {dtype!r}") + + +def pack_two_fp4(matrix, dtype): + row, col = matrix.shape + flat = fp4_nibbles(matrix, dtype).reshape(-1) + high = flat[::2] & 0x0F + low = (flat[1::2] & 0x0F) << 4 + return (low | high).reshape(row, col // 2) + + # --------------------------------------------------------------------------- # Terminal styling # --------------------------------------------------------------------------- diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tgemv_mx/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tgemv_mx/cases.py index dd9c7a1fed..f9d13a02e1 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tgemv_mx/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tgemv_mx/cases.py @@ -12,12 +12,12 @@ import numpy as np import ml_dtypes -import en_dtypes +from st_common import FP4_E1M2, FP4_E2M1 fp8_e4m3fn = ml_dtypes.float8_e4m3fn fp8_e5m2 = ml_dtypes.float8_e5m2 -fp4_e1m2x2 = en_dtypes.float4_e1m2 -fp4_e2m1x2 = en_dtypes.float4_e2m1 +fp4_e1m2x2 = FP4_E1M2 +fp4_e2m1x2 = FP4_E2M1 CASES = [ {"name": "gemv_mx_fp4_e1m2_1x128x62", "atype": fp4_e1m2x2, "btype": fp4_e1m2x2, "m": 1, "k": 128, "n": 62, "m_padded": 16, "n_storage": 64, "n_padded": 64, "is_bias": False, "is_fp4": True, "is_split_k": False, "eps": 1e-3}, diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tgemv_mx/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/tgemv_mx/gen_data.py index 3dc8e01f1b..cc2b9f4bc6 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tgemv_mx/gen_data.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tgemv_mx/gen_data.py @@ -13,12 +13,18 @@ import math import numpy as np import ml_dtypes -import en_dtypes +from st_common import ( + FP4_E1M2, + FP4_E2M1, + make_fp4_random, + zeros_fp4, + pack_two_fp4 as mx_pack_two_fp4, +) fp8_e4m3fn = ml_dtypes.float8_e4m3fn fp8_e5m2 = ml_dtypes.float8_e5m2 -fp4_e1m2x2 = en_dtypes.float4_e1m2 -fp4_e2m1x2 = en_dtypes.float4_e2m1 +fp4_e1m2x2 = FP4_E1M2 +fp4_e2m1x2 = FP4_E2M1 np.random.seed(19) @@ -159,26 +165,26 @@ def gen_golden(case): k_aligned = ceil_align(k, 64) if atype == fp4_e2m1x2: - x1 = np.random.randint(-6, 6, [m, k]).astype(atype) + x1 = make_fp4_random(atype, [m, k]) elif atype == fp4_e1m2x2: - x1 = np.random.randint(-1, 2, [m, k]).astype(atype) + x1 = make_fp4_random(atype, [m, k]) else: x1 = np.random.randint(-10, 10, [m, k]).astype(atype) if btype == fp4_e2m1x2: - x2 = np.random.randint(-6, 6, [k, n]).astype(btype) + x2 = make_fp4_random(btype, [k, n]) elif btype == fp4_e1m2x2: - x2 = np.random.randint(-1, 2, [k, n]).astype(btype) + x2 = make_fp4_random(btype, [k, n]) else: x2 = np.random.randint(-10, 10, [k, n]).astype(btype) if is_fp4: - x1_padded = np.zeros([m_padded, k_aligned], dtype=atype) + x1_padded = zeros_fp4(atype, [m_padded, k_aligned]) x1_padded[:m, :k] = x1 - x2_padded = np.zeros([k_aligned, n_storage], dtype=btype) + x2_padded = zeros_fp4(btype, [k_aligned, n_storage]) x2_padded[:k, :n] = x2 - x1_bin = pack_two_fp4(x1_padded) - x2_bin = pack_two_fp4(x2_padded) + x1_bin = mx_pack_two_fp4(x1_padded, atype) + x2_bin = mx_pack_two_fp4(x2_padded, btype) else: x1_padded = np.zeros([m_padded, k_aligned], dtype=atype) x1_padded[:m, :k] = x1 diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tload_mat/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/tload_mat/gen_data.py index 0da417f4cd..481fa2dd6c 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tload_mat/gen_data.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tload_mat/gen_data.py @@ -28,6 +28,12 @@ def bf16_to_uint16(arr): return (f32_view >> 16).astype(np.uint16) +def bf16_to_float32(arr): + """Round float32 values to bfloat16 precision and return float32 values.""" + u16 = bf16_to_uint16(arr) + return (u16.astype(np.uint32) << 16).view(np.float32) + + for case in CASES: setup_case_rng(case) @@ -41,20 +47,30 @@ def bf16_to_uint16(arr): x1_f32 = np.random.uniform(-1, 1, size=(M, K)).astype(np.float32) x2_f32 = np.random.uniform(-1, 1, size=(K, N)).astype(np.float32) - # Golden = matmul result in f32 (ACC output is always f32 for float matmul) - golden_f32 = np.matmul(x1_f32, x2_f32) + dtype_raw = case.get("dtype_raw", None) + if dtype_raw == "bf16": + x1_compute = bf16_to_float32(x1_f32) + x2_compute = bf16_to_float32(x2_f32) + elif case["dtype"] == np.float16: + x1_compute = x1_f32.astype(np.float16).astype(np.float32) + x2_compute = x2_f32.astype(np.float16).astype(np.float32) + else: + x1_compute = x1_f32 + x2_compute = x2_f32 + + # Golden = matmul result in f32 using the exact input precision consumed by the kernel. + golden_f32 = np.matmul(x1_compute, x2_compute) - # For DN2NZ layout, input data must be stored in DN (col-major/transposed) format - # in GM. The kernel's tload handles the DN→NZ conversion internally. + # For DN2NZ cases, x1 is stored in DN (col-major/transposed) format in GM. + # x2 stays ND so the matmul pipeline isolates the DN load path under test. if layout == "dn2nz": x1_gm_f32 = x1_f32.T # [K, M] transposed for DN format - x2_gm_f32 = x2_f32.T # [N, K] transposed for DN format + x2_gm_f32 = x2_f32 else: x1_gm_f32 = x1_f32 x2_gm_f32 = x2_f32 # Prepare input data in source dtype (using DN-layout data if needed) - dtype_raw = case.get("dtype_raw", None) if dtype_raw == "bf16": x1_bin = bf16_to_uint16(x1_gm_f32) x2_bin = bf16_to_uint16(x2_gm_f32) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tload_mat/tload_mat.pto b/test/tilelang_st/npu/a5/src/st/testcase/tload_mat/tload_mat.pto index 1edd5e8326..c53d091c81 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tload_mat/tload_mat.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tload_mat/tload_mat.pto @@ -15,7 +15,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind NZ MAT) func.func @TLOAD_MAT_f16_nd2nz(%dst_gm: !pto.ptr, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -26,15 +26,21 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view %sv_x2 = pto.partition_view %tv_x2, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xf16> %mat_b = pto.alloc_tile : !pto.tile_buf pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -48,7 +54,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind NZ MAT) func.func @TLOAD_MAT_bf16_nd2nz(%dst_gm: !pto.ptr, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -59,15 +65,21 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xbf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view %sv_x2 = pto.partition_view %tv_x2, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xbf16> %mat_b = pto.alloc_tile : !pto.tile_buf pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xbf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -81,7 +93,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind NZ MAT) func.func @TLOAD_MAT_f32_nd2nz(%dst_gm: !pto.ptr, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -92,15 +104,21 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xf32>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view %sv_x2 = pto.partition_view %tv_x2, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xf32> %mat_b = pto.alloc_tile : !pto.tile_buf pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xf32>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -115,28 +133,33 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind NZ MAT) // DN format: GM tensor has shape [K, M] with strides [M, 1] (transposed) func.func @TLOAD_MAT_f16_dn2nz(%dst_gm: !pto.ptr, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index %c32 = arith.constant 32 : index - // DN2NZ: x1 source is col-major [16, 16] with DN strides [16, 1] - %tv_x1 = pto.make_tensor_view %x1_gm, shape = [%c16, %c16], strides = [%c16, %c1] : !pto.tensor_view + // DN2NZ: logical [16,16] source with DN strides over transposed GM storage. + %tv_x1 = pto.make_tensor_view %x1_gm, shape = [%c16, %c16], strides = [%c1, %c16] {layout = #pto.layout} : !pto.tensor_view %sv_x1 = pto.partition_view %tv_x1, offsets = [%c0, %c0], sizes = [%c16, %c16] : !pto.tensor_view -> !pto.partition_tensor_view<16x16xf16> %mat_a = pto.alloc_tile : !pto.tile_buf pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) - // DN2NZ: x2 source is col-major [N, K]=[32,16] with DN strides [16, 1] (transposed) - %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c32, %c16], strides = [%c16, %c1] : !pto.tensor_view - %sv_x2 = pto.partition_view %tv_x2, offsets = [%c0, %c0], sizes = [%c32, %c16] : !pto.tensor_view -> !pto.partition_tensor_view<32x16xf16> + %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view + %sv_x2 = pto.partition_view %tv_x2, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xf16> %mat_b = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%sv_x2 : !pto.partition_tensor_view<32x16xf16>) + pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -150,27 +173,32 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind NZ MAT) func.func @TLOAD_MAT_bf16_dn2nz(%dst_gm: !pto.ptr, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index %c32 = arith.constant 32 : index - %tv_x1 = pto.make_tensor_view %x1_gm, shape = [%c16, %c16], strides = [%c16, %c1] : !pto.tensor_view + %tv_x1 = pto.make_tensor_view %x1_gm, shape = [%c16, %c16], strides = [%c1, %c16] {layout = #pto.layout} : !pto.tensor_view %sv_x1 = pto.partition_view %tv_x1, offsets = [%c0, %c0], sizes = [%c16, %c16] : !pto.tensor_view -> !pto.partition_tensor_view<16x16xbf16> %mat_a = pto.alloc_tile : !pto.tile_buf pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xbf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) - // DN2NZ: x2 source is col-major [N, K]=[32,16] with DN strides [16, 1] (transposed) - %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c32, %c16], strides = [%c16, %c1] : !pto.tensor_view - %sv_x2 = pto.partition_view %tv_x2, offsets = [%c0, %c0], sizes = [%c32, %c16] : !pto.tensor_view -> !pto.partition_tensor_view<32x16xbf16> + %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view + %sv_x2 = pto.partition_view %tv_x2, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xbf16> %mat_b = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%sv_x2 : !pto.partition_tensor_view<32x16xbf16>) + pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xbf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -184,27 +212,32 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind NZ MAT) func.func @TLOAD_MAT_f32_dn2nz(%dst_gm: !pto.ptr, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index %c32 = arith.constant 32 : index - %tv_x1 = pto.make_tensor_view %x1_gm, shape = [%c16, %c16], strides = [%c16, %c1] : !pto.tensor_view + %tv_x1 = pto.make_tensor_view %x1_gm, shape = [%c16, %c16], strides = [%c1, %c16] {layout = #pto.layout} : !pto.tensor_view %sv_x1 = pto.partition_view %tv_x1, offsets = [%c0, %c0], sizes = [%c16, %c16] : !pto.tensor_view -> !pto.partition_tensor_view<16x16xf32> %mat_a = pto.alloc_tile : !pto.tile_buf pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xf32>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) - // DN2NZ: x2 source is col-major [N, K]=[32,16] with DN strides [16, 1] (transposed) - %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c32, %c16], strides = [%c16, %c1] : !pto.tensor_view - %sv_x2 = pto.partition_view %tv_x2, offsets = [%c0, %c0], sizes = [%c32, %c16] : !pto.tensor_view -> !pto.partition_tensor_view<32x16xf32> + %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view + %sv_x2 = pto.partition_view %tv_x2, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xf32> %mat_b = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%sv_x2 : !pto.partition_tensor_view<32x16xf32>) + pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xf32>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_bias_mx/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_bias_mx/cases.py index 9c41d0ca7f..ade454c91d 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_bias_mx/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_bias_mx/cases.py @@ -12,12 +12,12 @@ import numpy as np import ml_dtypes -import en_dtypes +from st_common import FP4_E1M2, FP4_E2M1 fp8_e4m3fn = ml_dtypes.float8_e4m3fn fp8_e5m2 = ml_dtypes.float8_e5m2 -fp4_e1m2x2 = en_dtypes.float4_e1m2 -fp4_e2m1x2 = en_dtypes.float4_e2m1 +fp4_e1m2x2 = FP4_E1M2 +fp4_e2m1x2 = FP4_E2M1 CASES = [ {"name": "bias_fp8_e5m2_e4m3_115x64x30", "atype": fp8_e5m2, "btype": fp8_e4m3fn, "m": 115, "k": 64, "n": 30, "m_padded": 128, "n_padded": 32, "is_bias": True, "is_fp4": False, "eps": 1e-3}, diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_bias_mx/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_bias_mx/gen_data.py index 3313565ed4..5a28e33e33 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_bias_mx/gen_data.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_bias_mx/gen_data.py @@ -13,12 +13,18 @@ import math import numpy as np import ml_dtypes -import en_dtypes +from st_common import ( + FP4_E1M2, + FP4_E2M1, + make_fp4_random, + zeros_fp4, + pack_two_fp4 as mx_pack_two_fp4, +) fp8_e4m3fn = ml_dtypes.float8_e4m3fn fp8_e5m2 = ml_dtypes.float8_e5m2 -fp4_e1m2x2 = en_dtypes.float4_e1m2 -fp4_e2m1x2 = en_dtypes.float4_e2m1 +fp4_e1m2x2 = FP4_E1M2 +fp4_e2m1x2 = FP4_E2M1 np.random.seed(19) @@ -124,28 +130,31 @@ def gen_golden(case): k_aligned = ceil_align(k, 64) if atype == fp4_e2m1x2: - x1 = np.random.randint(-6, 6, [m, k]).astype(atype) + x1 = make_fp4_random(atype, [m, k]) elif atype == fp4_e1m2x2: - x1 = np.random.randint(-1, 2, [m, k]).astype(atype) + x1 = make_fp4_random(atype, [m, k]) else: x1 = np.random.randint(-10, 10, [m, k]).astype(atype) if btype == fp4_e2m1x2: - x2 = np.random.randint(-6, 6, [k, n]).astype(btype) + x2 = make_fp4_random(btype, [k, n]) elif btype == fp4_e1m2x2: - x2 = np.random.randint(-1, 2, [k, n]).astype(btype) + x2 = make_fp4_random(btype, [k, n]) else: x2 = np.random.randint(-10, 10, [k, n]).astype(btype) - x1_padded = np.zeros([m_padded, k_aligned], dtype=atype) - x1_padded[:m, :k] = x1 - x2_padded = np.zeros([k_aligned, n_padded], dtype=btype) - x2_padded[:k, :n] = x2 - if is_fp4: - x1_bin = pack_two_fp4(x1_padded) - x2_bin = pack_two_fp4(x2_padded) + x1_padded = zeros_fp4(atype, [m_padded, k_aligned]) + x1_padded[:m, :k] = x1 + x2_padded = zeros_fp4(btype, [k_aligned, n_padded]) + x2_padded[:k, :n] = x2 + x1_bin = mx_pack_two_fp4(x1_padded, atype) + x2_bin = mx_pack_two_fp4(x2_padded, btype) else: + x1_padded = np.zeros([m_padded, k_aligned], dtype=atype) + x1_padded[:m, :k] = x1 + x2_padded = np.zeros([k_aligned, n_padded], dtype=btype) + x2_padded[:k, :n] = x2 if split_m_physical_rows is not None: x1_bin = pack_mx_lhs_fp8_fractal_chunks(x1_padded, split_m_physical_rows) else: diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_mx/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_mx/cases.py index b930f39190..774acabe6c 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_mx/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_mx/cases.py @@ -12,12 +12,12 @@ import numpy as np import ml_dtypes -import en_dtypes +from st_common import FP4_E1M2, FP4_E2M1 fp8_e4m3fn = ml_dtypes.float8_e4m3fn fp8_e5m2 = ml_dtypes.float8_e5m2 -fp4_e1m2x2 = en_dtypes.float4_e1m2 -fp4_e2m1x2 = en_dtypes.float4_e2m1 +fp4_e1m2x2 = FP4_E1M2 +fp4_e2m1x2 = FP4_E2M1 CASES = [ {"name": "fp8_e5m2_128x64x64", "atype": fp8_e5m2, "btype": fp8_e5m2, "m": 128, "k": 64, "n": 64, "m_padded": 128, "n_padded": 64, "is_bias": False, "is_fp4": False, "eps": 1e-3}, diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_mx/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_mx/gen_data.py index 4b8b740bee..ece96824f9 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_mx/gen_data.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmatmul_mx/gen_data.py @@ -13,12 +13,18 @@ import math import numpy as np import ml_dtypes -import en_dtypes +from st_common import ( + FP4_E1M2, + FP4_E2M1, + make_fp4_random, + zeros_fp4, + pack_two_fp4 as mx_pack_two_fp4, +) fp8_e4m3fn = ml_dtypes.float8_e4m3fn fp8_e5m2 = ml_dtypes.float8_e5m2 -fp4_e1m2x2 = en_dtypes.float4_e1m2 -fp4_e2m1x2 = en_dtypes.float4_e2m1 +fp4_e1m2x2 = FP4_E1M2 +fp4_e2m1x2 = FP4_E2M1 np.random.seed(19) @@ -103,28 +109,31 @@ def gen_golden(case): k_aligned = ceil_align(k, 64) if atype == fp4_e2m1x2: - x1 = np.random.randint(-6, 6, [m, k]).astype(atype) + x1 = make_fp4_random(atype, [m, k]) elif atype == fp4_e1m2x2: - x1 = np.random.randint(-1, 2, [m, k]).astype(atype) + x1 = make_fp4_random(atype, [m, k]) else: x1 = np.random.randint(-10, 10, [m, k]).astype(atype) if btype == fp4_e2m1x2: - x2 = np.random.randint(-6, 6, [k, n]).astype(btype) + x2 = make_fp4_random(btype, [k, n]) elif btype == fp4_e1m2x2: - x2 = np.random.randint(-1, 2, [k, n]).astype(btype) + x2 = make_fp4_random(btype, [k, n]) else: x2 = np.random.randint(-10, 10, [k, n]).astype(btype) - x1_padded = np.zeros([m_padded, k_aligned], dtype=atype) - x1_padded[:m, :k] = x1 - x2_padded = np.zeros([k_aligned, n_padded], dtype=btype) - x2_padded[:k, :n] = x2 - if is_fp4: - x1_bin = pack_two_fp4(x1_padded) - x2_bin = pack_two_fp4(x2_padded) + x1_padded = zeros_fp4(atype, [m_padded, k_aligned]) + x1_padded[:m, :k] = x1 + x2_padded = zeros_fp4(btype, [k_aligned, n_padded]) + x2_padded[:k, :n] = x2 + x1_bin = mx_pack_two_fp4(x1_padded, atype) + x2_bin = mx_pack_two_fp4(x2_padded, btype) else: + x1_padded = np.zeros([m_padded, k_aligned], dtype=atype) + x1_padded[:m, :k] = x1 + x2_padded = np.zeros([k_aligned, n_padded], dtype=btype) + x2_padded[:k, :n] = x2 if case["name"] == "fp8_e4m3_16x32x16": x1_bin = np.ascontiguousarray(x1_padded) x2_bin = np.ascontiguousarray(x2_padded) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/cases.py index 12d5786a35..837471349e 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/cases.py @@ -162,7 +162,7 @@ "N": 32, "K": 16, "quant_mode": 2, - "scaling_dtype": np.float16, + "scaling_dtype": np.float32, "eps": 1e-3, }, { @@ -174,10 +174,9 @@ "N": 32, "K": 16, "quant_mode": 2, - "scaling_dtype": None, + "scaling_dtype": np.float32, "eps": 1e-3, "src_dtype_raw": "bf16", "dst_dtype_raw": "bf16", - "scaling_dtype_raw": "bf16", }, ] diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/gen_data.py index a9f8266769..ac5fee8633 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/gen_data.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/gen_data.py @@ -31,10 +31,10 @@ def bf16_to_uint16(arr): - """Convert float32 array to bfloat16 stored as uint16.""" - # bfloat16 is the upper 16 bits of float32 + """Convert float32 array to rounded bfloat16 stored as uint16.""" f32_view = arr.view(np.uint32) - bf16_uint16 = (f32_view >> 16).astype(np.uint16) + rounding_bias = np.uint32(0x7FFF) + ((f32_view >> 16) & np.uint32(1)) + bf16_uint16 = ((f32_view + rounding_bias) >> 16).astype(np.uint16) return bf16_uint16 @@ -55,31 +55,36 @@ def uint16_to_bf16_as_f32(arr): src_dtype_raw = case.get("src_dtype_raw", None) dst_dtype_raw = case.get("dst_dtype_raw", None) - scaling_dtype_raw = case.get("scaling_dtype_raw", None) - dst_layout = case.get("dst_layout", "nz2nd") # default NZ2ND (row-major) # Generate input matrices as float32 for computation x1_f32 = np.random.uniform(-1, 1, size=(M, K)).astype(np.float32) x2_f32 = np.random.uniform(-1, 1, size=(K, N)).astype(np.float32) - # Compute golden in float32 - golden_f32 = np.matmul(x1_f32, x2_f32) - # Prepare input data in source dtype if src_dtype_raw == "bf16": x1_bin = bf16_to_uint16(x1_f32) x2_bin = bf16_to_uint16(x2_f32) + x1_compute = uint16_to_bf16_as_f32(x1_bin) + x2_compute = uint16_to_bf16_as_f32(x2_bin) elif case["src_dtype"] == np.float16: x1_bin = x1_f32.astype(np.float16) x2_bin = x2_f32.astype(np.float16) + x1_compute = x1_bin.astype(np.float32) + x2_compute = x2_bin.astype(np.float32) elif case["src_dtype"] == np.int8: x1_bin = np.random.randint(-5, 5, size=(M, K)).astype(np.int8) x2_bin = np.random.randint(-5, 5, size=(K, N)).astype(np.int8) - golden_f32 = np.matmul(x1_bin.astype(np.float32), x2_bin.astype(np.float32)) + x1_compute = x1_bin.astype(np.float32) + x2_compute = x2_bin.astype(np.float32) else: x1_bin = x1_f32 x2_bin = x2_f32 + x1_compute = x1_bin.astype(np.float32) + x2_compute = x2_bin.astype(np.float32) + + # Hardware accumulates the rounded source dtype values. + golden_f32 = np.matmul(x1_compute, x2_compute).astype(np.float32) # Prepare golden in destination dtype if dst_dtype_raw == "bf16": @@ -107,15 +112,11 @@ def uint16_to_bf16_as_f32(arr): # Use non-trivial quantization values to actually test scaling functionality. # Values near 1.0 keep results within dtype range while being distinct from 1.0. quant_vector_f32 = np.random.uniform(0.5, 1.5, size=(1, N)).astype(np.float32) - if scaling_dtype_raw == "bf16": - quant_bin = bf16_to_uint16(quant_vector_f32) - elif case.get("scaling_dtype") == np.float16: - quant_bin = quant_vector_f32.astype(np.float16) - else: - quant_bin = bf16_to_uint16(quant_vector_f32) - data_dict["quant_vector"] = quant_bin - # Golden with quantization: result * quant_vector - golden_quant = golden_f32 * quant_vector_f32 + quant_vector_gm = quant_vector_f32.view(np.uint32).astype(np.uint64) + quant_bits = np.bitwise_and(quant_vector_f32.view(np.uint32).copy(), 0xFFFFE000) + quant_compute = quant_bits.view(np.float32) + data_dict["quant_vector"] = quant_vector_gm + golden_quant = golden_f32 * quant_compute if dst_dtype_raw == "bf16": golden_bin = bf16_to_uint16(golden_quant) elif case["dst_dtype"] == np.float16: diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/launch.cpp index 2f53b72257..26749ba2d9 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/launch.cpp @@ -103,18 +103,18 @@ void LaunchTSTORE_ACC2GM_bf16_f32_bf16_nz2nz(void *dst, void *x1, void *x2, void // Case 5: f16 scaling, f16 dst (4 void* args) extern "C" __global__ AICORE void TSTORE_ACC2GM_f16_f32_f16_vec( - __gm__ half *dst, __gm__ half *x1, __gm__ half *x2, __gm__ half *quant); + __gm__ half *dst, __gm__ half *x1, __gm__ half *x2, __gm__ uint64_t *quant); void LaunchTSTORE_ACC2GM_f16_f32_f16_vec(void *dst, void *x1, void *x2, void *quant, void *stream) { TSTORE_ACC2GM_f16_f32_f16_vec<<<1, nullptr, stream>>>( - (__gm__ half *)dst, (__gm__ half *)x1, (__gm__ half *)x2, (__gm__ half *)quant); + (__gm__ half *)dst, (__gm__ half *)x1, (__gm__ half *)x2, (__gm__ uint64_t *)quant); } // Case 6: bf16 scaling, bf16 dst (4 void* args) extern "C" __global__ AICORE void TSTORE_ACC2GM_bf16_f32_bf16_vec( - __gm__ uint16_t *dst, __gm__ uint16_t *x1, __gm__ uint16_t *x2, __gm__ uint16_t *quant); + __gm__ uint16_t *dst, __gm__ uint16_t *x1, __gm__ uint16_t *x2, __gm__ uint64_t *quant); void LaunchTSTORE_ACC2GM_bf16_f32_bf16_vec(void *dst, void *x1, void *x2, void *quant, void *stream) { TSTORE_ACC2GM_bf16_f32_bf16_vec<<<1, nullptr, stream>>>( - (__gm__ uint16_t *)dst, (__gm__ uint16_t *)x1, (__gm__ uint16_t *)x2, (__gm__ uint16_t *)quant); + (__gm__ uint16_t *)dst, (__gm__ uint16_t *)x1, (__gm__ uint16_t *)x2, (__gm__ uint64_t *)quant); } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/main.cpp index 6d873576b2..97853a2286 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/main.cpp @@ -74,8 +74,8 @@ static const TestCase kCases[] = { {"f16_f32_f32_nz2nz", (void*)LaunchTSTORE_ACC2GM_f16_f32_f32_nz2nz, QUANT_NONE, 16, 32, 16, 2, 2, 4, 0, true}, {"bf16_f32_bf16_nz2nz",(void*)LaunchTSTORE_ACC2GM_bf16_f32_bf16_nz2nz,QUANT_NONE, 16, 32, 16, 2, 2, 2, 0, false}, // TSTORE_FP cases (vector quant, NZ2ND) - {"f16_f32_f16_vec", (void*)LaunchTSTORE_ACC2GM_f16_f32_f16_vec, QUANT_VECTOR, 16, 32, 16, 2, 2, 2, 2, false}, - {"bf16_f32_bf16_vec", (void*)LaunchTSTORE_ACC2GM_bf16_f32_bf16_vec, QUANT_VECTOR, 16, 32, 16, 2, 2, 2, 2, false}, + {"f16_f32_f16_vec", (void*)LaunchTSTORE_ACC2GM_f16_f32_f16_vec, QUANT_VECTOR, 16, 32, 16, 2, 2, 2, 8, false}, + {"bf16_f32_bf16_vec", (void*)LaunchTSTORE_ACC2GM_bf16_f32_bf16_vec, QUANT_VECTOR, 16, 32, 16, 2, 2, 2, 8, false}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); @@ -84,7 +84,7 @@ static int RunCase(const TestCase &tc, int deviceId, aclrtStream stream) { const size_t x1Count = tc.M * tc.K; const size_t x2Count = tc.K * tc.N; const size_t dstCount = tc.M * tc.N; - const size_t quantCount = (tc.quant_mode == QUANT_VECTOR) ? 1 * tc.N : 0; + const size_t quantCount = (tc.quant_mode == QUANT_VECTOR) ? tc.N : 0; const size_t x1Size = x1Count * tc.x1ElemSize; const size_t x2Size = x2Count * tc.x2ElemSize; diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/tstore_acc2gm.pto b/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/tstore_acc2gm.pto index 6b6015f085..6534bc81ef 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/tstore_acc2gm.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tstore_acc2gm/tstore_acc2gm.pto @@ -15,7 +15,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -28,6 +28,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) // TLOAD.MAT ND2NZ: x2 [16,32] from GM -> MAT %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -35,10 +38,13 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) // TMATMUL: MAT -> ACC [16,32] f32 %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) // TSTORE.ACC NZ2ND: ACC -> GM [16,32] f32 @@ -53,7 +59,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -65,6 +71,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) // TLOAD.MAT ND2NZ: x2 [16,32] %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -72,10 +81,13 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) // TMATMUL: MAT -> ACC f32 %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) // TSTORE.ACC NZ2ND: ACC f32 -> GM f16 (dtype conversion) @@ -90,7 +102,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -102,6 +114,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xbf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) // TLOAD.MAT ND2NZ: x2 [16,32] bf16 %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -109,10 +124,13 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xbf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) // TMATMUL: MAT bf16 -> ACC f32 %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) // TSTORE.ACC NZ2ND: ACC f32 -> GM f32 @@ -127,7 +145,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -139,6 +157,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xbf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) // TLOAD.MAT ND2NZ: x2 [16,32] bf16 %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -146,10 +167,13 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xbf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) // TMATMUL: MAT bf16 -> ACC f32 %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) // TSTORE.ACC NZ2ND: ACC f32 -> GM bf16 (dtype conversion) @@ -164,7 +188,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -173,20 +197,26 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind %sv_x1 = pto.partition_view %tv_x1, offsets = [%c0, %c0], sizes = [%c16, %c16] : !pto.tensor_view -> !pto.partition_tensor_view<16x16xi8> - %mat_a = pto.alloc_tile : !pto.tile_buf + %mat_a = pto.alloc_tile : !pto.tile_buf pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xi8>) - outs(%mat_a : !pto.tile_buf) + outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) // TLOAD.MAT ND2NZ: x2 [16,32] i8 %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view %sv_x2 = pto.partition_view %tv_x2, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xi8> - %mat_b = pto.alloc_tile : !pto.tile_buf + %mat_b = pto.alloc_tile : !pto.tile_buf pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xi8>) - outs(%mat_b : !pto.tile_buf) + outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) // TMATMUL: MAT i8 -> ACC i32 %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) // TSTORE.ACC NZ2ND: ACC i32 -> GM i32 @@ -202,7 +232,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -214,6 +244,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) // TLOAD.MAT ND2NZ: x2 [16,32] f16 %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -221,14 +254,17 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) // TMATMUL: MAT f16 -> ACC f32 %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) // TSTORE.ACC NZ2DN: ACC f32 -> GM f32 (DN col-major format [N, M]) - %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c32, %c16], strides = [%c16, %c1] : !pto.tensor_view + %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c32, %c16], strides = [%c1, %c32] {layout = #pto.layout} : !pto.tensor_view %sv_dst = pto.partition_view %tv_dst, offsets = [%c0, %c0], sizes = [%c32, %c16] : !pto.tensor_view -> !pto.partition_tensor_view<32x16xf32> pto.tstore ins(%acc_c : !pto.tile_buf) outs(%sv_dst : !pto.partition_tensor_view<32x16xf32>) @@ -239,7 +275,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -251,6 +287,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xbf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) // TLOAD.MAT ND2NZ: x2 [16,32] bf16 %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -258,14 +297,17 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xbf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) // TMATMUL: MAT bf16 -> ACC f32 %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) // TSTORE.ACC NZ2DN: ACC f32 -> GM bf16 (DN col-major format [N, M]) - %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c32, %c16], strides = [%c16, %c1] : !pto.tensor_view + %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c32, %c16], strides = [%c1, %c32] {layout = #pto.layout} : !pto.tensor_view %sv_dst = pto.partition_view %tv_dst, offsets = [%c0, %c0], sizes = [%c32, %c16] : !pto.tensor_view -> !pto.partition_tensor_view<32x16xbf16> pto.tstore ins(%acc_c : !pto.tile_buf) outs(%sv_dst : !pto.partition_tensor_view<32x16xbf16>) @@ -277,7 +319,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -289,6 +331,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) // TLOAD.MAT ND2NZ: x2 [16,32] f16 %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -296,10 +341,13 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) // TMATMUL: MAT f16 -> ACC f32 %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) // TSTORE.ACC NZ2NZ: ACC f32 -> GM f32 (NZ fractal format) @@ -314,7 +362,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -326,6 +374,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xbf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) // TLOAD.MAT ND2NZ: x2 [16,32] bf16 %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -333,10 +384,13 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xbf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) // TMATMUL: MAT bf16 -> ACC f32 %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) // TSTORE.ACC NZ2NZ: ACC f32 -> GM bf16 (NZ fractal format) @@ -350,8 +404,8 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr, %quant_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @TSTORE_ACC2GM_f16_f32_f16_vec(%dst_gm: !pto.ptr, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr, %quant_gm: !pto.ptr) + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -363,6 +417,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) // TLOAD.MAT ND2NZ: x2 [16,32] f16 %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -370,23 +427,29 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) // TMATMUL: MAT f16 -> ACC f32 %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) - // TLOAD quant_vector from GM to scaling buffer (1x32 f16) for TSTORE_FP - %tv_quant = pto.make_tensor_view %quant_gm, shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view - %sv_quant = pto.partition_view %tv_quant, offsets = [%c0, %c0], sizes = [%c1, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf16> - %fp_tile = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%sv_quant : !pto.partition_tensor_view<1x32xf16>) - outs(%fp_tile : !pto.tile_buf) + // TLOAD quant_vector from GM to MAT, then TMOV MAT -> scaling for TSTORE_FP + %tv_quant = pto.make_tensor_view %quant_gm, shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view + %sv_quant = pto.partition_view %tv_quant, offsets = [%c0, %c0], sizes = [%c1, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<1x32xui64> + %mat_fp = pto.alloc_tile : !pto.tile_buf + pto.tload ins(%sv_quant : !pto.partition_tensor_view<1x32xui64>) + outs(%mat_fp : !pto.tile_buf) + %fp_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_fp : !pto.tile_buf) + outs(%fp_tile : !pto.tile_buf) // TSTORE_FP: ACC f32 + scaling f16 -> GM f16 (qf322f16_pre_vec) %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view %sv_dst = pto.partition_view %tv_dst, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xf16> - pto.tstore_fp ins(%acc_c, %fp_tile : !pto.tile_buf, !pto.tile_buf) + pto.tstore_fp ins(%acc_c, %fp_tile : !pto.tile_buf, !pto.tile_buf) outs(%sv_dst : !pto.partition_tensor_view<16x32xf16>) pto.barrier #pto.pipe @@ -394,8 +457,8 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr, %quant_gm: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @TSTORE_ACC2GM_bf16_f32_bf16_vec(%dst_gm: !pto.ptr, %x1_gm: !pto.ptr, %x2_gm: !pto.ptr, %quant_gm: !pto.ptr) + attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -407,6 +470,9 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x1 : !pto.partition_tensor_view<16x16xbf16>) outs(%mat_a : !pto.tile_buf) + %l0a_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_a : !pto.tile_buf) + outs(%l0a_tile : !pto.tile_buf) // TLOAD.MAT ND2NZ: x2 [16,32] bf16 %tv_x2 = pto.make_tensor_view %x2_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -414,23 +480,29 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind pto.tload ins(%sv_x2 : !pto.partition_tensor_view<16x32xbf16>) outs(%mat_b : !pto.tile_buf) + %l0b_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_b : !pto.tile_buf) + outs(%l0b_tile : !pto.tile_buf) // TMATMUL: MAT bf16 -> ACC f32 %acc_c = pto.alloc_tile : !pto.tile_buf - pto.tmatmul ins(%mat_a, %mat_b : !pto.tile_buf, !pto.tile_buf) + pto.tmatmul ins(%l0a_tile, %l0b_tile : !pto.tile_buf, !pto.tile_buf) outs(%acc_c : !pto.tile_buf) - // TLOAD quant_vector from GM to scaling buffer (1x32 bf16) for TSTORE_FP - %tv_quant = pto.make_tensor_view %quant_gm, shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view - %sv_quant = pto.partition_view %tv_quant, offsets = [%c0, %c0], sizes = [%c1, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<1x32xbf16> - %fp_tile = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%sv_quant : !pto.partition_tensor_view<1x32xbf16>) - outs(%fp_tile : !pto.tile_buf) + // TLOAD quant_vector from GM to MAT, then TMOV MAT -> scaling for TSTORE_FP + %tv_quant = pto.make_tensor_view %quant_gm, shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view + %sv_quant = pto.partition_view %tv_quant, offsets = [%c0, %c0], sizes = [%c1, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<1x32xui64> + %mat_fp = pto.alloc_tile : !pto.tile_buf + pto.tload ins(%sv_quant : !pto.partition_tensor_view<1x32xui64>) + outs(%mat_fp : !pto.tile_buf) + %fp_tile = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%mat_fp : !pto.tile_buf) + outs(%fp_tile : !pto.tile_buf) // TSTORE_FP: ACC f32 + scaling bf16 -> GM bf16 (qf322bf16_pre_vec) %tv_dst = pto.make_tensor_view %dst_gm, shape = [%c16, %c32], strides = [%c32, %c1] : !pto.tensor_view %sv_dst = pto.partition_view %tv_dst, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xbf16> - pto.tstore_fp ins(%acc_c, %fp_tile : !pto.tile_buf, !pto.tile_buf) + pto.tstore_fp ins(%acc_c, %fp_tile : !pto.tile_buf, !pto.tile_buf) outs(%sv_dst : !pto.partition_tensor_view<16x32xbf16>) pto.barrier #pto.pipe diff --git a/tilelang-dsl/docs/user_guide/12-cube-operations.md b/tilelang-dsl/docs/user_guide/12-cube-operations.md index b8cedd4712..6c2c79dc85 100644 --- a/tilelang-dsl/docs/user_guide/12-cube-operations.md +++ b/tilelang-dsl/docs/user_guide/12-cube-operations.md @@ -643,8 +643,9 @@ result rectangle to write. **Constraints**: - Clause order is canonical: `unit_flag -> pre_quant -> pre_relu -> layout -> loop3 -> sat`. - `pre_quant` requires payload and mode together. -- Vector `pre_quant` modes require an `fb` pointer with `f16`, `bf16`, or `f32` - element type. +- Vector `pre_quant` modes require an `fb` pointer. The payload may use + `f16`, `bf16`, or `f32` elements, or `ui64` when it is already packed as + 8-byte Quant_PRE parameter entries. - Scalar `pre_quant` modes require an `f16`, `bf16`, or `f32` scalar payload. - `pre_quant` source element type must be `f32` or `i32`, and the selected mode must be compatible with source and destination element types. diff --git a/tilelang-dsl/python/tilelang_dsl/pybind_renderer.py b/tilelang-dsl/python/tilelang_dsl/pybind_renderer.py index 8b9ad249c0..68fe6fdff5 100644 --- a/tilelang-dsl/python/tilelang_dsl/pybind_renderer.py +++ b/tilelang-dsl/python/tilelang_dsl/pybind_renderer.py @@ -1527,6 +1527,7 @@ def _get_element_size(self, semantic_type: SemanticType) -> int: name = dtype.name sizes = { "i1": 1, "i8": 1, "i16": 2, "i32": 4, "i64": 8, + "si64": 8, "ui64": 8, "f16": 2, "bf16": 2, "f32": 4, "hif8": 1, "f8e4m3": 1, "f8e5m2": 1, "f4e1m2x2": 1, "f4e2m1x2": 1, diff --git a/tilelang-dsl/python/tilelang_dsl/semantic.py b/tilelang-dsl/python/tilelang_dsl/semantic.py index 2ff13fb8d4..79cbb0586d 100644 --- a/tilelang-dsl/python/tilelang_dsl/semantic.py +++ b/tilelang-dsl/python/tilelang_dsl/semantic.py @@ -4672,8 +4672,10 @@ def _require_fixpipe_scalar_payload(self, expr: SemanticExpr, context: str) -> N def _require_fixpipe_vector_payload(self, expr: SemanticExpr, context: str) -> None: ptr = self._require_pointer_expr(expr, context, memory_space="scaling") - if ptr.type.element_dtype not in {f16, bf16, f32}: - raise TypeError(f"{context} must be an fb pointer with f16, bf16, or f32 elements in TileLang DSL v1") + if ptr.type.element_dtype not in {f16, bf16, f32, ui64}: + raise TypeError( + f"{context} must be an fb pointer with f16, bf16, f32 elements or packed ui64 Quant_PRE entries in TileLang DSL v1" + ) def _validate_fixpipe_payload( self,