diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index a72e21c7b4..f807c6a3cb 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -6164,6 +6164,82 @@ def TRowSumOp: PTO_TOp<"trowsum", [ }]; } +def TInterleaveOp: PTO_TOp<"tinterleave", [ + PTO_DpsInitOpInterface, + OpPipeInterface, + DeclareOpInterfaceMethods +]> { + let summary = "TINTERLEAVE: Interleave two source tiles into two destination tiles."; + let description = [{ + Interleaves src0 and src1 and splits the result into dst0 and dst1. + The operation uses the same valid shape for all four tiles. + }]; + + let arguments = (ins + PTODpsType:$src0, + PTODpsType:$src1, + PTODpsType:$dst0, + PTODpsType:$dst1 + ); + + let results = (outs); + + let hasVerifier = 1; + + let assemblyFormat = [{ + `ins` `(` $src0 `,` $src1 `:` + qualified(type($src0)) `,` qualified(type($src1)) `)` + `outs` `(` $dst0 `,` $dst1 `:` + qualified(type($dst0)) `,` qualified(type($dst1)) `)` + attr-dict + }]; + + let extraClassDeclaration = [{ + ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } + ::mlir::MutableOperandRange getDpsInitsMutable() { + return ::mlir::MutableOperandRange(getOperation(), 2, 2); + } + }]; +} + +def TDeInterleaveOp: PTO_TOp<"tdeinterleave", [ + AttrSizedOperandSegments, + PTO_DpsInitOpInterface, + OpPipeInterface, + DeclareOpInterfaceMethods +]> { + let summary = "TDEINTERLEAVE: De-interleave one or two source tiles into two destination tiles."; + let description = [{ + The two-source form de-interleaves two source tiles with matching valid + shapes. The single-source form de-interleaves one source tile into two + destination tiles whose valid column count is half the source count. + }]; + + let arguments = (ins + Variadic:$srcs, + Variadic:$dsts + ); + + let results = (outs); + + let hasVerifier = 1; + + let hasCustomAssemblyFormat = 1; + + let extraClassDeclaration = [{ + ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } + ::mlir::MutableOperandRange getDpsInitsMutable() { + return getDstsMutable(); + } + ::mlir::Value getSrc0() { return getSrcs().front(); } + ::mlir::Value getSrc1() { + return getSrcs().size() == 2 ? getSrcs()[1] : ::mlir::Value(); + } + ::mlir::Value getDst0() { return getDsts().front(); } + ::mlir::Value getDst1() { return getDsts()[1]; } + }]; +} + def TRowProdOp: PTO_TOp<"trowprod", [ PTO_DpsInitOpInterface, OpPipeInterface, diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 9191a3f0d3..a4106478e7 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -11147,6 +11147,101 @@ ParseResult mlir::pto::TCvtOp::parse(OpAsmParser &parser, OperationState &result return success(); } +void mlir::pto::TDeInterleaveOp::print(OpAsmPrinter &p) { + p << " ins("; + llvm::interleaveComma(getSrcs(), p, [&](Value src) { p << src; }); + p << " : "; + llvm::interleaveComma(getSrcs().getTypes(), p, [&](Type type) { p << type; }); + p << ") outs(" << getDst0() << ", " << getDst1() << " : " + << getDst0().getType() << ", " << getDst1().getType() << ")"; + p.printOptionalAttrDict((*this)->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); +} + +ParseResult mlir::pto::TDeInterleaveOp::parse(OpAsmParser &parser, + OperationState &result) { + bool invalidHeader = parser.parseKeyword("ins") || parser.parseLParen(); + if (invalidHeader) { + return failure(); + } + + SmallVector srcs; + OpAsmParser::UnresolvedOperand src; + bool invalidSrc = failed(parser.parseOperand(src)); + if (invalidSrc) { + return failure(); + } + srcs.push_back(src); + while (succeeded(parser.parseOptionalComma())) { + bool invalidNextSrc = failed(parser.parseOperand(src)); + if (invalidNextSrc) { + return failure(); + } + srcs.push_back(src); + } + + SmallVector srcTypes; + Type srcType; + bool invalidSrcType = failed(parser.parseColonType(srcType)); + if (invalidSrcType) { + return failure(); + } + srcTypes.push_back(srcType); + while (succeeded(parser.parseOptionalComma())) { + bool invalidNextType = failed(parser.parseType(srcType)); + if (invalidNextType) { + return failure(); + } + srcTypes.push_back(srcType); + } + bool invalidOutsHeader = + parser.parseRParen() || parser.parseKeyword("outs") || parser.parseLParen(); + if (invalidOutsHeader) { + return failure(); + } + + SmallVector dsts; + OpAsmParser::UnresolvedOperand dst; + bool invalidDst = parser.parseOperand(dst) || parser.parseComma(); + if (invalidDst) { + return failure(); + } + dsts.push_back(dst); + bool invalidSecondDst = failed(parser.parseOperand(dst)); + if (invalidSecondDst) { + return failure(); + } + dsts.push_back(dst); + Type dst0Ty; + Type dst1Ty; + bool invalidDstTypes = parser.parseColonType(dst0Ty) || + parser.parseComma() || parser.parseType(dst1Ty) || + parser.parseRParen(); + if (invalidDstTypes) { + return failure(); + } + bool invalidSourceCount = srcs.size() < 1 || srcs.size() > 2; + if (invalidSourceCount) { + return parser.emitError(parser.getCurrentLocation(), + "tdeinterleave expects one or two source operands"); + } + + bool unresolvedSources = + failed(parser.resolveOperands(srcs, srcTypes, parser.getCurrentLocation(), + result.operands)); + bool unresolvedDsts = + failed(parser.resolveOperand(dsts[0], dst0Ty, result.operands)) || + failed(parser.resolveOperand(dsts[1], dst1Ty, result.operands)); + if (unresolvedSources || unresolvedDsts) { + return failure(); + } + result.addAttribute( + "operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr( + {static_cast(srcs.size()), 2})); + return parser.parseOptionalAttrDict(result.attributes); +} + void mlir::pto::TMrgSortOp::print(OpAsmPrinter &p) { if (isFormat1()) { p << " ins(" << getSrc() << ", " << getBlockLen() << " : " << getSrc().getType() @@ -14974,6 +15069,170 @@ mlir::LogicalResult mlir::pto::TRowSumOp::verify() { return dispatchVerifierByArch(getOperation(), verifyByArch, verifyByArch); } +mlir::LogicalResult mlir::pto::TInterleaveOp::verify() { + auto verifyA2A3 = [&]() -> LogicalResult { + return emitOpError("tinterleave is only supported on A5 targets"); + }; + + auto verifyA5 = [&]() -> LogicalResult { + Type src0Ty = getSrc0().getType(); + Type src1Ty = getSrc1().getType(); + Type dst0Ty = getDst0().getType(); + Type dst1Ty = getDst1().getType(); + + bool invalidTile = + failed(verifyVecTileCommon(*this, src0Ty, "src0")) || + failed(verifyVecTileCommon(*this, src1Ty, "src1")) || + failed(verifyVecTileCommon(*this, dst0Ty, "dst0")) || + failed(verifyVecTileCommon(*this, dst1Ty, "dst1")); + if (invalidTile) { + return failure(); + } + + bool mismatchedElementTypes = + failed(verifyTileBufSameElemType(*this, src0Ty, src1Ty, "src0", "src1")) || + failed(verifyTileBufSameElemType(*this, src0Ty, dst0Ty, "src0", "dst0")) || + failed(verifyTileBufSameElemType(*this, src0Ty, dst1Ty, "src0", "dst1")); + if (mismatchedElementTypes) { + return failure(); + } + if (!isSupportedVecElemType(getElemTy(src0Ty), /*allowBf16=*/true, + /*allowInt8=*/true)) { + return emitOpError("expects vec tile element types to be supported"); + } + + bool mismatchedValidShapes = + failed(verifyTileBufSameValidShape(*this, src0Ty, src1Ty, "src0", "src1")) || + failed(verifyTileBufSameValidShape(*this, src0Ty, dst0Ty, "src0", "dst0")) || + failed(verifyTileBufSameValidShape(*this, src0Ty, dst1Ty, "src0", "dst1")); + if (mismatchedValidShapes) { + return failure(); + } + + auto validShape = getValidShapeVec(dst0Ty); + bool hasInvalidRank = validShape.size() != 2; + if (hasInvalidRank) { + return emitOpError("expects src0, src1, dst0, and dst1 to have rank-2 valid_shape"); + } + bool hasOddValidColumns = + validShape[1] != ShapedType::kDynamic && (validShape[1] & 1) != 0; + if (hasOddValidColumns) { + return emitOpError("expects valid_shape[1] to be even"); + } + + return success(); + }; + + return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); +} + +mlir::LogicalResult mlir::pto::TDeInterleaveOp::verify() { + auto verifyA2A3 = [&]() -> LogicalResult { + return emitOpError("tdeinterleave is only supported on A5 targets"); + }; + + auto verifyA5 = [&]() -> LogicalResult { + Type src0Ty = getSrc0().getType(); + Type dst0Ty = getDst0().getType(); + Type dst1Ty = getDst1().getType(); + + bool invalidTile = + failed(verifyVecTileCommon(*this, src0Ty, "src0")) || + failed(verifyVecTileCommon(*this, dst0Ty, "dst0")) || + failed(verifyVecTileCommon(*this, dst1Ty, "dst1")); + if (invalidTile) { + return failure(); + } + bool mismatchedElementTypes = + failed(verifyTileBufSameElemType(*this, src0Ty, dst0Ty, "src0", "dst0")) || + failed(verifyTileBufSameElemType(*this, src0Ty, dst1Ty, "src0", "dst1")); + if (mismatchedElementTypes) { + return failure(); + } + if (!isSupportedVecElemType(getElemTy(src0Ty), /*allowBf16=*/true, + /*allowInt8=*/true)) { + return emitOpError("expects vec tile element types to be supported"); + } + bool hasNonRowMajorTile = + !isRowMajorTileBuf(src0Ty) || !isRowMajorTileBuf(dst0Ty) || + !isRowMajorTileBuf(dst1Ty); + if (hasNonRowMajorTile) { + return emitOpError("expects src and dst tiles to use row-major layout"); + } + + auto src0Valid = getValidShapeVec(src0Ty); + auto dst0Valid = getValidShapeVec(dst0Ty); + auto dst1Valid = getValidShapeVec(dst1Ty); + bool hasInvalidRank = + src0Valid.size() != 2 || dst0Valid.size() != 2 || dst1Valid.size() != 2; + if (hasInvalidRank) { + return emitOpError("expects src and dst tiles to have rank-2 valid_shape"); + } + + bool hasSecondSource = getSrcs().size() == 2; + if (hasSecondSource) { + Type src1Ty = getSrc1().getType(); + bool invalidSource = + failed(verifyVecTileCommon(*this, src1Ty, "src1")) || + failed(verifyTileBufSameElemType(*this, src0Ty, src1Ty, "src0", "src1")) || + failed(verifyTileBufSameValidShape(*this, src0Ty, src1Ty, "src0", "src1")) || + failed(verifyTileBufSameValidShape(*this, src0Ty, dst0Ty, "src0", "dst0")) || + failed(verifyTileBufSameValidShape(*this, src0Ty, dst1Ty, "src0", "dst1")); + if (invalidSource) { + return failure(); + } + if (!isRowMajorTileBuf(src1Ty)) { + return emitOpError("expects src1 to use row-major layout"); + } + bool hasOddValidColumns = + src0Valid[1] != ShapedType::kDynamic && (src0Valid[1] & 1) != 0; + if (hasOddValidColumns) { + return emitOpError("expects two-source valid_shape[1] to be even"); + } + return success(); + } + + bool hasRowMismatchWithDst0 = + src0Valid[0] != ShapedType::kDynamic && + dst0Valid[0] != ShapedType::kDynamic && + src0Valid[0] != dst0Valid[0]; + if (hasRowMismatchWithDst0) { + return emitOpError("expects src0 and dst0 to have the same valid_shape[0]"); + } + bool hasRowMismatchWithDst1 = + src0Valid[0] != ShapedType::kDynamic && + dst1Valid[0] != ShapedType::kDynamic && + src0Valid[0] != dst1Valid[0]; + if (hasRowMismatchWithDst1) { + return emitOpError("expects src0 and dst1 to have the same valid_shape[0]"); + } + bool hasOddValidColumns = + src0Valid[1] != ShapedType::kDynamic && (src0Valid[1] & 1) != 0; + if (hasOddValidColumns) { + return emitOpError("expects single-source valid_shape[1] to be even"); + } + bool hasColumnMismatchWithDst0 = + src0Valid[1] != ShapedType::kDynamic && + dst0Valid[1] != ShapedType::kDynamic && + dst0Valid[1] != src0Valid[1] / 2; + if (hasColumnMismatchWithDst0) { + return emitOpError( + "expects dst0 valid_shape[1] to be half of src0 valid_shape[1]"); + } + bool hasColumnMismatchWithDst1 = + src0Valid[1] != ShapedType::kDynamic && + dst1Valid[1] != ShapedType::kDynamic && + dst1Valid[1] != src0Valid[1] / 2; + if (hasColumnMismatchWithDst1) { + return emitOpError( + "expects dst1 valid_shape[1] to be half of src0 valid_shape[1]"); + } + return success(); + }; + + return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); +} + mlir::LogicalResult mlir::pto::TRowProdOp::verify() { auto verifyA2A3 = [&]() -> LogicalResult { if (!getTmp()) { @@ -17624,6 +17883,22 @@ PTO_DEFINE_UNARY_EFFECTS(TOrSOp, getSrcMutable(), getDstMutable()) PTO_DEFINE_BINARY_EFFECTS(TPartAddOp, getSrc0Mutable(), getSrc1Mutable(), getDstMutable()) PTO_DEFINE_BINARY_EFFECTS(TPartMaxOp, getSrc0Mutable(), getSrc1Mutable(), getDstMutable()) PTO_DEFINE_BINARY_EFFECTS(TPartMinOp, getSrc0Mutable(), getSrc1Mutable(), getDstMutable()) +void TInterleaveOp::getEffects( + SmallVectorImpl> &effects) { + PTO_ADD_READ(getSrc0Mutable()); + PTO_ADD_READ(getSrc1Mutable()); + PTO_ADD_WRITE(getDst0Mutable()); + PTO_ADD_WRITE(getDst1Mutable()); +} +void TDeInterleaveOp::getEffects( + SmallVectorImpl> &effects) { + for (auto &operand : getSrcsMutable()) { + PTO_ADD_READ(operand); + } + for (auto &operand : getDstsMutable()) { + PTO_ADD_WRITE(operand); + } +} void TPartArgMaxOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrc0Mutable()); diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 1a35699d48..1c016243f7 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -11800,6 +11800,48 @@ struct PTORowSumToEmitC : public OpConversionPattern { } }; +struct PTOTInterleaveToEmitC + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite( + pto::TInterleaveOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + createLastUseAwareOpaqueCall( + rewriter, op.getOperation(), TypeRange{}, "TINTERLEAVE", + ValueRange{adaptor.getDst1(), adaptor.getDst0(), adaptor.getSrc1(), + adaptor.getSrc0()}); + rewriter.eraseOp(op); + return success(); + } +}; + +struct PTOTDeInterleaveToEmitC + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite( + pto::TDeInterleaveOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Value dst1 = adaptor.getDsts()[1]; + Value dst0 = adaptor.getDsts()[0]; + Value src0 = adaptor.getSrcs()[0]; + bool hasSecondSource = adaptor.getSrcs().size() == 2; + if (hasSecondSource) { + Value src1 = adaptor.getSrcs()[1]; + createLastUseAwareOpaqueCall( + rewriter, op.getOperation(), TypeRange{}, "TDEINTERLEAVE", + ValueRange{dst1, dst0, src1, src0}); + } else { + createLastUseAwareOpaqueCall( + rewriter, op.getOperation(), TypeRange{}, "TDEINTERLEAVE", + ValueRange{dst1, dst0, src0}); + } + rewriter.eraseOp(op); + return success(); + } +}; + struct PTORowProdToEmitC : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -13592,6 +13634,8 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); + patterns.add(typeConverter, ctx); + patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); diff --git a/test/samples/DeInterleave/tdeinterleave-pto.pto b/test/samples/DeInterleave/tdeinterleave-pto.pto new file mode 100644 index 0000000000..b25cfaecff --- /dev/null +++ b/test/samples/DeInterleave/tdeinterleave-pto.pto @@ -0,0 +1,34 @@ +// 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. Please make sure you comply with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS 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"} { + func.func private @tdeinterleave_single_source( + %src: !pto.tile_buf, + %dst0: !pto.tile_buf, + %dst1: !pto.tile_buf) { + pto.tdeinterleave ins(%src : !pto.tile_buf) + outs(%dst0, %dst1 : + !pto.tile_buf, + !pto.tile_buf) + return + } + + func.func private @tdeinterleave_two_sources( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst0: !pto.tile_buf, + %dst1: !pto.tile_buf) { + pto.tdeinterleave ins(%src0, %src1 : + !pto.tile_buf, + !pto.tile_buf) + outs(%dst0, %dst1 : + !pto.tile_buf, + !pto.tile_buf) + return + } +} diff --git a/test/samples/DeInterleave/tdeinterleave_single_runtime.py b/test/samples/DeInterleave/tdeinterleave_single_runtime.py new file mode 100644 index 0000000000..d590e00e93 --- /dev/null +++ b/test/samples/DeInterleave/tdeinterleave_single_runtime.py @@ -0,0 +1,94 @@ +# 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 + +from ptoas.mlir.ir import ( + Context, + InsertionPoint, + IndexType, + Location, + Module, + StringAttr, + UnitAttr, +) +from ptoas.mlir.dialects import arith, func, pto +from ptoas.mlir.ir import F32Type + + +def build(): + with Context() as ctx: + pto.register_dialect(ctx, load=True) + with Location.unknown(ctx): + module = Module.create() + arch = os.environ.get("PTOAS_SAMPLE_ARCH", "a5") + module.operation.attributes["pto.target_arch"] = StringAttr.get(arch) + f32 = F32Type.get(ctx) + ptr_f32 = pto.PtrType.get(f32, ctx) + tensor_view = pto.TensorViewType.get(2, f32, ctx) + src_partition_view = pto.PartitionTensorViewType.get([16, 128], f32, ctx) + dst_partition_view = pto.PartitionTensorViewType.get([16, 64], f32, ctx) + vec = pto.AddressSpaceAttr.get(pto.AddressSpace.VEC, ctx) + config = pto.TileBufConfigAttr.get( + pto.BLayoutAttr.get(pto.BLayout.RowMajor, ctx), + pto.SLayoutAttr.get(pto.SLayout.NoneBox, ctx), + pto.TileConfig.fractalABSize, + pto.PadValueAttr.get(pto.PadValue.Null, ctx), + ctx, + ) + src_tile_type = pto.TileBufType.get([16, 128], f32, vec, [16, 128], config, ctx) + dst_tile_type = pto.TileBufType.get([16, 64], f32, vec, [16, 64], config, ctx) + + function_type = func.FunctionType.get([ptr_f32] * 3, []) + with InsertionPoint(module.body): + function = func.FuncOp("tdeinterleave_single_runtime_kernel", function_type) + function.operation.attributes["pto.entry"] = UnitAttr.get(ctx) + entry = function.add_entry_block() + + with InsertionPoint(entry): + c0 = arith.ConstantOp(IndexType.get(ctx), 0).result + c1 = arith.ConstantOp(IndexType.get(ctx), 1).result + c16 = arith.ConstantOp(IndexType.get(ctx), 16).result + c64 = arith.ConstantOp(IndexType.get(ctx), 64).result + c128 = arith.ConstantOp(IndexType.get(ctx), 128).result + src_ptr, dst0_ptr, dst1_ptr = entry.arguments + + src_view = pto.MakeTensorViewOp( + tensor_view, src_ptr, [c16, c128], [c128, c1] + ).result + dst0_view = pto.MakeTensorViewOp( + tensor_view, dst0_ptr, [c16, c64], [c64, c1] + ).result + dst1_view = pto.MakeTensorViewOp( + tensor_view, dst1_ptr, [c16, c64], [c64, c1] + ).result + src_partition = pto.PartitionViewOp( + src_partition_view, src_view, offsets=[c0, c0], sizes=[c16, c128] + ).result + dst0_partition = pto.PartitionViewOp( + dst_partition_view, dst0_view, offsets=[c0, c0], sizes=[c16, c64] + ).result + dst1_partition = pto.PartitionViewOp( + dst_partition_view, dst1_view, offsets=[c0, c0], sizes=[c16, c64] + ).result + + src_tile = pto.AllocTileOp(src_tile_type).result + dst0_tile = pto.AllocTileOp(dst_tile_type).result + dst1_tile = pto.AllocTileOp(dst_tile_type).result + pto.TLoadOp(None, src_partition, src_tile) + pto.TDeInterleaveOp([src_tile], [dst0_tile, dst1_tile]) + pto.TStoreOp(None, dst0_tile, dst0_partition) + pto.TStoreOp(None, dst1_tile, dst1_partition) + func.ReturnOp([]) + + module.operation.verify() + return module + + +if __name__ == "__main__": + print(build()) diff --git a/test/samples/DeInterleave/tdeinterleave_single_runtime_compare.py b/test/samples/DeInterleave/tdeinterleave_single_runtime_compare.py new file mode 100644 index 0000000000..7cea88eedd --- /dev/null +++ b/test/samples/DeInterleave/tdeinterleave_single_runtime_compare.py @@ -0,0 +1,15 @@ +#!/usr/bin/env 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 numpy as np +from validation_runtime import compare_outputs + + +if __name__ == "__main__": + compare_outputs(np.float32, atol=0.001) diff --git a/test/samples/DeInterleave/tdeinterleave_single_runtime_golden.py b/test/samples/DeInterleave/tdeinterleave_single_runtime_golden.py new file mode 100644 index 0000000000..ce9ced8eb7 --- /dev/null +++ b/test/samples/DeInterleave/tdeinterleave_single_runtime_golden.py @@ -0,0 +1,37 @@ +#!/usr/bin/env 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 numpy as np +from pathlib import Path +import sys + +for search_root in (Path(__file__).resolve().parent, Path(__file__).resolve().parents[1]): + if (search_root / "validation_runtime.py").is_file(): + sys.path.insert(0, str(search_root)) + break + +from validation_runtime import default_buffers, float_values, load_case_meta, matrix32, rng, write_buffers, write_golden + + +def main(): + meta = load_case_meta() + [src_name] = meta.inputs + src = matrix32( + float_values(rng(), meta.elem_counts[src_name], style="signed"), + rows=16, + cols=128, + ) + buffers = default_buffers(meta) + buffers[src_name] = src.reshape(-1) + write_buffers(meta, buffers) + write_golden(meta, {meta.outputs[0]: src[:, 0::2].reshape(-1), meta.outputs[1]: src[:, 1::2].reshape(-1)}) + + +if __name__ == "__main__": + main() diff --git a/test/samples/DeInterleave/tdeinterleave_two_runtime.py b/test/samples/DeInterleave/tdeinterleave_two_runtime.py new file mode 100644 index 0000000000..e201889d22 --- /dev/null +++ b/test/samples/DeInterleave/tdeinterleave_two_runtime.py @@ -0,0 +1,83 @@ +# 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 + +from ptoas.mlir.ir import ( + Context, + InsertionPoint, + IndexType, + Location, + Module, + StringAttr, + UnitAttr, +) +from ptoas.mlir.dialects import arith, func, pto +from ptoas.mlir.ir import F32Type + + +def build(): + with Context() as ctx: + pto.register_dialect(ctx, load=True) + with Location.unknown(ctx): + module = Module.create() + arch = os.environ.get("PTOAS_SAMPLE_ARCH", "a5") + module.operation.attributes["pto.target_arch"] = StringAttr.get(arch) + f32 = F32Type.get(ctx) + ptr_f32 = pto.PtrType.get(f32, ctx) + tensor_view = pto.TensorViewType.get(2, f32, ctx) + partition_view = pto.PartitionTensorViewType.get([16, 64], f32, ctx) + vec = pto.AddressSpaceAttr.get(pto.AddressSpace.VEC, ctx) + config = pto.TileBufConfigAttr.get( + pto.BLayoutAttr.get(pto.BLayout.RowMajor, ctx), + pto.SLayoutAttr.get(pto.SLayout.NoneBox, ctx), + pto.TileConfig.fractalABSize, + pto.PadValueAttr.get(pto.PadValue.Null, ctx), + ctx, + ) + tile_type = pto.TileBufType.get([16, 64], f32, vec, [16, 64], config, ctx) + + function_type = func.FunctionType.get([ptr_f32] * 4, []) + with InsertionPoint(module.body): + function = func.FuncOp("tdeinterleave_two_runtime_kernel", function_type) + function.operation.attributes["pto.entry"] = UnitAttr.get(ctx) + entry = function.add_entry_block() + + with InsertionPoint(entry): + c0 = arith.ConstantOp(IndexType.get(ctx), 0).result + c1 = arith.ConstantOp(IndexType.get(ctx), 1).result + c16 = arith.ConstantOp(IndexType.get(ctx), 16).result + c64 = arith.ConstantOp(IndexType.get(ctx), 64).result + src0_ptr, src1_ptr, dst0_ptr, dst1_ptr = entry.arguments + views = [ + pto.MakeTensorViewOp(tensor_view, ptr, [c16, c64], [c64, c1]).result + for ptr in (src0_ptr, src1_ptr, dst0_ptr, dst1_ptr) + ] + partitions = [ + pto.PartitionViewOp( + partition_view, view, offsets=[c0, c0], sizes=[c16, c64] + ).result + for view in views + ] + src0_tile = pto.AllocTileOp(tile_type).result + src1_tile = pto.AllocTileOp(tile_type).result + dst0_tile = pto.AllocTileOp(tile_type).result + dst1_tile = pto.AllocTileOp(tile_type).result + pto.TLoadOp(None, partitions[0], src0_tile) + pto.TLoadOp(None, partitions[1], src1_tile) + pto.TDeInterleaveOp([src0_tile, src1_tile], [dst0_tile, dst1_tile]) + pto.TStoreOp(None, dst0_tile, partitions[2]) + pto.TStoreOp(None, dst1_tile, partitions[3]) + func.ReturnOp([]) + + module.operation.verify() + return module + + +if __name__ == "__main__": + print(build()) diff --git a/test/samples/DeInterleave/tdeinterleave_two_runtime_compare.py b/test/samples/DeInterleave/tdeinterleave_two_runtime_compare.py new file mode 100644 index 0000000000..7cea88eedd --- /dev/null +++ b/test/samples/DeInterleave/tdeinterleave_two_runtime_compare.py @@ -0,0 +1,15 @@ +#!/usr/bin/env 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 numpy as np +from validation_runtime import compare_outputs + + +if __name__ == "__main__": + compare_outputs(np.float32, atol=0.001) diff --git a/test/samples/DeInterleave/tdeinterleave_two_runtime_golden.py b/test/samples/DeInterleave/tdeinterleave_two_runtime_golden.py new file mode 100644 index 0000000000..1be94d3ba5 --- /dev/null +++ b/test/samples/DeInterleave/tdeinterleave_two_runtime_golden.py @@ -0,0 +1,57 @@ +#!/usr/bin/env 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 numpy as np +from pathlib import Path +import sys + +for search_root in (Path(__file__).resolve().parent, Path(__file__).resolve().parents[1]): + if (search_root / "validation_runtime.py").is_file(): + sys.path.insert(0, str(search_root)) + break + +from validation_runtime import default_buffers, float_values, load_case_meta, matrix32, rng, write_buffers, write_golden + + +def main(): + meta = load_case_meta() + src0_name, src1_name = meta.inputs + src0 = matrix32( + float_values(rng(), meta.elem_counts[src0_name], style="signed"), + rows=16, + cols=64, + ) + src1 = matrix32( + float_values(rng(), meta.elem_counts[src1_name], style="signed"), + rows=16, + cols=64, + ) + rows, cols = src0.shape + half = cols // 2 + dst0 = np.zeros_like(src0) + dst1 = np.zeros_like(src1) + dst0[:, :half] = src0[:, 0::2] + dst1[:, :half] = src0[:, 1::2] + dst0[:, half:] = src1[:, 0::2] + dst1[:, half:] = src1[:, 1::2] + buffers = default_buffers(meta) + buffers[src0_name] = src0.reshape(-1) + buffers[src1_name] = src1.reshape(-1) + write_buffers(meta, buffers) + write_golden( + meta, + { + meta.outputs[0]: dst0.reshape(-1), + meta.outputs[1]: dst1.reshape(-1), + }, + ) + + +if __name__ == "__main__": + main() diff --git a/test/samples/Interleave/tinterleave-pto.pto b/test/samples/Interleave/tinterleave-pto.pto new file mode 100644 index 0000000000..d739a5ef87 --- /dev/null +++ b/test/samples/Interleave/tinterleave-pto.pto @@ -0,0 +1,23 @@ +// 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. Please make sure you comply with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS 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"} { + func.func private @tinterleave_sample( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst0: !pto.tile_buf, + %dst1: !pto.tile_buf) { + pto.tinterleave ins(%src0, %src1 : + !pto.tile_buf, + !pto.tile_buf) + outs(%dst0, %dst1 : + !pto.tile_buf, + !pto.tile_buf) + return + } +} diff --git a/test/samples/Interleave/tinterleave_runtime.py b/test/samples/Interleave/tinterleave_runtime.py new file mode 100644 index 0000000000..4dd5fbf40c --- /dev/null +++ b/test/samples/Interleave/tinterleave_runtime.py @@ -0,0 +1,101 @@ +# 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 + +from ptoas.mlir.ir import ( + Context, + InsertionPoint, + IndexType, + Location, + Module, + StringAttr, + UnitAttr, +) +from ptoas.mlir.dialects import arith, func, pto +from ptoas.mlir.ir import F32Type + + +def build(): + with Context() as ctx: + pto.register_dialect(ctx, load=True) + with Location.unknown(ctx): + module = Module.create() + arch = os.environ.get("PTOAS_SAMPLE_ARCH", "a5") + module.operation.attributes["pto.target_arch"] = StringAttr.get(arch) + f32 = F32Type.get(ctx) + ptr_f32 = pto.PtrType.get(f32, ctx) + tensor_view = pto.TensorViewType.get(2, f32, ctx) + partition_view = pto.PartitionTensorViewType.get([16, 64], f32, ctx) + vec = pto.AddressSpaceAttr.get(pto.AddressSpace.VEC, ctx) + config = pto.TileBufConfigAttr.get( + pto.BLayoutAttr.get(pto.BLayout.RowMajor, ctx), + pto.SLayoutAttr.get(pto.SLayout.NoneBox, ctx), + pto.TileConfig.fractalABSize, + pto.PadValueAttr.get(pto.PadValue.Null, ctx), + ctx, + ) + tile_type = pto.TileBufType.get([16, 64], f32, vec, [16, 64], config, ctx) + + function_type = func.FunctionType.get([ptr_f32] * 4, []) + with InsertionPoint(module.body): + function = func.FuncOp("tinterleave_runtime_kernel", function_type) + function.operation.attributes["pto.entry"] = UnitAttr.get(ctx) + entry = function.add_entry_block() + + with InsertionPoint(entry): + c0 = arith.ConstantOp(IndexType.get(ctx), 0).result + c1 = arith.ConstantOp(IndexType.get(ctx), 1).result + c16 = arith.ConstantOp(IndexType.get(ctx), 16).result + c64 = arith.ConstantOp(IndexType.get(ctx), 64).result + src0_ptr, src1_ptr, dst0_ptr, dst1_ptr = entry.arguments + + src0_view = pto.MakeTensorViewOp( + tensor_view, src0_ptr, [c16, c64], [c64, c1] + ).result + src1_view = pto.MakeTensorViewOp( + tensor_view, src1_ptr, [c16, c64], [c64, c1] + ).result + dst0_view = pto.MakeTensorViewOp( + tensor_view, dst0_ptr, [c16, c64], [c64, c1] + ).result + dst1_view = pto.MakeTensorViewOp( + tensor_view, dst1_ptr, [c16, c64], [c64, c1] + ).result + + src0_partition = pto.PartitionViewOp( + partition_view, src0_view, offsets=[c0, c0], sizes=[c16, c64] + ).result + src1_partition = pto.PartitionViewOp( + partition_view, src1_view, offsets=[c0, c0], sizes=[c16, c64] + ).result + dst0_partition = pto.PartitionViewOp( + partition_view, dst0_view, offsets=[c0, c0], sizes=[c16, c64] + ).result + dst1_partition = pto.PartitionViewOp( + partition_view, dst1_view, offsets=[c0, c0], sizes=[c16, c64] + ).result + + src0_tile = pto.AllocTileOp(tile_type).result + src1_tile = pto.AllocTileOp(tile_type).result + dst0_tile = pto.AllocTileOp(tile_type).result + dst1_tile = pto.AllocTileOp(tile_type).result + + pto.TLoadOp(None, src0_partition, src0_tile) + pto.TLoadOp(None, src1_partition, src1_tile) + pto.TInterleaveOp(src0_tile, src1_tile, dst0_tile, dst1_tile) + pto.TStoreOp(None, dst0_tile, dst0_partition) + pto.TStoreOp(None, dst1_tile, dst1_partition) + func.ReturnOp([]) + + module.operation.verify() + return module + + +if __name__ == "__main__": + print(build()) diff --git a/test/samples/Interleave/tinterleave_runtime_compare.py b/test/samples/Interleave/tinterleave_runtime_compare.py new file mode 100644 index 0000000000..7cea88eedd --- /dev/null +++ b/test/samples/Interleave/tinterleave_runtime_compare.py @@ -0,0 +1,15 @@ +#!/usr/bin/env 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 numpy as np +from validation_runtime import compare_outputs + + +if __name__ == "__main__": + compare_outputs(np.float32, atol=0.001) diff --git a/test/samples/Interleave/tinterleave_runtime_golden.py b/test/samples/Interleave/tinterleave_runtime_golden.py new file mode 100644 index 0000000000..00a9f3f933 --- /dev/null +++ b/test/samples/Interleave/tinterleave_runtime_golden.py @@ -0,0 +1,52 @@ +#!/usr/bin/env 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 numpy as np +from pathlib import Path +import sys + +for search_root in (Path(__file__).resolve().parent, Path(__file__).resolve().parents[1]): + if (search_root / "validation_runtime.py").is_file(): + sys.path.insert(0, str(search_root)) + break + +from validation_runtime import default_buffers, float_values, load_case_meta, matrix32, rng, write_buffers, write_golden + + +def main(): + meta = load_case_meta() + src0_name, src1_name = meta.inputs + generator = rng() + src0 = float_values(generator, meta.elem_counts[src0_name], style="signed") + src1 = float_values(generator, meta.elem_counts[src1_name], style="signed") + src0_matrix = matrix32(src0, rows=16, cols=64) + src1_matrix = matrix32(src1, rows=16, cols=64) + rows, cols = src0_matrix.shape + half = cols // 2 + dst0 = np.empty_like(src0_matrix) + dst1 = np.empty_like(src1_matrix) + dst0[:, 0::2] = src0_matrix[:, :half] + dst0[:, 1::2] = src1_matrix[:, :half] + dst1[:, 0::2] = src0_matrix[:, half:] + dst1[:, 1::2] = src1_matrix[:, half:] + buffers = default_buffers(meta) + buffers[src0_name] = src0 + buffers[src1_name] = src1 + write_buffers(meta, buffers) + write_golden( + meta, + { + meta.outputs[0]: dst0.reshape(-1), + meta.outputs[1]: dst1.reshape(-1), + }, + ) + + +if __name__ == "__main__": + main() diff --git a/test/samples/runop.sh b/test/samples/runop.sh index 4ec92ab597..0fa8990a04 100755 --- a/test/samples/runop.sh +++ b/test/samples/runop.sh @@ -34,7 +34,7 @@ for model_path in "${BASE_DIR}"/Qwen* "${BASE_DIR}"/Deepseek*; do ;; esac done -PTO_PTO_DIRS="${PTO_PTO_DIRS:-Sync${MODEL_PTO_DIRS} CommSync Prelu Rem Rems Gemvmx MatmulMxLowPrecision TquantMx TquantMxDn Movfp}" +PTO_PTO_DIRS="${PTO_PTO_DIRS:-Sync${MODEL_PTO_DIRS} CommSync Prelu Rem Rems Gemvmx MatmulMxLowPrecision TquantMx TquantMxDn Movfp Interleave DeInterleave PairReduceSum}" ENABLE_BC=0 usage() { @@ -80,7 +80,7 @@ sample_dir_arch() { case "$1" in TPipe|TAxpy|TColArgMax|TColArgMin|TConcatIdx|\ TRowArgMax|TRowArgMin|Qwen*A3|Deepseek*A3) printf 'a3\n' ;; - Qwen*A5|Deepseek*A5|TquantMx|TquantMxDn) printf 'a5\n' ;; + Qwen*A5|Deepseek*A5|TquantMx|TquantMxDn|Interleave|DeInterleave|PairReduceSum) printf 'a5\n' ;; esac }