From 75b2364db1d9dbc3ef0ef129ae887d80c7f83a12 Mon Sep 17 00:00:00 2001 From: peanutchan Date: Sat, 25 Jul 2026 10:51:21 +0800 Subject: [PATCH 1/3] feat(ptodsl): expose store atomic configuration Add the minimal PTODSL and VPTO lowering needed for TileLang UB-to-GM atomic add, max, and min copies. Co-authored-by: Cursor --- include/PTO/IR/VPTOOps.td | 1 + lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp | 8 ++++++++ lib/PTO/Transforms/VPTOLLVMEmitter.cpp | 8 ++++++++ ptodsl/ptodsl/_ops.py | 9 +++++++++ ptodsl/ptodsl/pto.py | 1 + ptodsl/tests/test_jit_compile.py | 6 ++++++ 6 files changed, 33 insertions(+) diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index 6b5411e78..6d4a0d7fc 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -506,6 +506,7 @@ def PTO_GetVms4SrOp : PTO_Op<"get_vms4_sr", [Pure]> { }]; } def PTO_SetCtrlOp : PTO_UnaryI64ConfigOp<"set_ctrl">; +def PTO_SetStoreAtomicCfgOp : PTO_UnaryI64ConfigOp<"set_store_atomic_cfg">; def PTO_StoreVfSimtInfoOp : PTO_SimtOp<"store_vfsimt_info"> { let summary = "Configure SIMT VF launch dimensions"; let arguments = (ins I32:$dimZ, I32:$dimY, I32:$dimX); diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index c432b313c..5ddecf92c 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -4060,6 +4060,12 @@ StringRef buildUnaryConfigCallee(MLIRContext *context) { return StringAttr::get(context, "llvm.hivm.SET.FPC").getValue(); } +template <> +StringRef buildUnaryConfigCallee( + MLIRContext *context) { + return StringAttr::get(context, "llvm.hivm.SET.ST.ATOMIC.CFG").getValue(); +} + template <> StringRef buildNullaryConfigCallee(MLIRContext *context) { return StringAttr::get(context, "llvm.hivm.SET.ATOMIC.S32").getValue(); @@ -10425,6 +10431,7 @@ static void populateVPTOOpLoweringPatterns(VPTOTypeConverter &typeConverter, LowerUnaryI64ConfigOpPattern, LowerUnaryI64ConfigOpPattern, LowerUnaryI64ConfigOpPattern, + LowerUnaryI64ConfigOpPattern, LowerNullaryConfigOpPattern, LowerNullaryConfigOpPattern, LowerPipeEventSyncOpPattern, @@ -10538,6 +10545,7 @@ static void configureVPTOOpLoweringTarget(ConversionTarget &target, pto::SetLoopSizeOutToL1Op, pto::SetMte2NzParaOp, pto::SetPadValOutToL1Op, pto::SetReluAlphaOp, pto::SetFixClipReluOp, pto::SetFpcOp, + pto::SetStoreAtomicCfgOp, pto::SetAtomicS32Op, pto::SetAtomicS8Op, pto::SetCtrlOp, pto::StoreVfSimtInfoOp, pto::SetMovPadValOp, pto::SetQuantPreOp>(); diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 27cf34c93..d8fdda4ee 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -4105,6 +4105,12 @@ StringRef buildUnaryConfigCallee(MLIRContext *context) { return StringAttr::get(context, "llvm.hivm.SET.FPC").getValue(); } +template <> +StringRef buildUnaryConfigCallee( + MLIRContext *context) { + return StringAttr::get(context, "llvm.hivm.SET.ST.ATOMIC.CFG").getValue(); +} + template <> StringRef buildNullaryConfigCallee(MLIRContext *context) { return StringAttr::get(context, "llvm.hivm.SET.ATOMIC.S32").getValue(); @@ -11105,6 +11111,7 @@ static void populateVPTOOpLoweringPatterns(VPTOTypeConverter &typeConverter, LowerUnaryI64ConfigOpPattern, LowerUnaryI64ConfigOpPattern, LowerUnaryI64ConfigOpPattern, + LowerUnaryI64ConfigOpPattern, LowerNullaryConfigOpPattern, LowerNullaryConfigOpPattern, LowerPipeEventSyncOpPattern, @@ -11284,6 +11291,7 @@ static void configureVPTOOpLoweringTarget(ConversionTarget &target, pto::SetLoopSizeOutToL1Op, pto::SetMte2NzParaOp, pto::SetPadValOutToL1Op, pto::SetReluAlphaOp, pto::SetFixClipReluOp, pto::SetFpcOp, + pto::SetStoreAtomicCfgOp, pto::SetAtomicS32Op, pto::SetAtomicS8Op, pto::SetCtrlOp, pto::StoreVfSimtInfoOp, pto::SetMovPadValOp, pto::SetQuantPreOp>(); diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index c3349f545..fcd020276 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -4892,6 +4892,14 @@ def mte_gm_ub(source, destination, l2_cache_ctl, len_burst, *, nburst, loops=Non ) +@_explicit_mode_only("pto.set_store_atomic_cfg(...)") +def set_store_atomic_cfg(config): + """Configure store atomic mode for subsequent UB-to-GM DMA stores.""" + _pto.SetStoreAtomicCfgOp( + _coerce_i64(config, context="set_store_atomic_cfg config") + ) + + @_explicit_mode_only("pto.mte_ub_gm(...)") def mte_ub_gm( source, @@ -6285,6 +6293,7 @@ def import_reserved_buffer(name, *, peer_func): "as_ptr", "mte_load", "mte_store", "mte_gm_ub", "mte_ub_gm", "mte_ub_ub", "mte_ub_l1", "mte_gm_l1", "mte_l1_ub", "mte_gm_l1_frac", "mte_l1_bt", "mte_l1_fb", "mem_bar", + "set_store_atomic_cfg", "mte_l1_l0a", "mte_l1_l0b", "mte_l1_l0a_mx", "mte_l1_l0b_mx", "mte_l0c_l1", "mte_l0c_gm", "mte_l0c_ub", "mad", "mad_acc", "mad_bias", "mad_mx", "mad_mx_acc", "mad_mx_bias", diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index 9c6279d7c..4ae39bdfe 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -122,6 +122,7 @@ tsort32, tmrgsort, tgather, mte_load, mte_store, mte_gm_ub, mte_ub_gm, mte_ub_ub, mte_ub_l1, mte_gm_l1, mte_l1_ub, mte_gm_l1_frac, mte_l1_bt, mte_l1_fb, mem_bar, + set_store_atomic_cfg, mte_l1_l0a, mte_l1_l0b, mte_l1_l0a_mx, mte_l1_l0b_mx, mte_l0c_l1, mte_l0c_gm, mte_l0c_ub, mad, mad_acc, mad_bias, mad_mx, mad_mx_acc, mad_mx_bias, diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 7047c31b1..7873b1d2a 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -3086,7 +3086,9 @@ def public_data_movement_surface_probe(): pto.mte_gm_ub(gm_src, ub_dst, 0, 256, nburst=(8, 256, 256), loops=[(4, 2048, 2048)]) pto.mte_gm_ub(gm_src, ub_dst, 0, 200, nburst=(64, 200, 256), pad=(0.0, 0, 0)) pto.mte_ub_gm(ub_src, gm_dst, 256, nburst=(64, 256, 1024)) + pto.set_store_atomic_cfg(4) pto.mte_ub_gm(ub_src, gm_dst, 128, nburst=(1, 128, 128), l2_cache="nared") + pto.set_store_atomic_cfg(0) pto.mte_ub_gm(ub_src, gm_dst, 64, nburst=(1, 64, 64), l2_cache="wtsred") pto.mte_ub_ub(ub_src, ub_dst, 8, nburst=(16, 0, 4)) pto.mte_ub_l1(ub_src, l1_dst, 8, nburst=(16, 0, 4)) @@ -6504,6 +6506,10 @@ def _enter_inline_simt_with_resource_attr(): expect("pto.sync.wait , 31" in sync_surface_text, "wait_intra_flag(Pipe.MTE3, 31) should lower the static physical event id through pto.sync.wait") expect(data_movement_surface_text.count("pto.mte_gm_ub") == 2, "public grouped GM->UB wrappers should lower to pto.mte_gm_ub") expect("pto.mte_ub_gm" in data_movement_surface_text, "public grouped UB->GM wrapper should lower to pto.mte_ub_gm") + expect( + data_movement_surface_text.count("pto.set_store_atomic_cfg") == 2, + "set_store_atomic_cfg should lower the requested config and reset", + ) expect( re.search(r"pto\.mte_ub_gm [^\n]+ nburst\([^)]+\) l2_cache_ctl\(%c7[^)\n]*\)", data_movement_surface_text) is not None, "mte_ub_gm(..., l2_cache='nared') should map to the l2_cache_ctl group", From 59c780428ef068d353e3145b50a79fcc5c0fce1d Mon Sep 17 00:00:00 2001 From: peanutchan Date: Sat, 25 Jul 2026 17:21:47 +0800 Subject: [PATCH 2/3] fix(vpto): lower MTE atomics through CTRL Expose CCE-style atomic configuration operations and preserve ordered CTRL reads so dtype and operation state reaches asynchronous MTE stores. Co-authored-by: Cursor --- include/PTO/IR/VPTOOps.td | 16 +++++- lib/PTO/Transforms/VPTOExpandWrapperOps.cpp | 60 +++++++++++++++++++++ ptodsl/ptodsl/_ops.py | 29 +++++++++- ptodsl/ptodsl/pto.py | 3 ++ ptodsl/tests/test_jit_compile.py | 26 +++++++++ test/lit/vpto/store_atomic_ctrl_expand.pto | 40 ++++++++++++++ 6 files changed, 172 insertions(+), 2 deletions(-) create mode 100644 test/lit/vpto/store_atomic_ctrl_expand.pto diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index 6d4a0d7fc..dd526ab12 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -496,7 +496,13 @@ def PTO_SetLoop3ParaOp : PTO_BinaryI64ConfigOp<"set_loop3_para">; def PTO_SetChannelParaOp : PTO_BinaryI64ConfigOp<"set_channel_para">; def PTO_SetMte2NzParaOp : PTO_UnaryI64ConfigOp<"set_mte2_nz_para">; def PTO_SetPadValOutToL1Op : PTO_UnaryI64ConfigOp<"set_pad_val_outtol1">; -def PTO_GetCtrlOp : PTO_NullaryI64PureOp<"get_ctrl">; +// Reading CTRL is ordered with respect to set_ctrl. It must not be modeled as +// Pure, otherwise CSE can reuse a value read before an intervening write. +def PTO_GetCtrlOp : PTO_Op<"get_ctrl"> { + let arguments = (ins); + let results = (outs I64:$result); + let assemblyFormat = [{ attr-dict `:` type($result) }]; +} def PTO_GetVms4SrOp : PTO_Op<"get_vms4_sr", [Pure]> { let arguments = (ins); let results = (outs I16:$list0, I16:$list1, I16:$list2, I16:$list3); @@ -543,7 +549,15 @@ def PTO_SetQuantPreOp : PTO_UnaryI64ConfigOp<"set_quant_pre">; def PTO_SetReluAlphaOp : PTO_UnaryI64ConfigOp<"set_relu_alpha">; def PTO_SetFixClipReluOp : PTO_UnaryI64ConfigOp<"set_fix_clip_relu">; def PTO_SetFpcOp : PTO_UnaryI64ConfigOp<"set_fpc">; +def PTO_SetAtomicAddOp : PTO_NullaryConfigOp<"set_atomic_add">; +def PTO_SetAtomicMaxOp : PTO_NullaryConfigOp<"set_atomic_max">; +def PTO_SetAtomicMinOp : PTO_NullaryConfigOp<"set_atomic_min">; +def PTO_SetAtomicNoneOp : PTO_NullaryConfigOp<"set_atomic_none">; +def PTO_SetAtomicF32Op : PTO_NullaryConfigOp<"set_atomic_f32">; +def PTO_SetAtomicF16Op : PTO_NullaryConfigOp<"set_atomic_f16">; +def PTO_SetAtomicBF16Op : PTO_NullaryConfigOp<"set_atomic_bf16">; def PTO_SetAtomicS32Op : PTO_NullaryConfigOp<"set_atomic_s32">; +def PTO_SetAtomicS16Op : PTO_NullaryConfigOp<"set_atomic_s16">; def PTO_SetAtomicS8Op : PTO_NullaryConfigOp<"set_atomic_s8">; def PTO_GetTidXOp : PTO_DirectSimtNullaryI32PureOp<"get_tid_x">; def PTO_GetTidYOp : PTO_DirectSimtNullaryI32PureOp<"get_tid_y">; diff --git a/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp b/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp index c455ce73e..7e66172cf 100644 --- a/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp +++ b/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp @@ -1908,6 +1908,56 @@ struct ExpandSimtLaunchPattern : public OpRewritePattern { } }; +struct AtomicCtrlUpdate { + uint64_t mask; + uint64_t value; +}; + +template +static AtomicCtrlUpdate getAtomicCtrlUpdate(); + +#define DEFINE_ATOMIC_CTRL_UPDATE(OpTy, Mask, Value) \ + template <> \ + AtomicCtrlUpdate getAtomicCtrlUpdate() { \ + return {Mask, Value}; \ + } + +// CCE set_atomic_* configures CTRL[10:6]. Dtype occupies [8:6] and the +// reduction operation occupies [10:9]. This matches the structured L0C-to-GM +// FIXP atomic CTRL encoding used by configureAccStoreCtrl above. +DEFINE_ATOMIC_CTRL_UPDATE(SetAtomicAddOp, 0x3ULL << 9, 0x0ULL << 9) +DEFINE_ATOMIC_CTRL_UPDATE(SetAtomicMaxOp, 0x3ULL << 9, 0x1ULL << 9) +DEFINE_ATOMIC_CTRL_UPDATE(SetAtomicMinOp, 0x3ULL << 9, 0x2ULL << 9) +DEFINE_ATOMIC_CTRL_UPDATE(SetAtomicNoneOp, 0x7ULL << 6, 0) +DEFINE_ATOMIC_CTRL_UPDATE(SetAtomicF32Op, 0x7ULL << 6, 0x1ULL << 6) +DEFINE_ATOMIC_CTRL_UPDATE(SetAtomicF16Op, 0x7ULL << 6, 0x2ULL << 6) +DEFINE_ATOMIC_CTRL_UPDATE(SetAtomicS16Op, 0x7ULL << 6, 0x3ULL << 6) +DEFINE_ATOMIC_CTRL_UPDATE(SetAtomicS32Op, 0x7ULL << 6, 0x4ULL << 6) +DEFINE_ATOMIC_CTRL_UPDATE(SetAtomicS8Op, 0x7ULL << 6, 0x5ULL << 6) +DEFINE_ATOMIC_CTRL_UPDATE(SetAtomicBF16Op, 0x7ULL << 6, 0x6ULL << 6) + +#undef DEFINE_ATOMIC_CTRL_UPDATE + +template +struct ExpandAtomicConfigPattern + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(AtomicConfigOp op, + PatternRewriter &rewriter) const override { + Location loc = op.getLoc(); + AtomicCtrlUpdate update = getAtomicCtrlUpdate(); + Value ctrl = rewriter.create(loc); + Value clearMask = getI64Constant(loc, rewriter, ~update.mask); + Value value = getI64Constant(loc, rewriter, update.value); + Value updated = rewriter.create(loc, ctrl, clearMask); + updated = rewriter.create(loc, updated, value); + rewriter.create(loc, updated); + rewriter.eraseOp(op); + return success(); + } +}; + struct VPTOExpandWrapperOpsPass : public pto::impl::VPTOExpandWrapperOpsBase { using pto::impl::VPTOExpandWrapperOpsBase< @@ -1938,6 +1988,16 @@ struct VPTOExpandWrapperOpsPass ExpandAccStoreGmPattern, ExpandAccStoreUbPattern, ExpandSimtLaunchPattern, + ExpandAtomicConfigPattern, + ExpandAtomicConfigPattern, + ExpandAtomicConfigPattern, + ExpandAtomicConfigPattern, + ExpandAtomicConfigPattern, + ExpandAtomicConfigPattern, + ExpandAtomicConfigPattern, + ExpandAtomicConfigPattern, + ExpandAtomicConfigPattern, + ExpandAtomicConfigPattern, ExpandMadSemanticPattern, ExpandMadSemanticPattern, ExpandMadSemanticPattern, diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index fcd020276..acff16879 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -4894,12 +4894,36 @@ def mte_gm_ub(source, destination, l2_cache_ctl, len_burst, *, nburst, loops=Non @_explicit_mode_only("pto.set_store_atomic_cfg(...)") def set_store_atomic_cfg(config): - """Configure store atomic mode for subsequent UB-to-GM DMA stores.""" + """Configure scalar ST atomic mode via ST_ATOMIC_CFG (SU path).""" _pto.SetStoreAtomicCfgOp( _coerce_i64(config, context="set_store_atomic_cfg config") ) +def _nullary_atomic_config(op_cls, api_name): + @_explicit_mode_only(f"pto.{api_name}(...)") + def _impl(): + op_cls() + + _impl.__name__ = api_name + _impl.__doc__ = ( + f"``pto.{api_name}`` – configure MTE/FIXP store-atomic CTRL state." + ) + return _impl + + +set_atomic_add = _nullary_atomic_config(_pto.SetAtomicAddOp, "set_atomic_add") +set_atomic_max = _nullary_atomic_config(_pto.SetAtomicMaxOp, "set_atomic_max") +set_atomic_min = _nullary_atomic_config(_pto.SetAtomicMinOp, "set_atomic_min") +set_atomic_none = _nullary_atomic_config(_pto.SetAtomicNoneOp, "set_atomic_none") +set_atomic_f32 = _nullary_atomic_config(_pto.SetAtomicF32Op, "set_atomic_f32") +set_atomic_f16 = _nullary_atomic_config(_pto.SetAtomicF16Op, "set_atomic_f16") +set_atomic_bf16 = _nullary_atomic_config(_pto.SetAtomicBF16Op, "set_atomic_bf16") +set_atomic_s32 = _nullary_atomic_config(_pto.SetAtomicS32Op, "set_atomic_s32") +set_atomic_s16 = _nullary_atomic_config(_pto.SetAtomicS16Op, "set_atomic_s16") +set_atomic_s8 = _nullary_atomic_config(_pto.SetAtomicS8Op, "set_atomic_s8") + + @_explicit_mode_only("pto.mte_ub_gm(...)") def mte_ub_gm( source, @@ -6294,6 +6318,9 @@ def import_reserved_buffer(name, *, peer_func): "mte_load", "mte_store", "mte_gm_ub", "mte_ub_gm", "mte_ub_ub", "mte_ub_l1", "mte_gm_l1", "mte_l1_ub", "mte_gm_l1_frac", "mte_l1_bt", "mte_l1_fb", "mem_bar", "set_store_atomic_cfg", + "set_atomic_add", "set_atomic_max", "set_atomic_min", "set_atomic_none", + "set_atomic_f32", "set_atomic_f16", "set_atomic_bf16", + "set_atomic_s32", "set_atomic_s16", "set_atomic_s8", "mte_l1_l0a", "mte_l1_l0b", "mte_l1_l0a_mx", "mte_l1_l0b_mx", "mte_l0c_l1", "mte_l0c_gm", "mte_l0c_ub", "mad", "mad_acc", "mad_bias", "mad_mx", "mad_mx_acc", "mad_mx_bias", diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index 4ae39bdfe..97af343d1 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -123,6 +123,9 @@ mte_load, mte_store, mte_gm_ub, mte_ub_gm, mte_ub_ub, mte_ub_l1, mte_gm_l1, mte_l1_ub, mte_gm_l1_frac, mte_l1_bt, mte_l1_fb, mem_bar, set_store_atomic_cfg, + set_atomic_add, set_atomic_max, set_atomic_min, set_atomic_none, + set_atomic_f32, set_atomic_f16, set_atomic_bf16, + set_atomic_s32, set_atomic_s16, set_atomic_s8, mte_l1_l0a, mte_l1_l0b, mte_l1_l0a_mx, mte_l1_l0b_mx, mte_l0c_l1, mte_l0c_gm, mte_l0c_ub, mad, mad_acc, mad_bias, mad_mx, mad_mx_acc, mad_mx_bias, diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 7873b1d2a..d31c77bf1 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -3089,6 +3089,16 @@ def public_data_movement_surface_probe(): pto.set_store_atomic_cfg(4) pto.mte_ub_gm(ub_src, gm_dst, 128, nburst=(1, 128, 128), l2_cache="nared") pto.set_store_atomic_cfg(0) + pto.set_atomic_s32() + pto.set_atomic_add() + pto.set_atomic_max() + pto.set_atomic_min() + pto.set_atomic_f32() + pto.set_atomic_f16() + pto.set_atomic_bf16() + pto.set_atomic_s16() + pto.set_atomic_s8() + pto.set_atomic_none() pto.mte_ub_gm(ub_src, gm_dst, 64, nburst=(1, 64, 64), l2_cache="wtsred") pto.mte_ub_ub(ub_src, ub_dst, 8, nburst=(16, 0, 4)) pto.mte_ub_l1(ub_src, l1_dst, 8, nburst=(16, 0, 4)) @@ -6510,6 +6520,22 @@ def _enter_inline_simt_with_resource_attr(): data_movement_surface_text.count("pto.set_store_atomic_cfg") == 2, "set_store_atomic_cfg should lower the requested config and reset", ) + for atomic_api in ( + "set_atomic_s32", + "set_atomic_add", + "set_atomic_max", + "set_atomic_min", + "set_atomic_f32", + "set_atomic_f16", + "set_atomic_bf16", + "set_atomic_s16", + "set_atomic_s8", + "set_atomic_none", + ): + expect( + f"pto.{atomic_api}" in data_movement_surface_text, + f"{atomic_api} should lower through the explicit PTODSL surface", + ) expect( re.search(r"pto\.mte_ub_gm [^\n]+ nburst\([^)]+\) l2_cache_ctl\(%c7[^)\n]*\)", data_movement_surface_text) is not None, "mte_ub_gm(..., l2_cache='nared') should map to the l2_cache_ctl group", diff --git a/test/lit/vpto/store_atomic_ctrl_expand.pto b/test/lit/vpto/store_atomic_ctrl_expand.pto new file mode 100644 index 000000000..de43cfabb --- /dev/null +++ b/test/lit/vpto/store_atomic_ctrl_expand.pto @@ -0,0 +1,40 @@ +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-before=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ROUNDTRIP +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=EXPAND + +module attributes {"pto.target_arch" = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @store_atomic_ctrl_probe() { + pto.set_atomic_s32 + pto.set_atomic_add + pto.set_atomic_max + pto.set_atomic_min + pto.set_atomic_f32 + pto.set_atomic_f16 + pto.set_atomic_bf16 + pto.set_atomic_s16 + pto.set_atomic_s8 + pto.set_atomic_none + return + } +} + +// ROUNDTRIP-LABEL: func.func @store_atomic_ctrl_probe +// ROUNDTRIP: pto.set_atomic_s32 +// ROUNDTRIP: pto.set_atomic_add +// ROUNDTRIP: pto.set_atomic_max +// ROUNDTRIP: pto.set_atomic_min +// ROUNDTRIP: pto.set_atomic_none + +// EXPAND-LABEL: func.func @store_atomic_ctrl_probe +// EXPAND-NOT: pto.set_atomic_ +// EXPAND: pto.get_ctrl +// EXPAND: arith.andi +// EXPAND: arith.ori +// EXPAND: pto.set_ctrl +// EXPAND: pto.get_ctrl +// EXPAND: arith.andi +// EXPAND: arith.ori +// EXPAND: pto.set_ctrl +// EXPAND: pto.get_ctrl +// EXPAND: arith.andi +// EXPAND: arith.ori +// EXPAND: pto.set_ctrl From c3958f5d0e386c8f294ddab2680bef461a1cc702 Mon Sep 17 00:00:00 2001 From: mouliangyu <21963576+mouliangyu@users.noreply.github.com> Date: Wed, 29 Jul 2026 20:20:06 +0800 Subject: [PATCH 3/3] test(vpto): cover UB-to-GM atomic add on sim --- .../mte-ub-gm-atomic-add-s32/compare.py | 35 ++++++++ .../mte-ub-gm-atomic-add-s32/golden.py | 37 +++++++++ .../mte-ub-gm-atomic-add-s32/kernel.pto | 40 ++++++++++ .../mte-ub-gm-atomic-add-s32/launch.cpp | 33 ++++++++ .../mte-ub-gm-atomic-add-s32/main.cpp | 80 +++++++++++++++++++ 5 files changed, 225 insertions(+) create mode 100644 test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/compare.py create mode 100644 test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/golden.py create mode 100644 test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/kernel.pto create mode 100644 test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/launch.cpp create mode 100644 test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/main.cpp diff --git a/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/compare.py b/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/compare.py new file mode 100644 index 000000000..80fc7c143 --- /dev/null +++ b/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/compare.py @@ -0,0 +1,35 @@ +#!/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 sys + +import numpy as np + + +def main() -> None: + golden = np.fromfile("golden_v2.bin", dtype=np.int32) + output = np.fromfile("v2.bin", dtype=np.int32) + if golden.shape == output.shape and np.array_equal(golden, output): + print("[INFO] compare passed") + return + if golden.shape != output.shape: + print(f"[ERROR] shape mismatch golden={golden.shape} output={output.shape}") + sys.exit(2) + diff = np.nonzero(golden != output)[0] + idx = int(diff[0]) if diff.size else -1 + print( + f"[ERROR] compare failed idx={idx} " + f"golden={golden[idx] if idx >= 0 else 'n/a'} " + f"output={output[idx] if idx >= 0 else 'n/a'}" + ) + sys.exit(2) + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/golden.py b/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/golden.py new file mode 100644 index 000000000..b49ec27d6 --- /dev/null +++ b/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/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 argparse +from pathlib import Path + +import numpy as np + +ELEMENTS = 64 + + +def generate(output_dir: Path) -> None: + src = np.full(ELEMENTS, 3, dtype=np.int32) + dst = np.full(ELEMENTS, 10, dtype=np.int32) + golden = np.full(ELEMENTS, 13, dtype=np.int32) + + output_dir.mkdir(parents=True, exist_ok=True) + src.tofile(output_dir / "v1.bin") + dst.tofile(output_dir / "v2.bin") + golden.tofile(output_dir / "golden_v2.bin") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, default=Path(".")) + args = parser.parse_args() + generate(args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/kernel.pto b/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/kernel.pto new file mode 100644 index 000000000..fa191dece --- /dev/null +++ b/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/kernel.pto @@ -0,0 +1,40 @@ +// 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. + +// ----------------------------------------------------------------------------- +// case: micro-op/vector-load-store/mte-ub-gm-atomic-add-s32 +// family: vector-load-store +// target_ops: pto.set_atomic_s32, pto.set_atomic_add, pto.mte_ub_gm, +// pto.set_atomic_none +// scenarios: ub-to-gm-atomic-add, s32, nonzero-gm-initial-value +// ----------------------------------------------------------------------------- + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @mte_ub_gm_atomic_add_s32_kernel(%src_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c256_i64 = arith.constant 256 : i64 + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c256_i64 + nburst(%c1_i64, %c256_i64, %c256_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.set_flag["PIPE_MTE2", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE3", "EVENT_ID0"] + + pto.set_atomic_s32 + pto.set_atomic_add + pto.mte_ub_gm %ub_src, %dst_gm, %c256_i64 + nburst(%c1_i64, %c256_i64, %c256_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.set_atomic_none + pto.barrier #pto.pipe + return + } +} diff --git a/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/launch.cpp b/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/launch.cpp new file mode 100644 index 000000000..3096c142f --- /dev/null +++ b/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/launch.cpp @@ -0,0 +1,33 @@ +// 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. + +#ifndef __VEC_SCOPE__ +#define __VEC_SCOPE__ +#endif +#include +#if !defined(__CCE_AICORE__) && !defined(TMRGSORT_HPP) +struct MrgSortExecutedNumList { + uint16_t mrgSortList0; + uint16_t mrgSortList1; + uint16_t mrgSortList2; + uint16_t mrgSortList3; +}; +#endif +#ifndef __CPU_SIM +#include "acl/acl.h" +#endif + +extern "C" __global__ [aicore] void +mte_ub_gm_atomic_add_s32_kernel(__gm__ int32_t *src, + __gm__ int32_t *dst); + +void LaunchMteUbGmAtomicAddS32Kernel(int32_t *src, int32_t *dst, + void *stream) { + mte_ub_gm_atomic_add_s32_kernel<<<1, nullptr, stream>>>( + (__gm__ int32_t *)src, (__gm__ int32_t *)dst); +} diff --git a/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/main.cpp b/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/main.cpp new file mode 100644 index 000000000..c96967492 --- /dev/null +++ b/test/vpto/cases/micro-op/vector-load-store/mte-ub-gm-atomic-add-s32/main.cpp @@ -0,0 +1,80 @@ +// 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 + +using namespace PtoTestCommon; + +#define ACL_CHECK(expr) \ + do { \ + const aclError _ret = (expr); \ + if (_ret != ACL_SUCCESS) { \ + std::fprintf(stderr, "[ERROR] %s failed: %d (%s:%d)\n", #expr, \ + (int)_ret, __FILE__, __LINE__); \ + rc = 1; \ + goto cleanup; \ + } \ + } while (0) + +void LaunchMteUbGmAtomicAddS32Kernel(int32_t *src, int32_t *dst, + void *stream); + +int main() { + constexpr size_t kElements = 64; + size_t bytes = kElements * sizeof(int32_t); + int32_t *srcHost = nullptr; + int32_t *dstHost = nullptr; + int32_t *srcDevice = nullptr; + int32_t *dstDevice = nullptr; + int rc = 0; + bool aclInited = false; + bool deviceSet = false; + int deviceId = 0; + aclrtStream stream = nullptr; + + ACL_CHECK(aclInit(nullptr)); + aclInited = true; + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) + deviceId = std::atoi(envDevice); + ACL_CHECK(aclrtSetDevice(deviceId)); + deviceSet = true; + ACL_CHECK(aclrtCreateStream(&stream)); + ACL_CHECK(aclrtMallocHost((void **)(&srcHost), bytes)); + ACL_CHECK(aclrtMallocHost((void **)(&dstHost), bytes)); + ACL_CHECK(aclrtMalloc((void **)&srcDevice, bytes, ACL_MEM_MALLOC_HUGE_FIRST)); + ACL_CHECK(aclrtMalloc((void **)&dstDevice, bytes, ACL_MEM_MALLOC_HUGE_FIRST)); + + ReadFile("./v1.bin", bytes, srcHost, bytes); + ReadFile("./v2.bin", bytes, dstHost, bytes); + ACL_CHECK(aclrtMemcpy(srcDevice, bytes, srcHost, bytes, + ACL_MEMCPY_HOST_TO_DEVICE)); + ACL_CHECK(aclrtMemcpy(dstDevice, bytes, dstHost, bytes, + ACL_MEMCPY_HOST_TO_DEVICE)); + LaunchMteUbGmAtomicAddS32Kernel(srcDevice, dstDevice, stream); + ACL_CHECK(aclrtSynchronizeStream(stream)); + ACL_CHECK(aclrtMemcpy(dstHost, bytes, dstDevice, bytes, + ACL_MEMCPY_DEVICE_TO_HOST)); + WriteFile("./v2.bin", dstHost, bytes); + +cleanup: + aclrtFree(srcDevice); + aclrtFree(dstDevice); + aclrtFreeHost(srcHost); + aclrtFreeHost(dstHost); + if (stream) + aclrtDestroyStream(stream); + if (deviceSet) + aclrtResetDevice(deviceId); + if (aclInited) + aclFinalize(); + return rc; +}