From c43c8e4ce26da65c8c2b01a1eb4b4eaecf3b996a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=86=AF=E7=91=9E?= Date: Wed, 23 Jul 2025 01:17:08 +0000 Subject: [PATCH 1/7] support atomic or/xor/xchg/cas MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 冯瑞 --- ascend/backend/compiler.py | 22 ++- ascend/backend/driver.py | 23 ++- .../TritonToLinalg/LoadStoreConverter.h | 19 +++ .../lib/TritonToLinalg/LoadStoreConverter.cpp | 161 ++++++++++++++++++ .../lib/TritonToLinalg/TritonToLinalgPass.cpp | 21 ++- .../lib/TritonToLinalg/UseAnalysis.cpp | 11 ++ 6 files changed, 248 insertions(+), 9 deletions(-) diff --git a/ascend/backend/compiler.py b/ascend/backend/compiler.py index c2ee7ab..7555cf0 100644 --- a/ascend/backend/compiler.py +++ b/ascend/backend/compiler.py @@ -249,12 +249,26 @@ def linalg_to_bin_enable_npu_compile(linalg: str, metadata, opt): ret = subprocess.run(cmd_list, capture_output=True, check=True) if Path(callback_path).is_file(): lib = ctypes.CDLL(callback_path) - callback_func = getattr( + ws_callback_func = getattr( lib, metadata["kernel_name"] + "_infer_workspace_shape_function" ) - callback_func.restype = ctypes.c_int64 - callback_func.argtypes = [] - metadata["workspace_size"] = callback_func() + ws_callback_func.restype = ctypes.c_int64 + ws_callback_func.argtypes = [] + metadata["workspace_size"] = ws_callback_func() + + lock_num_callback_func = getattr( + lib, metadata["kernel_name"] + "_infer_sync_block_lock_num_function " + ) + lock_num_callback_func.restype = ctypes.c_int64 + lock_num_callback_func.argtypes = [] + metadata["lock_num"] = lock_num_callback_func() + + lock_init_val_callback_func = getattr( + lib, metadata["kernel_name"] + "_infer_sync_block_lock_num_function" + ) + lock_init_val_callback_func.restype = ctypes.c_int64 + lock_init_val_callback_func.argtypes = [] + metadata["lock_init_val"] = lock_init_val_callback_func() return Path(bin_path).read_bytes() diff --git a/ascend/backend/driver.py b/ascend/backend/driver.py index 151161f..0f98079 100644 --- a/ascend/backend/driver.py +++ b/ascend/backend/driver.py @@ -66,13 +66,18 @@ class NPULauncher(object): debug_mode = metadata.debug workspace_size = int(metadata.workspace_size) \ if hasattr(metadata, 'workspace_size') else -1 + lock_init_value = int(metadata.lock_init_value) \ + if hasattr(metadata, 'lock_init_value') else 0 + lock_num = int(metadata.lock_num) \ + if hasattr(metadata, 'lock_num') else -1 constants = src.constants if hasattr(src, "constants") else dict() cst_key = lambda i: src.fn.arg_names.index(i) if isinstance(i, str) else i constants = {cst_key(key): value for key, value in constants.items()} signature = {cst_key(key): value for key, value in src.signature.items()} mix_mode = metadata.mix_mode wrapper_src = generate_npu_wrapper_src(constants, signature, \ - workspace_size, mix_mode) + workspace_size, mix_mode, \ + lock_num, lock_init_value) so_launcher_path = make_npu_launcher_stub(wrapper_src, debug_mode) # initialize launcher import importlib.util @@ -207,7 +212,7 @@ def make_npu_launcher_stub(src, debug=False): # the template is from triton-adapter HEAD. Wrapping the generated kernel binary into a python module -def generate_npu_wrapper_src(constants, signature, workspace_size, mix_mode): +def generate_npu_wrapper_src(constants, signature, workspace_size, mix_mode, lock_num, lock_ini_val): import os def _ty_to_cpp(ty): if ty[0] == '*': @@ -467,9 +472,21 @@ static void _launch(const char* kernelName, const void* func, rtStream_t stream, return {'ret' if enable_taskqueue else ''}; }} // stub argument for workspace + void *syncBlockLock = NULL; void *workspace_addr = NULL; {f''' uint16_t ModuleId = 0; + uint64_t syncBlockLockSize = {lock_num} * sizeof(int64_t); + ret = rtMalloc(reinterpret_cast(&syncBlockLock), + syncBlockLockSize, RT_MEMORY_HBM, 0); + std::vector lockInitData({lock_num}, {lock_ini_val}); + ret = rtMemcpy(syncBlockLock, syncBlockLockSize, reinterpret_cast(lockInitData.data()), + syncBlockLockSize, RT_MEMCPY_HOST_TO_DEVICE); + if (ret != RT_ERROR_NONE) {{ + return {'ret' if enable_taskqueue else ''}; + }} + ''' if lock_num > 0 else ''} + {f''' uint64_t totalWorkSpaceSize = {workspace_size} * blockNum; ret = rtMalloc(reinterpret_cast(&workspace_addr), totalWorkSpaceSize, RT_MEMORY_HBM, ModuleId); @@ -479,12 +496,14 @@ static void _launch(const char* kernelName, const void* func, rtStream_t stream, ''' if workspace_size > 0 else ''} struct __attribute__((packed)) {{ void* ffts_addr __attribute__((aligned(8))); + void* syncBlockLock __attribute__((aligned(8))); void* workspace_addr __attribute__((aligned(8))); {' '.join(f'{_ty_to_cpp(ty)} arg{i} __attribute__((aligned({4 if ty[0] != "*" and ty[-2:] != "64" else 8})));' for i, ty in signature.items() if i not in constants)} {' '.join(f'{_ty_to_cpp(ty)} grid{mark} __attribute__((aligned(4)));' for mark, ty in grid_info.items())} {'void* DTData __attribute__((aligned(8)));' if enable_device_print else ''} }} args = {{ static_cast(ffts_addr), + static_cast(syncBlockLock), static_cast(workspace_addr), {', '.join(f'static_cast<{_ty_to_cpp(ty)}>(arg{i})' for i, ty in signature.items() if i not in constants)}, {', '.join(f'static_cast<{_ty_to_cpp(ty)}>(grid{mark})' for mark, ty in grid_info.items())} diff --git a/ascend/triton-adapter/include/TritonToLinalg/LoadStoreConverter.h b/ascend/triton-adapter/include/TritonToLinalg/LoadStoreConverter.h index 7bb438e..58eb3bf 100644 --- a/ascend/triton-adapter/include/TritonToLinalg/LoadStoreConverter.h +++ b/ascend/triton-adapter/include/TritonToLinalg/LoadStoreConverter.h @@ -123,6 +123,23 @@ class ScalarAtomicRMWCanonicalizer PatternRewriter &rewriter) const override; }; +class ScalarAtomicCASCanonicalizer + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(triton::AtomicCASOp op, + PatternRewriter &rewriter) const override; +}; + +class AtomicCASConverter : public OpConversionPattern { +public: + explicit AtomicCASConverter(MLIRContext *context) : OpConversionPattern(context) {} + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::AtomicCASOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + class AtomicRMWConverter : public OpConversionPattern { private: Value createAtomicBinaryOps(OpBuilder &builder, Location loc, @@ -157,6 +174,8 @@ private: } else { binaryOp = builder.create(loc, lhs, rhs); } + } else if (rmwOp == triton::RMWOp::XCHG) { + binaryOp = rhs; } else { op.emitOpError("unsupported atomic RMW operation: "); llvm_unreachable( diff --git a/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp b/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp index 6f66cd7..abd2282 100644 --- a/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp +++ b/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp @@ -516,6 +516,141 @@ AtomicRMWConverter::matchAndRewrite(triton::AtomicRMWOp op, OpAdaptor adaptor, return success(); } +LogicalResult +AtomicCASConverter::matchAndRewrite(triton::AtomicCASOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + // If the result of AtomicCASOp is not used, we don't need to load the old + // data stored at the ptr + auto ptr = adaptor.getPtr(); + auto cmp = op.getCmp(); + auto val = op.getVal(); + auto loc = op.getLoc(); + + auto resType = dyn_cast(op.getResult().getType()); + if (!resType) { + return rewriter.notifyMatchFailure( + op, "atomicCASConverter: scalar will be handled by " + "ScalarAtomicCASCanonicalizer"); + } + + // 1. Simple case where no mask is used. + auto type = dyn_cast(ptr.getType()); + if (!type) { + // Seen when implicit broadcasting is done late in a chain of + // operations. The workaround is to broadcast the pointers early in the + // address calculation. A proper fix is complicated, but at least we can + // provide a better error message. + return rewriter.notifyMatchFailure( + op, "AtomicCASOp expects a memref, not a memref of pointers"); + } + + auto dstMemref = ptr; + // Well, linalg structure op wouldn't support mixed tensor/buffer semantics + // any more in latest LLVM(triton LLVM dependency has involed this), so we + // need to convert tensor to buffer early. + auto dstType = dstMemref.getType(); + Value inputMemref = + rewriter.create(loc, dstType, val); + + Value cmpMemref = + rewriter.create(loc, dstType, cmp); + + // 3. If needed, handle the return value of atomic op + // + // tt.atomicRMW op has two part of feature + // 1. load the old data at the ptr + // 2. atomically store the data on ub to the ptr + // at the same time it perform the action it has been assigned + // So we lower this op to load + atomically store + // + // The first part is not necessary when the returned value of atomic op + // is not used, it will be deleted cause it's meaningless + // Here, we preemptively determine whether it will be used + // and decide whether it is necessary to create the load process based on + // this assessment. + // + // logic of handling is copied + // TODO: decoupling the logic of load, put it in the Utils + if (!op.getResult().use_empty()) { + auto tensorType = + RankedTensorType::get(type.getShape(), type.getElementType()); + auto alloc = rewriter.create( + loc, MemRefType::get(type.getShape(), type.getElementType())); + + // For the return value, don't need to care about mask for now + // this op don't support other, so we best not fill it + rewriter.create(loc, ptr, alloc); + Value tensor = rewriter.create( + loc, tensorType, alloc, true /* restrict */, true /* writable */); + rewriter.replaceOp(op, tensor); + } + + // create element-wise map + int64_t rank = type.getRank(); + SmallVector inputDims; + auto context = rewriter.getContext(); + + for (int i = 0; i < rank; i++) { + inputDims.push_back(getAffineDimExpr(i, context)); + } + + SmallVector indexingMaps; + // As mask has been erased for now + // the number of input must be 2 + // the input memref is also the output memref + // Thus, there are a total of three inputs and outputs. + // so here we have 3 map to create + for (int i = 0; i < 4; i++) { + indexingMaps.push_back(AffineMap::get(rank, 0, inputDims, context)); + } + + auto linalgOp = rewriter.create( + loc, /* operands */ ValueRange{dstMemref, cmpMemref, inputMemref}, + mlir::ValueRange{dstMemref}, indexingMaps, + mlir::ConverterUtils::getNParallelLoopsAttrs(rank), + [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) { + Value lhs = blockArgs[0]; + Value rhs = blockArgs[1]; + Value setValue = blockArgs[2]; + Value cond = nestedBuilder.create(nestedLoc, + arith::CmpIPredicate::eq, + lhs, rhs); + auto ifOp = nestedBuilder.create(nestedLoc, TypeRange{setValue.getType()}, cond, true); + { + OpBuilder::InsertionGuard guard(nestedBuilder); + nestedBuilder.setInsertionPointToEnd(&ifOp.getThenRegion().front()); + nestedBuilder.create(nestedLoc, setValue); + } + { + OpBuilder::InsertionGuard guard(nestedBuilder); + nestedBuilder.setInsertionPointToEnd(&ifOp.getElseRegion().front()); + nestedBuilder.create(nestedLoc, lhs); + } + nestedBuilder.setInsertionPointToEnd(nestedBuilder.getBlock()); + nestedBuilder.create(nestedLoc, ifOp.getResult(0)); + }); + + const StringRef genericAtomicRMW = "GenericAtomicRMW"; + const StringRef memSemantic = "MemSemantic"; + const StringRef memSyncScope = "MemSyncScope"; + auto attr = mlir::StringAttr::get(context, "cas"); + + linalgOp->setAttr(genericAtomicRMW, attr); + linalgOp->setAttr(memSemantic, + rewriter.getStringAttr(stringifyEnum(op.getSem()))); + linalgOp->setAttr(memSyncScope, + rewriter.getStringAttr(stringifyEnum(op.getScope()))); + + linalgOp->setAttr("Software", rewriter.getUnitAttr()); + + // if the result hasn't been replace by load + // we need to erase it here + if (op.getResult().use_empty()) { + rewriter.eraseOp(op); + } + return success(); +} + LogicalResult ScalarStoreCanonicalizer::matchAndRewrite(triton::StoreOp op, PatternRewriter &rewriter) const { @@ -564,6 +699,32 @@ ScalarAtomicRMWCanonicalizer::matchAndRewrite(triton::AtomicRMWOp op, return success(); } +LogicalResult +ScalarAtomicCASCanonicalizer::matchAndRewrite(triton::AtomicCASOp op, + PatternRewriter &rewriter) const { + + if (!op.getVal().getType().isIntOrIndexOrFloat() && !op.getCmp().getType().isIntOrIndexOrFloat()) { + return rewriter.notifyMatchFailure( + op, "ScalarAtomicCASCanonicalizer handles scalar atomic cas op scene!"); + } + + auto ptr = op.getPtr(); + auto ptrTy = RankedTensorType::get({(int64_t)1}, ptr.getType()); + auto ptrSplat = rewriter.create(op.getLoc(), ptrTy, ptr); + auto cmpTy = RankedTensorType::get({(int64_t)1}, op.getCmp().getType()); + auto cmpSplat = + rewriter.create(op.getLoc(), cmpTy, op.getCmp()); + auto valTy = RankedTensorType::get({(int64_t)1}, op.getVal().getType()); + auto valSplat = + rewriter.create(op.getLoc(), valTy, op.getVal()); + + auto newAtomicOp = rewriter.create( + op.getLoc(), valTy, ptrSplat, cmpSplat, valSplat, + op.getSem(), op.getScope()); + rewriter.replaceOp(op, newAtomicOp); + return success(); +} + // The atomic max op with float input will be devided into // two atomic max ops with integer input // One handles the part of the tensor greater than zero diff --git a/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp b/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp index a0a1fe6..4fe8303 100644 --- a/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp +++ b/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp @@ -374,11 +374,13 @@ void TritonToLinalgPass::populateTritonToLinalgCanonicalizationPatterns(RewriteP patterns.add(patterns.getContext()); patterns.add, LoadStoreConverter::LoadStoreCanonicalizer, - LoadStoreConverter::LoadStoreCanonicalizer>(patterns.getContext()); + LoadStoreConverter::LoadStoreCanonicalizer, + LoadStoreConverter::LoadStoreCanonicalizer>(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add< TTOpConverters::ScalarMathCanonicalizer, @@ -438,6 +440,7 @@ void TritonToLinalgPass::populateTritonToLinalgConversionPatterns( patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); @@ -604,7 +607,19 @@ void TritonToLinalgPass::runOnOperation() { continue; auto context = func.getContext(); - constexpr int64_t workspaceArgIdx = 0; + constexpr int64_t syncBlockLockArgIdx = 0; + NamedAttribute syncBlockLockArgAttr(StringAttr::get(context, "syncBlockLock"), + UnitAttr::get(context)); + MemRefType syncBlockLockArgType = + MemRefType::get(SmallVector(1, ShapedType::kDynamic), + IntegerType::get(context, 8)); + func.insertArgument(/*argIndex*/ syncBlockLockArgIdx, + /*argType*/ syncBlockLockArgType, + /*dicAttr*/ nullptr, func->getLoc()); + func->setAttr("SyncBlockLockArgType", + IntegerAttr::get(IntegerType::get(&getContext(), 64), 0)); + + constexpr int64_t workspaceArgIdx = 1; MemRefType workspaceArgType = MemRefType::get(SmallVector(1, ShapedType::kDynamic), IntegerType::get(context, 8)); @@ -615,7 +630,7 @@ void TritonToLinalgPass::runOnOperation() { /*argType*/ workspaceArgType, /*dicAttr*/ nullptr, func->getLoc()); func->setAttr("WorkspaceArgIdx", - IntegerAttr::get(IntegerType::get(&getContext(), 64), 0)); + IntegerAttr::get(IntegerType::get(&getContext(), 64), 1)); } // Fix the Location info diff --git a/ascend/triton-adapter/lib/TritonToLinalg/UseAnalysis.cpp b/ascend/triton-adapter/lib/TritonToLinalg/UseAnalysis.cpp index 4b09631..0580b7c 100644 --- a/ascend/triton-adapter/lib/TritonToLinalg/UseAnalysis.cpp +++ b/ascend/triton-adapter/lib/TritonToLinalg/UseAnalysis.cpp @@ -89,6 +89,12 @@ void triton::UseAnalysis::visitOperation(Operation *op, propagateUse(operands[2], UseType::MetaUse); } }) + .Case([&](auto atomicOp) { + propagateUse(operands[0], UseType::MetaUse); + propagateUse(operands[1], UseType::DataUse); + propagateUse(operands[2], UseType::DataUse); + auto value = atomicOp.getVal(); + }) .Case([&](auto dot) { propagateResults(operands[0], results); propagateResults(operands[1], results); @@ -212,6 +218,11 @@ LogicalResult triton::runUseAnalysis(triton::FuncOp &funcOp) { if (result == ptr || result == mask) metaUsers.insert(user); }) + .Case([&](auto atomicOp) { + auto ptr = atomicOp.getPtr(); + if (result == ptr) + metaUsers.insert(user); + }) .Case([&](auto dot) { auto opc = dot.getC(); triton::SplatOp splat; -- Gitee From fcb0998ca77eb2ea0ca7981904aa05ee2c40fffa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=86=AF=E7=91=9E?= Date: Wed, 23 Jul 2025 12:04:26 +0800 Subject: [PATCH 2/7] fix codecheck --- .../include/TritonToLinalg/LoadStoreConverter.h | 3 ++- .../lib/TritonToLinalg/LoadStoreConverter.cpp | 12 ++++-------- .../lib/TritonToLinalg/TritonToLinalgPass.cpp | 10 +++++----- 3 files changed, 11 insertions(+), 14 deletions(-) diff --git a/ascend/triton-adapter/include/TritonToLinalg/LoadStoreConverter.h b/ascend/triton-adapter/include/TritonToLinalg/LoadStoreConverter.h index 58eb3bf..c1d2dbc 100644 --- a/ascend/triton-adapter/include/TritonToLinalg/LoadStoreConverter.h +++ b/ascend/triton-adapter/include/TritonToLinalg/LoadStoreConverter.h @@ -132,7 +132,8 @@ class ScalarAtomicCASCanonicalizer class AtomicCASConverter : public OpConversionPattern { public: - explicit AtomicCASConverter(MLIRContext *context) : OpConversionPattern(context) {} + explicit AtomicCASConverter(MLIRContext *context) : + OpConversionPattern(context) {} using OpConversionPattern::OpConversionPattern; LogicalResult diff --git a/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp b/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp index abd2282..6ff6fba 100644 --- a/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp +++ b/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp @@ -570,7 +570,6 @@ AtomicCASConverter::matchAndRewrite(triton::AtomicCASOp op, OpAdaptor adaptor, // this assessment. // // logic of handling is copied - // TODO: decoupling the logic of load, put it in the Utils if (!op.getResult().use_empty()) { auto tensorType = RankedTensorType::get(type.getShape(), type.getElementType()); @@ -598,14 +597,14 @@ AtomicCASConverter::matchAndRewrite(triton::AtomicCASOp op, OpAdaptor adaptor, // As mask has been erased for now // the number of input must be 2 // the input memref is also the output memref - // Thus, there are a total of three inputs and outputs. - // so here we have 3 map to create - for (int i = 0; i < 4; i++) { + // Thus, there are a total of four inputs and outputs. + // so here we have 4 map to create + for (int i = 0; i < 4; i++) { // 4: 3 input and 1 output indexingMaps.push_back(AffineMap::get(rank, 0, inputDims, context)); } auto linalgOp = rewriter.create( - loc, /* operands */ ValueRange{dstMemref, cmpMemref, inputMemref}, + loc, ValueRange{dstMemref, cmpMemref, inputMemref}, mlir::ValueRange{dstMemref}, indexingMaps, mlir::ConverterUtils::getNParallelLoopsAttrs(rank), [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) { @@ -654,7 +653,6 @@ AtomicCASConverter::matchAndRewrite(triton::AtomicCASOp op, OpAdaptor adaptor, LogicalResult ScalarStoreCanonicalizer::matchAndRewrite(triton::StoreOp op, PatternRewriter &rewriter) const { - if (!op.getValue().getType().isIntOrIndexOrFloat()) { return rewriter.notifyMatchFailure( op, "ScalarStoreCanonicalizer handles scalar store scene!"); @@ -676,7 +674,6 @@ ScalarStoreCanonicalizer::matchAndRewrite(triton::StoreOp op, LogicalResult ScalarAtomicRMWCanonicalizer::matchAndRewrite(triton::AtomicRMWOp op, PatternRewriter &rewriter) const { - if (!op.getVal().getType().isIntOrIndexOrFloat()) { return rewriter.notifyMatchFailure( op, "ScalarAtomicRMWCanonicalizer handles scalar atomic rmw op scene!"); @@ -702,7 +699,6 @@ ScalarAtomicRMWCanonicalizer::matchAndRewrite(triton::AtomicRMWOp op, LogicalResult ScalarAtomicCASCanonicalizer::matchAndRewrite(triton::AtomicCASOp op, PatternRewriter &rewriter) const { - if (!op.getVal().getType().isIntOrIndexOrFloat() && !op.getCmp().getType().isIntOrIndexOrFloat()) { return rewriter.notifyMatchFailure( op, "ScalarAtomicCASCanonicalizer handles scalar atomic cas op scene!"); diff --git a/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp b/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp index 4fe8303..2240cf4 100644 --- a/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp +++ b/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp @@ -613,11 +613,11 @@ void TritonToLinalgPass::runOnOperation() { MemRefType syncBlockLockArgType = MemRefType::get(SmallVector(1, ShapedType::kDynamic), IntegerType::get(context, 8)); - func.insertArgument(/*argIndex*/ syncBlockLockArgIdx, - /*argType*/ syncBlockLockArgType, - /*dicAttr*/ nullptr, func->getLoc()); + func.insertArgument(syncBlockLockArgIdx, // argIndex + syncBlockLockArgType, // argType + nullptr, func->getLoc()); // dicAttr func->setAttr("SyncBlockLockArgType", - IntegerAttr::get(IntegerType::get(&getContext(), 64), 0)); + IntegerAttr::get(IntegerType::get(&getContext(), 64), 0)); // 64: 64位整型 constexpr int64_t workspaceArgIdx = 1; MemRefType workspaceArgType = @@ -630,7 +630,7 @@ void TritonToLinalgPass::runOnOperation() { /*argType*/ workspaceArgType, /*dicAttr*/ nullptr, func->getLoc()); func->setAttr("WorkspaceArgIdx", - IntegerAttr::get(IntegerType::get(&getContext(), 64), 1)); + IntegerAttr::get(IntegerType::get(&getContext(), 64), 1)); // 64: 64位整型 } // Fix the Location info -- Gitee From 41ed77d0456cc6a5781ac37d62b68f2f651d4d23 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=86=AF=E7=91=9E?= Date: Wed, 23 Jul 2025 12:47:16 +0800 Subject: [PATCH 3/7] add ret value check --- ascend/backend/driver.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ascend/backend/driver.py b/ascend/backend/driver.py index 0f98079..acfc5ad 100644 --- a/ascend/backend/driver.py +++ b/ascend/backend/driver.py @@ -479,6 +479,9 @@ static void _launch(const char* kernelName, const void* func, rtStream_t stream, uint64_t syncBlockLockSize = {lock_num} * sizeof(int64_t); ret = rtMalloc(reinterpret_cast(&syncBlockLock), syncBlockLockSize, RT_MEMORY_HBM, 0); + if (ret != RT_ERROR_NONE) {{ + return {'ret' if enable_taskqueue else ''}; + }} std::vector lockInitData({lock_num}, {lock_ini_val}); ret = rtMemcpy(syncBlockLock, syncBlockLockSize, reinterpret_cast(lockInitData.data()), syncBlockLockSize, RT_MEMCPY_HOST_TO_DEVICE); -- Gitee From 974ee79f8917c14c0d5727dbcccb27d61d0d2384 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=86=AF=E7=91=9E?= Date: Thu, 24 Jul 2025 11:17:47 +0800 Subject: [PATCH 4/7] fix bug of lock init --- ascend/backend/compiler.py | 34 +++++++------------ ascend/backend/driver.py | 2 +- .../lib/TritonToLinalg/TritonToLinalgPass.cpp | 2 +- 3 files changed, 14 insertions(+), 24 deletions(-) diff --git a/ascend/backend/compiler.py b/ascend/backend/compiler.py index 7555cf0..b640422 100644 --- a/ascend/backend/compiler.py +++ b/ascend/backend/compiler.py @@ -199,6 +199,15 @@ def llir_to_cpuasm(llir: str, metadata, opt): return Path(dst_path).read_text() +def __get_metadata_attr_by_callback(lib, postfix: str, metadata, meta_key: str): + func_symbol = metadata["kernel_name"] + postfix + if hasattr(lib, func_symbol): + callback_func = getattr(lib, func_symbol) + callback_func.restype = ctypes.c_int64 + callback_func.argtypes = [] + metadata[meta_key] = callback_func() + + def linalg_to_bin_enable_npu_compile(linalg: str, metadata, opt): # Note: Compiled Kernel requires to estimate size of shared memory to occupy # Currently, NPU backend does not limit on shared memory @@ -249,28 +258,9 @@ def linalg_to_bin_enable_npu_compile(linalg: str, metadata, opt): ret = subprocess.run(cmd_list, capture_output=True, check=True) if Path(callback_path).is_file(): lib = ctypes.CDLL(callback_path) - ws_callback_func = getattr( - lib, metadata["kernel_name"] + "_infer_workspace_shape_function" - ) - ws_callback_func.restype = ctypes.c_int64 - ws_callback_func.argtypes = [] - metadata["workspace_size"] = ws_callback_func() - - lock_num_callback_func = getattr( - lib, metadata["kernel_name"] + "_infer_sync_block_lock_num_function " - ) - lock_num_callback_func.restype = ctypes.c_int64 - lock_num_callback_func.argtypes = [] - metadata["lock_num"] = lock_num_callback_func() - - lock_init_val_callback_func = getattr( - lib, metadata["kernel_name"] + "_infer_sync_block_lock_num_function" - ) - lock_init_val_callback_func.restype = ctypes.c_int64 - lock_init_val_callback_func.argtypes = [] - metadata["lock_init_val"] = lock_init_val_callback_func() - - return Path(bin_path).read_bytes() + __get_metadata_attr_by_callback(lib, "_infer_workspace_shape_function", metadata, "workspace_size") + __get_metadata_attr_by_callback(lib, "_infer_sync_block_lock_num_function", metadata, "lock_num") + __get_metadata_attr_by_callback(lib, "_infer_sync_block_lock_init_function", metadata, "lock_init_val") @dataclass(frozen=True) diff --git a/ascend/backend/driver.py b/ascend/backend/driver.py index acfc5ad..ab7901f 100644 --- a/ascend/backend/driver.py +++ b/ascend/backend/driver.py @@ -474,8 +474,8 @@ static void _launch(const char* kernelName, const void* func, rtStream_t stream, // stub argument for workspace void *syncBlockLock = NULL; void *workspace_addr = NULL; - {f''' uint16_t ModuleId = 0; + {f''' uint64_t syncBlockLockSize = {lock_num} * sizeof(int64_t); ret = rtMalloc(reinterpret_cast(&syncBlockLock), syncBlockLockSize, RT_MEMORY_HBM, 0); diff --git a/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp b/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp index 2240cf4..c29f900 100644 --- a/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp +++ b/ascend/triton-adapter/lib/TritonToLinalg/TritonToLinalgPass.cpp @@ -616,7 +616,7 @@ void TritonToLinalgPass::runOnOperation() { func.insertArgument(syncBlockLockArgIdx, // argIndex syncBlockLockArgType, // argType nullptr, func->getLoc()); // dicAttr - func->setAttr("SyncBlockLockArgType", + func->setAttr("SyncBlockLockArgIdx", IntegerAttr::get(IntegerType::get(&getContext(), 64), 0)); // 64: 64位整型 constexpr int64_t workspaceArgIdx = 1; -- Gitee From 91b806289eb6576724643519b2c07324df8e5ab8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=86=AF=E7=91=9E?= Date: Thu, 24 Jul 2025 17:56:37 +0800 Subject: [PATCH 5/7] update docs --- README.md | 12 ++++++++---- docs/sources/python-api/outline.md | 12 ++++++++---- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index eb468eb..5893f56 100644 --- a/README.md +++ b/README.md @@ -339,12 +339,12 @@ triton autotune性能配置说明参考本仓库的 docs\sources\getting-started | | gather | × | × | × | × | × | ✓ | ✓ | ✓ | × | | Atomic Ops | atomic_add | ✓ | ✓ | ✓ | × | × | ✓ | ✓ | ✓ | × | | | atomic_and | × | × | × | × | × | × | × | × | × | -| | atomic_cas | × | × | × | × | × | × | × | × | × | +| | atomic_cas | ✓ | ✓ | ✓ | × | ✓ | ✓ | ✓ | × | × | | | atomic_max | ✓ | ✓ | ✓ | × | × | ✓ | ✓ | ✓ | × | | | atomic_min | ✓ | ✓ | ✓ | × | × | ✓ | ✓ | ✓ | × | -| | atomic_or | × | × | × | × | × | × | × | × | × | -| | atomic_xchg | × | × | × | × | × | × | × | × | × | -| | atomic_xor | × | × | × | × | × | × | × | × | × | +| | atomic_or | ✓ | ✓ | ✓ | × | ✓ | × | × | × | × | +| | atomic_xchg | ✓ | ✓ | ✓ | × | ✓ | ✓ | ✓ | × | × | +| | atomic_xor | ✓ | ✓ | ✓ | × | ✓ | × | × | × | × | | Random Number Generation | randint4x | × | × | ✓ | × | × | × | × | × | × | | | randint | × | × | ✓ | × | × | × | × | × | × | | | rand | × | × | × | × | × | × | ✓ | × | × | @@ -379,6 +379,10 @@ triton autotune性能配置说明参考本仓库的 docs\sources\getting-started - atomic_min: 不支持标量（包括长度为1的tensor）访存 +- atomic_or: sem只支持默认值"acq_rel"模式，其他值均按默认值处理；scope只支持默认值"gpu"，其他值均按默认值处理 + +- atomic_xor: sem只支持默认值"acq_rel"模式，其他值均按默认值处理；scope只支持默认值"gpu"，其他值均按默认值处理 + - permute: 不支持不相邻轴转置，如`(0, 1, 2) -> (2, 1, 0)` - trans: 不支持不相邻轴转置，如`(0, 1, 2) -> (2, 1, 0)` diff --git a/docs/sources/python-api/outline.md b/docs/sources/python-api/outline.md index 0ac7381..40a4ccc 100644 --- a/docs/sources/python-api/outline.md +++ b/docs/sources/python-api/outline.md @@ -90,12 +90,12 @@ | | gather | × | × | × | × | × | ✓ | ✓ | ✓ | × | | Atomic Ops | atomic_add | ✓ | ✓ | ✓ | × | × | ✓ | ✓ | ✓ | × | | | atomic_and | × | × | × | × | × | × | × | × | × | -| | atomic_cas | × | × | × | × | × | × | × | × | × | +| | atomic_cas | ✓ | ✓ | ✓ | × | ✓ | ✓ | ✓ | × | × | | | atomic_max | ✓ | ✓ | ✓ | × | × | ✓ | ✓ | ✓ | × | | | atomic_min | ✓ | ✓ | ✓ | × | × | ✓ | ✓ | ✓ | × | -| | atomic_or | × | × | × | × | × | × | × | × | × | -| | atomic_xchg | × | × | × | × | × | × | × | × | × | -| | atomic_xor | × | × | × | × | × | × | × | × | × | +| | atomic_or | ✓ | ✓ | ✓ | × | ✓ | × | × | × | × | +| | atomic_xchg | ✓ | ✓ | ✓ | × | ✓ | ✓ | ✓ | × | × | +| | atomic_xor | ✓ | ✓ | ✓ | × | ✓ | × | × | × | × | | Random Number Generation | randint4x | × | × | ✓ | × | × | × | × | × | × | | | randint | × | × | ✓ | × | × | × | × | × | × | | | rand | × | × | × | × | × | × | ✓ | × | × | @@ -130,6 +130,10 @@ - atomic_min: 不支持标量（包括长度为1的tensor）访存 +- atomic_or: sem只支持默认值"acq_rel"模式，其他值均按默认值处理；scope只支持默认值"gpu"，其他值均按默认值处理 + +- atomic_xor: sem只支持默认值"acq_rel"模式，其他值均按默认值处理；scope只支持默认值"gpu"，其他值均按默认值处理 + - permute: 不支持不相邻轴转置，如`(0, 1, 2) -> (2, 1, 0)` - trans: 不支持不相邻轴转置，如`(0, 1, 2) -> (2, 1, 0)` -- Gitee From 050739709d0da6a45bb0473de32b448e40923957 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=86=AF=E7=91=9E?= Date: Fri, 25 Jul 2025 19:06:34 +0800 Subject: [PATCH 6/7] fix bug --- ascend/backend/compiler.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ascend/backend/compiler.py b/ascend/backend/compiler.py index b640422..df04dd5 100644 --- a/ascend/backend/compiler.py +++ b/ascend/backend/compiler.py @@ -261,6 +261,8 @@ def linalg_to_bin_enable_npu_compile(linalg: str, metadata, opt): __get_metadata_attr_by_callback(lib, "_infer_workspace_shape_function", metadata, "workspace_size") __get_metadata_attr_by_callback(lib, "_infer_sync_block_lock_num_function", metadata, "lock_num") __get_metadata_attr_by_callback(lib, "_infer_sync_block_lock_init_function", metadata, "lock_init_val") + + return Path(bin_path).read_bytes() @dataclass(frozen=True) -- Gitee From e1706f71c4a99fc863cf68045009ed64c0b810fb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=86=AF=E7=91=9E?= Date: Sat, 26 Jul 2025 07:58:31 +0000 Subject: [PATCH 7/7] atomic_cas support float input MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 冯瑞 --- .../lib/TritonToLinalg/LoadStoreConverter.cpp | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp b/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp index 6ff6fba..9ba73de 100644 --- a/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp +++ b/ascend/triton-adapter/lib/TritonToLinalg/LoadStoreConverter.cpp @@ -611,9 +611,16 @@ AtomicCASConverter::matchAndRewrite(triton::AtomicCASOp op, OpAdaptor adaptor, Value lhs = blockArgs[0]; Value rhs = blockArgs[1]; Value setValue = blockArgs[2]; - Value cond = nestedBuilder.create(nestedLoc, - arith::CmpIPredicate::eq, - lhs, rhs); + Value cond; + if (mlir::isa(lhs.getType())) { + cond = nestedBuilder.create(nestedLoc, + arith::CmpFPredicate::UEQ, + lhs, rhs); + } else { + cond = nestedBuilder.create(nestedLoc, + arith::CmpIPredicate::eq, + lhs, rhs); + } auto ifOp = nestedBuilder.create(nestedLoc, TypeRange{setValue.getType()}, cond, true); { OpBuilder::InsertionGuard guard(nestedBuilder); -- Gitee