diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index 3de31a89bf..01a1f91ed6 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -107,6 +107,7 @@ LogicalResult validateVPTOEmissionIR(ModuleOp module, llvm::raw_ostream *diagOS = nullptr); std::unique_ptr createPTOValidateVPTOIRPass(); std::unique_ptr createPTOValidateVPTOEmissionIRPass(); +std::unique_ptr createPTOInsertVMemBarPass(); std::unique_ptr createExpandTileOpPass(); std::unique_ptr createExpandTileOpPass(const ExpandTileOpOptions &options); std::unique_ptr createFoldTileBufIntrinsicsPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 04165718b9..f6ea8913a7 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -941,4 +941,20 @@ def VPTOPtrCastCleanup "mlir::memref::MemRefDialect"]; } +def PTOInsertVMemBar + : Pass<"pto-insert-v-membar", "mlir::func::FuncOp"> { + let summary = + "Insert VPTO V-pipeline mem_bar for overlapping UB tile hazards"; + let description = [{ + Analyzes structured PIPE_V tile operations before TileLang expansion and + inserts `pto.mem_bar` for same-block UB RAW, WAR, and WAW hazards. The pass + uses DPS init operands as writes, other tile operands as reads, and constant + `pto.alloc_tile addr` ranges when available. + }]; + let constructor = "mlir::pto::createPTOInsertVMemBarPass()"; + let dependentDialects = ["mlir::pto::PTODialect", + "mlir::arith::ArithDialect", + "mlir::func::FuncDialect"]; +} + #endif // MLIR_DIALECT_PTO_PASSES diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index 3a6c04aed0..32c2842993 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -40,6 +40,7 @@ add_mlir_dialect_library(PTOTransforms InsertSync/PTOInsertSync.cpp PTOInjectBarrierAllSync.cpp InsertSync/InsertSyncDebug.cpp + PTOInsertVMemBar.cpp PTOViewToMemref.cpp PTOValidateIntToPtrUses.cpp ExpandTileOp.cpp diff --git a/lib/PTO/Transforms/PTOInsertVMemBar.cpp b/lib/PTO/Transforms/PTOInsertVMemBar.cpp new file mode 100644 index 0000000000..b467c6fae5 --- /dev/null +++ b/lib/PTO/Transforms/PTOInsertVMemBar.cpp @@ -0,0 +1,310 @@ +// 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. + +//===- PTOInsertVMemBar.cpp - VPTO V-pipeline mem_bar insertion -----------===// +// +// Structured tile ops can later be expanded and inlined into adjacent VPTO +// vlds/vsts sequences in the same vecscope. A5 V-pipeline memory operations are +// not a sufficient ordering guarantee for overlapping UB RAW/WAR/WAW hazards, +// so this pass inserts pto.mem_bar before TileLang expansion while tile operand +// read/write roles and alloc_tile addresses are still explicit. +// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/IR/PTOTypeUtils.h" +#include "PTO/Transforms/Passes.h" + +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Matchers.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" + +#include + +using namespace mlir; +using namespace mlir::pto; + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_PTOINSERTVMEMBAR +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +namespace { + +struct AddrRange { + uint64_t base = 0; + uint64_t size = 0; + bool valid = false; +}; + +struct TileAccess { + AddrRange range; + bool isWrite = false; +}; + +struct HazardSet { + bool raw = false; // previous write -> current read + bool war = false; // previous read -> current write + bool waw = false; // previous write -> current write + + bool empty() const { return !raw && !war && !waw; } + unsigned count() const { + return static_cast(raw) + static_cast(war) + + static_cast(waw); + } + + void merge(const HazardSet &rhs) { + raw |= rhs.raw; + war |= rhs.war; + waw |= rhs.waw; + } +}; + +static bool overlaps(const AddrRange &lhs, const AddrRange &rhs) { + if (!lhs.valid || !rhs.valid) + return true; + return lhs.base < rhs.base + rhs.size && rhs.base < lhs.base + lhs.size; +} + +static std::optional getConstantUInt(Value value) { + APInt intValue; + if (!value || !matchPattern(value, m_ConstantInt(&intValue)) || + intValue.isNegative()) + return std::nullopt; + return intValue.getZExtValue(); +} + +static uint64_t getTileByteSize(TileBufType tileTy) { + uint64_t elemBytes = getPTOStorageElemByteSize(tileTy.getElementType()); + if (elemBytes == 0) + return 0; + + uint64_t elems = 1; + for (int64_t dim : tileTy.getShape()) { + if (dim <= 0) + return 0; + elems *= static_cast(dim); + } + return elems * elemBytes; +} + +static std::optional resolveTileBase(Value tileVal, + unsigned depth = 0) { + if (!tileVal || depth > 8) + return std::nullopt; + + if (auto alloc = tileVal.getDefiningOp()) + return getConstantUInt(alloc.getAddr()); + + if (auto subview = tileVal.getDefiningOp()) { + auto parentBase = resolveTileBase(subview.getSource(), depth + 1); + if (!parentBase) + return std::nullopt; + + auto sourceTy = dyn_cast(subview.getSource().getType()); + if (!sourceTy || sourceTy.getShape().size() < 2) + return std::nullopt; + int64_t sourceCols = sourceTy.getShape()[1]; + if (sourceCols <= 0) + return std::nullopt; + + auto offsets = subview.getOffsets(); + if (offsets.size() < 2) + return std::nullopt; + auto rowOffset = getConstantUInt(offsets[0]); + auto colOffset = getConstantUInt(offsets[1]); + if (!rowOffset || !colOffset) + return std::nullopt; + + uint64_t elemBytes = getPTOStorageElemByteSize(sourceTy.getElementType()); + if (elemBytes == 0) + return std::nullopt; + + uint64_t linearElemOffset = + (*rowOffset) * static_cast(sourceCols) + (*colOffset); + return *parentBase + linearElemOffset * elemBytes; + } + + return std::nullopt; +} + +static AddrRange resolveTileAddrRange(Value tileVal) { + AddrRange range; + auto tileTy = dyn_cast(tileVal.getType()); + if (!tileTy) + return range; + + range.size = getTileByteSize(tileTy); + auto base = resolveTileBase(tileVal); + if (!base || range.size == 0) + return range; + + range.base = *base; + range.valid = true; + return range; +} + +static SmallVector collectTileAccesses(Operation *op) { + SmallVector accesses; + auto dps = dyn_cast(op); + if (!dps) + return accesses; + + SmallPtrSet writeOperands; + for (OpOperand &init : dps.getDpsInitsMutable()) { + if (isa(init.get().getType())) + writeOperands.insert(&init); + } + + for (OpOperand &operand : op->getOpOperands()) { + if (!isa(operand.get().getType())) + continue; + accesses.push_back( + TileAccess{resolveTileAddrRange(operand.get()), + writeOperands.contains(&operand)}); + } + return accesses; +} + +static HazardSet classifyHazards(const SmallVector &previous, + const SmallVector ¤t) { + HazardSet hazards; + for (const TileAccess &prev : previous) { + for (const TileAccess &cur : current) { + if (!overlaps(prev.range, cur.range)) + continue; + if (prev.isWrite && !cur.isWrite) + hazards.raw = true; + if (!prev.isWrite && cur.isWrite) + hazards.war = true; + if (prev.isWrite && cur.isWrite) + hazards.waw = true; + } + } + return hazards; +} + +static bool memBarCovers(MemBarKind kind, const HazardSet &hazards) { + if (hazards.empty()) + return true; + if (kind == MemBarKind::VV_ALL) + return true; + if (hazards.raw && kind != MemBarKind::VST_VLD) + return false; + if (hazards.war && kind != MemBarKind::VLD_VST) + return false; + if (hazards.waw && kind != MemBarKind::VST_VST) + return false; + return hazards.count() == 1; +} + +static MemBarKind getMemBarKindFor(const HazardSet &hazards) { + if (hazards.count() != 1) + return MemBarKind::VV_ALL; + if (hazards.raw) + return MemBarKind::VST_VLD; + if (hazards.war) + return MemBarKind::VLD_VST; + return MemBarKind::VST_VST; +} + +static bool hasCoveringMemBarBetween(Operation *previous, Operation *current, + const HazardSet &hazards) { + for (Operation *op = previous->getNextNode(); op && op != current; + op = op->getNextNode()) { + if (auto memBar = dyn_cast(op)) { + if (memBarCovers(memBar.getKind().getKind(), hazards)) + return true; + } + } + return false; +} + +static bool isStructuredVTileOp(Operation *op) { + auto pipe = dyn_cast(op); + if (!pipe || pipe.getPipe() != pto::PIPE::PIPE_V) + return false; + return isa(op); +} + +static bool hasUnknownRange(ArrayRef accesses) { + return llvm::any_of(accesses, [](const TileAccess &access) { + return !access.range.valid; + }); +} + +struct VTileOpInfo { + Operation *op = nullptr; + SmallVector accesses; +}; + +struct PTOInsertVMemBarPass + : public mlir::pto::impl::PTOInsertVMemBarBase { + void runOnOperation() override { + func::FuncOp func = getOperation(); + if (!pto::isTargetArchA5(func.getOperation())) + return; + + SmallVector, 16> insertions; + bool sawUnknownRange = false; + + for (Block &block : func.getBody()) { + SmallVector vOps; + for (Operation &op : block) { + if (!isStructuredVTileOp(&op)) + continue; + auto accesses = collectTileAccesses(&op); + if (accesses.empty()) + continue; + sawUnknownRange |= hasUnknownRange(accesses); + + HazardSet uncovered; + for (const VTileOpInfo &previous : vOps) { + HazardSet hazards = classifyHazards(previous.accesses, accesses); + if (hazards.empty()) + continue; + if (!hasCoveringMemBarBetween(previous.op, &op, hazards)) + uncovered.merge(hazards); + } + + if (!uncovered.empty()) + insertions.push_back({&op, getMemBarKindFor(uncovered)}); + vOps.push_back(VTileOpInfo{&op, std::move(accesses)}); + } + } + + OpBuilder builder(func.getContext()); + for (auto [op, kind] : insertions) { + builder.setInsertionPoint(op); + builder.create(op->getLoc(), + pto::MemBarAttr::get(func.getContext(), + kind)); + } + + if (sawUnknownRange) { + func.emitRemark() + << "pto-insert-v-membar conservatively treats unresolved VPTO tile " + "UB ranges as may-alias"; + } + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createPTOInsertVMemBarPass() { + return std::make_unique(); +} diff --git a/test/lit/vpto/vmembar_tile_hazards.pto b/test/lit/vpto/vmembar_tile_hazards.pto new file mode 100644 index 0000000000..550ced4832 --- /dev/null +++ b/test/lit/vpto/vmembar_tile_hazards.pto @@ -0,0 +1,126 @@ +// 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 --pto-level=level3 --emit-vpto \ +// RUN: --mlir-print-ir-after=pto-insert-v-membar %s -o /dev/null 2>&1 || true; } | FileCheck %s + +!tile = !pto.tile_buf + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmembar_tile_hazards() attributes {pto.kernel_kind = #pto.kernel_kind} { + %c0 = arith.constant 0 : i64 + %c256 = arith.constant 256 : i64 + %c512 = arith.constant 512 : i64 + %c768 = arith.constant 768 : i64 + %c1024 = arith.constant 1024 : i64 + %c1280 = arith.constant 1280 : i64 + %c1536 = arith.constant 1536 : i64 + %c4096 = arith.constant 4096 : i64 + %c4352 = arith.constant 4352 : i64 + %c4608 = arith.constant 4608 : i64 + %c8192 = arith.constant 8192 : i64 + %c8448 = arith.constant 8448 : i64 + %c8704 = arith.constant 8704 : i64 + %c8960 = arith.constant 8960 : i64 + %c12288 = arith.constant 12288 : i64 + %c12544 = arith.constant 12544 : i64 + %c12800 = arith.constant 12800 : i64 + %c16384 = arith.constant 16384 : i64 + %c16640 = arith.constant 16640 : i64 + %c16896 = arith.constant 16896 : i64 + %c17152 = arith.constant 17152 : i64 + %c20480 = arith.constant 20480 : i64 + %c20736 = arith.constant 20736 : i64 + %c20992 = arith.constant 20992 : i64 + %c21248 = arith.constant 21248 : i64 + + // #866 shape: two producers write temporary UB tiles; tsub reads them. + %k_lo = pto.alloc_tile addr = %c0 : !tile + %k_hi = pto.alloc_tile addr = %c256 : !tile + %cos = pto.alloc_tile addr = %c512 : !tile + %sin = pto.alloc_tile addr = %c768 : !tile + %prod_lo = pto.alloc_tile addr = %c1024 : !tile + %prod_hi = pto.alloc_tile addr = %c1280 : !tile + %out0 = pto.alloc_tile addr = %c1536 : !tile + pto.tcolexpandmul ins(%k_lo, %cos : !tile, !tile) outs(%prod_lo : !tile) + pto.tcolexpandmul ins(%k_hi, %sin : !tile, !tile) outs(%prod_hi : !tile) + pto.tsub ins(%prod_lo, %prod_hi : !tile, !tile) outs(%out0 : !tile) + + // WAW: two ops write the same UB tile range. + %waw_a = pto.alloc_tile addr = %c4096 : !tile + %waw_b = pto.alloc_tile addr = %c4352 : !tile + %waw_out0 = pto.alloc_tile addr = %c4608 : !tile + %waw_out1 = pto.alloc_tile addr = %c4608 : !tile + pto.tadd ins(%waw_a, %waw_b : !tile, !tile) outs(%waw_out0 : !tile) + pto.tsub ins(%waw_a, %waw_b : !tile, !tile) outs(%waw_out1 : !tile) + + // WAR: previous op reads a range that the next op overwrites. + %war_a = pto.alloc_tile addr = %c8192 : !tile + %war_b = pto.alloc_tile addr = %c8448 : !tile + %war_tmp = pto.alloc_tile addr = %c8704 : !tile + %war_c = pto.alloc_tile addr = %c8960 : !tile + %war_overwrite_a = pto.alloc_tile addr = %c8192 : !tile + pto.tadd ins(%war_a, %war_b : !tile, !tile) outs(%war_tmp : !tile) + pto.tsub ins(%war_c, %war_b : !tile, !tile) outs(%war_overwrite_a : !tile) + + // RAW + WAW together must become VV_ALL, not just VST_VLD. + %combo_a = pto.alloc_tile addr = %c12288 : !tile + %combo_b = pto.alloc_tile addr = %c12544 : !tile + %combo_tmp = pto.alloc_tile addr = %c12800 : !tile + %combo_out_alias = pto.alloc_tile addr = %c12800 : !tile + pto.tadd ins(%combo_a, %combo_b : !tile, !tile) outs(%combo_tmp : !tile) + pto.tsub ins(%combo_tmp, %combo_b : !tile, !tile) outs(%combo_out_alias : !tile) + + // Existing covering barrier must be respected and not duplicated. + %cover_a = pto.alloc_tile addr = %c16384 : !tile + %cover_b = pto.alloc_tile addr = %c16640 : !tile + %cover_tmp = pto.alloc_tile addr = %c16896 : !tile + %cover_out = pto.alloc_tile addr = %c17152 : !tile + pto.tadd ins(%cover_a, %cover_b : !tile, !tile) outs(%cover_tmp : !tile) + pto.mem_bar "VST_VLD" + pto.tsub ins(%cover_tmp, %cover_b : !tile, !tile) outs(%cover_out : !tile) + + // Existing non-covering barrier must not suppress the required RAW barrier. + %noncover_a = pto.alloc_tile addr = %c20480 : !tile + %noncover_b = pto.alloc_tile addr = %c20736 : !tile + %noncover_tmp = pto.alloc_tile addr = %c20992 : !tile + %noncover_out = pto.alloc_tile addr = %c21248 : !tile + pto.tadd ins(%noncover_a, %noncover_b : !tile, !tile) outs(%noncover_tmp : !tile) + pto.mem_bar "VLD_VST" + pto.tsub ins(%noncover_tmp, %noncover_b : !tile, !tile) outs(%noncover_out : !tile) + + return + } +} + +// CHECK-LABEL: func.func @vmembar_tile_hazards +// CHECK: pto.tcolexpandmul +// CHECK-NEXT: pto.tcolexpandmul +// CHECK-NEXT: pto.mem_bar "VST_VLD" +// CHECK-NEXT: pto.tsub + +// CHECK: pto.tadd +// CHECK-NEXT: pto.mem_bar "VST_VST" +// CHECK-NEXT: pto.tsub + +// CHECK: pto.tadd +// CHECK-NEXT: pto.mem_bar "VLD_VST" +// CHECK-NEXT: pto.tsub + +// CHECK: pto.tadd +// CHECK-NEXT: pto.mem_bar "VV_ALL" +// CHECK-NEXT: pto.tsub + +// CHECK: pto.tadd +// CHECK-NEXT: pto.mem_bar "VST_VLD" +// CHECK-NEXT: pto.tsub + +// CHECK: pto.tadd +// CHECK-NEXT: pto.mem_bar "VLD_VST" +// CHECK-NEXT: pto.mem_bar "VST_VLD" +// CHECK-NEXT: pto.tsub diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index ca58c85a7d..93e80ae7d3 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -346,6 +346,12 @@ static llvm::cl::opt enableInjectBarrierAllSync( "pto.barrier PIPE_ALL before memory-effecting PTO pipe ops"), llvm::cl::init(false)); +static llvm::cl::opt disableVMembar( + "disable-v-membar", + llvm::cl::desc("Disable automatic VPTO V-pipeline mem_bar insertion " + "(debug only; disabling may expose UB RAW/WAR/WAW hazards)"), + llvm::cl::init(false)); + static llvm::cl::opt enableGraphSyncSolver( "enable-graph-sync-solver", llvm::cl::desc("Enable the graph-based intra-core sync solver " @@ -2869,6 +2875,12 @@ int mlir::pto::compilePTOASModule( pm.addNestedPass(pto::createPTOA5NormalizeTMovPass()); pm.addNestedPass( pto::createPTOValidateIntToPtrUsesPass()); + if (effectiveBackend == PTOBackend::VPTO && !disableVMembar) { + // Insert V-pipeline mem_bar while tile ops are still structured and before + // VPTO fusion plans contiguous tile spans. The inserted barriers then act + // as normal tile-native scheduling boundaries. + pm.addNestedPass(pto::createPTOInsertVMemBarPass()); + } // Keep frontend fusion on tile-native PTO IR and annotate last_use directly // on scheduled block-local spans before the shared mainline lowers tiles.