feat(ptodsl): expose store atomic configuration

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);
@@ -506,6 +512,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);
@@ -542,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">;

lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp

@@ -4060,6 +4060,12 @@ StringRef buildUnaryConfigCallee<pto::SetFpcOp>(MLIRContext *context) {
   return StringAttr::get(context, "llvm.hivm.SET.FPC").getValue();
 }
 
+template <>
+StringRef buildUnaryConfigCallee<pto::SetStoreAtomicCfgOp>(
+    MLIRContext *context) {
+  return StringAttr::get(context, "llvm.hivm.SET.ST.ATOMIC.CFG").getValue();
+}
+
 template <>
 StringRef buildNullaryConfigCallee<pto::SetAtomicS32Op>(MLIRContext *context) {
   return StringAttr::get(context, "llvm.hivm.SET.ATOMIC.S32").getValue();
@@ -10425,6 +10431,7 @@ static void populateVPTOOpLoweringPatterns(VPTOTypeConverter &typeConverter,
                LowerUnaryI64ConfigOpPattern<pto::SetMte2NzParaOp>,
                LowerUnaryI64ConfigOpPattern<pto::SetPadValOutToL1Op>,
                LowerUnaryI64ConfigOpPattern<pto::SetFpcOp>,
+               LowerUnaryI64ConfigOpPattern<pto::SetStoreAtomicCfgOp>,
                LowerNullaryConfigOpPattern<pto::SetAtomicS32Op>,
                LowerNullaryConfigOpPattern<pto::SetAtomicS8Op>,
                LowerPipeEventSyncOpPattern<pto::SetFlagOp>,
@@ -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>();

lib/PTO/Transforms/VPTOExpandWrapperOps.cpp

@@ -1908,6 +1908,56 @@ struct ExpandSimtLaunchPattern : public OpRewritePattern<pto::SimtLaunchOp> {
   }
 };
 
+struct AtomicCtrlUpdate {
+  uint64_t mask;
+  uint64_t value;
+};
+
+template <typename AtomicConfigOp>
+static AtomicCtrlUpdate getAtomicCtrlUpdate();
+
+#define DEFINE_ATOMIC_CTRL_UPDATE(OpTy, Mask, Value)                            \
+  template <>                                                                   \
+  AtomicCtrlUpdate getAtomicCtrlUpdate<pto::OpTy>() {                           \
+    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 <typename AtomicConfigOp>
+struct ExpandAtomicConfigPattern
+    : public OpRewritePattern<AtomicConfigOp> {
+  using OpRewritePattern<AtomicConfigOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AtomicConfigOp op,
+                                PatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    AtomicCtrlUpdate update = getAtomicCtrlUpdate<AtomicConfigOp>();
+    Value ctrl = rewriter.create<pto::GetCtrlOp>(loc);
+    Value clearMask = getI64Constant(loc, rewriter, ~update.mask);
+    Value value = getI64Constant(loc, rewriter, update.value);
+    Value updated = rewriter.create<arith::AndIOp>(loc, ctrl, clearMask);
+    updated = rewriter.create<arith::OrIOp>(loc, updated, value);
+    rewriter.create<pto::SetCtrlOp>(loc, updated);
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
 struct VPTOExpandWrapperOpsPass
     : public pto::impl::VPTOExpandWrapperOpsBase<VPTOExpandWrapperOpsPass> {
   using pto::impl::VPTOExpandWrapperOpsBase<
@@ -1938,6 +1988,16 @@ struct VPTOExpandWrapperOpsPass
                  ExpandAccStoreGmPattern,
                  ExpandAccStoreUbPattern,
                  ExpandSimtLaunchPattern,
+                 ExpandAtomicConfigPattern<pto::SetAtomicAddOp>,
+                 ExpandAtomicConfigPattern<pto::SetAtomicMaxOp>,
+                 ExpandAtomicConfigPattern<pto::SetAtomicMinOp>,
+                 ExpandAtomicConfigPattern<pto::SetAtomicNoneOp>,
+                 ExpandAtomicConfigPattern<pto::SetAtomicF32Op>,
+                 ExpandAtomicConfigPattern<pto::SetAtomicF16Op>,
+                 ExpandAtomicConfigPattern<pto::SetAtomicBF16Op>,
+                 ExpandAtomicConfigPattern<pto::SetAtomicS32Op>,
+                 ExpandAtomicConfigPattern<pto::SetAtomicS16Op>,
+                 ExpandAtomicConfigPattern<pto::SetAtomicS8Op>,
                  ExpandMadSemanticPattern<pto::MadOp>,
                  ExpandMadSemanticPattern<pto::MadAccOp>,
                  ExpandMadSemanticPattern<pto::MadBiasOp>,

lib/PTO/Transforms/VPTOLLVMEmitter.cpp

@@ -4105,6 +4105,12 @@ StringRef buildUnaryConfigCallee<pto::SetFpcOp>(MLIRContext *context) {
   return StringAttr::get(context, "llvm.hivm.SET.FPC").getValue();
 }
 
+template <>
+StringRef buildUnaryConfigCallee<pto::SetStoreAtomicCfgOp>(
+    MLIRContext *context) {
+  return StringAttr::get(context, "llvm.hivm.SET.ST.ATOMIC.CFG").getValue();
+}
+
 template <>
 StringRef buildNullaryConfigCallee<pto::SetAtomicS32Op>(MLIRContext *context) {
   return StringAttr::get(context, "llvm.hivm.SET.ATOMIC.S32").getValue();
@@ -11105,6 +11111,7 @@ static void populateVPTOOpLoweringPatterns(VPTOTypeConverter &typeConverter,
                LowerUnaryI64ConfigOpPattern<pto::SetMte2NzParaOp>,
                LowerUnaryI64ConfigOpPattern<pto::SetPadValOutToL1Op>,
                LowerUnaryI64ConfigOpPattern<pto::SetFpcOp>,
+               LowerUnaryI64ConfigOpPattern<pto::SetStoreAtomicCfgOp>,
                LowerNullaryConfigOpPattern<pto::SetAtomicS32Op>,
                LowerNullaryConfigOpPattern<pto::SetAtomicS8Op>,
                LowerPipeEventSyncOpPattern<pto::SetFlagOp>,
@@ -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>();

ptodsl/ptodsl/_ops.py

@@ -4892,6 +4892,38 @@ 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 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,
@@ -6285,6 +6317,10 @@ 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",
+    "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",

ptodsl/ptodsl/pto.py

@@ -122,6 +122,10 @@
     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,
+    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,

ptodsl/tests/test_jit_compile.py

@@ -3086,7 +3086,19 @@ 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.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))
@@ -6504,6 +6516,26 @@ def _enter_inline_simt_with_resource_attr():
     expect("pto.sync.wait <PIPE_MTE3>, 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",
+    )
+    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",

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<vector>} {
+  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

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()

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()