From d28b0bcc34c14b5f7f219f265cc8c7dafa00c0e4 Mon Sep 17 00:00:00 2001 From: Alan Wang Date: Wed, 22 Jul 2026 15:12:51 +0000 Subject: [PATCH 1/4] Initial commit for removing tilelang dsl --- CMakeLists.txt | 20 - _ptoas_build_backend.py | 5 +- lib/TileOps/__init__.py | 137 +- .../templates => lib/TileOps}/a5/__init__.py | 0 .../templates => lib/TileOps}/a5/_col_arg.py | 0 .../templates => lib/TileOps}/a5/_common.py | 0 .../templates => lib/TileOps}/a5/_cube.py | 0 .../TileOps}/a5/_elementwise.py | 0 .../TileOps}/a5/_expand_binary.py | 0 .../templates => lib/TileOps}/a5/_fillpad.py | 0 .../TileOps}/a5/_load_store.py | 0 .../templates => lib/TileOps}/a5/_part.py | 0 .../TileOps}/a5/_reductions.py | 0 .../TileOps}/a5/_remainder.py | 0 .../templates => lib/TileOps}/a5/_row_arg.py | 0 .../TileOps}/a5/_row_reductions.py | 0 .../templates => lib/TileOps}/a5/div_hp.py | 0 .../templates => lib/TileOps}/a5/tabs.py | 0 .../templates => lib/TileOps}/a5/tadd.py | 0 .../templates => lib/TileOps}/a5/tadds.py | 0 .../templates => lib/TileOps}/a5/tand.py | 0 .../templates => lib/TileOps}/a5/tands.py | 0 .../templates => lib/TileOps}/a5/tci.py | 0 .../templates => lib/TileOps}/a5/tcmp.py | 0 .../templates => lib/TileOps}/a5/tcmps.py | 0 .../TileOps}/a5/tcolargmax.py | 0 .../TileOps}/a5/tcolargmin.py | 0 .../TileOps}/a5/tcolexpand.py | 0 .../TileOps}/a5/tcolexpandadd.py | 0 .../TileOps}/a5/tcolexpanddiv.py | 0 .../TileOps}/a5/tcolexpandexpdif.py | 0 .../TileOps}/a5/tcolexpandmax.py | 0 .../TileOps}/a5/tcolexpandmin.py | 0 .../TileOps}/a5/tcolexpandmul.py | 0 .../TileOps}/a5/tcolexpandsub.py | 0 .../templates => lib/TileOps}/a5/tcolmax.py | 0 .../templates => lib/TileOps}/a5/tcolmin.py | 0 .../templates => lib/TileOps}/a5/tcolprod.py | 0 .../templates => lib/TileOps}/a5/tcolsum.py | 0 .../templates => lib/TileOps}/a5/tconcat.py | 0 .../templates => lib/TileOps}/a5/tcvt.py | 0 .../templates => lib/TileOps}/a5/tdequant.py | 0 .../templates => lib/TileOps}/a5/tdiv.py | 0 .../templates => lib/TileOps}/a5/tdivs.py | 0 .../templates => lib/TileOps}/a5/texp.py | 0 .../templates => lib/TileOps}/a5/texpand.py | 0 .../templates => lib/TileOps}/a5/textract.py | 0 .../TileOps}/a5/textract_fp.py | 0 .../templates => lib/TileOps}/a5/tfillpad.py | 0 .../TileOps}/a5/tfillpad_expand.py | 0 .../TileOps}/a5/tfillpad_inplace.py | 0 .../templates => lib/TileOps}/a5/tfmod.py | 0 .../templates => lib/TileOps}/a5/tfmods.py | 0 .../templates => lib/TileOps}/a5/tgatherb.py | 0 .../templates => lib/TileOps}/a5/tgemv.py | 0 .../templates => lib/TileOps}/a5/tgemv_acc.py | 0 .../TileOps}/a5/tgemv_bias.py | 0 .../templates => lib/TileOps}/a5/tgemv_mx.py | 0 .../TileOps}/a5/thistogram.py | 0 .../templates => lib/TileOps}/a5/tinsert.py | 0 .../templates => lib/TileOps}/a5/tload.py | 0 .../templates => lib/TileOps}/a5/tlog.py | 0 .../templates => lib/TileOps}/a5/tlrelu.py | 0 .../templates => lib/TileOps}/a5/tmatmul.py | 0 .../TileOps}/a5/tmatmul_acc.py | 0 .../TileOps}/a5/tmatmul_bias.py | 0 .../TileOps}/a5/tmatmul_mx.py | 0 .../templates => lib/TileOps}/a5/tmax.py | 0 .../templates => lib/TileOps}/a5/tmaxs.py | 0 .../templates => lib/TileOps}/a5/tmin.py | 0 .../templates => lib/TileOps}/a5/tmins.py | 0 .../templates => lib/TileOps}/a5/tmov.py | 0 .../templates => lib/TileOps}/a5/tmov2bias.py | 0 .../templates => lib/TileOps}/a5/tmov2left.py | 0 .../TileOps}/a5/tmov2right.py | 0 .../TileOps}/a5/tmov2scale.py | 0 .../templates => lib/TileOps}/a5/tmov2vec.py | 0 .../templates => lib/TileOps}/a5/tmov_fp.py | 0 .../templates => lib/TileOps}/a5/tmrgsort.py | 0 .../templates => lib/TileOps}/a5/tmul.py | 0 .../templates => lib/TileOps}/a5/tmuls.py | 0 .../templates => lib/TileOps}/a5/tneg.py | 0 .../templates => lib/TileOps}/a5/tnot.py | 0 .../templates => lib/TileOps}/a5/tor.py | 0 .../templates => lib/TileOps}/a5/tors.py | 0 .../templates => lib/TileOps}/a5/tpartadd.py | 0 .../templates => lib/TileOps}/a5/tpartmax.py | 0 .../templates => lib/TileOps}/a5/tpartmin.py | 0 .../templates => lib/TileOps}/a5/tpartmul.py | 0 .../templates => lib/TileOps}/a5/tprelu.py | 0 .../templates => lib/TileOps}/a5/trandom.py | 0 .../templates => lib/TileOps}/a5/trecip.py | 0 .../templates => lib/TileOps}/a5/trelu.py | 0 .../templates => lib/TileOps}/a5/trem.py | 0 .../templates => lib/TileOps}/a5/trems.py | 0 .../TileOps}/a5/trowargmax.py | 0 .../TileOps}/a5/trowargmin.py | 0 .../TileOps}/a5/trowexpand.py | 0 .../TileOps}/a5/trowexpandadd.py | 0 .../TileOps}/a5/trowexpanddiv.py | 0 .../TileOps}/a5/trowexpandexpdif.py | 0 .../TileOps}/a5/trowexpandmax.py | 0 .../TileOps}/a5/trowexpandmin.py | 0 .../TileOps}/a5/trowexpandmul.py | 0 .../TileOps}/a5/trowexpandsub.py | 0 .../templates => lib/TileOps}/a5/trowmax.py | 0 .../templates => lib/TileOps}/a5/trowmin.py | 0 .../templates => lib/TileOps}/a5/trowprod.py | 0 .../templates => lib/TileOps}/a5/trowsum.py | 0 .../templates => lib/TileOps}/a5/trsqrt.py | 0 .../templates => lib/TileOps}/a5/tsel.py | 0 .../templates => lib/TileOps}/a5/tsels.py | 0 .../templates => lib/TileOps}/a5/tshl.py | 0 .../templates => lib/TileOps}/a5/tshls.py | 0 .../templates => lib/TileOps}/a5/tshr.py | 0 .../templates => lib/TileOps}/a5/tshrs.py | 0 .../templates => lib/TileOps}/a5/tsort32.py | 0 .../templates => lib/TileOps}/a5/tsqrt.py | 0 .../templates => lib/TileOps}/a5/tstore.py | 0 .../templates => lib/TileOps}/a5/tsub.py | 0 .../templates => lib/TileOps}/a5/tsubs.py | 0 .../templates => lib/TileOps}/a5/ttrans.py | 0 .../templates => lib/TileOps}/a5/ttri.py | 0 .../templates => lib/TileOps}/a5/txor.py | 0 .../templates => lib/TileOps}/a5/txors.py | 0 lib/TileOps/div_hp.py | 455 ----- lib/TileOps/exp_hp.py | 30 - lib/TileOps/math.py | 584 ------- lib/TileOps/render_template_mlir.py | 378 ----- lib/TileOps/sqrt_hp.py | 83 - lib/TileOps/tabs_template.py | 29 - lib/TileOps/tadd_template.py | 32 - lib/TileOps/tadds_template.py | 31 - lib/TileOps/tand_template.py | 32 - lib/TileOps/tands_template.py | 40 - lib/TileOps/tcmp_template.py | 130 -- lib/TileOps/tcmps_template.py | 137 -- lib/TileOps/tcolargmax_template.py | 213 --- lib/TileOps/tcolargmin_template.py | 213 --- lib/TileOps/tcolexpand_template.py | 29 - lib/TileOps/tcolexpandadd_template.py | 31 - lib/TileOps/tcolexpanddiv_template.py | 54 - lib/TileOps/tcolexpandexpdif_template.py | 57 - lib/TileOps/tcolexpandmax_template.py | 31 - lib/TileOps/tcolexpandmin_template.py | 31 - lib/TileOps/tcolexpandmul_template.py | 31 - lib/TileOps/tcolexpandsub_template.py | 31 - lib/TileOps/tcolmax_template.py | 56 - lib/TileOps/tcolmin_template.py | 56 - lib/TileOps/tcolprod_template.py | 56 - lib/TileOps/tcolsum_template.py | 57 - lib/TileOps/tcvt_template.py | 1348 --------------- lib/TileOps/tdiv_template.py | 50 - lib/TileOps/tdivs_template.py | 92 -- lib/TileOps/texp_template.py | 33 - lib/TileOps/texpand_template.py | 72 - lib/TileOps/textract_template.py | 413 ----- lib/TileOps/tfillpad_expand_template.py | 167 -- lib/TileOps/tfillpad_inplace_template.py | 165 -- lib/TileOps/tfillpad_template.py | 173 -- lib/TileOps/tfmod_template.py | 71 - lib/TileOps/tfmods_template.py | 111 -- lib/TileOps/tgemv_acc_template.py | 27 - lib/TileOps/tgemv_bias_template.py | 28 - lib/TileOps/tgemv_mx_template.py | 96 -- lib/TileOps/tgemv_template.py | 33 - lib/TileOps/tinsert_template.py | 1463 ----------------- lib/TileOps/tload_template.py | 528 ------ lib/TileOps/tlog_template.py | 66 - lib/TileOps/tlrelu_template.py | 46 - lib/TileOps/tmatmul_acc_template.py | 27 - lib/TileOps/tmatmul_bias_template.py | 33 - lib/TileOps/tmatmul_mx_template.py | 100 -- lib/TileOps/tmatmul_template.py | 33 - lib/TileOps/tmax_template.py | 32 - lib/TileOps/tmaxs_template.py | 31 - lib/TileOps/tmin_template.py | 32 - lib/TileOps/tmins_template.py | 31 - lib/TileOps/tmov2bias_template.py | 92 -- lib/TileOps/tmov2left_template.py | 78 - lib/TileOps/tmov2right_template.py | 79 - lib/TileOps/tmov2scale_template.py | 96 -- lib/TileOps/tmov2vec_template.py | 85 - lib/TileOps/tmov_fp_template.py | 145 -- lib/TileOps/tmov_template.py | 100 -- lib/TileOps/tmrgsort_template.py | 273 --- lib/TileOps/tmul_template.py | 32 - lib/TileOps/tmuls_template.py | 31 - lib/TileOps/tneg_template.py | 29 - lib/TileOps/tnot_template.py | 30 - lib/TileOps/tor_template.py | 32 - lib/TileOps/tors_template.py | 40 - lib/TileOps/tpartadd_template.py | 79 - lib/TileOps/tpartmax_template.py | 54 - lib/TileOps/tpartmin_template.py | 54 - lib/TileOps/tpartmul_template.py | 79 - lib/TileOps/tprelu_template.py | 45 - lib/TileOps/trandom_template.py | 285 ---- lib/TileOps/trecip_template.py | 61 - lib/TileOps/trelu_template.py | 43 - lib/TileOps/trem_template.py | 66 - lib/TileOps/trems_template.py | 96 -- lib/TileOps/trowargmax_template.py | 74 - lib/TileOps/trowargmin_template.py | 74 - lib/TileOps/trowexpand_template.py | 48 - lib/TileOps/trowexpandadd_template.py | 52 - lib/TileOps/trowexpanddiv_template.py | 131 -- lib/TileOps/trowexpandexpdif_template.py | 78 - lib/TileOps/trowexpandmax_template.py | 52 - lib/TileOps/trowexpandmin_template.py | 52 - lib/TileOps/trowexpandmul_template.py | 76 - lib/TileOps/trowexpandsub_template.py | 76 - lib/TileOps/trowmax_template.py | 61 - lib/TileOps/trowmin_template.py | 63 - lib/TileOps/trowprod_template.py | 71 - lib/TileOps/trowsum_template.py | 75 - lib/TileOps/trsqrt_template.py | 36 - lib/TileOps/tsel_template.py | 99 -- lib/TileOps/tsels_template.py | 121 -- lib/TileOps/tshl_template.py | 32 - lib/TileOps/tshls_template.py | 31 - lib/TileOps/tshr_template.py | 32 - lib/TileOps/tshrs_template.py | 31 - lib/TileOps/tsort32_template.py | 241 --- lib/TileOps/tsqrt_template.py | 33 - lib/TileOps/tstore_template.py | 662 -------- lib/TileOps/tsub_template.py | 32 - lib/TileOps/tsubs_template.py | 38 - lib/TileOps/txor_template.py | 32 - lib/TileOps/txors_template.py | 39 - ptodsl/CMakeLists.txt | 18 +- ptodsl/README.md | 4 +- .../tilelib-template-authoring.md | 2 +- ptodsl/ptodsl/tilelib/registry.py | 2 +- ptodsl/ptodsl/tilelib/render.py | 2 +- ptodsl/ptodsl/tilelib/serving/daemon.py | 2 +- ptodsl/ptodsl/tilelib/templates/__init__.py | 144 +- ptodsl/tests/test_tilelib_render.py | 2 +- scripts/ptoas_env.sh | 2 + .../npu/a5/src/st/testcase/CMakeLists.txt | 1 - tools/ptoas/CMakeLists.txt | 3 +- tools/ptoas/ptoas.cpp | 157 +- 242 files changed, 264 insertions(+), 12891 deletions(-) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/__init__.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/_col_arg.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/_common.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/_cube.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/_elementwise.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/_expand_binary.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/_fillpad.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/_load_store.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/_part.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/_reductions.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/_remainder.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/_row_arg.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/_row_reductions.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/div_hp.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tabs.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tadd.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tadds.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tand.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tands.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tci.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcmp.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcmps.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolargmax.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolargmin.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolexpand.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolexpandadd.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolexpanddiv.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolexpandexpdif.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolexpandmax.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolexpandmin.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolexpandmul.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolexpandsub.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolmax.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolmin.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolprod.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcolsum.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tconcat.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tcvt.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tdequant.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tdiv.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tdivs.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/texp.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/texpand.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/textract.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/textract_fp.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tfillpad.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tfillpad_expand.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tfillpad_inplace.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tfmod.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tfmods.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tgatherb.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tgemv.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tgemv_acc.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tgemv_bias.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tgemv_mx.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/thistogram.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tinsert.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tload.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tlog.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tlrelu.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmatmul.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmatmul_acc.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmatmul_bias.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmatmul_mx.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmax.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmaxs.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmin.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmins.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmov.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmov2bias.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmov2left.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmov2right.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmov2scale.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmov2vec.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmov_fp.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmrgsort.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmul.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tmuls.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tneg.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tnot.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tor.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tors.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tpartadd.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tpartmax.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tpartmin.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tpartmul.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tprelu.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trandom.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trecip.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trelu.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trem.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trems.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowargmax.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowargmin.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowexpand.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowexpandadd.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowexpanddiv.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowexpandexpdif.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowexpandmax.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowexpandmin.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowexpandmul.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowexpandsub.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowmax.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowmin.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowprod.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trowsum.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/trsqrt.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tsel.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tsels.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tshl.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tshls.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tshr.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tshrs.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tsort32.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tsqrt.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tstore.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tsub.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/tsubs.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/ttrans.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/ttri.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/txor.py (100%) rename {ptodsl/ptodsl/tilelib/templates => lib/TileOps}/a5/txors.py (100%) delete mode 100644 lib/TileOps/div_hp.py delete mode 100644 lib/TileOps/exp_hp.py delete mode 100644 lib/TileOps/math.py delete mode 100644 lib/TileOps/render_template_mlir.py delete mode 100644 lib/TileOps/sqrt_hp.py delete mode 100644 lib/TileOps/tabs_template.py delete mode 100644 lib/TileOps/tadd_template.py delete mode 100644 lib/TileOps/tadds_template.py delete mode 100644 lib/TileOps/tand_template.py delete mode 100644 lib/TileOps/tands_template.py delete mode 100644 lib/TileOps/tcmp_template.py delete mode 100644 lib/TileOps/tcmps_template.py delete mode 100644 lib/TileOps/tcolargmax_template.py delete mode 100644 lib/TileOps/tcolargmin_template.py delete mode 100644 lib/TileOps/tcolexpand_template.py delete mode 100644 lib/TileOps/tcolexpandadd_template.py delete mode 100644 lib/TileOps/tcolexpanddiv_template.py delete mode 100644 lib/TileOps/tcolexpandexpdif_template.py delete mode 100644 lib/TileOps/tcolexpandmax_template.py delete mode 100644 lib/TileOps/tcolexpandmin_template.py delete mode 100644 lib/TileOps/tcolexpandmul_template.py delete mode 100644 lib/TileOps/tcolexpandsub_template.py delete mode 100644 lib/TileOps/tcolmax_template.py delete mode 100644 lib/TileOps/tcolmin_template.py delete mode 100644 lib/TileOps/tcolprod_template.py delete mode 100644 lib/TileOps/tcolsum_template.py delete mode 100644 lib/TileOps/tcvt_template.py delete mode 100644 lib/TileOps/tdiv_template.py delete mode 100644 lib/TileOps/tdivs_template.py delete mode 100644 lib/TileOps/texp_template.py delete mode 100644 lib/TileOps/texpand_template.py delete mode 100644 lib/TileOps/textract_template.py delete mode 100644 lib/TileOps/tfillpad_expand_template.py delete mode 100644 lib/TileOps/tfillpad_inplace_template.py delete mode 100644 lib/TileOps/tfillpad_template.py delete mode 100644 lib/TileOps/tfmod_template.py delete mode 100644 lib/TileOps/tfmods_template.py delete mode 100644 lib/TileOps/tgemv_acc_template.py delete mode 100644 lib/TileOps/tgemv_bias_template.py delete mode 100644 lib/TileOps/tgemv_mx_template.py delete mode 100644 lib/TileOps/tgemv_template.py delete mode 100644 lib/TileOps/tinsert_template.py delete mode 100644 lib/TileOps/tload_template.py delete mode 100644 lib/TileOps/tlog_template.py delete mode 100644 lib/TileOps/tlrelu_template.py delete mode 100644 lib/TileOps/tmatmul_acc_template.py delete mode 100644 lib/TileOps/tmatmul_bias_template.py delete mode 100644 lib/TileOps/tmatmul_mx_template.py delete mode 100644 lib/TileOps/tmatmul_template.py delete mode 100644 lib/TileOps/tmax_template.py delete mode 100644 lib/TileOps/tmaxs_template.py delete mode 100644 lib/TileOps/tmin_template.py delete mode 100644 lib/TileOps/tmins_template.py delete mode 100644 lib/TileOps/tmov2bias_template.py delete mode 100644 lib/TileOps/tmov2left_template.py delete mode 100644 lib/TileOps/tmov2right_template.py delete mode 100644 lib/TileOps/tmov2scale_template.py delete mode 100644 lib/TileOps/tmov2vec_template.py delete mode 100644 lib/TileOps/tmov_fp_template.py delete mode 100644 lib/TileOps/tmov_template.py delete mode 100644 lib/TileOps/tmrgsort_template.py delete mode 100644 lib/TileOps/tmul_template.py delete mode 100644 lib/TileOps/tmuls_template.py delete mode 100644 lib/TileOps/tneg_template.py delete mode 100644 lib/TileOps/tnot_template.py delete mode 100644 lib/TileOps/tor_template.py delete mode 100644 lib/TileOps/tors_template.py delete mode 100644 lib/TileOps/tpartadd_template.py delete mode 100644 lib/TileOps/tpartmax_template.py delete mode 100644 lib/TileOps/tpartmin_template.py delete mode 100644 lib/TileOps/tpartmul_template.py delete mode 100644 lib/TileOps/tprelu_template.py delete mode 100644 lib/TileOps/trandom_template.py delete mode 100644 lib/TileOps/trecip_template.py delete mode 100644 lib/TileOps/trelu_template.py delete mode 100644 lib/TileOps/trem_template.py delete mode 100644 lib/TileOps/trems_template.py delete mode 100644 lib/TileOps/trowargmax_template.py delete mode 100644 lib/TileOps/trowargmin_template.py delete mode 100644 lib/TileOps/trowexpand_template.py delete mode 100644 lib/TileOps/trowexpandadd_template.py delete mode 100644 lib/TileOps/trowexpanddiv_template.py delete mode 100644 lib/TileOps/trowexpandexpdif_template.py delete mode 100644 lib/TileOps/trowexpandmax_template.py delete mode 100644 lib/TileOps/trowexpandmin_template.py delete mode 100644 lib/TileOps/trowexpandmul_template.py delete mode 100644 lib/TileOps/trowexpandsub_template.py delete mode 100644 lib/TileOps/trowmax_template.py delete mode 100644 lib/TileOps/trowmin_template.py delete mode 100644 lib/TileOps/trowprod_template.py delete mode 100644 lib/TileOps/trowsum_template.py delete mode 100644 lib/TileOps/trsqrt_template.py delete mode 100644 lib/TileOps/tsel_template.py delete mode 100644 lib/TileOps/tsels_template.py delete mode 100644 lib/TileOps/tshl_template.py delete mode 100644 lib/TileOps/tshls_template.py delete mode 100644 lib/TileOps/tshr_template.py delete mode 100644 lib/TileOps/tshrs_template.py delete mode 100644 lib/TileOps/tsort32_template.py delete mode 100644 lib/TileOps/tsqrt_template.py delete mode 100644 lib/TileOps/tstore_template.py delete mode 100644 lib/TileOps/tsub_template.py delete mode 100644 lib/TileOps/tsubs_template.py delete mode 100644 lib/TileOps/txor_template.py delete mode 100644 lib/TileOps/txors_template.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 7fb9465edb..41617bb313 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -247,7 +247,6 @@ if(PTO_ENABLE_PYTHON_BINDING) add_subdirectory(python) add_subdirectory(ptodsl) - add_subdirectory(tilelang-dsl) endif() # ========================================================= @@ -283,25 +282,6 @@ if(BUILD_TESTING) ) add_subdirectory(ptodsl/tests) - add_test( - NAME tilelang_dsl_import - COMMAND "${Python3_EXECUTABLE}" - "${CMAKE_CURRENT_SOURCE_DIR}/tilelang-dsl/tests/import_tilelang_dsl.py" - ) - set_tests_properties(tilelang_dsl_import PROPERTIES - LABELS "PTODSL" - ENVIRONMENT "${_pto_python_test_env}" - ) - add_test( - NAME tilelang_dsl_unittest - COMMAND "${Python3_EXECUTABLE}" -m unittest discover - -s "${CMAKE_CURRENT_SOURCE_DIR}/tilelang-dsl/tests" - -p "test_*.py" - ) - set_tests_properties(tilelang_dsl_unittest PROPERTIES - LABELS "PTODSL" - ENVIRONMENT "${_pto_python_test_env}" - ) add_test( NAME pto_memory_consistency_bindings COMMAND "${Python3_EXECUTABLE}" diff --git a/_ptoas_build_backend.py b/_ptoas_build_backend.py index 1dfea891db..80f3895d51 100644 --- a/_ptoas_build_backend.py +++ b/_ptoas_build_backend.py @@ -63,6 +63,7 @@ def _find_llvm_dir(): ) _BUILD_DIR = Path(os.environ.get("PTO_BUILD_DIR", str(_REPO / "build"))) _PTODSL_SOURCE_ROOT = _REPO / "ptodsl" +_TILEOPS_SOURCE_ROOT = _REPO / "lib" _MLIR_PY_PKG = None if "MLIR_PYTHON_PACKAGE_DIR" in os.environ: _MLIR_PY_PKG = Path(os.environ["MLIR_PYTHON_PACKAGE_DIR"]) @@ -303,9 +304,11 @@ def build_editable(wheel_directory, config_settings=None, metadata_directory=Non # Prefer the repository PTODSL sources so editable installs pick up # local Python edits instead of staged/install-tree copies. str(_PTODSL_SOURCE_ROOT), + # PTODSL TileLib templates live in the top-level TileOps package. + str(_TILEOPS_SOURCE_ROOT), # Handwritten Python sources (pto/dialects/pto.py, etc.). str(_REPO / "python"), - # Installed PTOAS runtime overlay (TileOps/resources when present). + # Installed PTOAS runtime overlay. str(_PTO_INSTALL_DIR), # Generated files (_pto.so, _pto_ops_gen.py) under mlir/ namespace. str(_BUILD_DIR / "python"), diff --git a/lib/TileOps/__init__.py b/lib/TileOps/__init__.py index 34437fbeea..77ad910701 100644 --- a/lib/TileOps/__init__.py +++ b/lib/TileOps/__init__.py @@ -5,5 +5,140 @@ # 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. +"""Lazy loader for ported, per-architecture TileLib templates.""" -"""Shared TileOps template helpers.""" +from functools import lru_cache +from importlib import import_module + + +_TEMPLATE_MODULES = { + ("a5", "pto.tabs"): ".a5.tabs", + ("a5", "pto.tadd"): ".a5.tadd", + ("a5", "pto.tadds"): ".a5.tadds", + ("a5", "pto.tand"): ".a5.tand", + ("a5", "pto.tands"): ".a5.tands", + ("a5", "pto.tcmp"): ".a5.tcmp", + ("a5", "pto.tcmps"): ".a5.tcmps", + ("a5", "pto.tcolargmax"): ".a5.tcolargmax", + ("a5", "pto.tcolargmin"): ".a5.tcolargmin", + ("a5", "pto.tcolmax"): ".a5.tcolmax", + ("a5", "pto.tcolexpand"): ".a5.tcolexpand", + ("a5", "pto.tcolexpandadd"): ".a5.tcolexpandadd", + ("a5", "pto.tcolexpanddiv"): ".a5.tcolexpanddiv", + ("a5", "pto.tcolexpandexpdif"): ".a5.tcolexpandexpdif", + ("a5", "pto.tcolexpandmax"): ".a5.tcolexpandmax", + ("a5", "pto.tcolexpandmin"): ".a5.tcolexpandmin", + ("a5", "pto.tcolexpandmul"): ".a5.tcolexpandmul", + ("a5", "pto.tcolexpandsub"): ".a5.tcolexpandsub", + ("a5", "pto.tcolmin"): ".a5.tcolmin", + ("a5", "pto.tcolprod"): ".a5.tcolprod", + ("a5", "pto.tcolsum"): ".a5.tcolsum", + ("a5", "pto.tdiv"): ".a5.tdiv", + ("a5", "pto.tdivs"): ".a5.tdivs", + ("a5", "pto.tcvt"): ".a5.tcvt", + ("a5", "pto.texp"): ".a5.texp", + ("a5", "pto.texpands"): ".a5.texpand", + ("a5", "pto.textract"): ".a5.textract", + ("a5", "pto.textract_fp"): ".a5.textract_fp", + ("a5", "pto.tfmod"): ".a5.tfmod", + ("a5", "pto.tfmods"): ".a5.tfmods", + ("a5", "pto.tfillpad"): ".a5.tfillpad", + ("a5", "pto.tfillpad_expand"): ".a5.tfillpad_expand", + ("a5", "pto.tfillpad_inplace"): ".a5.tfillpad_inplace", + ("a5", "pto.tgatherb"): ".a5.tgatherb", + ("a5", "pto.tgemv"): ".a5.tgemv", + ("a5", "pto.tgemv.acc"): ".a5.tgemv_acc", + ("a5", "pto.tgemv.bias"): ".a5.tgemv_bias", + ("a5", "pto.tgemv.mx"): ".a5.tgemv_mx", + ("a5", "pto.tgemv.mx.acc"): ".a5.tgemv_mx", + ("a5", "pto.tgemv.mx.bias"): ".a5.tgemv_mx", + ("a5", "pto.tinsert"): ".a5.tinsert", + ("a5", "pto.tlrelu"): ".a5.tlrelu", + ("a5", "pto.tlog"): ".a5.tlog", + ("a5", "pto.tload"): ".a5.tload", + ("a5", "pto.tmax"): ".a5.tmax", + ("a5", "pto.tmaxs"): ".a5.tmaxs", + ("a5", "pto.tmin"): ".a5.tmin", + ("a5", "pto.tmins"): ".a5.tmins", + ("a5", "pto.tmatmul"): ".a5.tmatmul", + ("a5", "pto.tmatmul.acc"): ".a5.tmatmul_acc", + ("a5", "pto.tmatmul.bias"): ".a5.tmatmul_bias", + ("a5", "pto.tmatmul.mx"): ".a5.tmatmul_mx", + ("a5", "pto.tmatmul.mx.acc"): ".a5.tmatmul_mx", + ("a5", "pto.tmatmul.mx.bias"): ".a5.tmatmul_mx", + ("a5", "pto.tmov"): ( + ".a5.tmov", + ".a5.tmov2bias", + ".a5.tmov2left", + ".a5.tmov2right", + ".a5.tmov2scale", + ".a5.tmov2vec", + ".a5.tmov_fp", + ), + ("a5", "pto.tmul"): ".a5.tmul", + ("a5", "pto.tmuls"): ".a5.tmuls", + ("a5", "pto.tneg"): ".a5.tneg", + ("a5", "pto.tnot"): ".a5.tnot", + ("a5", "pto.tor"): ".a5.tor", + ("a5", "pto.tors"): ".a5.tors", + ("a5", "pto.tpartadd"): ".a5.tpartadd", + ("a5", "pto.tpartmax"): ".a5.tpartmax", + ("a5", "pto.tpartmin"): ".a5.tpartmin", + ("a5", "pto.tpartmul"): ".a5.tpartmul", + ("a5", "pto.tprelu"): ".a5.tprelu", + ("a5", "pto.trelu"): ".a5.trelu", + ("a5", "pto.trecip"): ".a5.trecip", + ("a5", "pto.trem"): ".a5.trem", + ("a5", "pto.trems"): ".a5.trems", + ("a5", "pto.trandom"): ".a5.trandom", + ("a5", "pto.trsqrt"): ".a5.trsqrt", + ("a5", "pto.trowargmax"): ".a5.trowargmax", + ("a5", "pto.trowargmin"): ".a5.trowargmin", + ("a5", "pto.trowexpand"): ".a5.trowexpand", + ("a5", "pto.trowexpandadd"): ".a5.trowexpandadd", + ("a5", "pto.trowexpanddiv"): ".a5.trowexpanddiv", + ("a5", "pto.trowexpandexpdif"): ".a5.trowexpandexpdif", + ("a5", "pto.trowexpandmax"): ".a5.trowexpandmax", + ("a5", "pto.trowexpandmin"): ".a5.trowexpandmin", + ("a5", "pto.trowexpandmul"): ".a5.trowexpandmul", + ("a5", "pto.trowexpandsub"): ".a5.trowexpandsub", + ("a5", "pto.trowmax"): ".a5.trowmax", + ("a5", "pto.trowmin"): ".a5.trowmin", + ("a5", "pto.trowprod"): ".a5.trowprod", + ("a5", "pto.trowsum"): ".a5.trowsum", + ("a5", "pto.tsel"): ".a5.tsel", + ("a5", "pto.tsels"): ".a5.tsels", + ("a5", "pto.tshl"): ".a5.tshl", + ("a5", "pto.tshls"): ".a5.tshls", + ("a5", "pto.tshr"): ".a5.tshr", + ("a5", "pto.tshrs"): ".a5.tshrs", + ("a5", "pto.tmrgsort"): ".a5.tmrgsort", + ("a5", "pto.tsort32"): ".a5.tsort32", + ("a5", "pto.tstore"): ".a5.tstore", + ("a5", "pto.tstore_fp"): ".a5.tstore", + ("a5", "pto.tsub"): ".a5.tsub", + ("a5", "pto.tsubs"): ".a5.tsubs", + ("a5", "pto.tsqrt"): ".a5.tsqrt", + ("a5", "pto.txor"): ".a5.txor", + ("a5", "pto.txors"): ".a5.txors", +} + + +@lru_cache(maxsize=None) +def load_template(op: str, target: str) -> bool: + """Import and register only the template module for ``(target, op)``. + + Both this cache and Python's module cache make repeated requests no-ops. + Returns ``False`` when this TileLib has no module for the requested pair. + """ + module_names = _TEMPLATE_MODULES.get((target, op)) + if module_names is None: + return False + if isinstance(module_names, str): + module_names = (module_names,) + for module_name in module_names: + import_module(module_name, package=__name__) + return True + + +__all__ = ["load_template"] diff --git a/ptodsl/ptodsl/tilelib/templates/a5/__init__.py b/lib/TileOps/a5/__init__.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/__init__.py rename to lib/TileOps/a5/__init__.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_col_arg.py b/lib/TileOps/a5/_col_arg.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/_col_arg.py rename to lib/TileOps/a5/_col_arg.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_common.py b/lib/TileOps/a5/_common.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/_common.py rename to lib/TileOps/a5/_common.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_cube.py b/lib/TileOps/a5/_cube.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/_cube.py rename to lib/TileOps/a5/_cube.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_elementwise.py b/lib/TileOps/a5/_elementwise.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/_elementwise.py rename to lib/TileOps/a5/_elementwise.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_expand_binary.py b/lib/TileOps/a5/_expand_binary.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/_expand_binary.py rename to lib/TileOps/a5/_expand_binary.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_fillpad.py b/lib/TileOps/a5/_fillpad.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/_fillpad.py rename to lib/TileOps/a5/_fillpad.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_load_store.py b/lib/TileOps/a5/_load_store.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/_load_store.py rename to lib/TileOps/a5/_load_store.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_part.py b/lib/TileOps/a5/_part.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/_part.py rename to lib/TileOps/a5/_part.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_reductions.py b/lib/TileOps/a5/_reductions.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/_reductions.py rename to lib/TileOps/a5/_reductions.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_remainder.py b/lib/TileOps/a5/_remainder.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/_remainder.py rename to lib/TileOps/a5/_remainder.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_row_arg.py b/lib/TileOps/a5/_row_arg.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/_row_arg.py rename to lib/TileOps/a5/_row_arg.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_row_reductions.py b/lib/TileOps/a5/_row_reductions.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/_row_reductions.py rename to lib/TileOps/a5/_row_reductions.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/div_hp.py b/lib/TileOps/a5/div_hp.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/div_hp.py rename to lib/TileOps/a5/div_hp.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tabs.py b/lib/TileOps/a5/tabs.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tabs.py rename to lib/TileOps/a5/tabs.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tadd.py b/lib/TileOps/a5/tadd.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tadd.py rename to lib/TileOps/a5/tadd.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tadds.py b/lib/TileOps/a5/tadds.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tadds.py rename to lib/TileOps/a5/tadds.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tand.py b/lib/TileOps/a5/tand.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tand.py rename to lib/TileOps/a5/tand.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tands.py b/lib/TileOps/a5/tands.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tands.py rename to lib/TileOps/a5/tands.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tci.py b/lib/TileOps/a5/tci.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tci.py rename to lib/TileOps/a5/tci.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcmp.py b/lib/TileOps/a5/tcmp.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcmp.py rename to lib/TileOps/a5/tcmp.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcmps.py b/lib/TileOps/a5/tcmps.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcmps.py rename to lib/TileOps/a5/tcmps.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolargmax.py b/lib/TileOps/a5/tcolargmax.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolargmax.py rename to lib/TileOps/a5/tcolargmax.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolargmin.py b/lib/TileOps/a5/tcolargmin.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolargmin.py rename to lib/TileOps/a5/tcolargmin.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolexpand.py b/lib/TileOps/a5/tcolexpand.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolexpand.py rename to lib/TileOps/a5/tcolexpand.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolexpandadd.py b/lib/TileOps/a5/tcolexpandadd.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolexpandadd.py rename to lib/TileOps/a5/tcolexpandadd.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolexpanddiv.py b/lib/TileOps/a5/tcolexpanddiv.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolexpanddiv.py rename to lib/TileOps/a5/tcolexpanddiv.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolexpandexpdif.py b/lib/TileOps/a5/tcolexpandexpdif.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolexpandexpdif.py rename to lib/TileOps/a5/tcolexpandexpdif.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolexpandmax.py b/lib/TileOps/a5/tcolexpandmax.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolexpandmax.py rename to lib/TileOps/a5/tcolexpandmax.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolexpandmin.py b/lib/TileOps/a5/tcolexpandmin.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolexpandmin.py rename to lib/TileOps/a5/tcolexpandmin.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolexpandmul.py b/lib/TileOps/a5/tcolexpandmul.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolexpandmul.py rename to lib/TileOps/a5/tcolexpandmul.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolexpandsub.py b/lib/TileOps/a5/tcolexpandsub.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolexpandsub.py rename to lib/TileOps/a5/tcolexpandsub.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolmax.py b/lib/TileOps/a5/tcolmax.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolmax.py rename to lib/TileOps/a5/tcolmax.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolmin.py b/lib/TileOps/a5/tcolmin.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolmin.py rename to lib/TileOps/a5/tcolmin.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolprod.py b/lib/TileOps/a5/tcolprod.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolprod.py rename to lib/TileOps/a5/tcolprod.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcolsum.py b/lib/TileOps/a5/tcolsum.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcolsum.py rename to lib/TileOps/a5/tcolsum.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tconcat.py b/lib/TileOps/a5/tconcat.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tconcat.py rename to lib/TileOps/a5/tconcat.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tcvt.py b/lib/TileOps/a5/tcvt.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tcvt.py rename to lib/TileOps/a5/tcvt.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tdequant.py b/lib/TileOps/a5/tdequant.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tdequant.py rename to lib/TileOps/a5/tdequant.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tdiv.py b/lib/TileOps/a5/tdiv.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tdiv.py rename to lib/TileOps/a5/tdiv.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tdivs.py b/lib/TileOps/a5/tdivs.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tdivs.py rename to lib/TileOps/a5/tdivs.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/texp.py b/lib/TileOps/a5/texp.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/texp.py rename to lib/TileOps/a5/texp.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/texpand.py b/lib/TileOps/a5/texpand.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/texpand.py rename to lib/TileOps/a5/texpand.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/textract.py b/lib/TileOps/a5/textract.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/textract.py rename to lib/TileOps/a5/textract.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/textract_fp.py b/lib/TileOps/a5/textract_fp.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/textract_fp.py rename to lib/TileOps/a5/textract_fp.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tfillpad.py b/lib/TileOps/a5/tfillpad.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tfillpad.py rename to lib/TileOps/a5/tfillpad.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tfillpad_expand.py b/lib/TileOps/a5/tfillpad_expand.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tfillpad_expand.py rename to lib/TileOps/a5/tfillpad_expand.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tfillpad_inplace.py b/lib/TileOps/a5/tfillpad_inplace.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tfillpad_inplace.py rename to lib/TileOps/a5/tfillpad_inplace.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tfmod.py b/lib/TileOps/a5/tfmod.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tfmod.py rename to lib/TileOps/a5/tfmod.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tfmods.py b/lib/TileOps/a5/tfmods.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tfmods.py rename to lib/TileOps/a5/tfmods.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tgatherb.py b/lib/TileOps/a5/tgatherb.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tgatherb.py rename to lib/TileOps/a5/tgatherb.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tgemv.py b/lib/TileOps/a5/tgemv.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tgemv.py rename to lib/TileOps/a5/tgemv.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tgemv_acc.py b/lib/TileOps/a5/tgemv_acc.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tgemv_acc.py rename to lib/TileOps/a5/tgemv_acc.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tgemv_bias.py b/lib/TileOps/a5/tgemv_bias.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tgemv_bias.py rename to lib/TileOps/a5/tgemv_bias.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tgemv_mx.py b/lib/TileOps/a5/tgemv_mx.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tgemv_mx.py rename to lib/TileOps/a5/tgemv_mx.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/thistogram.py b/lib/TileOps/a5/thistogram.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/thistogram.py rename to lib/TileOps/a5/thistogram.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tinsert.py b/lib/TileOps/a5/tinsert.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tinsert.py rename to lib/TileOps/a5/tinsert.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tload.py b/lib/TileOps/a5/tload.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tload.py rename to lib/TileOps/a5/tload.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tlog.py b/lib/TileOps/a5/tlog.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tlog.py rename to lib/TileOps/a5/tlog.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tlrelu.py b/lib/TileOps/a5/tlrelu.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tlrelu.py rename to lib/TileOps/a5/tlrelu.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmatmul.py b/lib/TileOps/a5/tmatmul.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmatmul.py rename to lib/TileOps/a5/tmatmul.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmatmul_acc.py b/lib/TileOps/a5/tmatmul_acc.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmatmul_acc.py rename to lib/TileOps/a5/tmatmul_acc.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmatmul_bias.py b/lib/TileOps/a5/tmatmul_bias.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmatmul_bias.py rename to lib/TileOps/a5/tmatmul_bias.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmatmul_mx.py b/lib/TileOps/a5/tmatmul_mx.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmatmul_mx.py rename to lib/TileOps/a5/tmatmul_mx.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmax.py b/lib/TileOps/a5/tmax.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmax.py rename to lib/TileOps/a5/tmax.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmaxs.py b/lib/TileOps/a5/tmaxs.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmaxs.py rename to lib/TileOps/a5/tmaxs.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmin.py b/lib/TileOps/a5/tmin.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmin.py rename to lib/TileOps/a5/tmin.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmins.py b/lib/TileOps/a5/tmins.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmins.py rename to lib/TileOps/a5/tmins.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmov.py b/lib/TileOps/a5/tmov.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmov.py rename to lib/TileOps/a5/tmov.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmov2bias.py b/lib/TileOps/a5/tmov2bias.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmov2bias.py rename to lib/TileOps/a5/tmov2bias.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmov2left.py b/lib/TileOps/a5/tmov2left.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmov2left.py rename to lib/TileOps/a5/tmov2left.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmov2right.py b/lib/TileOps/a5/tmov2right.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmov2right.py rename to lib/TileOps/a5/tmov2right.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmov2scale.py b/lib/TileOps/a5/tmov2scale.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmov2scale.py rename to lib/TileOps/a5/tmov2scale.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmov2vec.py b/lib/TileOps/a5/tmov2vec.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmov2vec.py rename to lib/TileOps/a5/tmov2vec.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmov_fp.py b/lib/TileOps/a5/tmov_fp.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmov_fp.py rename to lib/TileOps/a5/tmov_fp.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmrgsort.py b/lib/TileOps/a5/tmrgsort.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmrgsort.py rename to lib/TileOps/a5/tmrgsort.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmul.py b/lib/TileOps/a5/tmul.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmul.py rename to lib/TileOps/a5/tmul.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tmuls.py b/lib/TileOps/a5/tmuls.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tmuls.py rename to lib/TileOps/a5/tmuls.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tneg.py b/lib/TileOps/a5/tneg.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tneg.py rename to lib/TileOps/a5/tneg.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tnot.py b/lib/TileOps/a5/tnot.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tnot.py rename to lib/TileOps/a5/tnot.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tor.py b/lib/TileOps/a5/tor.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tor.py rename to lib/TileOps/a5/tor.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tors.py b/lib/TileOps/a5/tors.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tors.py rename to lib/TileOps/a5/tors.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tpartadd.py b/lib/TileOps/a5/tpartadd.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tpartadd.py rename to lib/TileOps/a5/tpartadd.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tpartmax.py b/lib/TileOps/a5/tpartmax.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tpartmax.py rename to lib/TileOps/a5/tpartmax.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tpartmin.py b/lib/TileOps/a5/tpartmin.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tpartmin.py rename to lib/TileOps/a5/tpartmin.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tpartmul.py b/lib/TileOps/a5/tpartmul.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tpartmul.py rename to lib/TileOps/a5/tpartmul.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tprelu.py b/lib/TileOps/a5/tprelu.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tprelu.py rename to lib/TileOps/a5/tprelu.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trandom.py b/lib/TileOps/a5/trandom.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trandom.py rename to lib/TileOps/a5/trandom.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trecip.py b/lib/TileOps/a5/trecip.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trecip.py rename to lib/TileOps/a5/trecip.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trelu.py b/lib/TileOps/a5/trelu.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trelu.py rename to lib/TileOps/a5/trelu.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trem.py b/lib/TileOps/a5/trem.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trem.py rename to lib/TileOps/a5/trem.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trems.py b/lib/TileOps/a5/trems.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trems.py rename to lib/TileOps/a5/trems.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowargmax.py b/lib/TileOps/a5/trowargmax.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowargmax.py rename to lib/TileOps/a5/trowargmax.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowargmin.py b/lib/TileOps/a5/trowargmin.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowargmin.py rename to lib/TileOps/a5/trowargmin.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowexpand.py b/lib/TileOps/a5/trowexpand.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowexpand.py rename to lib/TileOps/a5/trowexpand.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowexpandadd.py b/lib/TileOps/a5/trowexpandadd.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowexpandadd.py rename to lib/TileOps/a5/trowexpandadd.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowexpanddiv.py b/lib/TileOps/a5/trowexpanddiv.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowexpanddiv.py rename to lib/TileOps/a5/trowexpanddiv.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowexpandexpdif.py b/lib/TileOps/a5/trowexpandexpdif.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowexpandexpdif.py rename to lib/TileOps/a5/trowexpandexpdif.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowexpandmax.py b/lib/TileOps/a5/trowexpandmax.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowexpandmax.py rename to lib/TileOps/a5/trowexpandmax.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowexpandmin.py b/lib/TileOps/a5/trowexpandmin.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowexpandmin.py rename to lib/TileOps/a5/trowexpandmin.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowexpandmul.py b/lib/TileOps/a5/trowexpandmul.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowexpandmul.py rename to lib/TileOps/a5/trowexpandmul.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowexpandsub.py b/lib/TileOps/a5/trowexpandsub.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowexpandsub.py rename to lib/TileOps/a5/trowexpandsub.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowmax.py b/lib/TileOps/a5/trowmax.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowmax.py rename to lib/TileOps/a5/trowmax.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowmin.py b/lib/TileOps/a5/trowmin.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowmin.py rename to lib/TileOps/a5/trowmin.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowprod.py b/lib/TileOps/a5/trowprod.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowprod.py rename to lib/TileOps/a5/trowprod.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trowsum.py b/lib/TileOps/a5/trowsum.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trowsum.py rename to lib/TileOps/a5/trowsum.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/trsqrt.py b/lib/TileOps/a5/trsqrt.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/trsqrt.py rename to lib/TileOps/a5/trsqrt.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tsel.py b/lib/TileOps/a5/tsel.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tsel.py rename to lib/TileOps/a5/tsel.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tsels.py b/lib/TileOps/a5/tsels.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tsels.py rename to lib/TileOps/a5/tsels.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tshl.py b/lib/TileOps/a5/tshl.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tshl.py rename to lib/TileOps/a5/tshl.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tshls.py b/lib/TileOps/a5/tshls.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tshls.py rename to lib/TileOps/a5/tshls.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tshr.py b/lib/TileOps/a5/tshr.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tshr.py rename to lib/TileOps/a5/tshr.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tshrs.py b/lib/TileOps/a5/tshrs.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tshrs.py rename to lib/TileOps/a5/tshrs.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tsort32.py b/lib/TileOps/a5/tsort32.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tsort32.py rename to lib/TileOps/a5/tsort32.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tsqrt.py b/lib/TileOps/a5/tsqrt.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tsqrt.py rename to lib/TileOps/a5/tsqrt.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tstore.py b/lib/TileOps/a5/tstore.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tstore.py rename to lib/TileOps/a5/tstore.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tsub.py b/lib/TileOps/a5/tsub.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tsub.py rename to lib/TileOps/a5/tsub.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tsubs.py b/lib/TileOps/a5/tsubs.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/tsubs.py rename to lib/TileOps/a5/tsubs.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/ttrans.py b/lib/TileOps/a5/ttrans.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/ttrans.py rename to lib/TileOps/a5/ttrans.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/ttri.py b/lib/TileOps/a5/ttri.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/ttri.py rename to lib/TileOps/a5/ttri.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/txor.py b/lib/TileOps/a5/txor.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/txor.py rename to lib/TileOps/a5/txor.py diff --git a/ptodsl/ptodsl/tilelib/templates/a5/txors.py b/lib/TileOps/a5/txors.py similarity index 100% rename from ptodsl/ptodsl/tilelib/templates/a5/txors.py rename to lib/TileOps/a5/txors.py diff --git a/lib/TileOps/div_hp.py b/lib/TileOps/div_hp.py deleted file mode 100644 index 66c7c76434..0000000000 --- a/lib/TileOps/div_hp.py +++ /dev/null @@ -1,455 +0,0 @@ -# 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. - -"""Shared IEEE 754 high-precision division algorithms for pto.tdiv and pto.tdivs - -This module provides inline_proc functions that implement IEEE 754 compliant -division with improved accuracy for: -- Precision-sensitive values (1/7, 7/3, etc.) -- Subnormal numbers (denormals) -- Overflow/underflow boundary cases -- NaN propagation - -Reference: pto-isa/include/pto/npu/a5/custom/Div754.hpp -""" - -import tilelang_dsl as pto - - -@pto.inline_proc -def _div_three_candidate_search_f32(lhs, rhs, mask): - """Three-candidate search core algorithm for IEEE 754 division accuracy improvement. - - Corresponds to DivPrecisionImpl in pto-isa/include/pto/npu/a5/custom/Div754.hpp:16-62 - - Algorithm: Computes three candidates (z, z-1, z+1) and selects the one with smallest - residual |lhs - z*rhs|, improving accuracy for values like 1/7 that have infinite - binary representation. - """ - - # IEEE 754 Float32 bit patterns (corresponds to Div754.hpp:18-19) - inf_bound_u32 = pto.ui32(0x7f800000) # Infinity bound: sign=0, exp=255, mant=0 - sign_bit_u32 = pto.ui32(0x80000000) # Sign bit mask: bit31=1, others=0 - zero_f32 = pto.f32(0.0) - one_f32 = pto.f32(1.0) - neg_one_f32 = pto.f32(-1.0) - - z = pto.vdiv(lhs, rhs, mask) - z_init = z - - z_u32 = pto.vbitcast(z, pto.ui32) - z_or_sign = pto.vor(z_u32, pto.vbr(sign_bit_u32), mask) - is_inf_nan = pto.vcmp(z_or_sign, pto.vbr(inf_bound_u32), mask, pto.CmpMode.GE) - - is_zero = pto.vcmp(z, pto.vbr(zero_f32), mask, pto.CmpMode.EQ) - - special_mask = pto.por(is_inf_nan, is_zero, mask) - - y = pto.vmuls(rhs, neg_one_f32, mask) - r = pto.vmula(lhs, z, y, mask) - - z_pre = pto.vadds(z, neg_one_f32, mask) - z_next = pto.vadds(z, one_f32, mask) - - r_pre = pto.vmula(lhs, z_pre, y, mask) - r_next = pto.vmula(lhs, z_next, y, mask) - - r_abs = pto.vabs(r, mask) - r_pre_abs = pto.vabs(r_pre, mask) - r_next_abs = pto.vabs(r_next, mask) - - better_pre = pto.vcmp(r_pre_abs, r_abs, mask, pto.CmpMode.LT) - z_best = pto.vsel(z_pre, z, better_pre) - r_best_abs = pto.vsel(r_pre_abs, r_abs, better_pre) - - better_next = pto.vcmp(r_next_abs, r_best_abs, mask, pto.CmpMode.LT) - z_best = pto.vsel(z_next, z_best, better_next) - - divided = pto.vsel(z_init, z_best, special_mask) - - return divided - - -@pto.inline_proc -def _div_ieee754_f32_impl(src0, src1, mask): - """Complete IEEE 754 float32 high-precision division with subnormal and overflow handling. - - Corresponds to DivIEEE754FloatImpl in pto-isa/include/pto/npu/a5/custom/Div754.hpp:65-288 - - Key improvements over pto-isa: - - Subnormal detection uses LT (line 94) instead of EQ (Div754.hpp:159) - Rationale: Covers entire subnormal range [2^-149, 2^-126), not just max subnormal - """ - - # IEEE 754 Float32 bit masks and constants (corresponds to Div754.hpp:69-81) - F32_INF = pto.ui32(0x7f800000) # +Infinity: sign=0, exp=255, mant=0 - sign_extractor = pto.ui32(0x80000000) # Sign bit mask (bit31) - exponent_extractor = pto.ui32(0x807FFFFF) # Clear exponent bits [30:23] - exponent_normalizer = pto.ui32(0x3F800000) # Bias 127: 1.0f reference - subnormal_threshold = pto.ui32(0x007FFFFF) # Max subnormal: (1-2^-23)*2^-126 - nan_value = pto.ui32(0x7fc00000) # Quiet NaN: exp=255, mant=0x400000 - min_denormal = pto.ui32(0x1) # Smallest positive: 2^-149 - - # Subnormal normalization factors (corresponds to Div754.hpp:86-89) - normalize_scale_enlarge = pto.f32(8388608.0) # 2^23: shifts subnormals to normal range - normalize_scale_reduce = pto.f32(1.1920928955078125e-07) # 2^-23: inverse for result compensation - - src0_abs = pto.vabs(src0, mask) - src1_abs = pto.vabs(src1, mask) - - src0_abs_u32 = pto.vbitcast(src0_abs, pto.ui32) - src1_abs_u32 = pto.vbitcast(src1_abs, pto.ui32) - - mask_inf_src0 = pto.vcmp(src0_abs_u32, pto.vbr(F32_INF), mask, pto.CmpMode.EQ) - mask_inf_src1 = pto.vcmp(src1_abs_u32, pto.vbr(F32_INF), mask, pto.CmpMode.EQ) - mask_invalid = pto.por(mask_inf_src0, mask_inf_src1, mask) - - mask_zero_src0 = pto.vcmp(src0_abs_u32, pto.vbr(pto.ui32(0)), mask, pto.CmpMode.EQ) - mask_invalid = pto.por(mask_invalid, mask_zero_src0, mask) - mask_zero_src1 = pto.vcmp(src1_abs_u32, pto.vbr(pto.ui32(0)), mask, pto.CmpMode.EQ) - mask_invalid = pto.por(mask_invalid, mask_zero_src1, mask) - - mask_valid = pto.pnot(mask_invalid, mask) - - # Detect subnormal numbers (denormals) - # NOTE: Uses EQ/LT comparison matching pto-isa Div754.hpp asymmetry: - # - src0: EQ comparison (Div754.hpp:159) - detects exact max subnormal - # - src1: LT comparison (Div754.hpp:166) - covers entire subnormal range - mask_src0_subnormal = pto.vcmp(src0_abs_u32, pto.vbr(subnormal_threshold), mask, pto.CmpMode.EQ) - mask_src0_normal = pto.pnot(mask_src0_subnormal, mask) - src0_subnormal = pto.vmuls(src0, normalize_scale_enlarge, mask_src0_subnormal) - - mask_src1_subnormal = pto.vcmp(src1_abs_u32, pto.vbr(subnormal_threshold), mask, pto.CmpMode.LT) - mask_src1_normal = pto.pnot(mask_src1_subnormal, mask) - src1_subnormal = pto.vmuls(src1, normalize_scale_enlarge, mask_src1_subnormal) - - src0_all = pto.vsel(src0, src0_subnormal, mask_src0_normal) - src1_all = pto.vsel(src1, src1_subnormal, mask_src1_normal) - - src0_all_u32 = pto.vbitcast(src0_all, pto.ui32) - src1_all_u32 = pto.vbitcast(src1_all, pto.ui32) - - src0_norm_u32 = pto.vand(src0_all_u32, pto.vbr(exponent_extractor), mask_valid) - src1_norm_u32 = pto.vand(src1_all_u32, pto.vbr(exponent_extractor), mask_valid) - - src0_norm_u32 = pto.vadd(src0_norm_u32, pto.vbr(exponent_normalizer), mask_valid) - src1_norm_u32 = pto.vadd(src1_norm_u32, pto.vbr(exponent_normalizer), mask_valid) - - src0_norm_f32 = pto.vbitcast(src0_norm_u32, pto.f32) - src1_norm_f32 = pto.vbitcast(src1_norm_u32, pto.f32) - src0_norm = pto.vsel(src0_norm_f32, src0_all, mask_valid) - src1_norm = pto.vsel(src1_norm_f32, src1_all, mask_valid) - - dst = _div_three_candidate_search_f32(src0_norm, src1_norm, mask_valid) - - mask0 = pto.pand(mask_src0_subnormal, mask_src1_normal, mask) - z1 = pto.vmuls(dst, normalize_scale_reduce, mask0) - dst = pto.vsel(z1, dst, mask0) - - mask0 = pto.pand(mask_src0_normal, mask_src1_subnormal, mask) - z1 = pto.vmuls(dst, normalize_scale_enlarge, mask0) - dst = pto.vsel(z1, dst, mask0) - - dst_u32 = pto.vbitcast(dst, pto.ui32) - dst_sign = pto.vand(dst_u32, pto.vbr(sign_extractor), mask) - - src0_exponent = pto.vand(src0_all_u32, pto.vbr(F32_INF), mask) - src1_exponent = pto.vand(src1_all_u32, pto.vbr(F32_INF), mask) - - src0_exp_shifted = pto.vshrs(src0_exponent, pto.i16(23), mask) - src1_exp_shifted = pto.vshrs(src1_exponent, pto.i16(23), mask) - - src0_exp_i32 = pto.vbitcast(src0_exp_shifted, pto.si32) - src1_exp_i32 = pto.vbitcast(src1_exp_shifted, pto.si32) - - scale = pto.vsub(src0_exp_i32, src1_exp_i32, mask) - scale = pto.vadds(scale, pto.si32(127), mask) - - neg23 = pto.si32(-23) - mask_underflow1 = pto.vcmp(scale, pto.vbr(neg23), mask, pto.CmpMode.EQ) - mask_underflow1 = pto.pand(mask_underflow1, mask_valid, mask) - - z1_u32 = pto.vadd(dst_sign, pto.vbr(min_denormal), mask_underflow1) - z2_u32 = pto.vadd(dst_sign, pto.vbr(pto.ui32(0)), mask_underflow1) - - src0_norm_abs = pto.vabs(src0_norm, mask_valid) - src1_norm_abs = pto.vabs(src1_norm, mask_valid) - mask_norm = pto.vcmp(src0_norm_abs, src1_norm_abs, mask_valid, pto.CmpMode.LE) - - z1_sel = pto.vsel(z2_u32, z1_u32, mask_norm) - dst_u32_temp = pto.vsel(z1_sel, dst_u32, mask_underflow1) - - mask_underflow1_not = pto.pnot(mask_underflow1, mask) - mask_valid_temp = pto.pand(mask_underflow1_not, mask_valid, mask) - - mask_underflow2 = pto.vcmp(scale, pto.vbr(neg23), mask, pto.CmpMode.LT) - mask_underflow2 = pto.pand(mask_underflow2, mask_valid_temp, mask) - - z1_u32 = pto.vadd(dst_sign, pto.vbr(pto.ui32(0)), mask_underflow2) - dst_u32_temp = pto.vsel(z1_u32, dst_u32_temp, mask_underflow2) - - mask_underflow2_not = pto.pnot(mask_underflow2, mask) - mask_valid_temp = pto.pand(mask_underflow2_not, mask_valid_temp, mask) - - max_exp = pto.si32(255) - mask_overflow1 = pto.vcmp(scale, pto.vbr(max_exp), mask, pto.CmpMode.EQ) - mask_overflow1 = pto.pand(mask_overflow1, mask_valid_temp, mask) - - scale_adj = pto.vadds(scale, pto.si32(-1), mask_overflow1) - scale = pto.vsel(scale_adj, scale, mask_overflow1) - - dst_f32_temp = pto.vbitcast(dst_u32_temp, pto.f32) - z1_f32 = pto.vmuls(dst_f32_temp, pto.f32(2.0), mask_overflow1) - dst_f32_temp = pto.vsel(z1_f32, dst_f32_temp, mask_overflow1) - - mask_overflow2 = pto.vcmp(scale, pto.vbr(max_exp), mask, pto.CmpMode.GT) - mask_overflow2 = pto.pand(mask_overflow2, mask_valid_temp, mask) - - z1_u32 = pto.vadd(dst_sign, pto.vbr(F32_INF), mask_overflow2) - dst_u32_temp = pto.vbitcast(dst_f32_temp, pto.ui32) - dst_u32_temp = pto.vsel(z1_u32, dst_u32_temp, mask_overflow2) - - mask_overflow2_not = pto.pnot(mask_overflow2, mask) - mask_valid_final = pto.pand(mask_overflow2_not, mask_valid_temp, mask) - - zero_exp = pto.si32(0) - mask_pos_exp = pto.vcmp(scale, pto.vbr(zero_exp), mask_valid_final, pto.CmpMode.GT) - - scale_u32 = pto.vbitcast(scale, pto.ui32) - exp_shifted = pto.vshls(scale_u32, pto.i16(23), mask_pos_exp) - exp_factor_f32 = pto.vbitcast(exp_shifted, pto.f32) - - dst_f32_temp = pto.vbitcast(dst_u32_temp, pto.f32) - z1_f32 = pto.vmul(dst_f32_temp, exp_factor_f32, mask_pos_exp) - dst_f32_temp = pto.vsel(z1_f32, dst_f32_temp, mask_pos_exp) - - mask_pos_exp_not = pto.pnot(mask_pos_exp, mask_valid_final) - - # Handle negative exponent (underflow scenarios) - # Corresponds to Div754.hpp:275 - # Value 0x00400000 = Float32 with exp=0, mantissa bit22=1 (used for shift calculation) - four_million = pto.ui32(4194304) # Normal float 1.0 in bit representation for exponent manipulation - scale_abs = pto.vabs(scale, mask_pos_exp_not) - - shr_base_vec = pto.vdup(four_million, mask_pos_exp_not) - shr_base_i32 = pto.vbitcast(shr_base_vec, pto.si32) - shr_factor_i32 = pto.vshr(shr_base_i32, scale_abs, mask_pos_exp_not) - shr_factor_f32 = pto.vbitcast(shr_factor_i32, pto.f32) - - z1_f32 = pto.vmul(dst_f32_temp, shr_factor_f32, mask_pos_exp_not) - dst_f32_temp = pto.vsel(z1_f32, dst_f32_temp, mask_pos_exp_not) - - mask_nan_src0 = pto.vcmp(src0_abs, src0_abs, mask, pto.CmpMode.NE) - mask_nan_src1 = pto.vcmp(src1_abs, src1_abs, mask, pto.CmpMode.NE) - mask_nan = pto.por(mask_nan_src0, mask_nan_src1, mask) - - nan_vec = pto.vbr(nan_value) - nan_f32_vec = pto.vbitcast(nan_vec, pto.f32) - dst_final = pto.vsel(nan_f32_vec, dst_f32_temp, mask_nan) - - return dst_final - - -@pto.inline_proc -def _div_ieee754_f16_impl(src0, src1, mask): - """Complete IEEE 754 float16 high-precision division with subnormal handling. - - Follows pto-isa Div754.hpp:291-502 (DivIEEE754HalfImpl). - - Key differences from F32 implementation: - - Uses LT for both src0/src1 subnormal detection (symmetric, not EQ/LT like F32) - - Normalization factor: 2^10 (not 2^23 for F32) - - Exponent bias: 15 (not 127 for F32) - - Exponent shift: 10 bits (not 23 for F32) - - Direct vdiv call (no three-candidate search) - """ - - # IEEE 754 Float16 bit masks and constants (corresponds to Div754.hpp:293-309) - F16_INF = pto.ui16(0x7C00) # +Infinity: sign=0, exp=31, mant=0 - exponent_extractor = pto.ui16(0x83FF) # Clear exponent bits [14:10] - exponent_normalizer = pto.ui16(0x3C00) # 1.0f16 reference (bias=15) - sign_extractor = pto.ui16(0x8000) # Sign bit mask (bit15) - subnormal_threshold = pto.ui16(0x03FF) # Max subnormal: (1-2^-10)*2^-14 - nan_value = pto.ui16(0x7E00) # Quiet NaN: exp=31, mant=0x200 - min_denormal = pto.ui16(0x1) # Smallest positive: 2^-24 - - # Subnormal normalization factors (corresponds to Div754.hpp:306-309) - normalize_scale_enlarge = pto.f16(1024.0) # 2^10: shifts subnormals to normal range - normalize_scale_reduce = pto.f16(0.0009765625) # 2^-10: inverse for result compensation - - src0_abs = pto.vabs(src0, mask) - src1_abs = pto.vabs(src1, mask) - - src0_abs_u16 = pto.vbitcast(src0_abs, pto.ui16) - src1_abs_u16 = pto.vbitcast(src1_abs, pto.ui16) - - # Detect Infinity values - mask_inf_src0 = pto.vcmp(src0_abs_u16, pto.vbr(F16_INF), mask, pto.CmpMode.EQ) - mask_inf_src1 = pto.vcmp(src1_abs_u16, pto.vbr(F16_INF), mask, pto.CmpMode.EQ) - mask_invalid = pto.por(mask_inf_src0, mask_inf_src1, mask) - - # Detect Zero values - mask_zero_src0 = pto.vcmp(src0_abs_u16, pto.vbr(pto.ui16(0)), mask, pto.CmpMode.EQ) - mask_invalid = pto.por(mask_invalid, mask_zero_src0, mask) - mask_zero_src1 = pto.vcmp(src1_abs_u16, pto.vbr(pto.ui16(0)), mask, pto.CmpMode.EQ) - mask_invalid = pto.por(mask_invalid, mask_zero_src1, mask) - - mask_valid = pto.pnot(mask_invalid, mask) - - # Detect subnormal numbers (denormals) - # NOTE: F16 uses LT for BOTH src0 and src1 (symmetric detection) - # Different from F32's asymmetric EQ/LT pattern - mask_src0_subnormal = pto.vcmp(src0_abs_u16, pto.vbr(subnormal_threshold), mask, pto.CmpMode.LT) - mask_src0_normal = pto.pnot(mask_src0_subnormal, mask) - src0_subnormal = pto.vmuls(src0, normalize_scale_enlarge, mask_src0_subnormal) - - mask_src1_subnormal = pto.vcmp(src1_abs_u16, pto.vbr(subnormal_threshold), mask, pto.CmpMode.LT) - mask_src1_normal = pto.pnot(mask_src1_subnormal, mask) - src1_subnormal = pto.vmuls(src1, normalize_scale_enlarge, mask_src1_subnormal) - - # Merge normalized subnormals with normal values - src0_all = pto.vsel(src0, src0_subnormal, mask_src0_normal) - src1_all = pto.vsel(src1, src1_subnormal, mask_src1_normal) - - src0_all_u16 = pto.vbitcast(src0_all, pto.ui16) - src1_all_u16 = pto.vbitcast(src1_all, pto.ui16) - - # Standardize exponent bits (corresponds to Div754.hpp:391-401) - src0_norm_u16 = pto.vand(src0_all_u16, pto.vbr(exponent_extractor), mask_valid) - src1_norm_u16 = pto.vand(src1_all_u16, pto.vbr(exponent_extractor), mask_valid) - - src0_norm_u16 = pto.vadd(src0_norm_u16, pto.vbr(exponent_normalizer), mask_valid) - src1_norm_u16 = pto.vadd(src1_norm_u16, pto.vbr(exponent_normalizer), mask_valid) - - src0_norm_f16 = pto.vbitcast(src0_norm_u16, pto.f16) - src1_norm_f16 = pto.vbitcast(src1_norm_u16, pto.f16) - src0_norm = pto.vsel(src0_norm_f16, src0_all, mask_valid) - src1_norm = pto.vsel(src1_norm_f16, src1_all, mask_valid) - - src0_norm_abs = pto.vabs(src0_norm, mask_valid) - src1_norm_abs = pto.vabs(src1_norm, mask_valid) - mask_norm = pto.vcmp(src0_norm_abs, src1_norm_abs, mask_valid, pto.CmpMode.LE) - - # Execute division directly (no three-candidate search for F16) - # Corresponds to Div754.hpp:406 - dst = pto.vdiv(src0_norm, src1_norm, mask) - - # Subnormal dividend, normal divisor: scale down result - # Corresponds to Div754.hpp:408-412 - mask0 = pto.pand(mask_src0_subnormal, mask_src1_normal, mask) - z1 = pto.vmuls(dst, normalize_scale_reduce, mask0) - dst = pto.vsel(z1, dst, mask0) - - # Normal dividend, subnormal divisor: scale up result - # Corresponds to Div754.hpp:414-419 - mask0 = pto.pand(mask_src0_normal, mask_src1_subnormal, mask) - z1 = pto.vmuls(dst, normalize_scale_enlarge, mask0) - dst = pto.vsel(z1, dst, mask0) - - # Preserve sign for overflow/underflow handling - dst_u16 = pto.vbitcast(dst, pto.ui16) - dst_sign = pto.vand(dst_u16, pto.vbr(sign_extractor), mask) - - # Extract exponent bits (corresponds to Div754.hpp:428-439) - src0_exponent = pto.vand(src0_all_u16, pto.vbr(F16_INF), mask) - src1_exponent = pto.vand(src1_all_u16, pto.vbr(F16_INF), mask) - - src0_exp_shifted = pto.vshrs(src0_exponent, pto.i16(10), mask) - src1_exp_shifted = pto.vshrs(src1_exponent, pto.i16(10), mask) - - src0_exp_i16 = pto.vbitcast(src0_exp_shifted, pto.si16) - src1_exp_i16 = pto.vbitcast(src1_exp_shifted, pto.si16) - - # Scale = src0_exp - src1_exp + bias(15) - scale = pto.vsub(src0_exp_i16, src1_exp_i16, mask) - scale = pto.vadds(scale, pto.si16(15), mask) - - # Underflow handling: scale == -9 (corresponds to Div754.hpp:443-453) - neg9 = pto.si16(-9) - mask_underflow1 = pto.vcmp(scale, pto.vbr(neg9), mask, pto.CmpMode.EQ) - mask_underflow1 = pto.pand(mask_underflow1, mask_valid, mask) - - z1_u16 = pto.vadd(dst_sign, pto.vbr(min_denormal), mask_underflow1) - z2_u16 = pto.vadd(dst_sign, pto.vbr(pto.ui16(0)), mask_underflow1) - - z1_sel = pto.vsel(z2_u16, z1_u16, mask_norm) - dst_u16_temp = pto.vsel(z1_sel, dst_u16, mask_underflow1) - - mask_underflow1_not = pto.pnot(mask_underflow1, mask) - mask_valid_temp = pto.pand(mask_underflow1_not, mask_valid, mask) - - # Underflow handling: scale < -9 (corresponds to Div754.hpp:456-463) - mask_underflow2 = pto.vcmp(scale, pto.vbr(neg9), mask, pto.CmpMode.LT) - mask_underflow2 = pto.pand(mask_underflow2, mask_valid_temp, mask) - - z1_u16 = pto.vadd(dst_sign, pto.vbr(pto.ui16(0)), mask_underflow2) - dst_u16_temp = pto.vsel(z1_u16, dst_u16_temp, mask_underflow2) - - mask_underflow2_not = pto.pnot(mask_underflow2, mask) - mask_valid_temp = pto.pand(mask_underflow2_not, mask_valid_temp, mask) - - # Overflow handling: scale == 31 (corresponds to Div754.hpp:465-472) - max_exp = pto.si16(31) - mask_overflow1 = pto.vcmp(scale, pto.vbr(max_exp), mask, pto.CmpMode.EQ) - mask_overflow1 = pto.pand(mask_overflow1, mask_valid_temp, mask) - - scale_adj = pto.vadds(scale, pto.si16(-1), mask_overflow1) - scale = pto.vsel(scale_adj, scale, mask_overflow1) - - dst_f16_temp = pto.vbitcast(dst_u16_temp, pto.f16) - z1_f16 = pto.vmuls(dst_f16_temp, pto.f16(2.0), mask_overflow1) - dst_f16_temp = pto.vsel(z1_f16, dst_f16_temp, mask_overflow1) - - # Overflow handling: scale > 31 (corresponds to Div754.hpp:474-480) - mask_overflow2 = pto.vcmp(scale, pto.vbr(max_exp), mask, pto.CmpMode.GT) - mask_overflow2 = pto.pand(mask_overflow2, mask_valid_temp, mask) - - z1_u16 = pto.vadd(dst_sign, pto.vbr(F16_INF), mask_overflow2) - dst_u16_temp = pto.vbitcast(dst_f16_temp, pto.ui16) - dst_u16_temp = pto.vsel(z1_u16, dst_u16_temp, mask_overflow2) - - mask_overflow2_not = pto.pnot(mask_overflow2, mask) - mask_valid_final = pto.pand(mask_overflow2_not, mask_valid_temp, mask) - - # Positive exponent handling (corresponds to Div754.hpp:482-486) - zero_exp = pto.si16(0) - mask_pos_exp = pto.vcmp(scale, pto.vbr(zero_exp), mask_valid_final, pto.CmpMode.GT) - - scale_u16 = pto.vbitcast(scale, pto.ui16) - exp_shifted = pto.vshls(scale_u16, pto.i16(10), mask_pos_exp) - exp_factor_f16 = pto.vbitcast(exp_shifted, pto.f16) - - dst_f16_temp = pto.vbitcast(dst_u16_temp, pto.f16) - z1_f16 = pto.vmul(dst_f16_temp, exp_factor_f16, mask_pos_exp) - dst_f16_temp = pto.vsel(z1_f16, dst_f16_temp, mask_pos_exp) - - # Negative exponent handling (corresponds to Div754.hpp:488-493) - mask_pos_exp_not = pto.pnot(mask_pos_exp, mask_valid_final) - - # Value 0x0200 = Float16 with exp=0, mantissa bit9=1 (used for shift calculation) - shr_base = pto.ui16(512) # 0x0200 - scale_abs = pto.vabs(scale, mask_pos_exp_not) - - shr_base_vec = pto.vdup(shr_base, mask_pos_exp_not) - shr_base_i16 = pto.vbitcast(shr_base_vec, pto.si16) - shr_factor_i16 = pto.vshr(shr_base_i16, scale_abs, mask_pos_exp_not) - shr_factor_f16 = pto.vbitcast(shr_factor_i16, pto.f16) - - z1_f16 = pto.vmul(dst_f16_temp, shr_factor_f16, mask_pos_exp_not) - dst_f16_temp = pto.vsel(z1_f16, dst_f16_temp, mask_pos_exp_not) - - # NaN propagation (corresponds to Div754.hpp:495-501) - mask_nan_src0 = pto.vcmp(src0_abs, src0_abs, mask, pto.CmpMode.NE) - mask_nan_src1 = pto.vcmp(src1_abs, src1_abs, mask, pto.CmpMode.NE) - mask_nan = pto.por(mask_nan_src0, mask_nan_src1, mask) - - nan_vec = pto.vbr(nan_value) - nan_f16_vec = pto.vbitcast(nan_vec, pto.f16) - dst_final = pto.vsel(nan_f16_vec, dst_f16_temp, mask_nan) - - return dst_final \ No newline at end of file diff --git a/lib/TileOps/exp_hp.py b/lib/TileOps/exp_hp.py deleted file mode 100644 index f967c12490..0000000000 --- a/lib/TileOps/exp_hp.py +++ /dev/null @@ -1,30 +0,0 @@ -# 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 tilelang_dsl as pto - -@pto.inline_proc -def _tl_exp_precision(src, mask, dtype): - if pto.constexpr(dtype == pto.f16): - subnormal_threshold = pto.f16("0x03ff") - two_val = pto.f16(2.0) - else: - subnormal_threshold = pto.f32("0x007FFFFF") - two_val = pto.f32(2.0) - - dst = pto.vexp(src, mask) - - subnormal_mask = pto.vcmps(dst, subnormal_threshold, mask, pto.CmpMode.LE) - - reg_two = pto.vbr(two_val) - tmp = pto.vdiv(src, reg_two, subnormal_mask) - tmp = pto.vexp(tmp, subnormal_mask) - tmp = pto.vmul(tmp, tmp, subnormal_mask) - - result = pto.vsel(tmp, dst, subnormal_mask) - return result diff --git a/lib/TileOps/math.py b/lib/TileOps/math.py deleted file mode 100644 index 83c4634f69..0000000000 --- a/lib/TileOps/math.py +++ /dev/null @@ -1,584 +0,0 @@ -# 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 tilelang_dsl as pto - - -@pto.inline_proc -def _tl_soft_vdiv_u8(vec, scalar_vec, mask): - zero = pto.ui8(0) - zero_q = pto.ui8(0xFF) - full_mask_b8 = pto.pset_b8(pto.PAT.ALL) - - zero_mask = pto.vcmps(scalar_vec, zero, mask, pto.CmpMode.EQ) - active_mask = pto.pnot(zero_mask, mask) - active_low = pto.punpack(active_mask, pto.PredicatePart.LOWER) - active_high = pto.punpack(active_mask, pto.PredicatePart.HIGHER) - - vec_low = pto.vzunpack(vec, 0) - vec_high = pto.vzunpack(vec, 1) - scalar_low = pto.vzunpack(scalar_vec, 0) - scalar_high = pto.vzunpack(scalar_vec, 1) - - q_low = _tl_soft_vdiv_u16(vec_low, scalar_low, active_low) - q_high = _tl_soft_vdiv_u16(vec_high, scalar_high, active_high) - packed_low = pto.vpack(q_low, pto.PredicatePart.LOWER) - packed_high = pto.vpack(q_high, pto.PredicatePart.HIGHER) - q = pto.vor(packed_low, packed_high, full_mask_b8) - return pto.vsel(pto.vbr(zero_q), q, zero_mask) - - -@pto.inline_proc -def _tl_soft_vdiv_i8(vec, scalar_vec, mask): - zero = pto.i8(0) - neg_one = pto.i8(-1) - full_mask_b8 = pto.pset_b8(pto.PAT.ALL) - - zero_mask = pto.vcmps(scalar_vec, zero, mask, pto.CmpMode.EQ) - active_mask = pto.pnot(zero_mask, mask) - active_low = pto.punpack(active_mask, pto.PredicatePart.LOWER) - active_high = pto.punpack(active_mask, pto.PredicatePart.HIGHER) - - vec_low = pto.vsunpack(vec, 0) - vec_high = pto.vsunpack(vec, 1) - scalar_low = pto.vsunpack(scalar_vec, 0) - scalar_high = pto.vsunpack(scalar_vec, 1) - - q_low = _tl_soft_vdiv_i16(vec_low, scalar_low, active_low) - q_high = _tl_soft_vdiv_i16(vec_high, scalar_high, active_high) - packed_low = pto.vpack(q_low, pto.PredicatePart.LOWER) - packed_high = pto.vpack(q_high, pto.PredicatePart.HIGHER) - q = pto.vbitcast(pto.vor(packed_low, packed_high, full_mask_b8), pto.i8) - return pto.vsel(pto.vbr(neg_one), q, zero_mask) - - -@pto.inline_proc -def _tl_soft_vdiv_u16(vec, scalar_vec, mask): - zero = pto.ui16(0) - one = pto.ui16(1) - fp32_one = pto.f32(1.0) - full_mask_b16 = pto.pset_b16(pto.PAT.ALL) - full_mask_b32 = pto.pset_b32(pto.PAT.ALL) - - zero_mask = pto.vcmps(scalar_vec, zero, mask, pto.CmpMode.EQ) - active_mask = pto.pnot(zero_mask, mask) - - zero_u16 = pto.vbr(zero) - vy_lower_u16, vy_higher_u16 = pto.vintlv(scalar_vec, zero_u16) - vy_lower_u32 = pto.vcvt(vy_lower_u16, pto.ui32, full_mask_b16, part=pto.VcvtPartMode.EVEN) - vy_higher_u32 = pto.vcvt(vy_higher_u16, pto.ui32, full_mask_b16, part=pto.VcvtPartMode.EVEN) - active_low = pto.vcmps(vy_lower_u32, pto.ui32(0), full_mask_b32, pto.CmpMode.NE) - active_high = pto.vcmps(vy_higher_u32, pto.ui32(0), full_mask_b32, pto.CmpMode.NE) - vy_lower_f32 = pto.vcvt(pto.vbitcast(vy_lower_u32, pto.i32), pto.f32, active_low, rnd=pto.VcvtRoundMode.F) - vy_higher_f32 = pto.vcvt(pto.vbitcast(vy_higher_u32, pto.i32), pto.f32, active_high, rnd=pto.VcvtRoundMode.F) - - vy_rec_lower = pto.vdiv(pto.vbr(fp32_one), vy_lower_f32, active_low) - vy_rec_higher = pto.vdiv(pto.vbr(fp32_one), vy_higher_f32, active_high) - vy_scale_lower = pto.vmul(vy_rec_lower, pto.vbr(pto.f32(65536.0)), active_low) - vy_scale_higher = pto.vmul(vy_rec_higher, pto.vbr(pto.f32(65536.0)), active_high) - - v_lower_i32 = pto.vcvt( - vy_scale_lower, - pto.i32, - active_low, - rnd=pto.VcvtRoundMode.F, - sat=pto.VcvtSatMode.NOSAT, - ) - v_higher_i32 = pto.vcvt( - vy_scale_higher, - pto.i32, - active_high, - rnd=pto.VcvtRoundMode.F, - sat=pto.VcvtSatMode.NOSAT, - ) - v_lower_u32 = pto.vbitcast(v_lower_i32, pto.ui32) - v_higher_u32 = pto.vbitcast(v_higher_i32, pto.ui32) - - vx_lower_u16, vx_higher_u16 = pto.vintlv(vec, zero_u16) - vx_lower_u32 = pto.vcvt(vx_lower_u16, pto.ui32, full_mask_b16, part=pto.VcvtPartMode.EVEN) - vx_higher_u32 = pto.vcvt(vx_higher_u16, pto.ui32, full_mask_b16, part=pto.VcvtPartMode.EVEN) - q_tmp_lower = pto.vmul(v_lower_u32, vx_lower_u32, active_low) - q_tmp_higher = pto.vmul(v_higher_u32, vx_higher_u32, active_high) - _q_lower, q_tmp = pto.vdintlv(pto.vbitcast(q_tmp_lower, pto.ui16), pto.vbitcast(q_tmp_higher, pto.ui16)) - - yq_tmp = pto.vmul(q_tmp, scalar_vec, active_mask) - r_tmp = pto.vsub(vec, yq_tmp, active_mask) - ge_mask = pto.vcmp(r_tmp, scalar_vec, active_mask, pto.CmpMode.GE) - refined_r = pto.vsub(r_tmp, scalar_vec, active_mask) - r_tmp = pto.vsel(refined_r, r_tmp, ge_mask) - q_inc = pto.vadds(q_tmp, one, active_mask) - q_tmp = pto.vsel(q_inc, q_tmp, ge_mask) - ge_mask = pto.vcmp(r_tmp, scalar_vec, active_mask, pto.CmpMode.GE) - refined_r = pto.vsub(r_tmp, scalar_vec, active_mask) - r_tmp = pto.vsel(refined_r, r_tmp, ge_mask) - q_inc = pto.vadds(q_tmp, one, active_mask) - q_tmp = pto.vsel(q_inc, q_tmp, ge_mask) - zero_q = pto.vbr(pto.ui16(0xFFFF)) - return pto.vsel(zero_q, q_tmp, zero_mask) - - -@pto.inline_proc -def _tl_soft_vdiv_i16(vec, scalar_vec, mask): - zero = pto.i16(0) - neg_one = pto.i16(-1) - - zero_mask = pto.vcmps(scalar_vec, zero, mask, pto.CmpMode.EQ) - active_mask = pto.pnot(zero_mask, mask) - - abs_x = pto.vbitcast(pto.vabs(vec, active_mask), pto.ui16) - abs_y = pto.vbitcast(pto.vabs(scalar_vec, active_mask), pto.ui16) - x_xor_y = pto.vxor(vec, scalar_vec, active_mask) - p_pos = pto.vcmps(x_xor_y, zero, active_mask, pto.CmpMode.GE) - - q_abs = _tl_soft_vdiv_u16(abs_x, abs_y, active_mask) - neg_q = pto.vneg(pto.vbitcast(q_abs, pto.i16), active_mask) - q = pto.vsel(pto.vbitcast(q_abs, pto.i16), neg_q, p_pos) - return pto.vsel(pto.vbr(neg_one), q, zero_mask) - - -@pto.inline_proc -def _tl_soft_vmod_u8(vec, scalar_vec, mask): - zero = pto.ui8(0) - zero_r = pto.ui8(0xFF) - full_mask_b8 = pto.pset_b8(pto.PAT.ALL) - - zero_mask = pto.vcmps(scalar_vec, zero, mask, pto.CmpMode.EQ) - active_mask = pto.pnot(zero_mask, mask) - active_low = pto.punpack(active_mask, pto.PredicatePart.LOWER) - active_high = pto.punpack(active_mask, pto.PredicatePart.HIGHER) - - vec_low = pto.vzunpack(vec, 0) - vec_high = pto.vzunpack(vec, 1) - scalar_low = pto.vzunpack(scalar_vec, 0) - scalar_high = pto.vzunpack(scalar_vec, 1) - - r_low = _tl_soft_vmod_u16(vec_low, scalar_low, active_low) - r_high = _tl_soft_vmod_u16(vec_high, scalar_high, active_high) - packed_low = pto.vpack(r_low, pto.PredicatePart.LOWER) - packed_high = pto.vpack(r_high, pto.PredicatePart.HIGHER) - r = pto.vor(packed_low, packed_high, full_mask_b8) - return pto.vsel(pto.vbr(zero_r), r, zero_mask) - - -@pto.inline_proc -def _tl_soft_vmod_i8(vec, scalar_vec, mask): - zero = pto.i8(0) - neg_one = pto.i8(-1) - full_mask_b8 = pto.pset_b8(pto.PAT.ALL) - - zero_mask = pto.vcmps(scalar_vec, zero, mask, pto.CmpMode.EQ) - active_mask = pto.pnot(zero_mask, mask) - active_low = pto.punpack(active_mask, pto.PredicatePart.LOWER) - active_high = pto.punpack(active_mask, pto.PredicatePart.HIGHER) - - vec_low = pto.vsunpack(vec, 0) - vec_high = pto.vsunpack(vec, 1) - scalar_low = pto.vsunpack(scalar_vec, 0) - scalar_high = pto.vsunpack(scalar_vec, 1) - - r_low = _tl_soft_vmod_i16(vec_low, scalar_low, active_low) - r_high = _tl_soft_vmod_i16(vec_high, scalar_high, active_high) - packed_low = pto.vpack(r_low, pto.PredicatePart.LOWER) - packed_high = pto.vpack(r_high, pto.PredicatePart.HIGHER) - r = pto.vbitcast(pto.vor(packed_low, packed_high, full_mask_b8), pto.i8) - return pto.vsel(pto.vbr(neg_one), r, zero_mask) - - -@pto.inline_proc -def _tl_soft_vmod_u16(vec, scalar_vec, mask): - zero = pto.ui16(0) - one = pto.ui16(1) - zero_r = pto.ui16(0xFFFF) - fp32_one = pto.f32(1.0) - full_mask_b16 = pto.pset_b16(pto.PAT.ALL) - full_mask_b32 = pto.pset_b32(pto.PAT.ALL) - - zero_mask = pto.vcmps(scalar_vec, zero, mask, pto.CmpMode.EQ) - active_mask = pto.pnot(zero_mask, mask) - - zero_u16 = pto.vbr(zero) - vy_lower_u16, vy_higher_u16 = pto.vintlv(scalar_vec, zero_u16) - vy_lower_u32 = pto.vcvt(vy_lower_u16, pto.ui32, full_mask_b16, part=pto.VcvtPartMode.EVEN) - vy_higher_u32 = pto.vcvt(vy_higher_u16, pto.ui32, full_mask_b16, part=pto.VcvtPartMode.EVEN) - active_low = pto.vcmps(vy_lower_u32, pto.ui32(0), full_mask_b32, pto.CmpMode.NE) - active_high = pto.vcmps(vy_higher_u32, pto.ui32(0), full_mask_b32, pto.CmpMode.NE) - vy_lower_f32 = pto.vcvt(pto.vbitcast(vy_lower_u32, pto.i32), pto.f32, active_low, rnd=pto.VcvtRoundMode.F) - vy_higher_f32 = pto.vcvt(pto.vbitcast(vy_higher_u32, pto.i32), pto.f32, active_high, rnd=pto.VcvtRoundMode.F) - - vy_rec_lower = pto.vdiv(pto.vbr(fp32_one), vy_lower_f32, active_low) - vy_rec_higher = pto.vdiv(pto.vbr(fp32_one), vy_higher_f32, active_high) - vy_scale_lower = pto.vmul(vy_rec_lower, pto.vbr(pto.f32(65536.0)), active_low) - vy_scale_higher = pto.vmul(vy_rec_higher, pto.vbr(pto.f32(65536.0)), active_high) - - v_lower_i32 = pto.vcvt( - vy_scale_lower, - pto.i32, - active_low, - rnd=pto.VcvtRoundMode.F, - sat=pto.VcvtSatMode.NOSAT, - ) - v_higher_i32 = pto.vcvt( - vy_scale_higher, - pto.i32, - active_high, - rnd=pto.VcvtRoundMode.F, - sat=pto.VcvtSatMode.NOSAT, - ) - v_lower_u32 = pto.vbitcast(v_lower_i32, pto.ui32) - v_higher_u32 = pto.vbitcast(v_higher_i32, pto.ui32) - - vx_lower_u16, vx_higher_u16 = pto.vintlv(vec, zero_u16) - vx_lower_u32 = pto.vcvt(vx_lower_u16, pto.ui32, full_mask_b16, part=pto.VcvtPartMode.EVEN) - vx_higher_u32 = pto.vcvt(vx_higher_u16, pto.ui32, full_mask_b16, part=pto.VcvtPartMode.EVEN) - q_tmp_lower = pto.vmul(v_lower_u32, vx_lower_u32, active_low) - q_tmp_higher = pto.vmul(v_higher_u32, vx_higher_u32, active_high) - _q_lower, q_tmp = pto.vdintlv(pto.vbitcast(q_tmp_lower, pto.ui16), pto.vbitcast(q_tmp_higher, pto.ui16)) - - yq_tmp = pto.vmul(q_tmp, scalar_vec, active_mask) - r_tmp = pto.vsub(vec, yq_tmp, active_mask) - ge_mask = pto.vcmp(r_tmp, scalar_vec, active_mask, pto.CmpMode.GE) - refined_r = pto.vsub(r_tmp, scalar_vec, active_mask) - r_tmp = pto.vsel(refined_r, r_tmp, ge_mask) - q_inc = pto.vadds(q_tmp, one, active_mask) - q_tmp = pto.vsel(q_inc, q_tmp, ge_mask) - ge_mask = pto.vcmp(r_tmp, scalar_vec, active_mask, pto.CmpMode.GE) - refined_r = pto.vsub(r_tmp, scalar_vec, active_mask) - r_tmp = pto.vsel(refined_r, r_tmp, ge_mask) - q_inc = pto.vadds(q_tmp, one, active_mask) - q_tmp = pto.vsel(q_inc, q_tmp, ge_mask) - return pto.vsel(pto.vbr(zero_r), r_tmp, zero_mask) - - -@pto.inline_proc -def _tl_soft_vdiv_u32(vec, scalar_vec, mask): - zero = pto.ui32(0) - one = pto.ui32(1) - zero_q = pto.ui32(0xFFFFFFFF) - fp32_one = pto.f32(1.0) - full_mask = pto.pset_b32(pto.PAT.ALL) - - zero_mask = pto.vcmps(scalar_vec, zero, mask, pto.CmpMode.EQ) - active_mask = pto.pnot(zero_mask, mask) - - zero_u32 = pto.vbr(zero) - zero_f32 = pto.vbr(pto.f32(0.0)) - vy_lower_u32, vy_higher_u32 = pto.vintlv(scalar_vec, zero_u32) - vy_lower_f32 = pto.vcvt(pto.vbitcast(vy_lower_u32, pto.i64), pto.f32, full_mask, rnd=pto.VcvtRoundMode.F, part=pto.VcvtPartMode.EVEN) - vy_higher_f32 = pto.vcvt(pto.vbitcast(vy_higher_u32, pto.i64), pto.f32, full_mask, rnd=pto.VcvtRoundMode.F, part=pto.VcvtPartMode.EVEN) - vy_float, _vy_waste = pto.vdintlv(vy_lower_f32, vy_higher_f32) - - vy_rec = pto.vdiv(pto.vbr(fp32_one), vy_float, full_mask) - vy_scale = pto.vmul(vy_rec, pto.vbr(pto.f32(4294966784.0)), full_mask) - - vy_scale_lower_f32, vy_scale_higher_f32 = pto.vintlv(vy_scale, zero_f32) - v_lower_i64 = pto.vcvt( - vy_scale_lower_f32, - pto.i64, - full_mask, - rnd=pto.VcvtRoundMode.F, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - v_higher_i64 = pto.vcvt( - vy_scale_higher_f32, - pto.i64, - full_mask, - rnd=pto.VcvtRoundMode.F, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - z, _z_waste = pto.vdintlv(pto.vbitcast(v_lower_i64, pto.ui32), pto.vbitcast(v_higher_i64, pto.ui32)) - - tmp_0 = pto.vmul(z, scalar_vec, full_mask) - tmp_0 = pto.vbitcast(pto.vneg(pto.vbitcast(tmp_0, pto.i32), full_mask), pto.ui32) - _z_lower, z_high = pto.vmull(z, tmp_0, full_mask) - z = pto.vadd(z, z_high, full_mask) - - _q_lower, q_tmp = pto.vmull(vec, z, full_mask) - yq_tmp = pto.vmul(q_tmp, scalar_vec, active_mask) - r_tmp = pto.vsub(vec, yq_tmp, active_mask) - ge_mask = pto.vcmp(r_tmp, scalar_vec, active_mask, pto.CmpMode.GE) - refined_r = pto.vsub(r_tmp, scalar_vec, active_mask) - r_tmp = pto.vsel(refined_r, r_tmp, ge_mask) - q_inc = pto.vadds(q_tmp, one, active_mask) - q_tmp = pto.vsel(q_inc, q_tmp, ge_mask) - ge_mask = pto.vcmp(r_tmp, scalar_vec, active_mask, pto.CmpMode.GE) - refined_r = pto.vsub(r_tmp, scalar_vec, active_mask) - r_tmp = pto.vsel(refined_r, r_tmp, ge_mask) - q_inc = pto.vadds(q_tmp, one, active_mask) - q_tmp = pto.vsel(q_inc, q_tmp, ge_mask) - return pto.vsel(pto.vbr(zero_q), q_tmp, zero_mask) - - -@pto.inline_proc -def _tl_soft_vmod_i16(vec, scalar_vec, mask): - zero = pto.i16(0) - neg_one = pto.i16(-1) - - zero_mask = pto.vcmps(scalar_vec, zero, mask, pto.CmpMode.EQ) - active_mask = pto.pnot(zero_mask, mask) - - abs_x = pto.vbitcast(pto.vabs(vec, active_mask), pto.ui16) - abs_y = pto.vbitcast(pto.vabs(scalar_vec, active_mask), pto.ui16) - x_xor_y = pto.vxor(vec, scalar_vec, active_mask) - p_pos = pto.vcmps(x_xor_y, zero, active_mask, pto.CmpMode.GE) - - q_abs = _tl_soft_vdiv_u16(abs_x, abs_y, active_mask) - neg_q = pto.vneg(pto.vbitcast(q_abs, pto.i16), active_mask) - q = pto.vsel(pto.vbitcast(q_abs, pto.i16), neg_q, p_pos) - - qy = pto.vmul(q, scalar_vec, active_mask) - remainder = pto.vsub(vec, qy, active_mask) - - nonzero_remainder = pto.vcmps(remainder, zero, active_mask, pto.CmpMode.NE) - sign_x = pto.vcmps(vec, zero, active_mask, pto.CmpMode.GE) - sign_y = pto.vcmps(scalar_vec, zero, active_mask, pto.CmpMode.GE) - sign_diff = pto.pxor(sign_x, sign_y, active_mask) - need_floor_fix = pto.pand(sign_diff, nonzero_remainder, active_mask) - amended_remainder = pto.vadd(scalar_vec, remainder, active_mask) - remainder = pto.vsel(amended_remainder, remainder, need_floor_fix) - return pto.vsel(pto.vbr(neg_one), remainder, zero_mask) - - -@pto.inline_proc -def _tl_soft_vdiv_i32(vec, scalar_vec, mask): - zero = pto.i32(0) - neg_one = pto.i32(-1) - fp32_one = pto.f32(1.0) - false_mask = pto.pset_b32(pto.PAT.ALLF) - - zero_mask = pto.vcmps(scalar_vec, zero, mask, pto.CmpMode.EQ) - active_mask = pto.pnot(zero_mask, mask) - - abs_x = pto.vbitcast(pto.vabs(vec, active_mask), pto.ui32) - abs_y = pto.vbitcast(pto.vabs(scalar_vec, active_mask), pto.ui32) - x_xor_y = pto.vxor(vec, scalar_vec, active_mask) - p_pos = pto.vcmps(x_xor_y, zero, active_mask, pto.CmpMode.GE) - - y_float = pto.vcvt(pto.vbitcast(abs_y, pto.i32), pto.f32, active_mask, rnd=pto.VcvtRoundMode.R) - y_rec = pto.vdiv(pto.vbr(fp32_one), y_float, active_mask) - f_z_tmp_bits = pto.vadds(pto.vbitcast(y_rec, pto.ui32), pto.ui32(0x0FFFFFFE), active_mask) - - low_mask, high_mask = pto.pintlv_b32(active_mask, false_mask) - lower_bits, higher_bits = pto.vintlv(f_z_tmp_bits, pto.vbr(pto.ui32(0))) - lower_i64 = pto.vcvt( - pto.vbitcast(lower_bits, pto.f32), - pto.i64, - low_mask, - rnd=pto.VcvtRoundMode.F, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - higher_i64 = pto.vcvt( - pto.vbitcast(higher_bits, pto.f32), - pto.i64, - high_mask, - rnd=pto.VcvtRoundMode.F, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - z, _z_waste = pto.vdintlv(pto.vbitcast(lower_i64, pto.ui32), pto.vbitcast(higher_i64, pto.ui32)) - active_mask, _waste_mask = pto.pdintlv_b32(low_mask, high_mask) - - fz_negative = pto.vcmps(pto.vbitcast(f_z_tmp_bits, pto.f32), pto.f32(0.0), active_mask, pto.CmpMode.LT) - z = pto.vsel(pto.vbr(pto.ui32(0)), z, fz_negative) - - tmp_0 = pto.vmul(z, abs_y, active_mask) - tmp_0 = pto.vbitcast(pto.vneg(pto.vbitcast(tmp_0, pto.i32), active_mask), pto.ui32) - _z_lower, z_high = pto.vmull(z, tmp_0, active_mask) - z = pto.vadd(z, z_high, active_mask) - - _q_lower, q_tmp = pto.vmull(abs_x, z, active_mask) - yq_tmp = pto.vmul(q_tmp, abs_y, active_mask) - r_tmp = pto.vsub(abs_x, yq_tmp, active_mask) - ge_mask = pto.vcmp(r_tmp, abs_y, active_mask, pto.CmpMode.GE) - refined_r = pto.vsub(r_tmp, abs_y, active_mask) - r_tmp = pto.vsel(refined_r, r_tmp, ge_mask) - q_inc = pto.vadds(q_tmp, pto.ui32(1), active_mask) - q_tmp = pto.vsel(q_inc, q_tmp, ge_mask) - ge_mask = pto.vcmp(r_tmp, abs_y, active_mask, pto.CmpMode.GE) - refined_r = pto.vsub(r_tmp, abs_y, active_mask) - r_tmp = pto.vsel(refined_r, r_tmp, ge_mask) - q_inc = pto.vadds(q_tmp, pto.ui32(1), active_mask) - q_tmp = pto.vsel(q_inc, q_tmp, ge_mask) - - neg_q = pto.vneg(pto.vbitcast(q_tmp, pto.i32), active_mask) - q = pto.vsel(pto.vbitcast(q_tmp, pto.i32), neg_q, p_pos) - return pto.vsel(pto.vbr(neg_one), q, zero_mask) - - -@pto.inline_proc -def _tl_soft_vmod_u32(vec, scalar_vec, mask): - zero = pto.ui32(0) - one = pto.ui32(1) - zero_r = pto.ui32(0xFFFFFFFF) - fp32_one = pto.f32(1.0) - full_mask = pto.pset_b32(pto.PAT.ALL) - - zero_mask = pto.vcmps(scalar_vec, zero, mask, pto.CmpMode.EQ) - active_mask = pto.pnot(zero_mask, mask) - - zero_u32 = pto.vbr(zero) - zero_f32 = pto.vbr(pto.f32(0.0)) - vy_lower_u32, vy_higher_u32 = pto.vintlv(scalar_vec, zero_u32) - vy_lower_f32 = pto.vcvt(pto.vbitcast(vy_lower_u32, pto.i64), pto.f32, full_mask, rnd=pto.VcvtRoundMode.F, part=pto.VcvtPartMode.EVEN) - vy_higher_f32 = pto.vcvt(pto.vbitcast(vy_higher_u32, pto.i64), pto.f32, full_mask, rnd=pto.VcvtRoundMode.F, part=pto.VcvtPartMode.EVEN) - vy_float, _vy_waste = pto.vdintlv(vy_lower_f32, vy_higher_f32) - - vy_rec = pto.vdiv(pto.vbr(fp32_one), vy_float, full_mask) - vy_scale = pto.vmul(vy_rec, pto.vbr(pto.f32(4294966784.0)), full_mask) - - vy_scale_lower_f32, vy_scale_higher_f32 = pto.vintlv(vy_scale, zero_f32) - v_lower_i64 = pto.vcvt( - vy_scale_lower_f32, - pto.i64, - full_mask, - rnd=pto.VcvtRoundMode.F, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - v_higher_i64 = pto.vcvt( - vy_scale_higher_f32, - pto.i64, - full_mask, - rnd=pto.VcvtRoundMode.F, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - z, _z_waste = pto.vdintlv(pto.vbitcast(v_lower_i64, pto.ui32), pto.vbitcast(v_higher_i64, pto.ui32)) - - tmp_0 = pto.vmul(z, scalar_vec, full_mask) - tmp_0 = pto.vbitcast(pto.vneg(pto.vbitcast(tmp_0, pto.i32), full_mask), pto.ui32) - _z_lower, z_high = pto.vmull(z, tmp_0, full_mask) - z = pto.vadd(z, z_high, full_mask) - - _q_lower, q_tmp = pto.vmull(vec, z, full_mask) - yq_tmp = pto.vmul(q_tmp, scalar_vec, active_mask) - r_tmp = pto.vsub(vec, yq_tmp, active_mask) - ge_mask = pto.vcmp(r_tmp, scalar_vec, active_mask, pto.CmpMode.GE) - refined_r = pto.vsub(r_tmp, scalar_vec, active_mask) - r_tmp = pto.vsel(refined_r, r_tmp, ge_mask) - q_inc = pto.vadds(q_tmp, one, active_mask) - q_tmp = pto.vsel(q_inc, q_tmp, ge_mask) - ge_mask = pto.vcmp(r_tmp, scalar_vec, active_mask, pto.CmpMode.GE) - refined_r = pto.vsub(r_tmp, scalar_vec, active_mask) - r_tmp = pto.vsel(refined_r, r_tmp, ge_mask) - q_inc = pto.vadds(q_tmp, one, active_mask) - q_tmp = pto.vsel(q_inc, q_tmp, ge_mask) - return pto.vsel(pto.vbr(zero_r), r_tmp, zero_mask) - - -@pto.inline_proc -def _tl_soft_vmod_i32(vec, scalar_vec, mask): - zero = pto.i32(0) - neg_one = pto.i32(-1) - fp32_one = pto.f32(1.0) - false_mask = pto.pset_b32(pto.PAT.ALLF) - - zero_mask = pto.vcmps(scalar_vec, zero, mask, pto.CmpMode.EQ) - active_mask = pto.pnot(zero_mask, mask) - - abs_x = pto.vbitcast(pto.vabs(vec, active_mask), pto.ui32) - abs_y = pto.vbitcast(pto.vabs(scalar_vec, active_mask), pto.ui32) - x_xor_y = pto.vxor(vec, scalar_vec, active_mask) - p_pos = pto.vcmps(x_xor_y, zero, active_mask, pto.CmpMode.GE) - - y_float = pto.vcvt(pto.vbitcast(abs_y, pto.i32), pto.f32, active_mask, rnd=pto.VcvtRoundMode.R) - y_rec = pto.vdiv(pto.vbr(fp32_one), y_float, active_mask) - f_z_tmp_bits = pto.vadds(pto.vbitcast(y_rec, pto.ui32), pto.ui32(0x0FFFFFFE), active_mask) - - low_mask, high_mask = pto.pintlv_b32(active_mask, false_mask) - lower_bits, higher_bits = pto.vintlv(f_z_tmp_bits, pto.vbr(pto.ui32(0))) - lower_i64 = pto.vcvt( - pto.vbitcast(lower_bits, pto.f32), - pto.i64, - low_mask, - rnd=pto.VcvtRoundMode.F, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - higher_i64 = pto.vcvt( - pto.vbitcast(higher_bits, pto.f32), - pto.i64, - high_mask, - rnd=pto.VcvtRoundMode.F, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - z, _z_waste = pto.vdintlv(pto.vbitcast(lower_i64, pto.ui32), pto.vbitcast(higher_i64, pto.ui32)) - active_mask, _waste_mask = pto.pdintlv_b32(low_mask, high_mask) - - fz_negative = pto.vcmps(pto.vbitcast(f_z_tmp_bits, pto.f32), pto.f32(0.0), active_mask, pto.CmpMode.LT) - z = pto.vsel(pto.vbr(pto.ui32(0)), z, fz_negative) - - tmp_0 = pto.vmul(z, abs_y, active_mask) - tmp_0 = pto.vbitcast(pto.vneg(pto.vbitcast(tmp_0, pto.i32), active_mask), pto.ui32) - _z_lower, z_high = pto.vmull(z, tmp_0, active_mask) - z = pto.vadd(z, z_high, active_mask) - - _q_lower, q_tmp = pto.vmull(abs_x, z, active_mask) - yq_tmp = pto.vmul(q_tmp, abs_y, active_mask) - r_tmp = pto.vsub(abs_x, yq_tmp, active_mask) - ge_mask = pto.vcmp(r_tmp, abs_y, active_mask, pto.CmpMode.GE) - refined_r = pto.vsub(r_tmp, abs_y, active_mask) - r_tmp = pto.vsel(refined_r, r_tmp, ge_mask) - q_inc = pto.vadds(q_tmp, pto.ui32(1), active_mask) - q_tmp = pto.vsel(q_inc, q_tmp, ge_mask) - ge_mask = pto.vcmp(r_tmp, abs_y, active_mask, pto.CmpMode.GE) - refined_r = pto.vsub(r_tmp, abs_y, active_mask) - r_tmp = pto.vsel(refined_r, r_tmp, ge_mask) - q_inc = pto.vadds(q_tmp, pto.ui32(1), active_mask) - q_tmp = pto.vsel(q_inc, q_tmp, ge_mask) - - neg_q = pto.vneg(pto.vbitcast(q_tmp, pto.i32), active_mask) - q = pto.vsel(pto.vbitcast(q_tmp, pto.i32), neg_q, p_pos) - - qy = pto.vmul(q, scalar_vec, active_mask) - remainder = pto.vsub(vec, qy, active_mask) - nonzero_remainder = pto.vcmps(pto.vbitcast(r_tmp, pto.i32), zero, active_mask, pto.CmpMode.NE) - sign_x = pto.vcmps(vec, zero, active_mask, pto.CmpMode.GE) - sign_y = pto.vcmps(scalar_vec, zero, active_mask, pto.CmpMode.GE) - sign_diff = pto.pxor(sign_x, sign_y, active_mask) - need_floor_fix = pto.pand(sign_diff, nonzero_remainder, active_mask) - amended_remainder = pto.vadd(scalar_vec, remainder, active_mask) - remainder = pto.vsel(amended_remainder, remainder, need_floor_fix) - return pto.vsel(pto.vbr(neg_one), remainder, zero_mask) - - -@pto.inline_proc -def _tl_soft_vmod(vec, scalar_vec, mask, dtype): - if pto.constexpr(dtype == pto.ui8): - result = _tl_soft_vmod_u8(vec, scalar_vec, mask) - elif pto.constexpr(dtype == pto.i8): - result = _tl_soft_vmod_i8(vec, scalar_vec, mask) - elif pto.constexpr(dtype == pto.ui16): - result = _tl_soft_vmod_u16(vec, scalar_vec, mask) - elif pto.constexpr(dtype == pto.i16): - result = _tl_soft_vmod_i16(vec, scalar_vec, mask) - elif pto.constexpr(dtype == pto.ui32): - result = _tl_soft_vmod_u32(vec, scalar_vec, mask) - else: - result = _tl_soft_vmod_i32(vec, scalar_vec, mask) - return result - - -@pto.inline_proc -def _tl_soft_vdiv(vec, scalar_vec, mask, dtype): - if pto.constexpr(dtype == pto.ui8): - result = _tl_soft_vdiv_u8(vec, scalar_vec, mask) - elif pto.constexpr(dtype == pto.i8): - result = _tl_soft_vdiv_i8(vec, scalar_vec, mask) - elif pto.constexpr(dtype == pto.ui16): - result = _tl_soft_vdiv_u16(vec, scalar_vec, mask) - elif pto.constexpr(dtype == pto.i16): - result = _tl_soft_vdiv_i16(vec, scalar_vec, mask) - elif pto.constexpr(dtype == pto.ui32): - result = _tl_soft_vdiv_u32(vec, scalar_vec, mask) - else: - result = _tl_soft_vdiv_i32(vec, scalar_vec, mask) - return result diff --git a/lib/TileOps/render_template_mlir.py b/lib/TileOps/render_template_mlir.py deleted file mode 100644 index b211189612..0000000000 --- a/lib/TileOps/render_template_mlir.py +++ /dev/null @@ -1,378 +0,0 @@ -# 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. - -#!/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. - -"""Materialize a TileLang DSL library template to authoring-form MLIR. - -Examples: - python3 lib/TileOps/render_template_mlir.py lib/TileOps/tload_template.py - python3 lib/TileOps/render_template_mlir.py lib/TileOps/tadd_template.py --tile dst=8x64@ub --tile src0=8x64@ub --tile src1=8x64@ub - python3 lib/TileOps/render_template_mlir.py lib/TileOps/tload_template.py --dtypes f16,f16 -o /tmp/tload.mlir -""" - -from __future__ import annotations - -import argparse -import importlib.util -import sys -from pathlib import Path -from types import ModuleType - - -REPO_ROOT = Path(__file__).resolve().parents[2] -TILELANG_PYTHON_DIR = REPO_ROOT / "tilelang-dsl" / "python" -if str(TILELANG_PYTHON_DIR) not in sys.path: - sys.path.insert(0, str(TILELANG_PYTHON_DIR)) - -import tilelang_dsl as pto - - -_DTYPE_BY_NAME = { - "i1": pto.i1, - "i8": pto.i8, - "i16": pto.i16, - "i32": pto.i32, - "i64": pto.i64, - "f16": pto.f16, - "bf16": pto.bf16, - "f32": pto.f32, -} -_MEMORY_SPACE_BY_NAME = { - "gm": pto.MemorySpace.GM, - "ub": pto.MemorySpace.UB, -} - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Load a TileLang DSL template file and emit its corresponding MLIR text.", - ) - parser.add_argument("template", help="Path to the template Python file") - parser.add_argument( - "--kernel", - help="Descriptor symbol name inside the module when the file defines multiple @pto.vkernel templates", - ) - parser.add_argument( - "--op", - help="Concrete op to bind when the descriptor matches multiple ops; defaults to the first match op", - ) - parser.add_argument( - "--dtypes", - help="Concrete operand dtypes as a comma-separated list, for example: f32,f32 or f16,f16,f16", - ) - parser.add_argument( - "--tile", - action="append", - default=[], - metavar="PARAM=SHAPE[@SPACE][:VALID]", - help=( - "Tile specialization override, for example: dst=16x32@ub or " - "dst=16x32@ub:8x32. May be repeated." - ), - ) - parser.add_argument( - "--default-tile-shape", - default="16x32", - help="Default shape for every bare Tile parameter when no --tile override is given", - ) - parser.add_argument( - "--default-tile-space", - default="ub", - choices=sorted(_MEMORY_SPACE_BY_NAME), - help="Default memory space for every bare Tile parameter", - ) - parser.add_argument( - "-o", - "--output", - help="Optional output path; defaults to stdout", - ) - return parser.parse_args() - - -def _load_module(template_path: Path) -> ModuleType: - template_parent = template_path.parent.parent - if str(template_parent) not in sys.path: - sys.path.insert(0, str(template_parent)) - module_name = f"_tileops_template_{template_path.stem}" - spec = importlib.util.spec_from_file_location(module_name, template_path) - if spec is None or spec.loader is None: - raise ValueError(f"failed to load Python module from {template_path}") - module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) - return module - - -def _find_descriptors(module: ModuleType) -> dict[str, pto.VKernelDescriptor]: - descriptors: dict[str, pto.VKernelDescriptor] = {} - for name, value in vars(module).items(): - if isinstance(value, pto.VKernelDescriptor): - descriptors[name] = value - return descriptors - - -def _select_descriptor( - descriptors: dict[str, pto.VKernelDescriptor], - kernel_name: str | None, -) -> tuple[str, pto.VKernelDescriptor]: - if not descriptors: - raise ValueError("no @pto.vkernel descriptor found in the template module") - if kernel_name is not None: - descriptor = descriptors.get(kernel_name) - if descriptor is None: - available = ", ".join(sorted(descriptors)) - raise ValueError( - f"kernel {kernel_name!r} was not found in the template module; available descriptors: {available}" - ) - return kernel_name, descriptor - if len(descriptors) == 1: - return next(iter(descriptors.items())) - available = ", ".join(sorted(descriptors)) - raise ValueError( - "the template module defines multiple @pto.vkernel descriptors; " - f"please pass --kernel. Available descriptors: {available}" - ) - - -def _parse_dtype_list(text: str) -> tuple[pto.ScalarType, ...]: - parts = [part.strip() for part in text.split(",") if part.strip()] - if not parts: - raise ValueError("--dtypes must contain at least one dtype") - try: - return tuple(_DTYPE_BY_NAME[part] for part in parts) - except KeyError as exc: - available = ", ".join(sorted(_DTYPE_BY_NAME)) - raise ValueError( - f"unsupported dtype {exc.args[0]!r}; available dtypes: {available}" - ) from exc - - -def _default_concrete_dtype(pattern: object) -> pto.ScalarType: - if isinstance(pattern, pto.ScalarType): - return pattern - if isinstance(pattern, pto.WildcardType): - if pattern.name in {"AnyType", "AnyFloat"}: - return pto.f32 - if pattern.name == "AnyInt": - return pto.i32 - if pattern.name == "AnyMask": - return pto.i1 - raise ValueError(f"unsupported wildcard dtype pattern {pattern!r}") - if isinstance(pattern, pto.TypeVariable): - return pto.f32 - raise ValueError(f"unsupported dtype pattern {pattern!r}") - - -def _default_parameter_dtype( - param_spec: object | None, - pattern: object, -) -> pto.ScalarType: - annotation = getattr(param_spec, "annotation", None) - if isinstance(annotation, pto.ScalarType): - return annotation - if isinstance(annotation, pto.WildcardType) and annotation.name != "AnyType": - return _default_concrete_dtype(annotation) - if isinstance(annotation, pto.MaskType): - return pto.i1 - return _default_concrete_dtype(pattern) - - -def _default_operand_types(descriptor: pto.VKernelDescriptor) -> tuple[pto.ScalarType, ...]: - if not descriptor.dtypes: - raise ValueError("descriptor does not declare any dtype signatures") - prototype = descriptor.dtypes[0] - parameter_specs = getattr(descriptor, "_parameter_specs", ()) - typevar_bindings: dict[str, pto.ScalarType] = {} - concrete: list[pto.ScalarType] = [] - for index, pattern in enumerate(prototype): - param_spec = parameter_specs[index] if index < len(parameter_specs) else None - if isinstance(pattern, pto.TypeVariable): - bound = typevar_bindings.get(pattern.name) - if bound is None: - bound = _default_parameter_dtype(param_spec, pattern) - typevar_bindings[pattern.name] = bound - concrete.append(bound) - continue - concrete.append(_default_parameter_dtype(param_spec, pattern)) - return tuple(concrete) - - -def _bind_descriptor( - descriptor: pto.VKernelDescriptor, - *, - op_name: str | None, - operand_types: tuple[pto.ScalarType, ...] | None, -) -> pto.VKernelDescriptor: - concrete_op = op_name - if concrete_op is None: - if descriptor.selected_op is not None: - concrete_op = descriptor.selected_op - elif len(descriptor.match_ops) == 1: - concrete_op = descriptor.match_ops[0] - else: - available = ", ".join(descriptor.match_ops) - raise ValueError( - f"descriptor matches multiple ops; pass --op. Available ops: {available}" - ) - - concrete_operand_types = operand_types - if concrete_operand_types is None: - if descriptor._selected_dtype_signature is not None: - concrete_operand_types = descriptor._selected_dtype_signature - else: - concrete_operand_types = _default_operand_types(descriptor) - - registry = pto.KernelRegistry((descriptor,)) - return pto.select_kernel( - target=descriptor.target, - op=concrete_op, - operand_types=concrete_operand_types, - registry=registry, - ) - - -def _parse_shape(text: str) -> tuple[int, ...]: - dims = [] - for part in text.split("x"): - part = part.strip() - if not part: - raise ValueError(f"invalid shape {text!r}") - value = int(part) - if value <= 0: - raise ValueError(f"shape dimensions must be positive integers, got {text!r}") - dims.append(value) - if not dims: - raise ValueError(f"invalid shape {text!r}") - return tuple(dims) - - -def _parse_tile_override(spec_text: str) -> tuple[str, pto.TileSpecialization]: - if "=" not in spec_text: - raise ValueError( - f"invalid --tile value {spec_text!r}; expected PARAM=SHAPE[@SPACE][:VALID]" - ) - param_name, payload = spec_text.split("=", 1) - param_name = param_name.strip() - payload = payload.strip() - if not param_name: - raise ValueError(f"invalid --tile value {spec_text!r}; missing parameter name") - - valid_shape = None - if ":" in payload: - payload, valid_text = payload.split(":", 1) - valid_shape = _parse_shape(valid_text.strip()) - - memory_space = pto.MemorySpace.UB - if "@" in payload: - shape_text, memory_space_text = payload.split("@", 1) - memory_space_key = memory_space_text.strip().lower() - try: - memory_space = _MEMORY_SPACE_BY_NAME[memory_space_key] - except KeyError as exc: - available = ", ".join(sorted(_MEMORY_SPACE_BY_NAME)) - raise ValueError( - f"unsupported memory space {memory_space_text!r}; available spaces: {available}" - ) from exc - else: - shape_text = payload - - shape = _parse_shape(shape_text.strip()) - if valid_shape is not None and len(valid_shape) != len(shape): - raise ValueError( - f"valid_shape rank {len(valid_shape)} does not match shape rank {len(shape)} for {param_name!r}" - ) - return ( - param_name, - pto.TileSpecialization( - shape=shape, - memory_space=memory_space, - valid_shape=valid_shape, - ), - ) - - -def _default_tile_specialization( - *, - shape: tuple[int, ...], - memory_space: pto.MemorySpace, -) -> pto.TileSpecialization: - return pto.TileSpecialization(shape=shape, memory_space=memory_space) - - -def _specialize_tiles( - descriptor: pto.VKernelDescriptor, - *, - tile_overrides: dict[str, pto.TileSpecialization], - default_shape: tuple[int, ...], - default_memory_space: pto.MemorySpace, -) -> pto.VKernelDescriptor: - if not descriptor.tile_parameters: - return descriptor - - specializations: dict[str, pto.TileSpecialization] = {} - for param in descriptor.tile_parameters: - specializations[param.name] = tile_overrides.get( - param.name, - _default_tile_specialization( - shape=default_shape, - memory_space=default_memory_space, - ), - ) - return descriptor.specialize(**specializations) - - -def _emit_output(text: str, output_path: str | None) -> None: - if output_path is None: - sys.stdout.write(text) - if not text.endswith("\n"): - sys.stdout.write("\n") - return - path = Path(output_path) - path.write_text(text, encoding="utf-8") - - -def main() -> int: - args = _parse_args() - template_path = Path(args.template).resolve() - if not template_path.is_file(): - print(f"error: template file not found: {template_path}", file=sys.stderr) - return 1 - - try: - module = _load_module(template_path) - _, descriptor = _select_descriptor(_find_descriptors(module), args.kernel) - operand_types = None if args.dtypes is None else _parse_dtype_list(args.dtypes) - bound = _bind_descriptor( - descriptor, - op_name=args.op, - operand_types=operand_types, - ) - tile_overrides = dict(_parse_tile_override(spec_text) for spec_text in args.tile) - specialized = _specialize_tiles( - bound, - tile_overrides=tile_overrides, - default_shape=_parse_shape(args.default_tile_shape), - default_memory_space=_MEMORY_SPACE_BY_NAME[args.default_tile_space], - ) - _emit_output(specialized.mlir_text(), args.output) - return 0 - except Exception as exc: - print(f"error: {exc}", file=sys.stderr) - return 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/lib/TileOps/sqrt_hp.py b/lib/TileOps/sqrt_hp.py deleted file mode 100644 index 11f0ea97da..0000000000 --- a/lib/TileOps/sqrt_hp.py +++ /dev/null @@ -1,83 +0,0 @@ -# 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 tilelang_dsl as pto - -@pto.inline_proc -def _tl_sqrt_precision_f16(src, mask): - multiply_factor0 = pto.f16("0x6c00") - multiply_factor1 = pto.f16("0x2400") - subnormal_threshold = pto.f16("0x03ff") - - subnormal_mask = pto.vcmps(src, subnormal_threshold, mask, pto.CmpMode.LT) - - tmp = pto.vmuls(src, multiply_factor0, subnormal_mask) - src_adjusted = pto.vsel(tmp, src, subnormal_mask) - - dst = pto.vsqrt(src_adjusted, mask) - - tmp = pto.vmuls(dst, multiply_factor1, subnormal_mask) - result = pto.vsel(tmp, dst, subnormal_mask) - - return result - - -@pto.inline_proc -def _tl_sqrt_precision_f32(src, mask): - multiply_factor0 = pto.f32(16777216.0) - multiply_factor1 = pto.f32(0.000244140625) - subnormal_bound = pto.f32(1.0) - half_factor = pto.f32(0.5) - neg_one = pto.f32(-1.0) - - subnormal_mask = pto.vcmps(src, subnormal_bound, mask, pto.CmpMode.LT) - - tmp = pto.vmuls(src, multiply_factor0, subnormal_mask) - src_adjusted = pto.vsel(tmp, src, subnormal_mask) - - reg_one = pto.vbr(pto.f32(1.0)) - tmp_sqrt = pto.vsqrt(src_adjusted, mask) - dst = pto.vdiv(reg_one, tmp_sqrt, mask) - - reg_neg_one = pto.vmuls(dst, neg_one, mask) - err = pto.vmul(dst, src_adjusted, mask) - reg_one_adj = pto.vmula(reg_one, err, reg_neg_one, mask) - tmp_half = pto.vmuls(dst, half_factor, mask) - dst = pto.vmula(dst, reg_one_adj, tmp_half, mask) - - res = pto.vmul(dst, src_adjusted, mask) - tmp_neg = pto.vmuls(res, neg_one, mask) - err = pto.vmula(src_adjusted, res, tmp_neg, mask) - tmp_half = pto.vmuls(dst, half_factor, mask) - tmp = pto.vmul(err, tmp_half, mask) - tmp = pto.vadd(tmp, res, mask) - - tmp_scaled = pto.vmuls(tmp, multiply_factor1, mask) - result = pto.vsel(tmp_scaled, tmp, subnormal_mask) - - pos_inf = pto.ui32(0x7f800000) - neg_zero = pto.ui32(0x80000000) - - src_as_u32 = pto.vbitcast(src_adjusted, pto.ui32) - is_inf_mask = pto.vcmps(src_as_u32, pos_inf, mask, pto.CmpMode.EQ) - src_with_sign = pto.vor(src_as_u32, pto.vbr(neg_zero), mask) - is_zero_mask = pto.vcmps(src_with_sign, neg_zero, mask, pto.CmpMode.EQ) - special_mask = pto.por(is_zero_mask, is_inf_mask, mask) - - result = pto.vsel(src_adjusted, result, special_mask) - - return result - - -@pto.inline_proc -def _tl_sqrt_precision(src, mask, dtype): - if pto.constexpr(dtype == pto.f16): - result = _tl_sqrt_precision_f16(src, mask) - else: - result = _tl_sqrt_precision_f32(src, mask) - return result \ No newline at end of file diff --git a/lib/TileOps/tabs_template.py b/lib/TileOps/tabs_template.py deleted file mode 100644 index 6c6802ae52..0000000000 --- a/lib/TileOps/tabs_template.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tabs""" - -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tabs" -) -def template_tabs(src: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vinput = pto.vlds(src[row, col:]) - result = pto.vabs(vinput, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tadd_template.py b/lib/TileOps/tadd_template.py deleted file mode 100644 index 8e247fd732..0000000000 --- a/lib/TileOps/tadd_template.py +++ /dev/null @@ -1,32 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tadd""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tadd" -) -def template_tadd(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - summed = pto.vadd(lhs, rhs, mask) - pto.vsts(summed, dst[row, col:], mask) - return diff --git a/lib/TileOps/tadds_template.py b/lib/TileOps/tadds_template.py deleted file mode 100644 index 7c3ddb06c7..0000000000 --- a/lib/TileOps/tadds_template.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tadds""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tadds", -) -def template_tadds(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - result = pto.vadds(vec, scalar, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tand_template.py b/lib/TileOps/tand_template.py deleted file mode 100644 index 6c14771976..0000000000 --- a/lib/TileOps/tand_template.py +++ /dev/null @@ -1,32 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tand""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tand" -) -def template_tand(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - result = pto.vand(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tands_template.py b/lib/TileOps/tands_template.py deleted file mode 100644 index 91258e7930..0000000000 --- a/lib/TileOps/tands_template.py +++ /dev/null @@ -1,40 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tands - -Note: A5 hardware implements tands as: - TEXPANDS_IMPL(dst, scalar); // broadcast scalar to dst - TAND_IMPL(dst, src, dst); // dst = src & dst - -This template uses vbr + vand to achieve element-wise bitwise AND. -Only supports tile, scalar order. -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tands", -) -def template_tands(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - scalar_vec = pto.vbr(scalar) - result = pto.vand(vec, scalar_vec, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tcmp_template.py b/lib/TileOps/tcmp_template.py deleted file mode 100644 index 6b6b2e88f0..0000000000 --- a/lib/TileOps/tcmp_template.py +++ /dev/null @@ -1,130 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tcmp - -Note: A5 hardware implements tcmp as packed comparison between two tiles: - dst = packed_mask(src0 cmp src1) - -Dst is i8 type with same shape as src (packed predicate mask bytes). -Uses vcmp + psts to produce packed predicate mask output. - -Implementation per TCmp.hpp: - - 32B types (f32, i32): uses TCmp_32B path with pdintlv_b8 + PK storage - - 16B types (f16, i16): uses TCmp_8B_16B path with PK storage - - 8B types (i8, u8): uses TCmp_8B_16B path with NORM storage - -Supported comparison modes (via cmp_mode attribute): - eq, ne, lt, gt, ge, le -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tcmp", - dtypes=[ - (pto.f32, pto.f32, pto.i8), - (pto.i32, pto.i32, pto.i8), - (pto.f16, pto.f16, pto.i8), - (pto.i16, pto.i16, pto.i8), - (pto.i8, pto.i8, pto.i8), - (pto.ui8, pto.ui8, pto.i8), - ], - advanced=True, -) -def template_tcmp(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - """src0 cmp src1 -> packed mask in dst (i8, same shape as src) - - TCmp.hpp structure: - - 32B: TCmp_32B with double iteration + pdintlv_b8 + PK - - 16B: TCmp_8B_16B with PK - - 8B: TCmp_8B_16B with NORM - """ - dtype = src0.element_type - valid_rows, valid_cols = src0.valid_shape - cmp_mode = pto.get_op_attr("cmp_mode", "eq") - - lanes = pto.get_lanes(dtype) - dst_ptr = dst.as_ptr() - dst_stride = dst.shape[1] - - if pto.constexpr(dtype == pto.f32 or dtype == pto.i32): - # 32B path: TCmp_32B implementation per TCmp.hpp - # repeatElm = CCE_VL / sizeof(uint32_t) = 64 - # repeatTimes = CeilDivision(validCol, repeatElm) + 1 - # iterations = repeatTimes // 2 - # Each iteration loads 2*64 elements (4 vlds), uses plt_b32 to split into - # two 32-lane comparisons, then pdintlv_b8 to interleave - repeat_times = (valid_cols + lanes - 1) // lanes + 1 - iterations = repeat_times // 2 - - for row in range(0, valid_rows, 1): - remained = valid_cols - for j in range(0, iterations, 1): - # Load 4 vector registers per TCmp.hpp structure - # src0Reg0, src0Reg1 from j*2*repeatElm offset - # src1Reg0, src1Reg1 from (j*2+1)*repeatElm offset - vec_src0_first = pto.vlds(src0[row, j * lanes * 2:]) - vec_src1_first = pto.vlds(src1[row, j * lanes * 2:]) - vec_src0_second = pto.vlds(src0[row, (j * 2 + 1) * lanes:]) - vec_src1_second = pto.vlds(src1[row, (j * 2 + 1) * lanes:]) - - # Use plt_b32 to create 32-lane masks (POST_UPDATE semantics) - mask_first, remained = pto.make_mask(dtype, remained) - cmp_first = pto.vcmp(vec_src0_first, vec_src1_first, mask_first, cmp_mode) - cmp_first_b8 = pto.pbitcast(cmp_first, pto.mask_b8) - - mask_second, remained = pto.make_mask(dtype, remained) - cmp_second = pto.vcmp(vec_src0_second, vec_src1_second, mask_second, cmp_mode) - cmp_second_b8 = pto.pbitcast(cmp_second, pto.mask_b8) - - # pdintlv_b8 interleave two mask_b8 results - packed_low, packed_high = pto.pdintlv_b8(cmp_first_b8, cmp_second_b8) - - # Store to dst: offset = (row * dstStride + j * 4) in uint32 units - # byte_offset = row * dst_stride + j * 16 - # For i8 dst, dstStride = RowStride / 4 = 16 uint32 units - byte_offset = row * dst_stride + j * 16 - pto.psts(packed_low, dst_ptr, byte_offset, pto.PredicateDist.PK) - - elif pto.constexpr(dtype == pto.f16 or dtype == pto.i16): - # 16B path: TCmp_8B_16B with PK - # vcmp returns mask_b16, cast to mask_b8 for psts PK - iters_per_row = (valid_cols + lanes - 1) // lanes - - for row in range(0, valid_rows, 1): - remained = valid_cols - for j in range(0, iters_per_row, 1): - mask, remained = pto.make_mask(dtype, remained) - vec0 = pto.vlds(src0[row, j * lanes:]) - vec1 = pto.vlds(src1[row, j * lanes:]) - cmp = pto.vcmp(vec0, vec1, mask, cmp_mode) - cmp_b8 = pto.pbitcast(cmp, pto.mask_b8) - byte_offset = row * dst_stride + j * 16 - pto.psts(cmp_b8, dst_ptr, byte_offset, pto.PredicateDist.PK) - - else: - # 8B path: TCmp_8B_16B with NORM - # vcmp returns mask_b8 directly, no cast needed - iters_per_row = (valid_cols + lanes - 1) // lanes - - for row in range(0, valid_rows, 1): - remained = valid_cols - for j in range(0, iters_per_row, 1): - mask, remained = pto.make_mask(dtype, remained) - vec0 = pto.vlds(src0[row, j * lanes:]) - vec1 = pto.vlds(src1[row, j * lanes:]) - cmp = pto.vcmp(vec0, vec1, mask, cmp_mode) - byte_offset = row * dst_stride + j * 32 - pto.psts(cmp, dst_ptr, byte_offset, pto.PredicateDist.NORM) - - return \ No newline at end of file diff --git a/lib/TileOps/tcmps_template.py b/lib/TileOps/tcmps_template.py deleted file mode 100644 index fad499c5cd..0000000000 --- a/lib/TileOps/tcmps_template.py +++ /dev/null @@ -1,137 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tcmps - -Note: A5 hardware implements tcmps as packed comparison with scalar: - dst = packed_mask(src cmp scalar) - -Uses vcmps + psts to produce packed predicate mask output. -Implementation: - - 32B types (f32, i32): 64 elements per repeat, 32 bytes per iteration (NORM mode). - - 16B types (f16, i16): 128 elements per repeat, 16 bytes per iteration (PK mode). - - 8B types (i8, u8): 256 elements per repeat, 32 bytes per iteration (NORM mode). - -Supported comparison modes (via cmp_mode op attribute): - eq, ne, lt, gt, ge, le -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tcmps", - dtypes=[ - (pto.f32, pto.f32, pto.ui8), (pto.i32, pto.i32, pto.ui8), - (pto.f16, pto.f16, pto.ui8), (pto.i16, pto.i16, pto.ui8), - (pto.i8, pto.i8, pto.ui8), (pto.ui8, pto.ui8, pto.ui8), - ], - advanced=True, -) -def template_tcmps(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): - """src cmp scalar -> packed mask in dst (ui8) - - - 32B: 1 repeat per iteration, 32 bytes/store (NORM mode, 1 bit per element) - - 16B: 1 repeat per iteration, 16 bytes/store (PK mode) - - 8B: 1 repeat per iteration, 32 bytes/store (NORM mode) - """ - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - cmp_mode = pto.get_op_attr("cmp_mode") - - lanes = pto.get_lanes(dtype) - dst_ptr = dst.as_ptr() - - if pto.constexpr(dtype == pto.f32 or dtype == pto.i32): - # 32B path: 2 vcmps + pbitcast + dintlv_b8 -> psts(PK) - # Use 2D slicing for safety, convert linear offset to (row, col) - bytes_per_iter = 16 - elem_size = 4 - total_elm = valid_rows * valid_cols - repeat_elm = lanes - - # Calculate repeat times matching ISA: CeilDivision + 1 - # But limit iterations to avoid complete out-of-bounds access - repeat_times = (total_elm + repeat_elm - 1) // repeat_elm + 1 - - # Safety: limit iterations to avoid elem_offset beyond total_elm + repeat_elm - # ISA allows one extra repeat for odd elements, but we need to protect DSL slicing - iterations_needed = repeat_times // 2 - - for i in range(0, iterations_needed, 1): - # Convert linear element offsets to (row, col) coordinates - elem_offset0 = i * 2 * repeat_elm - elem_offset1 = (i * 2 + 1) * repeat_elm - - row0 = elem_offset0 // valid_cols - col0 = elem_offset0 % valid_cols - row1 = elem_offset1 // valid_cols - col1 = elem_offset1 % valid_cols - - # Remaining elements for each position (clamp to >= 0) - # When remaining <= 0, make_mask returns all-zero mask (safe) - remaining0 = total_elm - elem_offset0 - if remaining0 < 0: - remaining0 = 0 - remaining1 = total_elm - elem_offset1 - if remaining1 < 0: - remaining1 = 0 - - # Predicate for each compare - mask0, _ = pto.make_mask(dtype, remaining0) - mask1, _ = pto.make_mask(dtype, remaining1) - - # Load using 2D slicing (safer than pointer+offset) - # When row/col exceeds valid_shape, mask ensures no invalid data is used - vec0 = pto.vlds(src[row0, col0:]) - vec1 = pto.vlds(src[row1, col1:]) - - cmp0 = pto.vcmps(vec0, scalar, mask0, cmp_mode) - cmp1 = pto.vcmps(vec1, scalar, mask1, cmp_mode) - - # Convert mask_b32 to mask_b8 and interleave - cmp0_b8 = pto.pbitcast(cmp0, pto.mask_b8) - cmp1_b8 = pto.pbitcast(cmp1, pto.mask_b8) - cmp_interleaved, _ = pto.pdintlv_b8(cmp0_b8, cmp1_b8) - - # Linear byte offset for output - byte_offset = i * bytes_per_iter - pto.psts(cmp_interleaved, dst_ptr, byte_offset, pto.PredicateDist.PK) - elif pto.constexpr(dtype == pto.f16 or dtype == pto.i16): - # 16B path: 128 elements per repeat, 16 bytes per iteration (PK mode). - # Each vcmps produces 128 bits; PK mode packs them into 16 bytes, - # achieving 1 bit per element. - bytes_per_iter = 16 - iters_per_row = (valid_cols + lanes - 1) // lanes - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - cmp = pto.vcmps(vec, scalar, mask, cmp_mode) - byte_offset = (row * iters_per_row + col // lanes) * bytes_per_iter - pto.psts(cmp, dst_ptr, byte_offset, pto.PredicateDist.PK) - else: - # 8B path: 256 elements per repeat, 32 bytes packed per iteration - bytes_per_iter = 32 - iters_per_row = (valid_cols + lanes - 1) // lanes - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - cmp = pto.vcmps(vec, scalar, mask, cmp_mode) - byte_offset = (row * iters_per_row + col // lanes) * bytes_per_iter - pto.psts(cmp, dst_ptr, byte_offset, pto.PredicateDist.NORM) - - return diff --git a/lib/TileOps/tcolargmax_template.py b/lib/TileOps/tcolargmax_template.py deleted file mode 100644 index e6781be272..0000000000 --- a/lib/TileOps/tcolargmax_template.py +++ /dev/null @@ -1,213 +0,0 @@ -# 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 -from pathlib import Path -import tilelang_dsl as pto - - -def _validate_tcolargmax( - src_shape=(), - src_valid_shape=(), - tmp_shape=(), - tmp_valid_shape=(), - dst_shape=(), - dst_valid_shape=(), - src_config=None, - tmp_config=None, - dst_config=None, - src_dtype=None, - tmp_dtype=None, - dst_dtype=None, -): - if src_config is None or tmp_config is None or dst_config is None: - return False - if src_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if tmp_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if dst_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if src_config.s_layout != pto.SLayout.NONE_BOX: - return False - if tmp_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_valid_shape[0] != 1: - return False - if src_dtype != tmp_dtype: - return False - return True - - -@pto.vkernel( - target="a5", - op="pto.tcolargmax", - dtypes=[ - (pto.ui8, pto.ui8, pto.i32), - (pto.i8, pto.i8, pto.i32), - ], - constraints=[_validate_tcolargmax], - advanced=True, -) -def template_tcolargmax_i8_to_i32(src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - src_valid_rows, src_valid_cols = src.valid_shape - src_dtype = src.element_type - lanes_i8 = pto.get_lanes(src_dtype) - lanes_i32 = pto.get_lanes(pto.i32) - - if pto.constexpr(src_dtype == pto.ui8): - intermediate_dtype = pto.ui16 - final_cvt_dtype = pto.ui32 - else: - intermediate_dtype = pto.i16 - final_cvt_dtype = pto.i32 - - lanes_intermediate = pto.get_lanes(intermediate_dtype) - - with pto.vecscope(): - all_mask_b8 = pto.make_mask(src_dtype, pto.PAT.ALL) - all_mask_intermediate = pto.make_mask(intermediate_dtype, pto.PAT.ALL) - - for col in range(0, src_valid_cols, lanes_i8): - remained = src_valid_cols - col - mask_i32_0, remained = pto.make_mask(pto.i32, remained) - mask_i32_1, remained = pto.make_mask(pto.i32, remained) - mask_i32_2, remained = pto.make_mask(pto.i32, remained) - mask_i32_3, _ = pto.make_mask(pto.i32, remained) - - index_old_even = pto.vdup(intermediate_dtype(0), all_mask_intermediate) - index_old_odd = pto.vdup(intermediate_dtype(0), all_mask_intermediate) - index_new_even = pto.vdup(intermediate_dtype(0), all_mask_intermediate) - index_new_odd = pto.vdup(intermediate_dtype(0), all_mask_intermediate) - - vreg_old = pto.vlds(src[0, col:]) - vreg_old_even = pto.vcvt(vreg_old, intermediate_dtype, all_mask_b8, part=pto.VcvtPartMode.EVEN) - vreg_old_odd = pto.vcvt(vreg_old, intermediate_dtype, all_mask_b8, part=pto.VcvtPartMode.ODD) - - for row in range(1, src_valid_rows, 1): - index_new_even = pto.vadds(index_new_even, intermediate_dtype(1), all_mask_intermediate) - index_new_odd = pto.vadds(index_new_odd, intermediate_dtype(1), all_mask_intermediate) - vreg_new = pto.vlds(src[row, col:]) - vreg_new_even = pto.vcvt(vreg_new, intermediate_dtype, all_mask_b8, part=pto.VcvtPartMode.EVEN) - vreg_new_odd = pto.vcvt(vreg_new, intermediate_dtype, all_mask_b8, part=pto.VcvtPartMode.ODD) - - select_even = pto.vcmp(vreg_new_even, vreg_old_even, all_mask_intermediate, "gt") - select_odd = pto.vcmp(vreg_new_odd, vreg_old_odd, all_mask_intermediate, "gt") - - index_old_even = pto.vsel(index_new_even, index_old_even, select_even) - index_old_odd = pto.vsel(index_new_odd, index_old_odd, select_odd) - - vreg_old_even = pto.vmax(vreg_old_even, vreg_new_even, all_mask_intermediate) - vreg_old_odd = pto.vmax(vreg_old_odd, vreg_new_odd, all_mask_intermediate) - - index_output_0, index_output_1 = pto.vintlv(index_old_even, index_old_odd) - output_even = pto.vcvt(index_output_0, final_cvt_dtype, all_mask_intermediate, part=pto.VcvtPartMode.EVEN) - output_odd = pto.vcvt(index_output_0, final_cvt_dtype, all_mask_intermediate, part=pto.VcvtPartMode.ODD) - output_0, output_1 = pto.vintlv(output_even, output_odd) - - output_0 = pto.vbitcast(output_0, pto.i32) - output_1 = pto.vbitcast(output_1, pto.i32) - - pto.vsts(output_0, dst[0, col:], mask_i32_0) - pto.vsts(output_1, dst[0, col + lanes_i32:], mask_i32_1) - - output_even = pto.vcvt(index_output_1, final_cvt_dtype, all_mask_intermediate, part=pto.VcvtPartMode.EVEN) - output_odd = pto.vcvt(index_output_1, final_cvt_dtype, all_mask_intermediate, part=pto.VcvtPartMode.ODD) - output_0, output_1 = pto.vintlv(output_even, output_odd) - - output_0 = pto.vbitcast(output_0, pto.i32) - output_1 = pto.vbitcast(output_1, pto.i32) - - pto.vsts(output_0, dst[0, col + 2 * lanes_i32:], mask_i32_2) - pto.vsts(output_1, dst[0, col + 3 * lanes_i32:], mask_i32_3) - - return - - -@pto.vkernel( - target="a5", - op="pto.tcolargmax", - dtypes=[ - (pto.f16, pto.f16, pto.i32), - (pto.ui16, pto.ui16, pto.i32), - ], - constraints=[_validate_tcolargmax], - advanced=True, -) -def template_tcolargmax_f16_to_i32(src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - src_valid_rows, src_valid_cols = src.valid_shape - lanes_f16 = pto.get_lanes(pto.f16) - lanes_i32 = pto.get_lanes(pto.i32) - - with pto.vecscope(): - all_mask = pto.make_mask(pto.f16, pto.PAT.ALL) - for col in range(0, src_valid_cols, lanes_f16): - remained = src_valid_cols - col - - mask_f16, _ = pto.make_mask(pto.f16, remained) - - mask_i32_0, remained = pto.make_mask(pto.i32, remained) - mask_i32_1, _ = pto.make_mask(pto.i32, remained) - - index_old = pto.vdup(pto.i16(0), mask_f16) - index_new = pto.vdup(pto.i16(0), mask_f16) - max_vals = pto.vlds(src[0, col:]) - - for row in range(1, src_valid_rows, 1): - index_new = pto.vadds(index_new, pto.i16(1), mask_f16) - new_vals = pto.vlds(src[row, col:]) - gt_mask = pto.vcmp(new_vals, max_vals, all_mask, "gt") - index_old = pto.vsel(index_new, index_old, gt_mask) - max_vals = pto.vmax(max_vals, new_vals, mask_f16) - - index_even = pto.vcvt(index_old, pto.i32, all_mask, part=pto.VcvtPartMode.EVEN) - index_odd = pto.vcvt(index_old, pto.i32, all_mask, part=pto.VcvtPartMode.ODD) - index_lo, index_hi = pto.vintlv(index_even, index_odd) - - pto.vsts(index_lo, dst[0, col:], mask_i32_0) - pto.vsts(index_hi, dst[0, col + lanes_i32:], mask_i32_1) - - return - - -@pto.vkernel( - target="a5", - op="pto.tcolargmax", - dtypes=[ - (pto.f32, pto.f32, pto.i32), - (pto.ui32, pto.ui32, pto.i32), - ], - constraints=[_validate_tcolargmax], - advanced=True, -) -def template_tcolargmax_f32_to_i32(src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - src_valid_rows, src_valid_cols = src.valid_shape - lanes = pto.get_lanes(pto.f32) - - with pto.vecscope(): - remained = src_valid_cols - for col in range(0, src_valid_cols, lanes): - all_mask = pto.make_mask(pto.f32, pto.PAT.ALL) - mask, remained = pto.make_mask(pto.f32, remained) - - index_old = pto.vdup(pto.i32(0), mask) - index_new = pto.vdup(pto.i32(0), mask) - max_vals = pto.vlds(src[0, col:]) - - for row in range(1, src_valid_rows, 1): - index_new = pto.vadds(index_new, pto.i32(1), mask) - new_vals = pto.vlds(src[row, col:]) - gt_mask = pto.vcmp(new_vals, max_vals, all_mask, "gt") - index_old = pto.vsel(index_new, index_old, gt_mask) - max_vals = pto.vmax(max_vals, new_vals, mask) - - pto.vsts(index_old, dst[0, col:], mask) - - return \ No newline at end of file diff --git a/lib/TileOps/tcolargmin_template.py b/lib/TileOps/tcolargmin_template.py deleted file mode 100644 index d5f6808ab1..0000000000 --- a/lib/TileOps/tcolargmin_template.py +++ /dev/null @@ -1,213 +0,0 @@ -# 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 -from pathlib import Path -import tilelang_dsl as pto - - -def _validate_tcolargmin( - src_shape=(), - src_valid_shape=(), - tmp_shape=(), - tmp_valid_shape=(), - dst_shape=(), - dst_valid_shape=(), - src_config=None, - tmp_config=None, - dst_config=None, - src_dtype=None, - tmp_dtype=None, - dst_dtype=None, -): - if src_config is None or tmp_config is None or dst_config is None: - return False - if src_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if tmp_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if dst_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if src_config.s_layout != pto.SLayout.NONE_BOX: - return False - if tmp_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_valid_shape[0] != 1: - return False - if src_dtype != tmp_dtype: - return False - return True - - -@pto.vkernel( - target="a5", - op="pto.tcolargmin", - dtypes=[ - (pto.ui8, pto.ui8, pto.i32), - (pto.i8, pto.i8, pto.i32), - ], - constraints=[_validate_tcolargmin], - advanced=True, -) -def template_tcolargmin_i8_to_i32(src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - src_valid_rows, src_valid_cols = src.valid_shape - src_dtype = src.element_type - lanes_i8 = pto.get_lanes(src_dtype) - lanes_i32 = pto.get_lanes(pto.i32) - - if pto.constexpr(src_dtype == pto.ui8): - intermediate_dtype = pto.ui16 - final_cvt_dtype = pto.ui32 - else: - intermediate_dtype = pto.i16 - final_cvt_dtype = pto.i32 - - lanes_intermediate = pto.get_lanes(intermediate_dtype) - - with pto.vecscope(): - all_mask_b8 = pto.make_mask(src_dtype, pto.PAT.ALL) - all_mask_intermediate = pto.make_mask(intermediate_dtype, pto.PAT.ALL) - - for col in range(0, src_valid_cols, lanes_i8): - remained = src_valid_cols - col - mask_i32_0, remained = pto.make_mask(pto.i32, remained) - mask_i32_1, remained = pto.make_mask(pto.i32, remained) - mask_i32_2, remained = pto.make_mask(pto.i32, remained) - mask_i32_3, _ = pto.make_mask(pto.i32, remained) - - index_old_even = pto.vdup(intermediate_dtype(0), all_mask_intermediate) - index_old_odd = pto.vdup(intermediate_dtype(0), all_mask_intermediate) - index_new_even = pto.vdup(intermediate_dtype(0), all_mask_intermediate) - index_new_odd = pto.vdup(intermediate_dtype(0), all_mask_intermediate) - - vreg_old = pto.vlds(src[0, col:]) - vreg_old_even = pto.vcvt(vreg_old, intermediate_dtype, all_mask_b8, part=pto.VcvtPartMode.EVEN) - vreg_old_odd = pto.vcvt(vreg_old, intermediate_dtype, all_mask_b8, part=pto.VcvtPartMode.ODD) - - for row in range(1, src_valid_rows, 1): - index_new_even = pto.vadds(index_new_even, intermediate_dtype(1), all_mask_intermediate) - index_new_odd = pto.vadds(index_new_odd, intermediate_dtype(1), all_mask_intermediate) - vreg_new = pto.vlds(src[row, col:]) - vreg_new_even = pto.vcvt(vreg_new, intermediate_dtype, all_mask_b8, part=pto.VcvtPartMode.EVEN) - vreg_new_odd = pto.vcvt(vreg_new, intermediate_dtype, all_mask_b8, part=pto.VcvtPartMode.ODD) - - select_even = pto.vcmp(vreg_new_even, vreg_old_even, all_mask_intermediate, "lt") - select_odd = pto.vcmp(vreg_new_odd, vreg_old_odd, all_mask_intermediate, "lt") - - index_old_even = pto.vsel(index_new_even, index_old_even, select_even) - index_old_odd = pto.vsel(index_new_odd, index_old_odd, select_odd) - - vreg_old_even = pto.vmin(vreg_old_even, vreg_new_even, all_mask_intermediate) - vreg_old_odd = pto.vmin(vreg_old_odd, vreg_new_odd, all_mask_intermediate) - - index_output_0, index_output_1 = pto.vintlv(index_old_even, index_old_odd) - output_even = pto.vcvt(index_output_0, final_cvt_dtype, all_mask_intermediate, part=pto.VcvtPartMode.EVEN) - output_odd = pto.vcvt(index_output_0, final_cvt_dtype, all_mask_intermediate, part=pto.VcvtPartMode.ODD) - output_0, output_1 = pto.vintlv(output_even, output_odd) - - output_0 = pto.vbitcast(output_0, pto.i32) - output_1 = pto.vbitcast(output_1, pto.i32) - - pto.vsts(output_0, dst[0, col:], mask_i32_0) - pto.vsts(output_1, dst[0, col + lanes_i32:], mask_i32_1) - - output_even = pto.vcvt(index_output_1, final_cvt_dtype, all_mask_intermediate, part=pto.VcvtPartMode.EVEN) - output_odd = pto.vcvt(index_output_1, final_cvt_dtype, all_mask_intermediate, part=pto.VcvtPartMode.ODD) - output_0, output_1 = pto.vintlv(output_even, output_odd) - - output_0 = pto.vbitcast(output_0, pto.i32) - output_1 = pto.vbitcast(output_1, pto.i32) - - pto.vsts(output_0, dst[0, col + 2 * lanes_i32:], mask_i32_2) - pto.vsts(output_1, dst[0, col + 3 * lanes_i32:], mask_i32_3) - - return - - -@pto.vkernel( - target="a5", - op="pto.tcolargmin", - dtypes=[ - (pto.f16, pto.f16, pto.i32), - (pto.ui16, pto.ui16, pto.i32), - ], - constraints=[_validate_tcolargmin], - advanced=True, -) -def template_tcolargmin_f16_to_i32(src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - src_valid_rows, src_valid_cols = src.valid_shape - lanes_f16 = pto.get_lanes(pto.f16) - lanes_i32 = pto.get_lanes(pto.i32) - - with pto.vecscope(): - all_mask = pto.make_mask(pto.f16, pto.PAT.ALL) - for col in range(0, src_valid_cols, lanes_f16): - remained = src_valid_cols - col - - mask_f16, _ = pto.make_mask(pto.f16, remained) - - mask_i32_0, remained = pto.make_mask(pto.i32, remained) - mask_i32_1, _ = pto.make_mask(pto.i32, remained) - - index_old = pto.vdup(pto.i16(0), mask_f16) - index_new = pto.vdup(pto.i16(0), mask_f16) - min_vals = pto.vlds(src[0, col:]) - - for row in range(1, src_valid_rows, 1): - index_new = pto.vadds(index_new, pto.i16(1), mask_f16) - new_vals = pto.vlds(src[row, col:]) - lt_mask = pto.vcmp(new_vals, min_vals, all_mask, "lt") - index_old = pto.vsel(index_new, index_old, lt_mask) - min_vals = pto.vmin(min_vals, new_vals, mask_f16) - - index_even = pto.vcvt(index_old, pto.i32, all_mask, part=pto.VcvtPartMode.EVEN) - index_odd = pto.vcvt(index_old, pto.i32, all_mask, part=pto.VcvtPartMode.ODD) - index_lo, index_hi = pto.vintlv(index_even, index_odd) - - pto.vsts(index_lo, dst[0, col:], mask_i32_0) - pto.vsts(index_hi, dst[0, col + lanes_i32:], mask_i32_1) - - return - - -@pto.vkernel( - target="a5", - op="pto.tcolargmin", - dtypes=[ - (pto.f32, pto.f32, pto.i32), - (pto.ui32, pto.ui32, pto.i32), - ], - constraints=[_validate_tcolargmin], - advanced=True, -) -def template_tcolargmin_f32_to_i32(src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - src_valid_rows, src_valid_cols = src.valid_shape - lanes = pto.get_lanes(pto.f32) - - with pto.vecscope(): - remained = src_valid_cols - for col in range(0, src_valid_cols, lanes): - all_mask = pto.make_mask(pto.f32, pto.PAT.ALL) - mask, remained = pto.make_mask(pto.f32, remained) - - index_old = pto.vdup(pto.i32(0), mask) - index_new = pto.vdup(pto.i32(0), mask) - min_vals = pto.vlds(src[0, col:]) - - for row in range(1, src_valid_rows, 1): - index_new = pto.vadds(index_new, pto.i32(1), mask) - new_vals = pto.vlds(src[row, col:]) - lt_mask = pto.vcmp(new_vals, min_vals, all_mask, "lt") - index_old = pto.vsel(index_new, index_old, lt_mask) - min_vals = pto.vmin(min_vals, new_vals, mask) - - pto.vsts(index_old, dst[0, col:], mask) - - return diff --git a/lib/TileOps/tcolexpand_template.py b/lib/TileOps/tcolexpand_template.py deleted file mode 100644 index 5d20fcbe1b..0000000000 --- a/lib/TileOps/tcolexpand_template.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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. -"""TileLang DSL template for pto.tcolexpand""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tcolexpand" -) -def template_tcolexpand(src0: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[0, col:]) - pto.vsts(lhs, dst[row, col:], mask) - return diff --git a/lib/TileOps/tcolexpandadd_template.py b/lib/TileOps/tcolexpandadd_template.py deleted file mode 100644 index 287be93dc7..0000000000 --- a/lib/TileOps/tcolexpandadd_template.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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. -"""TileLang DSL template for pto.tcolexpandadd""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tcolexpandadd" -) -def template_tcolexpandadd(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[0, col:]) - summed = pto.vadd(lhs, rhs, mask) - pto.vsts(summed, dst[row, col:], mask) - return diff --git a/lib/TileOps/tcolexpanddiv_template.py b/lib/TileOps/tcolexpanddiv_template.py deleted file mode 100644 index 7092ba2f51..0000000000 --- a/lib/TileOps/tcolexpanddiv_template.py +++ /dev/null @@ -1,54 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tcolexpanddiv with IEEE 754 high-precision support - -Divide each column of src0 by a per-column scalar from src1[0, col]. -Semantics: dst[row, col] = src0[row, col] / src1[0, col] -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - -# Import shared high-precision division algorithms -from div_hp import _div_ieee754_f32_impl, _div_ieee754_f16_impl - - -@pto.vkernel( - target="a5", - op="pto.tcolexpanddiv" -) -def template_tcolexpanddiv(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - """Template for pto.tcolexpanddiv with optional high-precision mode.""" - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - precision_type = pto.get_op_attr("precisionType", "default") - if pto.constexpr(precision_type == "high_precision"): - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[0, col:]) - if pto.constexpr(dtype == pto.f32): - result = _div_ieee754_f32_impl(lhs, rhs, mask) - else: # dtype == pto.f16 (guaranteed by MLIR validation) - result = _div_ieee754_f16_impl(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) - else: - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[0, col:]) - result = pto.vdiv(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tcolexpandexpdif_template.py b/lib/TileOps/tcolexpandexpdif_template.py deleted file mode 100644 index 0ae28ee0c6..0000000000 --- a/lib/TileOps/tcolexpandexpdif_template.py +++ /dev/null @@ -1,57 +0,0 @@ -# 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. -"""TileLang DSL template for pto.tcolexpandexpdif""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tcolexpandexpdif", - dtypes=[ - (pto.f16, pto.f16, pto.f16), - ], -) -def template_tcolexpandexpdif_f16(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[0, col:]) - diff = pto.vsub(lhs, rhs, mask) - result = pto.vexp(diff, mask) - pto.vsts(result, dst[row, col:], mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tcolexpandexpdif", - dtypes=[ - (pto.f32, pto.f32, pto.f32), - ], -) -def template_tcolexpandexpdif_f32(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[0, col:]) - result = pto.vexpdif(lhs, rhs, mask, pto.VcvtPartMode.ODD) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tcolexpandmax_template.py b/lib/TileOps/tcolexpandmax_template.py deleted file mode 100644 index 79f3699b71..0000000000 --- a/lib/TileOps/tcolexpandmax_template.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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. -"""TileLang DSL template for pto.tcolexpandmax""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tcolexpandmax" -) -def template_tcolexpandmax(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[0, col:]) - result = pto.vmax(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tcolexpandmin_template.py b/lib/TileOps/tcolexpandmin_template.py deleted file mode 100644 index 054b35dfc0..0000000000 --- a/lib/TileOps/tcolexpandmin_template.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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. -"""TileLang DSL template for pto.tcolexpandmin""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tcolexpandmin" -) -def template_tcolexpandmin(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[0, col:]) - result = pto.vmin(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tcolexpandmul_template.py b/lib/TileOps/tcolexpandmul_template.py deleted file mode 100644 index 5dcacfa912..0000000000 --- a/lib/TileOps/tcolexpandmul_template.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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. -"""TileLang DSL template for pto.tcolexpandmul""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tcolexpandmul" -) -def template_tcolexpandmul(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[0, col:]) - result = pto.vmul(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tcolexpandsub_template.py b/lib/TileOps/tcolexpandsub_template.py deleted file mode 100644 index f46bbaf110..0000000000 --- a/lib/TileOps/tcolexpandsub_template.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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. -"""TileLang DSL template for pto.tcolexpandsub""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tcolexpandsub" -) -def template_tcolexpandsub(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[0, col:]) - result = pto.vsub(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tcolmax_template.py b/lib/TileOps/tcolmax_template.py deleted file mode 100644 index f08df7edac..0000000000 --- a/lib/TileOps/tcolmax_template.py +++ /dev/null @@ -1,56 +0,0 @@ -# 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 -from pathlib import Path -import tilelang_dsl as pto - -def _validate_tcolmax( - src_shape=(), - src_valid_shape=(), - dst_shape=(), - dst_valid_shape=(), - src_config=None, - dst_config=None -): - if src_config is None or dst_config is None: - return False - if src_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if dst_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if src_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_valid_shape[0] != 1: - return False - return True - -@pto.vkernel( - target="a5", - op="pto.tcolmax", - constraints=[_validate_tcolmax] -) -def template_tcolmax(src: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = src.valid_shape - - lanes = pto.get_lanes(dtype) - remained = valid_cols - - for col_chunk in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - - acc = pto.vlds(src[0, col_chunk:]) - for row in range(1, valid_rows, 1): - row_vec = pto.vlds(src[row, col_chunk:]) - acc = pto.vmax(acc, row_vec, mask) - pto.vsts(acc, dst[0, col_chunk:], mask) - - return diff --git a/lib/TileOps/tcolmin_template.py b/lib/TileOps/tcolmin_template.py deleted file mode 100644 index 2a36dcdd5e..0000000000 --- a/lib/TileOps/tcolmin_template.py +++ /dev/null @@ -1,56 +0,0 @@ -# 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 -from pathlib import Path -import tilelang_dsl as pto - -def _validate_tcolmin( - src_shape=(), - src_valid_shape=(), - dst_shape=(), - dst_valid_shape=(), - src_config=None, - dst_config=None -): - if src_config is None or dst_config is None: - return False - if src_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if dst_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if src_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_valid_shape[0] != 1: - return False - return True - -@pto.vkernel( - target="a5", - op="pto.tcolmin", - constraints=[_validate_tcolmin] -) -def template_tcolmin(src: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = src.valid_shape - - lanes = pto.get_lanes(dtype) - remained = valid_cols - - for col_chunk in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - - acc = pto.vlds(src[0, col_chunk:]) - for row in range(1, valid_rows, 1): - row_vec = pto.vlds(src[row, col_chunk:]) - acc = pto.vmin(acc, row_vec, mask) - pto.vsts(acc, dst[0, col_chunk:], mask) - - return diff --git a/lib/TileOps/tcolprod_template.py b/lib/TileOps/tcolprod_template.py deleted file mode 100644 index 4ebb99f484..0000000000 --- a/lib/TileOps/tcolprod_template.py +++ /dev/null @@ -1,56 +0,0 @@ -# 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 -from pathlib import Path -import tilelang_dsl as pto - -def _validate_tcolprod( - src_shape=(), - src_valid_shape=(), - dst_shape=(), - dst_valid_shape=(), - src_config=None, - dst_config=None -): - if src_config is None or dst_config is None: - return False - if src_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if dst_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if src_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_valid_shape[0] != 1: - return False - return True - -@pto.vkernel( - target="a5", - op="pto.tcolprod", - constraints=[_validate_tcolprod] -) -def template_tcolprod(src: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = src.valid_shape - - lanes = pto.get_lanes(dtype) - remained = valid_cols - - for col_chunk in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - - acc = pto.vlds(src[0, col_chunk:]) - for row in range(1, valid_rows, 1): - row_vec = pto.vlds(src[row, col_chunk:]) - acc = pto.vmul(acc, row_vec, mask) - pto.vsts(acc, dst[0, col_chunk:], mask) - - return diff --git a/lib/TileOps/tcolsum_template.py b/lib/TileOps/tcolsum_template.py deleted file mode 100644 index b187b45a33..0000000000 --- a/lib/TileOps/tcolsum_template.py +++ /dev/null @@ -1,57 +0,0 @@ -# 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 -from pathlib import Path -import tilelang_dsl as pto - -def _validate_tcolsum( - src_shape=(), - src_valid_shape=(), - dst_shape=(), - dst_valid_shape=(), - src_config=None, - dst_config=None -): - if src_config is None or dst_config is None: - return False - if src_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if dst_config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if src_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_valid_shape[0] != 1: - return False - return True - -# Todo: This is the basic implementation. Later the binary colsum algorithm should be implemented also. -@pto.vkernel( - target="a5", - op="pto.tcolsum", - constraints=[_validate_tcolsum] -) -def template_tcolsum(src: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = src.valid_shape - - lanes = pto.get_lanes(dtype) - remained = valid_cols - - for col_chunk in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - - acc = pto.vlds(src[0, col_chunk:]) - for row in range(1, valid_rows, 1): - row_vec = pto.vlds(src[row, col_chunk:]) - acc = pto.vadd(acc, row_vec, mask) - pto.vsts(acc, dst[0, col_chunk:], mask) - - return diff --git a/lib/TileOps/tcvt_template.py b/lib/TileOps/tcvt_template.py deleted file mode 100644 index 9906e55e82..0000000000 --- a/lib/TileOps/tcvt_template.py +++ /dev/null @@ -1,1348 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tcvt.""" - -import tilelang_dsl as pto - - -def _config_value(config, name): - if config is None: - return None - if isinstance(config, dict): - return config.get(name) - return getattr(config, name, None) - - -def _matches_layout(value, expected, expected_name): - if value is None: - return False - return value == expected or value == expected_name or str(value).lower().endswith(expected_name) - - -def _supports_basic_rowwise_tcvt( - src=None, - dst=None, - src_shape=(), - src_valid_shape=(), - dst_shape=(), - dst_valid_shape=(), -): - if tuple(src_shape) != tuple(dst_shape): - return False - if tuple(src_valid_shape) != tuple(dst_valid_shape): - return False - if len(src_shape) != 2 or len(dst_shape) != 2: - return False - if src is None or dst is None: - return False - src_config = src.config - dst_config = dst.config - if src_config is None or dst_config is None: - return False - if not _matches_layout(_config_value(src_config, "b_layout"), pto.BLayout.ROW_MAJOR, "row_major"): - return False - if not _matches_layout(_config_value(dst_config, "b_layout"), pto.BLayout.ROW_MAJOR, "row_major"): - return False - if not _matches_layout(_config_value(src_config, "s_layout"), pto.SLayout.NONE_BOX, "none_box"): - return False - if not _matches_layout(_config_value(dst_config, "s_layout"), pto.SLayout.NONE_BOX, "none_box"): - return False - return True - - -def _supports_bf16_to_fp4_rowwise_tcvt( - src=None, - dst=None, - src_shape=(), - src_valid_shape=(), - dst_shape=(), - dst_valid_shape=(), -): - if len(src_shape) != 2 or len(dst_shape) != 2: - return False - if src_shape[0] != dst_shape[0] or src_shape[1] != dst_shape[1] * 2: - return False - if src_valid_shape[0] != dst_valid_shape[0] or src_valid_shape[1] != dst_valid_shape[1] * 2: - return False - if src is None or dst is None: - return False - src_config = src.config - dst_config = dst.config - if src_config is None or dst_config is None: - return False - if not _matches_layout(_config_value(src_config, "b_layout"), pto.BLayout.ROW_MAJOR, "row_major"): - return False - if not _matches_layout(_config_value(dst_config, "b_layout"), pto.BLayout.ROW_MAJOR, "row_major"): - return False - if not _matches_layout(_config_value(src_config, "s_layout"), pto.SLayout.NONE_BOX, "none_box"): - return False - if not _matches_layout(_config_value(dst_config, "s_layout"), pto.SLayout.NONE_BOX, "none_box"): - return False - return True - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f32, pto.f16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_f32_to_f16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - full_mask = pto.make_mask(pto.f32, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f32)): - store_mask, remained = pto.make_mask(pto.f32, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.f16, - full_mask, - rnd=rnd, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B32) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f32, pto.i32), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_f32_to_i32(src: pto.Tile, dst: pto.Tile): - dst_dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dst_dtype)): - mask, remained = pto.make_mask(dst_dtype, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - dst_dtype, - mask, - rnd=rnd, - sat=pto.VcvtSatMode.SAT, - ) - pto.vsts(converted, dst[row, col:], mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.i32, pto.f32), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_i32_to_f32(src: pto.Tile, dst: pto.Tile): - dst_dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dst_dtype)): - mask, remained = pto.make_mask(dst_dtype, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - dst_dtype, - mask, - rnd=rnd, - ) - pto.vsts(converted, dst[row, col:], mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f16, pto.f32), - (pto.bf16, pto.f32), - (pto.i16, pto.f32), - (pto.i16, pto.i32), - (pto.i16, pto.ui32), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_16_to_32(src: pto.Tile, dst: pto.Tile): - src_dtype = src.element_type - dst_dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(src_dtype, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dst_dtype)): - store_mask, remained = pto.make_mask(dst_dtype, remained) - vec = pto.vlds(src[row, col:], dist=pto.VLoadDist.UNPK_B16) - converted = pto.vcvt( - vec, - dst_dtype, - full_mask, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.i32, pto.i64), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_i32_to_i64(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(pto.i32, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols * 2 # i64 requires double the mask - for col in range(0, valid_cols, pto.get_lanes(pto.i64)): - store_mask, remained = pto.make_mask(pto.i64, remained) - vec = pto.vlds(src[row, col:], dist=pto.VLoadDist.UNPK_B32) - converted = pto.vcvt( - vec, - pto.i64, - full_mask, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.NORM_B32) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.ui8, pto.f16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_ui8_to_f16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(pto.ui8, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f16)): - store_mask, remained = pto.make_mask(pto.f16, remained) - vec = pto.vlds(src[row, col:], dist=pto.VLoadDist.UNPK_B8) - converted = pto.vcvt( - vec, - pto.f16, - full_mask, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.ui8, pto.ui16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_ui8_to_ui16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(pto.ui8, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.ui16)): - store_mask, remained = pto.make_mask(pto.ui16, remained) - vec = pto.vlds(src[row, col:], dist=pto.VLoadDist.UNPK_B8) - converted = pto.vcvt( - vec, - pto.ui16, - full_mask, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.si8, pto.f16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_si8_to_f16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(pto.si8, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f16)): - store_mask, remained = pto.make_mask(pto.f16, remained) - vec = pto.vlds(src[row, col:], dist=pto.VLoadDist.UNPK_B8) - converted = pto.vcvt( - vec, - pto.f16, - full_mask, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.si8, pto.si16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_si8_to_si16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(pto.si8, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.si16)): - store_mask, remained = pto.make_mask(pto.si16, remained) - vec = pto.vlds(src[row, col:], dist=pto.VLoadDist.UNPK_B8) - converted = pto.vcvt( - vec, - pto.si16, - full_mask, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.NORM_B16) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.si8, pto.i32), - ], - constraints=[_supports_basic_rowwise_tcvt], - advanced=True, -) -def template_tcvt_si8_to_i32(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - b8_mask = pto.make_mask(pto.ui8, pto.PAT.ALL) - v_zero = pto.vdup(pto.ui8(0), b8_mask) - lanes_i32 = pto.get_lanes(pto.i32) - for row in range(0, valid_rows, 1): - remained = valid_cols - next_remained = 0 - if valid_cols > lanes_i32: - next_remained = valid_cols - lanes_i32 - for col in range(0, valid_cols, pto.get_lanes(pto.i16)): - mask_b16_cur, remained = pto.make_mask(pto.i16, remained) - mask_b16_next, next_remained = pto.make_mask(pto.i16, next_remained) - mask_b32_cur = pto.punpack(mask_b16_cur, pto.PredicatePart.LOWER) - mask_b32_next = pto.punpack(mask_b16_next, pto.PredicatePart.LOWER) - vec_si8_0 = pto.vlds(src[row, col:], dist=pto.VLoadDist.UNPK_B8) - vec_ui8_0 = pto.vbitcast(vec_si8_0, pto.ui8) - vec_ui8_1, vec_ui8_2 = pto.vintlv(vec_ui8_0, v_zero) - vec_si8_1 = pto.vbitcast(vec_ui8_1, pto.si8) - vec_si8_2 = pto.vbitcast(vec_ui8_2, pto.si8) - output_0 = pto.vcvt(vec_si8_1, pto.i32, b8_mask, part=pto.VcvtPartMode.P0) - output_1 = pto.vcvt(vec_si8_2, pto.i32, b8_mask, part=pto.VcvtPartMode.P0) - pto.vsts(output_0, dst[row, col:], mask_b32_cur, dist=pto.VStoreDist.NORM_B32) - pto.vsts(output_1, dst[row, col + lanes_i32:], mask_b32_next, dist=pto.VStoreDist.NORM_B32) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f32, pto.f32), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_f32_to_f32(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f32)): - mask, remained = pto.make_mask(pto.f32, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vtrc(vec, mask, rnd=rnd) - pto.vsts(converted, dst[row, col:], mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f16, pto.i32), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_f16_to_i32(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - full_mask = pto.make_mask(pto.f16, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.i32)): - store_mask, remained = pto.make_mask(pto.i32, remained) - vec = pto.vlds(src[row, col:], dist=pto.VLoadDist.UNPK_B16) - converted = pto.vcvt( - vec, - pto.i32, - full_mask, - rnd=rnd, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.i16, pto.f16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_i16_to_f16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - full_mask = pto.make_mask(pto.i16, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.i16)): - store_mask, remained = pto.make_mask(pto.f16, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.f16, - full_mask, - rnd=rnd, - ) - pto.vsts(converted, dst[row, col:], store_mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.i64, pto.f32), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_i64_to_f32(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - for row in range(0, valid_rows, 1): - remained = valid_cols * 2 # i64 requires double the mask - full_mask, _ = pto.make_mask(pto.i64, remained) - for col in range(0, valid_cols, pto.get_lanes(pto.i64)): - store_mask, remained = pto.make_mask(pto.f32, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.f32, - full_mask, - rnd=rnd, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B64) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.i16, pto.ui8), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_i16_to_ui8(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(pto.i16, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.i16)): - store_mask, remained = pto.make_mask(pto.i16, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.ui8, - full_mask, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B16) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.i32, pto.i16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_i32_to_i16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(pto.i32, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.i32)): - store_mask, remained = pto.make_mask(pto.i32, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.i16, - full_mask, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B32) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.i32, pto.ui16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_i32_to_ui16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(pto.i32, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.i32)): - store_mask, remained = pto.make_mask(pto.i32, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.ui16, - full_mask, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B32) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.i32, pto.ui8), - ], - constraints=[_supports_basic_rowwise_tcvt], - advanced=True, -) -def template_tcvt_i32_to_ui8(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(pto.i32, pto.PAT.ALL) - idx_mask_b8 = pto.pset_b8(pto.PAT.ALL) - idx_mask_b16 = pto.pbitcast(idx_mask_b8, pto.mask_b16) - lanes_i32 = pto.get_lanes(pto.i32) - v_idx = pto.vci(pto.i8(0), pto.OrderMode.ASC) - v_idx_i16 = pto.vbitcast(v_idx, pto.i16) - v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) - v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) - for row in range(0, valid_rows, 1): - mask_len_tail = valid_cols % lanes_i32 - if valid_cols % lanes_i32 == 0: - mask_len_tail = lanes_i32 - for col in range(0, valid_cols, lanes_i32): - mask_len = lanes_i32 - if valid_cols < lanes_i32 or col == valid_cols - lanes_i32: - mask_len = mask_len_tail - store_mask, _ = pto.make_mask(pto.ui8, mask_len) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.ui8, - full_mask, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.P0, - ) - result = pto.vselr(converted, v_idx_ui8) - pto.mem_bar(pto.BarrierType.VST_VST) - pto.vsts(result, dst[row, col:], store_mask, dist=pto.VStoreDist.NORM_B8) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.ui32, pto.i16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_ui32_to_i16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(pto.ui32, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.ui32)): - store_mask, remained = pto.make_mask(pto.ui32, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.i16, - full_mask, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B32) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.ui32, pto.ui16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_ui32_to_ui16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(pto.ui32, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.ui32)): - store_mask, remained = pto.make_mask(pto.ui32, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.ui16, - full_mask, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B32) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.ui32, pto.ui8), - ], - constraints=[_supports_basic_rowwise_tcvt], - advanced=True, -) -def template_tcvt_ui32_to_ui8(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - full_mask = pto.make_mask(pto.ui32, pto.PAT.ALL) - idx_mask_b8 = pto.pset_b8(pto.PAT.ALL) - idx_mask_b16 = pto.pbitcast(idx_mask_b8, pto.mask_b16) - lanes_ui32 = pto.get_lanes(pto.ui32) - v_idx = pto.vci(pto.i8(0), pto.OrderMode.ASC) - v_idx_i16 = pto.vbitcast(v_idx, pto.i16) - v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) - v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) - for row in range(0, valid_rows, 1): - mask_len_tail = valid_cols % lanes_ui32 - if valid_cols % lanes_ui32 == 0: - mask_len_tail = lanes_ui32 - for col in range(0, valid_cols, lanes_ui32): - mask_len = lanes_ui32 - if valid_cols < lanes_ui32 or col == valid_cols - lanes_ui32: - mask_len = mask_len_tail - store_mask, _ = pto.make_mask(pto.ui8, mask_len) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.ui8, - full_mask, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.P0, - ) - result = pto.vselr(converted, v_idx_ui8) - pto.mem_bar(pto.BarrierType.VST_VST) - pto.vsts(result, dst[row, col:], store_mask, dist=pto.VStoreDist.NORM_B8) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.i64, pto.i32), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_i64_to_i32(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - for row in range(0, valid_rows, 1): - remained = valid_cols * 2 # i64 requires double the mask - full_mask, _ = pto.make_mask(pto.i64, remained) - for col in range(0, valid_cols, pto.get_lanes(pto.i64)): - store_mask, remained = pto.make_mask(pto.i32, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.i32, - full_mask, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B64) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f32, pto.bf16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_f32_to_bf16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - full_mask = pto.make_mask(pto.f32, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f32)): - store_mask, remained = pto.make_mask(pto.f32, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.bf16, - full_mask, - rnd=rnd, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B32) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f32, pto.i64), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_f32_to_i64(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - full_mask = pto.make_mask(pto.f32, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols * 2 # i64 requires double the mask - for col in range(0, valid_cols, pto.get_lanes(pto.i64)): - store_mask, remained = pto.make_mask(pto.i64, remained) - vec = pto.vlds(src[row, col:], dist=pto.VLoadDist.UNPK_B32) - converted = pto.vcvt( - vec, - pto.i64, - full_mask, - rnd=rnd, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.NORM_B32) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f16, pto.ui8), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_f16_to_ui8(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - full_mask = pto.make_mask(pto.f16, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f16)): - store_mask, remained = pto.make_mask(pto.f16, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.ui8, - full_mask, - rnd=rnd, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B16) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.bf16, pto.i32), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_bf16_to_i32(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - full_mask = pto.make_mask(pto.bf16, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.i32)): - store_mask, remained = pto.make_mask(pto.i32, remained) - vec = pto.vlds(src[row, col:], dist=pto.VLoadDist.UNPK_B16) - converted = pto.vcvt( - vec, - pto.i32, - full_mask, - rnd=rnd, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], store_mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.bf16, pto.f16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_bf16_to_f16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - full_mask = pto.make_mask(pto.bf16, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.bf16)): - store_mask, remained = pto.make_mask(pto.f16, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - pto.f16, - full_mask, - sat=pto.VcvtSatMode.SAT, - rnd=rnd, - ) - pto.vsts(converted, dst[row, col:], store_mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f32, pto.i16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_f32_to_i16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - full_mask = pto.make_mask(pto.f32, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f32)): - store_mask, remained = pto.make_mask(pto.f32, remained) - vec_f32 = pto.vlds(src[row, col:]) - # sat=OFF NonSatTorch - vec_i32 = pto.vcvt( - vec_f32, - pto.i32, - full_mask, - rnd=rnd, - sat=pto.VcvtSatMode.NOSAT, - ) - vec_i16 = pto.vcvt( - vec_i32, - pto.i16, - full_mask, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(vec_i16, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B32) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f16, pto.i16), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_f16_to_i16(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - full_mask_b16 = pto.make_mask(pto.f16, pto.PAT.ALL) - full_mask_b32 = pto.make_mask(pto.i32, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f32)): - store_mask, remained = pto.make_mask(pto.i32, remained) - vec_f16 = pto.vlds(src[row, col:], dist=pto.VLoadDist.UNPK_B16) - # sat=OFF NonSatTorch - vec_i32 = pto.vcvt( - vec_f16, - pto.i32, - full_mask_b16, - rnd=rnd, - part=pto.VcvtPartMode.EVEN, - ) - vec_i16 = pto.vcvt( - vec_i32, - pto.i16, - full_mask_b32, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(vec_i16, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B32) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f16, pto.si8), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_f16_to_si8(src: pto.Tile, dst: pto.Tile): - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - - lanes_f16 = pto.get_lanes(pto.f16) - pg = pto.make_mask(pto.f16, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes_f16): - full_mask, _ = pto.make_mask(pto.f16, lanes_f16) - store_mask, remained = pto.make_mask(pto.f16, remained) - vec_f16 = pto.vlds(src[row, col:]) - # sat=OFF NonSatTorch - vec_i16 = pto.vcvt( - vec_f16, - pto.i16, - full_mask, - rnd=rnd, - sat=pto.VcvtSatMode.NOSAT, - ) - v_mask = pto.vdup(pto.i16(255), pg) - vec_i16_and = pto.vand(vec_i16, v_mask, store_mask) - vec_f16_temp = pto.vcvt( - vec_i16_and, - pto.f16, - full_mask, - rnd=rnd, - ) - vec_si8 = pto.vcvt( - vec_f16_temp, - pto.si8, - full_mask, - rnd=rnd, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(vec_si8, dst[row, col:], store_mask, dist=pto.VStoreDist.PK_B16) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f32, pto.f8e4m3), - (pto.f32, pto.f8e5m2), - ], - constraints=[_supports_basic_rowwise_tcvt], - advanced=True, -) -def template_tcvt_f32_to_fp8(src: pto.Tile, dst: pto.Tile): - dst_dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - lanes_f32 = pto.get_lanes(pto.f32) - src_mask = pto.make_mask(pto.f32, pto.PAT.ALL) - idx_mask_b8 = pto.pset_b8(pto.PAT.ALL) - idx_mask_b16 = pto.pbitcast(idx_mask_b8, pto.mask_b16) - v_idx = pto.vci(pto.i8(0), pto.OrderMode.ASC) - v_idx_i16 = pto.vbitcast(v_idx, pto.i16) - v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) - v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) - for row in range(0, valid_rows, 1): - for col in range(0, valid_cols, lanes_f32): - mask_len = valid_cols - col - if mask_len > lanes_f32: - mask_len = lanes_f32 - dst_mask, _ = pto.make_mask(dst_dtype, mask_len) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - dst_dtype, - src_mask, - rnd=rnd, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.P0, - ) - result = pto.vselr(converted, v_idx_ui8) - pto.mem_bar(pto.BarrierType.VST_VST) - pto.vsts(result, dst[row, col:], dst_mask, dist=pto.VStoreDist.NORM_B8) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f32, pto.hif8), - ], - constraints=[_supports_basic_rowwise_tcvt], - advanced=True, -) -def template_tcvt_f32_to_hif8(src: pto.Tile, dst: pto.Tile): - dst_dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.A - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - lanes_f32 = pto.get_lanes(pto.f32) - src_mask = pto.make_mask(pto.f32, pto.PAT.ALL) - idx_mask_b8 = pto.pset_b8(pto.PAT.ALL) - idx_mask_b16 = pto.pbitcast(idx_mask_b8, pto.mask_b16) - v_idx = pto.vci(pto.i8(0), pto.OrderMode.ASC) - v_idx_i16 = pto.vbitcast(v_idx, pto.i16) - v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) - v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) - for row in range(0, valid_rows, 1): - for col in range(0, valid_cols, lanes_f32): - mask_len = valid_cols - col - if mask_len > lanes_f32: - mask_len = lanes_f32 - dst_mask, _ = pto.make_mask(dst_dtype, mask_len) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - dst_dtype, - src_mask, - rnd=rnd, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.P0, - ) - result = pto.vselr(converted, v_idx_ui8) - pto.mem_bar(pto.BarrierType.VST_VST) - pto.vsts(result, dst[row, col:], dst_mask, dist=pto.VStoreDist.NORM_B8) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.f16, pto.hif8), - ], - constraints=[_supports_basic_rowwise_tcvt], -) -def template_tcvt_f16_to_hif8(src: pto.Tile, dst: pto.Tile): - dst_dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.A - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - src_mask = pto.make_mask(pto.f16, pto.PAT.ALL) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f16)): - dst_mask, remained = pto.make_mask(pto.f16, remained) - vec = pto.vlds(src[row, col:]) - converted = pto.vcvt( - vec, - dst_dtype, - src_mask, - rnd=rnd, - sat=pto.VcvtSatMode.NOSAT, - part=pto.VcvtPartMode.EVEN, - ) - pto.vsts(converted, dst[row, col:], dst_mask, dist=pto.VStoreDist.PK_B16) - return - - -@pto.vkernel( - target="a5", - op="pto.tcvt", - dtypes=[ - (pto.bf16, pto.f4e1m2x2), - (pto.bf16, pto.f4e2m1x2), - ], - constraints=[_supports_bf16_to_fp4_rowwise_tcvt], - advanced=True, -) -def template_tcvt_bf16_to_fp4(src: pto.Tile, dst: pto.Tile): - dst_dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - round_mode = pto.get_op_attr("round_mode", "RINT") - rnd = pto.VcvtRoundMode.R - if pto.constexpr(round_mode == "ROUND"): - rnd = pto.VcvtRoundMode.A - elif pto.constexpr(round_mode == "FLOOR"): - rnd = pto.VcvtRoundMode.F - elif pto.constexpr(round_mode == "CEIL"): - rnd = pto.VcvtRoundMode.C - elif pto.constexpr(round_mode == "TRUNC"): - rnd = pto.VcvtRoundMode.Z - elif pto.constexpr(round_mode == "ODD"): - rnd = pto.VcvtRoundMode.O - lanes_bf16 = pto.get_lanes(pto.bf16) - dst_chunk_cols = lanes_bf16 // 2 - idx_mask_b8 = pto.pset_b8(pto.PAT.ALL) - idx_mask_b16 = pto.pbitcast(idx_mask_b8, pto.mask_b16) - v_idx = pto.vci(pto.i8(0), pto.OrderMode.ASC) - v_idx_i16 = pto.vbitcast(v_idx, pto.i16) - v_idx_i16 = pto.vmuls(v_idx_i16, pto.i16(4), idx_mask_b16) - v_idx_ui8 = pto.vbitcast(v_idx_i16, pto.ui8) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, dst_chunk_cols): - dst_mask, remained = pto.make_mask(dst_dtype, remained) - store_len = valid_cols - col - if store_len > dst_chunk_cols: - store_len = dst_chunk_cols - src_mask, _ = pto.make_mask(pto.bf16, store_len * 2) - vec = pto.vlds(src[row, col * 2:]) - converted = pto.vcvt( - vec, - dst_dtype, - src_mask, - rnd=rnd, - part=pto.VcvtPartMode.P0, - ) - result = pto.vselr(converted, v_idx_ui8) - pto.mem_bar(pto.BarrierType.VST_VST) - pto.vsts(result, dst[row, col:], dst_mask, dist=pto.VStoreDist.NORM_B8) - return diff --git a/lib/TileOps/tdiv_template.py b/lib/TileOps/tdiv_template.py deleted file mode 100644 index 4703915095..0000000000 --- a/lib/TileOps/tdiv_template.py +++ /dev/null @@ -1,50 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tdiv with IEEE 754 high-precision support""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - -# Import shared high-precision division algorithms -from div_hp import _div_ieee754_f32_impl, _div_ieee754_f16_impl - - -@pto.vkernel( - target="a5", - op="pto.tdiv" -) -def template_tdiv(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - """Element-wise division with optional high-precision mode""" - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - precision_type = pto.get_op_attr("precisionType", "default") - if pto.constexpr(precision_type == "high_precision"): - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - if pto.constexpr(dtype == pto.f32): - divided = _div_ieee754_f32_impl(lhs, rhs, mask) - else: # dtype == pto.f16 (guaranteed by MLIR validation) - divided = _div_ieee754_f16_impl(lhs, rhs, mask) - pto.vsts(divided, dst[row, col:], mask) - else: - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - divided = pto.vdiv(lhs, rhs, mask) - pto.vsts(divided, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tdivs_template.py b/lib/TileOps/tdivs_template.py deleted file mode 100644 index 76ac2cbd35..0000000000 --- a/lib/TileOps/tdivs_template.py +++ /dev/null @@ -1,92 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tdivs with IEEE 754 high-precision support - -Supports two operand orders: - 1. tdivs(src_tile, scalar, dst) -> src / scalar - 2. tdivs(scalar, src_tile, dst) -> scalar / src - -High-precision mode uses IEEE 754 compliant division algorithms from div_hp module -for improved accuracy with precision-sensitive, subnormal, and overflow boundary cases. -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - -# Import shared high-precision division algorithms -from div_hp import _div_ieee754_f32_impl, _div_ieee754_f16_impl - - -@pto.vkernel( - target="a5", - op="pto.tdivs", -) -def template_tdivs_tile_scalar(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): - """src / scalar with optional high-precision mode""" - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - precision_type = pto.get_op_attr("precisionType", "default") - if pto.constexpr(precision_type == "high_precision"): - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - scalar_vec = pto.vbr(scalar) - if pto.constexpr(dtype == pto.f32): - result = _div_ieee754_f32_impl(vec, scalar_vec, mask) - else: # dtype == pto.f16 (guaranteed by MLIR validation) - result = _div_ieee754_f16_impl(vec, scalar_vec, mask) - pto.vsts(result, dst[row, col:], mask) - else: - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - scalar_vec = pto.vbr(scalar) - result = pto.vdiv(vec, scalar_vec, mask) - pto.vsts(result, dst[row, col:], mask) - return - - -@pto.vkernel( - target="a5", - op="pto.tdivs", -) -def template_tdivs_scalar_tile(scalar: pto.AnyType, src: pto.Tile, dst: pto.Tile): - """scalar / src with optional high-precision mode""" - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - precision_type = pto.get_op_attr("precisionType", "default") - if pto.constexpr(precision_type == "high_precision"): - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - scalar_vec = pto.vbr(scalar) - if pto.constexpr(dtype == pto.f32): - result = _div_ieee754_f32_impl(scalar_vec, vec, mask) - else: # dtype == pto.f16 (guaranteed by MLIR validation) - result = _div_ieee754_f16_impl(scalar_vec, vec, mask) - pto.vsts(result, dst[row, col:], mask) - else: - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - scalar_vec = pto.vbr(scalar) - result = pto.vdiv(scalar_vec, vec, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/texp_template.py b/lib/TileOps/texp_template.py deleted file mode 100644 index 1ce9d8edbf..0000000000 --- a/lib/TileOps/texp_template.py +++ /dev/null @@ -1,33 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.texp""" - -import tilelang_dsl as pto -from exp_hp import _tl_exp_precision - -@pto.vkernel( - target="a5", - op="pto.texp" -) -def template_texp(src: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - precision_type = pto.get_op_attr("precisionType", "default") - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vinput = pto.vlds(src[row, col:]) - if pto.constexpr(precision_type == "high_precision"): - result = _tl_exp_precision(vinput, mask, dtype) - else: - result = pto.vexp(vinput, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/texpand_template.py b/lib/TileOps/texpand_template.py deleted file mode 100644 index d2e07360f9..0000000000 --- a/lib/TileOps/texpand_template.py +++ /dev/null @@ -1,72 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.texpands - -This template implements scalar broadcast expansion for location=VEC tiles. -It fills dst.valid_shape region with the broadcasted scalar value. - -Location constraint: - - This template is designed for tiles with location=VEC (vector buffer) - - In PTO-ISA, texpands has separate implementations for VEC and MAT locations - - For MAT location tiles, a different template/implementation path should be used - - Current tilelang_dsl MemorySpace only distinguishes GM and UB, where UB maps to - both VEC and MAT locations. The constraint checks memory_space=="ub" as a proxy - - Future enhancement: tilelang_dsl should support explicit location distinction - (e.g., MemorySpace.VEC vs MemorySpace.MAT) for more precise constraint matching - -Layout considerations: - - PTO-ISA has both rowmajor and colmajor expands implementations - - However, expands (scalar broadcast) is layout-agnostic: it simply fills - the tile with a scalar value using vector stores - - The vector store (vsts) writes data according to the tile's physical layout, - which is handled by the underlying DMA engine - - Therefore, this single template covers both rowmajor and colmajor cases -""" - -import tilelang_dsl as pto - - -def _texpands_vec_location_constraint(scalar, dst) -> bool: - """Constraint: dst tile must have location=VEC (represented as memory_space=ub). - - PTO-ISA defines texpands for both MAT and VEC locations: - - MAT location: expands matrix tiles (different implementation path, not supported here) - - VEC location: expands vector tiles (this template) - - Current tilelang_dsl limitation: - MemorySpace only has UB and GM. VEC and MAT both map to UB. - We check memory_space=="ub" as a proxy for VEC location. - MAT tiles should use a different op/template path and won't match here. - """ - # Check memory_space is "ub" (VEC/MAT location, not GM) - # In current tilelang_dsl, VEC location tiles have memory_space="ub" - ms = dst.memory_space - if isinstance(ms, str): - return ms == "ub" - return hasattr(ms, "value") and ms.value == "ub" - - -@pto.vkernel( - target="a5", - op="pto.texpands", - constraints=[_texpands_vec_location_constraint], -) -def template_texpands(scalar: pto.AnyType, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - # Use vdup for scalar broadcast - vec = pto.vdup(scalar, mask) - pto.vsts(vec, dst[row, col:], mask) - - return \ No newline at end of file diff --git a/lib/TileOps/textract_template.py b/lib/TileOps/textract_template.py deleted file mode 100644 index 2126e96d3b..0000000000 --- a/lib/TileOps/textract_template.py +++ /dev/null @@ -1,413 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.textract and pto.textract_fp (A5). - -TEXTRACT extracts a sub-tile window from src into dst at offset (indexRow, indexCol): - dst[i, j] = src[i + indexRow, j + indexCol] - -Supported data-flow directions (A5): - -Cube path (Mat -> Left/Right): - The mte_l1_l0a/l0b transpose flag is set based on src tile layout: - - Same fractal (src.col_major blayout + src.row_major slayout): - transpose=False — direct extraction (TExtractToA / TExtractToACompact) - - Cross fractal (src.row_major blayout + src.col_major slayout): - transpose=True — transposed extraction (TExtractToATransCompact / TExtractToA) - Offset extraction (indexRow/indexCol != 0) is supported via start_row/start_col - keywords on mte_l1_l0a/l0b (PR #469). - -Fix-pipe path (Acc -> Mat with FP quantization): - - textract_fp : PIPE_FIX, uses pto.mte_l0c_l1 with pre_quant keyword - Currently constrained to indexRow=0, indexCol=0. The index_row/index_col - parameters are accepted but unused in the lowering because the start - position for fix-pipe copy is implicit (whole-tile extraction from L0C). - Future work: pass index offset to mte_l0c_l1 when PTO-ISA supports it. - -Vector path (Vec -> Vec): - - Vec -> Vec ND : PIPE_V, uses pto.vlds/vsts subscript syntax - Note: Uses vector load/store instead of DMA (copy_ubuf_to_ubuf) because - textract requires arbitrary offset support (index_row/index_col), which - DMA block copy cannot provide. - -Blocked paths: - - Mat -> ScaleLeft/ScaleRight : mte_l1_l0a_mx/l0b_mx reject memory_space="scaling" dst - - Vec -> Mat : now available via pto.mte_ub_l1 (PR #405) -""" - -import tilelang_dsl as pto - -_FP_QUANT_MODE_MAP = { - (pto.f32, pto.si8): "qf322b8_pre_vec", - (pto.f32, pto.ui8): "qf322b8_pre_vec", - (pto.f32, pto.f16): "qf322f16_pre_vec", - (pto.f32, pto.bf16): "qf322bf16_pre_vec", - (pto.f32, pto.f32): "qf322f32_pre_vec", - (pto.si32, pto.si8): "req8_vec", - (pto.si32, pto.ui8): "req8_vec", - (pto.si32, pto.f16): "deqf16_vec", - (pto.si32, pto.bf16): "qs322bf16_pre_vec", -} - - -def _is_same_fractal_as_left(src) -> bool: - if src.config is None: - return True - return (src.config.b_layout != pto.BLayout.ROW_MAJOR - and src.config.s_layout == pto.SLayout.ROW_MAJOR) - - -def _is_cross_fractal_as_left(src) -> bool: - if src.config is None: - return False - return (src.config.b_layout == pto.BLayout.ROW_MAJOR - and src.config.s_layout == pto.SLayout.COL_MAJOR) - - -def _textract_cube_dst_is_left_same_fractal(src, index_row, index_col, dst) -> bool: - dst_ms = dst.memory_space - if not (dst_ms == "left" if isinstance(dst_ms, str) - else dst_ms.value == "left"): - return False - if not _is_same_fractal_as_left(src): - return False - return True - - -def _textract_cube_dst_is_left_cross_fractal(src, index_row, index_col, dst) -> bool: - dst_ms = dst.memory_space - if not (dst_ms == "left" if isinstance(dst_ms, str) - else dst_ms.value == "left"): - return False - if not _is_cross_fractal_as_left(src): - return False - return True - - -def _is_same_fractal_as_right(src) -> bool: - if src.config is None: - return True - return (src.config.b_layout != pto.BLayout.ROW_MAJOR - and src.config.s_layout == pto.SLayout.ROW_MAJOR) - - -def _is_cross_fractal_as_right(src) -> bool: - if src.config is None: - return False - return (src.config.b_layout == pto.BLayout.ROW_MAJOR - and src.config.s_layout == pto.SLayout.COL_MAJOR) - - -def _textract_cube_dst_is_right_same_fractal(src, index_row, index_col, dst) -> bool: - dst_ms = dst.memory_space - if not (dst_ms == "right" if isinstance(dst_ms, str) - else dst_ms.value == "right"): - return False - if not _is_same_fractal_as_right(src): - return False - return True - - -def _textract_cube_dst_is_right_cross_fractal(src, index_row, index_col, dst) -> bool: - dst_ms = dst.memory_space - if not (dst_ms == "right" if isinstance(dst_ms, str) - else dst_ms.value == "right"): - return False - if not _is_cross_fractal_as_right(src): - return False - return True - - -def _textract_vec2vec_nd_constraint(src, index_row, index_col, dst) -> bool: - src_ms = src.memory_space - dst_ms = dst.memory_space - src_is_ub = (src_ms == "ub" if isinstance(src_ms, str) - else src_ms.value == "ub") - dst_is_ub = (dst_ms == "ub" if isinstance(dst_ms, str) - else dst_ms.value == "ub") - if not (src_is_ub and dst_is_ub): - return False - if src.config is None or dst.config is None: - return False - if src.config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if src.config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst.config.b_layout != pto.BLayout.ROW_MAJOR: - return False - if dst.config.s_layout != pto.SLayout.NONE_BOX: - return False - if src.dtype != dst.dtype: - return False - return True - - -def _textract_fp_acc2mat_constraint(src, fp, index_row, index_col, dst) -> bool: - src_ms = src.memory_space - fp_ms = fp.memory_space - dst_ms = dst.memory_space - if not (src_ms == "acc" if isinstance(src_ms, str) - else src_ms.value == "acc"): - return False - if not (fp_ms == "scaling" if isinstance(fp_ms, str) - else fp_ms.value == "scaling"): - return False - if not (dst_ms == "mat" if isinstance(dst_ms, str) - else dst_ms.value == "mat"): - return False - return (src.dtype, dst.dtype) in _FP_QUANT_MODE_MAP - - -def _make_fp_constraint(src_dtype, dst_dtype): - def _fp_constraint(src, fp, index_row, index_col, dst) -> bool: - src_ms = src.memory_space - fp_ms = fp.memory_space - dst_ms = dst.memory_space - if not (src_ms == "acc" if isinstance(src_ms, str) - else src_ms.value == "acc"): - return False - if not (fp_ms == "scaling" if isinstance(fp_ms, str) - else fp_ms.value == "scaling"): - return False - if not (dst_ms == "mat" if isinstance(dst_ms, str) - else dst_ms.value == "mat"): - return False - return src.dtype == src_dtype and dst.dtype == dst_dtype - return _fp_constraint - - -@pto.ckernel( - target="a5", - op="pto.textract", - constraints=[_textract_cube_dst_is_left_same_fractal], -) -def template_textract_mat2left(src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - m, k = dst.valid_shape - pto.mte_l1_l0a(src.as_ptr(), dst.as_ptr(), m, k, - start_row=index_row, start_col=index_col) - return None - - -@pto.ckernel( - target="a5", - op="pto.textract", - constraints=[_textract_cube_dst_is_left_cross_fractal], -) -def template_textract_mat2left_trans(src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - m, k = dst.valid_shape - pto.mte_l1_l0a(src.as_ptr(), dst.as_ptr(), m, k, - start_row=index_row, start_col=index_col, transpose=True) - return None - - -@pto.ckernel( - target="a5", - op="pto.textract", - constraints=[_textract_cube_dst_is_right_same_fractal], -) -def template_textract_mat2right(src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - k, n = dst.valid_shape - pto.mte_l1_l0b(src.as_ptr(), dst.as_ptr(), k, n, - start_row=index_row, start_col=index_col) - return None - - -@pto.ckernel( - target="a5", - op="pto.textract", - constraints=[_textract_cube_dst_is_right_cross_fractal], -) -def template_textract_mat2right_trans(src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - k, n = dst.valid_shape - pto.mte_l1_l0b(src.as_ptr(), dst.as_ptr(), k, n, - start_row=index_row, start_col=index_col, transpose=True) - return None - - -@pto.vkernel( - target="a5", - op="pto.textract", - advanced=True, - constraints=[_textract_vec2vec_nd_constraint], -) -def template_textract_vec2vec_nd(src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - dtype = dst.element_type - lanes = pto.get_lanes(dtype) - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - data = pto.vlds(src[index_row + row, index_col + col:]) - pto.vsts(data, dst[row, col:], mask) - - return None - - -@pto.ckernel( - target="a5", - op="pto.textract_fp", - dtypes=((pto.f32, pto.f32, pto.i32, pto.i32, pto.si8),), - constraints=[_make_fp_constraint(pto.f32, pto.si8)], -) -def template_textract_fp_f32_si8(src: pto.Tile, fp: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - m, n = dst.valid_shape - src_stride = src.shape[0] - dst_stride = dst.shape[1] - pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, - pre_quant=(fp.as_ptr(), "qf322b8_pre_vec")) - return None - - -@pto.ckernel( - target="a5", - op="pto.textract_fp", - dtypes=((pto.f32, pto.f32, pto.i32, pto.i32, pto.ui8),), - constraints=[_make_fp_constraint(pto.f32, pto.ui8)], -) -def template_textract_fp_f32_ui8(src: pto.Tile, fp: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - m, n = dst.valid_shape - src_stride = src.shape[0] - dst_stride = dst.shape[1] - pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, - pre_quant=(fp.as_ptr(), "qf322b8_pre_vec")) - return None - - -@pto.ckernel( - target="a5", - op="pto.textract_fp", - dtypes=((pto.f32, pto.f32, pto.i32, pto.i32, pto.f16),), - constraints=[_make_fp_constraint(pto.f32, pto.f16)], -) -def template_textract_fp_f32_f16(src: pto.Tile, fp: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - m, n = dst.valid_shape - src_stride = src.shape[0] - dst_stride = dst.shape[1] - pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, - pre_quant=(fp.as_ptr(), "qf322f16_pre_vec")) - return None - - -@pto.ckernel( - target="a5", - op="pto.textract_fp", - dtypes=((pto.f32, pto.f32, pto.i32, pto.i32, pto.bf16),), - constraints=[_make_fp_constraint(pto.f32, pto.bf16)], -) -def template_textract_fp_f32_bf16(src: pto.Tile, fp: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - m, n = dst.valid_shape - src_stride = src.shape[0] - dst_stride = dst.shape[1] - pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, - pre_quant=(fp.as_ptr(), "qf322bf16_pre_vec")) - return None - - -@pto.ckernel( - target="a5", - op="pto.textract_fp", - dtypes=((pto.f32, pto.f32, pto.i32, pto.i32, pto.f32),), - constraints=[_make_fp_constraint(pto.f32, pto.f32)], -) -def template_textract_fp_f32_f32(src: pto.Tile, fp: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - m, n = dst.valid_shape - src_stride = src.shape[0] - dst_stride = dst.shape[1] - pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, - pre_quant=(fp.as_ptr(), "qf322f32_pre_vec")) - return None - - -@pto.ckernel( - target="a5", - op="pto.textract_fp", - dtypes=((pto.si32, pto.f32, pto.i32, pto.i32, pto.si8),), - constraints=[_make_fp_constraint(pto.si32, pto.si8)], -) -def template_textract_fp_si32_si8(src: pto.Tile, fp: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - m, n = dst.valid_shape - src_stride = src.shape[0] - dst_stride = dst.shape[1] - pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, - pre_quant=(fp.as_ptr(), "req8_vec")) - return None - - -@pto.ckernel( - target="a5", - op="pto.textract_fp", - dtypes=((pto.si32, pto.f32, pto.i32, pto.i32, pto.ui8),), - constraints=[_make_fp_constraint(pto.si32, pto.ui8)], -) -def template_textract_fp_si32_ui8(src: pto.Tile, fp: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - m, n = dst.valid_shape - src_stride = src.shape[0] - dst_stride = dst.shape[1] - pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, - pre_quant=(fp.as_ptr(), "req8_vec")) - return None - - -@pto.ckernel( - target="a5", - op="pto.textract_fp", - dtypes=((pto.si32, pto.f32, pto.i32, pto.i32, pto.f16),), - constraints=[_make_fp_constraint(pto.si32, pto.f16)], -) -def template_textract_fp_si32_f16(src: pto.Tile, fp: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - m, n = dst.valid_shape - src_stride = src.shape[0] - dst_stride = dst.shape[1] - pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, - pre_quant=(fp.as_ptr(), "deqf16_vec")) - return None - - -@pto.ckernel( - target="a5", - op="pto.textract_fp", - dtypes=((pto.si32, pto.f32, pto.i32, pto.i32, pto.bf16),), - constraints=[_make_fp_constraint(pto.si32, pto.bf16)], -) -def template_textract_fp_si32_bf16(src: pto.Tile, fp: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile): - m, n = dst.valid_shape - src_stride = src.shape[0] - dst_stride = dst.shape[1] - pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, - pre_quant=(fp.as_ptr(), "qs322bf16_pre_vec")) - return None \ No newline at end of file diff --git a/lib/TileOps/tfillpad_expand_template.py b/lib/TileOps/tfillpad_expand_template.py deleted file mode 100644 index 9fa3f435cd..0000000000 --- a/lib/TileOps/tfillpad_expand_template.py +++ /dev/null @@ -1,167 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tfillpad_expand - -Expand mode semantics: - - TFILLPAD_EXPAND: src rows may be less than dst rows - - Copy src.valid data to dst - - Fill cols from src.valid_cols to dst.valid_cols with FillPadVal - - Fill rows from src.rows to dst.rows with FillPadVal - -Strategy: - - Phase 1: Copy aligned valid blocks (cols 0 to aligned_col-1) - - Phase 2: Fill cols aligned_col to dst_valid_cols-1 with FillPadVal - - Phase 3: Copy tail valid lanes (cols aligned_col to src_valid_cols-1) - - Phase 4: Fill row expansion - -Address alignment and unaligned handling: - - vlds/vsts require 32-byte aligned base addresses - - Phase 1: col=0 is always aligned (tile base address is aligned), each iteration - accesses col + lanes which maintains alignment - - Phase 2/3/4: handle non-aligned lengths using make_mask() to control active lanes - - make_mask approach: simpler than vldus/vstus for isolated tail operations, no need - for alignment state management (vldas/vldus/vsta sequence) - - vldus/vstus is suitable for continuous unaligned streams; for single tail ops, - mask-controlled vlds/vsts is more direct and efficient -""" - -import tilelang_dsl as pto - -_NEG1_F32 = -1.0 - -# All supported dtype pairs for tfillpad_expand -_DTYPE_SIGNATURES = [ - (pto.f32, pto.f32), - (pto.i16, pto.i16), - (pto.si16, pto.si16), - (pto.ui16, pto.ui16), - (pto.i32, pto.i32), - (pto.si32, pto.si32), - (pto.ui32, pto.ui32), - (pto.i8, pto.i8), - (pto.si8, pto.si8), - (pto.ui8, pto.ui8), -] - - -@pto.vkernel( - target="a5", - op="pto.tfillpad_expand", - dtypes=_DTYPE_SIGNATURES, -) -def template_tfillpad_expand(src: pto.Tile, dst: pto.Tile): - """Unified tfillpad_expand template for all dtypes. - - Main logic is identical across dtypes; only PadValue handling differs: - - f32: ZERO + expansion uses -1.0 (special encoding), otherwise eval() or 0.0 - - integer families: eval() or dtype-specific zero constant - """ - dtype = dst.element_type - src_rows, _ = src.shape - src_valid_rows, src_valid_cols = src.valid_shape - dst_rows, _ = dst.shape - dst_valid_rows, dst_valid_cols = dst.valid_shape - - lanes = pto.get_lanes(dtype) - aligned_col = (src_valid_cols // lanes) * lanes - has_tail = src_valid_cols > aligned_col - has_valid_expansion = (src_valid_cols < dst_valid_cols) or (src_valid_rows < dst_valid_rows) - - # PadValue handling - dtype-specific - if pto.constexpr(dtype == pto.f32): - if pto.constexpr(dst.pad_value == pto.PadValue.ZERO): - if has_valid_expansion: - fill_scalar = pto.f32(_NEG1_F32) - else: - fill_scalar = pto.f32(0.0) - elif pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.f32(0.0) - elif pto.constexpr(dtype == pto.ui16): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.ui16(0) - elif pto.constexpr(dtype == pto.si16): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.si16(0) - elif pto.constexpr(dtype == pto.i16): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.i16(0) - elif pto.constexpr(dtype == pto.ui32): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.ui32(0) - elif pto.constexpr(dtype == pto.si32): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.si32(0) - elif pto.constexpr(dtype == pto.i32): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.i32(0) - elif pto.constexpr(dtype == pto.ui8): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.ui8(0) - elif pto.constexpr(dtype == pto.si8): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.si8(0) - elif pto.constexpr(dtype == pto.i8): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.i8(0) - - # Phase 1: Copy aligned valid blocks - for row in range(0, src_valid_rows, 1): - remained = aligned_col - for col in range(0, aligned_col, lanes): - mask, remained = pto.make_mask(dtype, remained) - data = pto.vlds(src[row, col:]) - pto.vsts(data, dst[row, col:], mask) - - # Phase 2: Fill col padding - if aligned_col < dst_valid_cols: - for row in range(0, dst_valid_rows, 1): - remained = dst_valid_cols - aligned_col - for col in range(aligned_col, dst_valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vdup(fill_scalar, mask) - pto.vsts(vec, dst[row, col:], mask) - - # Phase 3: Copy tail valid lanes - if has_tail: - for row in range(0, src_valid_rows, 1): - remained = src_valid_cols - aligned_col - mask_copy, remained = pto.make_mask(dtype, remained) - data = pto.vlds(src[row, aligned_col:]) - pto.vsts(data, dst[row, aligned_col:], mask_copy) - - # Phase 4: Fill row expansion - if pto.constexpr(src_rows < dst_rows): - for row in range(src_rows, dst_rows, 1): - remained = dst_valid_cols - for col in range(0, dst_valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vdup(fill_scalar, mask) - pto.vsts(vec, dst[row, col:], mask) - - return diff --git a/lib/TileOps/tfillpad_inplace_template.py b/lib/TileOps/tfillpad_inplace_template.py deleted file mode 100644 index 6214476c64..0000000000 --- a/lib/TileOps/tfillpad_inplace_template.py +++ /dev/null @@ -1,165 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tfillpad_inplace - -Semantic (based on TFillPad.hpp reference): - - TFILLPAD_INPLACE: same physical buffer (src == dst), skips copy phase, only fills expansion - - Inplace mode: src and dst share the same physical UB address - -Strategy (inplace mode): - - Skip Phase 1+3: Copy phases (data already in buffer) - - Phase 2: Fill cols from src_valid_cols to dst_valid_cols-1 with FillPadVal - - Phase 4: Fill rows from src_valid_rows to dst_valid_rows-1 with FillPadVal -""" - -import tilelang_dsl as pto - -_NEG1_F32 = -1.0 - -# All supported dtype pairs -_DTYPE_SIGNATURES = [ - (pto.f32, pto.f32), - (pto.i16, pto.i16), - (pto.si16, pto.si16), - (pto.ui16, pto.ui16), - (pto.i32, pto.i32), - (pto.si32, pto.si32), - (pto.ui32, pto.ui32), - (pto.i8, pto.i8), - (pto.si8, pto.si8), - (pto.ui8, pto.ui8), -] - - -@pto.vkernel( - target="a5", - op="pto.tfillpad_inplace", - dtypes=_DTYPE_SIGNATURES, - advanced=True, # Required for as_ptr() -) -def template_tfillpad_inplace(src: pto.Tile, dst: pto.Tile): - """tfillpad_inplace: skip copy phase, only fill expansion regions. - -Uses vstus+vstas for unaligned column fill, matching TFillPad.hpp. -""" - dtype = dst.element_type - _, cols = dst.shape - src_valid_rows, src_valid_cols = src.valid_shape - dst_valid_rows, dst_valid_cols = dst.valid_shape - - lanes = pto.get_lanes(dtype) - byte_width = pto.bytewidth(dtype) - has_valid_expansion = (src_valid_cols < dst_valid_cols) or (src_valid_rows < dst_valid_rows) - - # PadValue handling - same as tfillpad_template.py - # Note: dtype and pad_value are compile-time known, so constexpr is valid for those. - # has_valid_expansion is a runtime value derived from dynamic shapes, so split the condition. - if pto.constexpr(dtype == pto.f32): - if pto.constexpr(dst.pad_value == pto.PadValue.ZERO): - # For ZERO pad_value, use -1.0 encoding only when there's valid expansion - if has_valid_expansion: - fill_scalar = pto.f32(_NEG1_F32) - else: - fill_scalar = pto.f32(0.0) - elif pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.f32(0.0) - elif pto.constexpr(dtype == pto.ui16): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.ui16(0) - elif pto.constexpr(dtype == pto.si16): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.si16(0) - elif pto.constexpr(dtype == pto.i16): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.i16(0) - elif pto.constexpr(dtype == pto.ui32): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.ui32(0) - elif pto.constexpr(dtype == pto.si32): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.si32(0) - elif pto.constexpr(dtype == pto.i32): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.i32(0) - elif pto.constexpr(dtype == pto.ui8): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.ui8(0) - elif pto.constexpr(dtype == pto.si8): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.si8(0) - elif pto.constexpr(dtype == pto.i8): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.i8(0) - - # Phase 2: Fill cols from src_valid_cols to cols-1 (physical buffer end) - # Use vstus+vstas for unaligned starting column, matching TFillPad.hpp - pad_cols = cols - src_valid_cols # TileDataDst::Cols - srcValidCol - - # Create fill vector once (reused across all rows) - fill_vec = pto.vdup(fill_scalar, pto.make_mask(dtype, pto.PAT.ALL)) - - # Get base pointer to UB buffer - base_ptr = dst.as_ptr() - - for row in range(0, src_valid_rows, 1): - # Initialize align register for this row - ureg = pto.init_align() - - # Pointer to dst[row, src_valid_cols]: base_ptr + (row * cols + src_valid_cols) * byte_width - # Matching: dstPtr + i * dstStride + srcValidCol - row_offset = (row * cols + src_valid_cols) * byte_width - row_ptr = pto.addptr(base_ptr, row_offset) - - # Simple loop: iterate pad_cols times with step lanes - # Use vstus + addptr in each branch to simulate POST_UPDATE behavior - # ureg, remaining, row_ptr are all loop-carried, updated in every iteration - remaining = pad_cols - for _ in range(0, pad_cols, lanes): - if remaining >= lanes: - ureg = pto.vstus(ureg, lanes, fill_vec, row_ptr) - row_ptr = pto.addptr(row_ptr, lanes * byte_width) - remaining = remaining - lanes - else: - ureg = pto.vstus(ureg, remaining, fill_vec, row_ptr) - row_ptr = pto.addptr(row_ptr, remaining * byte_width) - remaining = 0 - - # vstas: flush buffered bytes - pto.vstas(ureg, row_ptr, 0) - - # Phase 4: Fill rows from src_valid_rows to dst_valid_rows-1 - # Fill entire physical rows (cols elements), matching: padRows * dstStride - for row in range(src_valid_rows, dst_valid_rows, 1): - remained = cols - for col in range(0, cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vdup(fill_scalar, mask) - pto.vsts(vec, dst[row, col:], mask) - - return \ No newline at end of file diff --git a/lib/TileOps/tfillpad_template.py b/lib/TileOps/tfillpad_template.py deleted file mode 100644 index 3378cb1983..0000000000 --- a/lib/TileOps/tfillpad_template.py +++ /dev/null @@ -1,173 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tfillpad - -Semantic (based on C++ TFillPad.hpp reference): - - TFILLPAD: copies src.valid data to dst, then fills dst expansion with FillPadVal - - TFILLPAD_INPLACE: same physical buffer (src == dst), skips copy phase, only fills expansion - -Key logic from C++: - if constexpr (!inplace) { - CopyValidElementsVec(dst, src, ...); // Phase 1+3: copy valid data - } - // Phase 2+4: fill expansion (always executed if dst has larger valid region) - FillExpansion(dst, padCols, padRows, padValue); - -Strategy: - - Phase 1: Copy aligned valid blocks (cols 0 to aligned_col-1) [only if !inplace] - - Phase 2: Fill cols aligned_col to dst_valid_cols-1 with FillPadVal - - Phase 3: Copy tail valid lanes (cols aligned_col to src_valid_cols-1) [only if !inplace] - - Phase 4: Fill row expansion - -Address alignment and unaligned handling: - - vlds/vsts require 32-byte aligned base addresses - - Phase 1: col=0 is always aligned (tile base address is aligned) - - Phase 2/4: handle non-aligned lengths using make_mask() to control active lanes - -Note: There is no separate pto.tfillpad_inplace operation in PTO IR. - In-place mode is expressed via tfillpad with src and dst being the same SSA value. - This template handles both cases by detecting if copy phase is needed. -""" - -import tilelang_dsl as pto - -_NEG1_F32 = -1.0 - -# All supported dtype pairs -_DTYPE_SIGNATURES = [ - (pto.f32, pto.f32), - (pto.i16, pto.i16), - (pto.si16, pto.si16), - (pto.ui16, pto.ui16), - (pto.i32, pto.i32), - (pto.si32, pto.si32), - (pto.ui32, pto.ui32), - (pto.i8, pto.i8), - (pto.si8, pto.si8), - (pto.ui8, pto.ui8), -] - - -@pto.vkernel( - target="a5", - op="pto.tfillpad", - dtypes=_DTYPE_SIGNATURES, -) -def template_tfillpad(src: pto.Tile, dst: pto.Tile): - """tfillpad: copy src.valid to dst and fill expansion regions. - - Based on C++ TFillPad.hpp reference: - - TFILLPAD (non-inplace): CopyValidElementsVec + FillExpansion - - tfillpad requires src.shape == dst.shape (same physical size). - If dst.valid > src.valid, fill the expansion regions. - """ - dtype = dst.element_type - src_rows, _ = src.shape - src_valid_rows, src_valid_cols = src.valid_shape - dst_rows, dst_cols = dst.shape - dst_valid_rows, dst_valid_cols = dst.valid_shape - - lanes = pto.get_lanes(dtype) - aligned_col = (src_valid_cols // lanes) * lanes - has_tail = src_valid_cols > aligned_col - has_valid_expansion = (src_valid_cols < dst_valid_cols) or (src_valid_rows < dst_valid_rows) - - # PadValue handling - dtype-specific (inline to avoid external call) - if pto.constexpr(dtype == pto.f32): - if pto.constexpr(dst.pad_value == pto.PadValue.ZERO): - if has_valid_expansion: - fill_scalar = pto.f32(_NEG1_F32) - else: - fill_scalar = pto.f32(0.0) - elif pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.f32(0.0) - elif pto.constexpr(dtype == pto.ui16): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.ui16(0) - elif pto.constexpr(dtype == pto.si16): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.si16(0) - elif pto.constexpr(dtype == pto.i16): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.i16(0) - elif pto.constexpr(dtype == pto.ui32): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.ui32(0) - elif pto.constexpr(dtype == pto.si32): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.si32(0) - elif pto.constexpr(dtype == pto.i32): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.i32(0) - elif pto.constexpr(dtype == pto.ui8): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.ui8(0) - elif pto.constexpr(dtype == pto.si8): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.si8(0) - elif pto.constexpr(dtype == pto.i8): - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - fill_scalar = dst.pad_value.eval() - else: - fill_scalar = pto.i8(0) - - # Phase 1: Copy aligned valid blocks - for row in range(0, src_valid_rows, 1): - remained = aligned_col - for col in range(0, aligned_col, lanes): - mask, remained = pto.make_mask(dtype, remained) - data = pto.vlds(src[row, col:]) - pto.vsts(data, dst[row, col:], mask) - - # Phase 2: Fill cols from aligned_col to dst_cols-1 - if aligned_col < dst_cols: - for row in range(0, src_valid_rows, 1): - remained = dst_cols - aligned_col - for col in range(aligned_col, dst_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vdup(fill_scalar, mask) - pto.vsts(vec, dst[row, col:], mask) - - # Phase 3: Copy tail valid lanes - if has_tail: - for row in range(0, src_valid_rows, 1): - remained = src_valid_cols - aligned_col - mask_copy, remained = pto.make_mask(dtype, remained) - data = pto.vlds(src[row, aligned_col:]) - pto.vsts(data, dst[row, aligned_col:], mask_copy) - - # Phase 4: Fill row expansion - if pto.constexpr(src_rows < dst_rows): - for row in range(src_rows, dst_rows, 1): - remained = dst_cols - for col in range(0, dst_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vdup(fill_scalar, mask) - pto.vsts(vec, dst[row, col:], mask) - - return diff --git a/lib/TileOps/tfmod_template.py b/lib/TileOps/tfmod_template.py deleted file mode 100644 index 32f7cb65f5..0000000000 --- a/lib/TileOps/tfmod_template.py +++ /dev/null @@ -1,71 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tfmod - -Aligned with pto-isa/include/pto/npu/a5/TFMod.hpp: -- float: vdiv -> vtrc(ROUND_Z) -> vmul -> vsub -- half: vcvt(half->float, PART_EVEN/ODD) -> vdiv -> vtrc(ROUND_Z) -> vmul -> vsub - -> vcvt(float->half, ROUND_Z, RS_ENABLE, PART_EVEN/ODD) -> vor -- other (i16/ui16): vdiv -> vmul -> vsub (no vtrc, integer div is trunc by nature) -""" - -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tfmod", - dtypes=[ - (pto.f32, pto.f32, pto.f32), - (pto.f16, pto.f16, pto.f16), - (pto.i16, pto.i16, pto.i16), - (pto.ui16, pto.ui16, pto.ui16), - ], -) -def template_tfmod(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - - if pto.constexpr(dtype == pto.f32): - quotient = pto.vdiv(lhs, rhs, mask) - quotient = pto.vtrc(quotient, mask, rnd=pto.VcvtRoundMode.Z) - truncated_mul = pto.vmul(quotient, rhs, mask) - result = pto.vsub(lhs, truncated_mul, mask) - elif pto.constexpr(dtype == pto.f16): - lhs_even = pto.vcvt(lhs, pto.f32, mask, part=pto.VcvtPartMode.EVEN) - rhs_even = pto.vcvt(rhs, pto.f32, mask, part=pto.VcvtPartMode.EVEN) - quotient_even = pto.vdiv(lhs_even, rhs_even, mask) - quotient_even = pto.vtrc(quotient_even, mask, rnd=pto.VcvtRoundMode.Z) - truncated_mul_even = pto.vmul(quotient_even, rhs_even, mask) - result_even = pto.vsub(lhs_even, truncated_mul_even, mask) - dst_even = pto.vcvt(result_even, pto.f16, mask, rnd=pto.VcvtRoundMode.Z, sat=pto.VcvtSatMode.SAT, part=pto.VcvtPartMode.EVEN) - - lhs_odd = pto.vcvt(lhs, pto.f32, mask, part=pto.VcvtPartMode.ODD) - rhs_odd = pto.vcvt(rhs, pto.f32, mask, part=pto.VcvtPartMode.ODD) - quotient_odd = pto.vdiv(lhs_odd, rhs_odd, mask) - quotient_odd = pto.vtrc(quotient_odd, mask, rnd=pto.VcvtRoundMode.Z) - truncated_mul_odd = pto.vmul(quotient_odd, rhs_odd, mask) - result_odd = pto.vsub(lhs_odd, truncated_mul_odd, mask) - dst_odd = pto.vcvt(result_odd, pto.f16, mask, rnd=pto.VcvtRoundMode.Z, sat=pto.VcvtSatMode.SAT, part=pto.VcvtPartMode.ODD) - - result = pto.vor(dst_even, dst_odd, mask) - else: - quotient = pto.vdiv(lhs, rhs, mask) - truncated_mul = pto.vmul(quotient, rhs, mask) - result = pto.vsub(lhs, truncated_mul, mask) - - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tfmods_template.py b/lib/TileOps/tfmods_template.py deleted file mode 100644 index 9d9606625c..0000000000 --- a/lib/TileOps/tfmods_template.py +++ /dev/null @@ -1,111 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tfmods - -Note: A5 hardware implements tfmods as: - - float: dst = src - trunc(src / scalar) * scalar (f32 precision) - - half: dst = src - trunc(src / scalar) * scalar (computed in f32 precision, then converted back to f16) - - integer: dst = src - (src / scalar) * scalar (integer division already truncates) - -f16 path: convert f16 to f32 (even/odd), compute in f32 with vtrc(ROUND_Z), -convert back to f16 with ROUND_Z, merge with vor. -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tfmods", - dtypes=[ - (pto.f32, pto.f32, pto.f32), - (pto.f16, pto.f16, pto.f16), - (pto.i32, pto.i32, pto.i32), - (pto.i16, pto.i16, pto.i16), - ], - advanced=True, -) -def template_tfmods(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): - """dst = src - trunc(src / scalar) * scalar""" - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - if pto.constexpr(dtype == pto.f32): - # f32 path: direct f32 computation - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f32)): - mask, remained = pto.make_mask(pto.f32, remained) - vec = pto.vlds(src[row, col:]) - scalar_vec = pto.vbr(scalar) - quotient = pto.vdiv(vec, scalar_vec, mask) - truncated = pto.vtrc(quotient, mask, rnd="Z") - product = pto.vmuls(truncated, scalar, mask) - result = pto.vsub(vec, product, mask) - pto.vsts(result, dst[row, col:], mask) - elif pto.constexpr(dtype == pto.f16): - # f16 path: compute in f32 precision, then convert back to f16 - full_mask_b16 = pto.make_mask(pto.f16, pto.PAT.ALL) - full_mask_b32 = pto.make_mask(pto.f32, pto.PAT.ALL) - scalar_vec_f16 = pto.vbr(scalar) - scalar_f32_vec = pto.vcvt(scalar_vec_f16, pto.f32, full_mask_b16, part=pto.VcvtPartMode.EVEN) - - for row in range(0, valid_rows, 1): - remained = valid_cols - remained_f32 = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f16)): - mask_f16, remained = pto.make_mask(pto.f16, remained) - mask_f32, remained_f32 = pto.make_mask(pto.f32, remained_f32) - vec = pto.vlds(src[row, col:]) - - # Convert f16 to f32 (even and odd parts) - vec_even = pto.vcvt(vec, pto.f32, full_mask_b16, part=pto.VcvtPartMode.EVEN) - vec_odd = pto.vcvt(vec, pto.f32, full_mask_b16, part=pto.VcvtPartMode.ODD) - - # Even part: f32 computation - scalar_f32 = pto.f32(scalar) - quotient_even = pto.vdiv(vec_even, scalar_f32_vec, mask_f32) - truncated_even = pto.vtrc(quotient_even, mask_f32, rnd="Z") - product_even = pto.vmuls(truncated_even, scalar_f32, mask_f32) - result_even = pto.vsub(vec_even, product_even, mask_f32) - - # Odd part: f32 computation - quotient_odd = pto.vdiv(vec_odd, scalar_f32_vec, mask_f32) - truncated_odd = pto.vtrc(quotient_odd, mask_f32, rnd="Z") - product_odd = pto.vmuls(truncated_odd, scalar_f32, mask_f32) - result_odd = pto.vsub(vec_odd, product_odd, mask_f32) - - # Convert f32 results back to f16 with ROUND_Z + saturation - result_f16_even = pto.vcvt(result_even, pto.f16, full_mask_b32, - rnd=pto.VcvtRoundMode.Z, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.EVEN) - result_f16_odd = pto.vcvt(result_odd, pto.f16, full_mask_b32, - rnd=pto.VcvtRoundMode.Z, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.ODD) - - # Merge even and odd parts - result_f16 = pto.vor(result_f16_even, result_f16_odd, mask_f16) - pto.vsts(result_f16, dst[row, col:], mask_f16) - else: - # Integer path: vdiv already truncates towards zero, no vtrc needed - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - scalar_vec = pto.vbr(scalar) - quotient = pto.vdiv(vec, scalar_vec, mask) - product = pto.vmuls(quotient, scalar, mask) - result = pto.vsub(vec, product, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tgemv_acc_template.py b/lib/TileOps/tgemv_acc_template.py deleted file mode 100644 index 0152578339..0000000000 --- a/lib/TileOps/tgemv_acc_template.py +++ /dev/null @@ -1,27 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tgemv.acc.""" - -import tilelang_dsl as pto - - -@pto.ckernel( - target="a5", - op="pto.tgemv.acc", - dtypes=[ - (pto.f32, pto.f16, pto.f16, pto.f32), - (pto.f32, pto.bf16, pto.bf16, pto.f32), - (pto.f32, pto.f32, pto.f32, pto.f32), - ], -) -def template_tgemv_acc(acc_in: pto.Tile, lhs: pto.Tile, rhs: pto.Tile, dst: pto.Tile): - _, k = lhs.valid_shape - _, n = rhs.valid_shape - pto.mad_acc(lhs.as_ptr(), rhs.as_ptr(), dst.as_ptr(), 1, n, k) - return None diff --git a/lib/TileOps/tgemv_bias_template.py b/lib/TileOps/tgemv_bias_template.py deleted file mode 100644 index 759c7abeec..0000000000 --- a/lib/TileOps/tgemv_bias_template.py +++ /dev/null @@ -1,28 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tgemv.bias.""" - -import tilelang_dsl as pto - - -@pto.ckernel( - target="a5", - op="pto.tgemv.bias", - dtypes=[ - (pto.f16, pto.f16, pto.f32, pto.f32), - (pto.bf16, pto.bf16, pto.f32, pto.f32), - (pto.f32, pto.f32, pto.f32, pto.f32), - (pto.i8, pto.i8, pto.i32, pto.i32), - ], -) -def template_tgemv_bias(lhs: pto.Tile, rhs: pto.Tile, bias: pto.Tile, dst: pto.Tile): - _, k = lhs.valid_shape - _, n = rhs.valid_shape - pto.mad_bias(lhs.as_ptr(), rhs.as_ptr(), dst.as_ptr(), bias.as_ptr(), 1, n, k) - return None diff --git a/lib/TileOps/tgemv_mx_template.py b/lib/TileOps/tgemv_mx_template.py deleted file mode 100644 index def0db88f5..0000000000 --- a/lib/TileOps/tgemv_mx_template.py +++ /dev/null @@ -1,96 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tgemv.mx.""" - -import tilelang_dsl as pto - -_LOW_PRECISION_PAIRS = ( - (pto.f8e4m3, pto.f8e4m3), - (pto.f8e4m3, pto.f8e5m2), - (pto.f8e5m2, pto.f8e4m3), - (pto.f8e5m2, pto.f8e5m2), - (pto.f4e1m2x2, pto.f4e1m2x2), - (pto.f4e1m2x2, pto.f4e2m1x2), - (pto.f4e2m1x2, pto.f4e1m2x2), - (pto.f4e2m1x2, pto.f4e2m1x2), -) - -_TGEMV_MX_DTYPES = tuple( - (lhs_dtype, pto.f16, rhs_dtype, pto.f16, pto.f32) - for lhs_dtype, rhs_dtype in _LOW_PRECISION_PAIRS -) - -_TGEMV_MX_ACC_DTYPES = tuple( - (pto.f32, lhs_dtype, pto.f16, rhs_dtype, pto.f16, pto.f32) - for lhs_dtype, rhs_dtype in _LOW_PRECISION_PAIRS -) - -_TGEMV_MX_BIAS_DTYPES = tuple( - (lhs_dtype, pto.f16, rhs_dtype, pto.f16, pto.f32, pto.f32) - for lhs_dtype, rhs_dtype in _LOW_PRECISION_PAIRS -) - - -@pto.ckernel( - target="a5", - op="pto.tgemv.mx", - dtypes=_TGEMV_MX_DTYPES, -) -def template_tgemv_mx( - lhs: pto.Tile, - lhs_scale: pto.Tile, - rhs: pto.Tile, - rhs_scale: pto.Tile, - dst: pto.Tile, -): - _, k = lhs.valid_shape - _, n = rhs.valid_shape - pto.mad_mx( - lhs.as_ptr(), rhs.as_ptr(), dst.as_ptr(), 1, n, k, sat="sat") - return None - - -@pto.ckernel( - target="a5", - op="pto.tgemv.mx.acc", - dtypes=_TGEMV_MX_ACC_DTYPES, -) -def template_tgemv_mx_acc( - acc_in: pto.Tile, - lhs: pto.Tile, - lhs_scale: pto.Tile, - rhs: pto.Tile, - rhs_scale: pto.Tile, - dst: pto.Tile, -): - _, k = lhs.valid_shape - _, n = rhs.valid_shape - pto.mad_mx_acc( - lhs.as_ptr(), rhs.as_ptr(), dst.as_ptr(), 1, n, k, sat="nosat") - return None - - -@pto.ckernel( - target="a5", - op="pto.tgemv.mx.bias", - dtypes=_TGEMV_MX_BIAS_DTYPES, -) -def template_tgemv_mx_bias( - lhs: pto.Tile, - lhs_scale: pto.Tile, - rhs: pto.Tile, - rhs_scale: pto.Tile, - bias: pto.Tile, - dst: pto.Tile, -): - _, k = lhs.valid_shape - _, n = rhs.valid_shape - pto.mad_mx_bias( - lhs.as_ptr(), rhs.as_ptr(), dst.as_ptr(), bias.as_ptr(), 1, n, k, sat="sat") - return None diff --git a/lib/TileOps/tgemv_template.py b/lib/TileOps/tgemv_template.py deleted file mode 100644 index f39d77ee5d..0000000000 --- a/lib/TileOps/tgemv_template.py +++ /dev/null @@ -1,33 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tgemv.""" - -import tilelang_dsl as pto - - -@pto.ckernel( - target="a5", - op="pto.tgemv", - dtypes=[ - (pto.f16, pto.f16, pto.f32), - (pto.bf16, pto.bf16, pto.f32), - (pto.f32, pto.f32, pto.f32), - (pto.i8, pto.i8, pto.i32), - (pto.ScalarType("f8E4M3FN"), pto.ScalarType("f8E4M3FN"), pto.f32), - (pto.ScalarType("f8E4M3FN"), pto.ScalarType("f8E5M2"), pto.f32), - (pto.ScalarType("f8E5M2"), pto.ScalarType("f8E4M3FN"), pto.f32), - (pto.ScalarType("f8E5M2"), pto.ScalarType("f8E5M2"), pto.f32), - (pto.ScalarType("hif8"), pto.ScalarType("hif8"), pto.f32), - ], -) -def template_tgemv(lhs: pto.Tile, rhs: pto.Tile, acc: pto.Tile): - _, k = lhs.valid_shape - _, n = rhs.valid_shape - pto.mad(lhs.as_ptr(), rhs.as_ptr(), acc.as_ptr(), 1, n, k) - return None diff --git a/lib/TileOps/tinsert_template.py b/lib/TileOps/tinsert_template.py deleted file mode 100644 index 8bb0b5da1b..0000000000 --- a/lib/TileOps/tinsert_template.py +++ /dev/null @@ -1,1463 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tinsert - tile data insertion - -Implements six data movement paths using @pto.ckernel: - - Acc→Mat NZ: pto.mte_l0c_l1 (copy_matrix_cc_to_cbuf) - - Acc→Vec ND/DN/NZ: pto.mte_l0c_ub (copy_matrix_cc_to_ub) - - Vec→Vec ND/NZ: pto.copy_ubuf_to_ubuf - - Vec→Mat NZ/ND: pto.mte_ub_l1 (copy_ubuf_to_cbuf) -""" - -import tilelang_dsl as pto - -BLOCK_BYTE_SIZE = 32 -BLOCK_BYTE_BITS = 256 -FRACTAL_NZ_ROW = 16 - - -# --------------------------------------------------------------------------- -# Constraint functions -# --------------------------------------------------------------------------- - - -def _acc_to_mat_constraint(src, dst) -> bool: - return ( - src.memory_space == "acc" - and dst.memory_space == "mat" - and dst.config.b_layout == pto.BLayout.COL_MAJOR - and dst.config.s_layout == pto.SLayout.ROW_MAJOR - ) - - -def _acc_to_vec_nd_constraint(src, dst) -> bool: - return ( - src.memory_space == "acc" - and dst.memory_space == "ub" - and dst.config.b_layout == pto.BLayout.ROW_MAJOR - and dst.config.s_layout == pto.SLayout.NONE_BOX - ) - - -def _acc_to_vec_dn_constraint(src, dst) -> bool: - return ( - src.memory_space == "acc" - and dst.memory_space == "ub" - and dst.config.b_layout == pto.BLayout.COL_MAJOR - and dst.config.s_layout == pto.SLayout.NONE_BOX - ) - - -def _acc_to_vec_nz_constraint(src, dst) -> bool: - return ( - src.memory_space == "acc" - and dst.memory_space == "ub" - and dst.config.b_layout == pto.BLayout.COL_MAJOR - and dst.config.s_layout == pto.SLayout.ROW_MAJOR - ) - - -def _vec_to_vec_nd_constraint(src, dst) -> bool: - return ( - src.memory_space == "ub" - and src.config.b_layout == pto.BLayout.ROW_MAJOR - and src.config.s_layout == pto.SLayout.NONE_BOX - and dst.memory_space == "ub" - and dst.config.b_layout == pto.BLayout.ROW_MAJOR - and dst.config.s_layout == pto.SLayout.NONE_BOX - and not (src.valid_shape[0] == 1 and src.valid_shape[1] == 1) - ) - - -def _vec_to_vec_nd_scalar_constraint(src, dst) -> bool: - return ( - src.memory_space == "ub" - and src.config.b_layout == pto.BLayout.ROW_MAJOR - and src.config.s_layout == pto.SLayout.NONE_BOX - and dst.memory_space == "ub" - and dst.config.b_layout == pto.BLayout.ROW_MAJOR - and dst.config.s_layout == pto.SLayout.NONE_BOX - and src.valid_shape[0] == 1 - and src.valid_shape[1] == 1 - ) - - -def _vec_to_vec_nz_constraint(src, dst) -> bool: - return ( - src.memory_space == "ub" - and src.config.b_layout == pto.BLayout.COL_MAJOR - and src.config.s_layout == pto.SLayout.ROW_MAJOR - and dst.memory_space == "ub" - and dst.config.b_layout == pto.BLayout.COL_MAJOR - and dst.config.s_layout == pto.SLayout.ROW_MAJOR - ) - - -def _vec_to_mat_nz_constraint(src, dst) -> bool: - return ( - src.memory_space == "ub" - and src.config.b_layout == pto.BLayout.COL_MAJOR - and src.config.s_layout == pto.SLayout.ROW_MAJOR - and dst.memory_space == "mat" - and dst.config.b_layout == pto.BLayout.COL_MAJOR - and dst.config.s_layout == pto.SLayout.ROW_MAJOR - ) - - -def _vec_to_mat_nd_constraint(src, dst) -> bool: - return ( - src.memory_space == "ub" - and src.config.b_layout == pto.BLayout.ROW_MAJOR - and src.config.s_layout == pto.SLayout.NONE_BOX - and dst.memory_space == "mat" - and dst.config.s_layout == pto.SLayout.NONE_BOX - and src.valid_shape[1] * pto.bytewidth(pto.ScalarType(src.dtype)) >= BLOCK_BYTE_SIZE - ) - - -def _pre_quant_vec_mode(src_dtype, dst_dtype): - if pto.constexpr(src_dtype == pto.f32): - if pto.constexpr(dst_dtype == pto.f16): - return "qf322f16_pre_vec" - elif pto.constexpr(dst_dtype == pto.bf16): - return "qf322bf16_pre_vec" - elif pto.constexpr(dst_dtype == pto.f32): - return "qf322f32_pre_vec" - else: - return "qf322b8_pre_vec" - else: - if pto.constexpr(dst_dtype == pto.f16): - return "deqf16_vec" - elif pto.constexpr(dst_dtype == pto.bf16): - return "qs322bf16_pre_vec" - elif pto.constexpr(dst_dtype == pto.i32): - return "deqs32_int_vec" - else: - return "req8_vec" - - -def _pre_quant_scalar_mode(src_dtype, dst_dtype): - if pto.constexpr(src_dtype == pto.f32): - if pto.constexpr(dst_dtype == pto.f16): - return "qf322f16_pre_scalar" - elif pto.constexpr(dst_dtype == pto.bf16): - return "qf322bf16_pre_scalar" - elif pto.constexpr(dst_dtype == pto.f32): - return "qf322f32_pre_scalar" - else: - return "qf322b8_pre_scalar" - else: - if pto.constexpr(dst_dtype == pto.f16): - return "deqf16_scalar" - elif pto.constexpr(dst_dtype == pto.bf16): - return "qs322bf16_pre_scalar" - elif pto.constexpr(dst_dtype == pto.i32): - return "deqs32_int_scalar" - else: - return "req8_scalar" - - -# --------------------------------------------------------------------------- -# Acc -> Mat -# --------------------------------------------------------------------------- - - -@pto.ckernel( - target="a5", - op="pto.tinsert", - dtypes=[ - (pto.f32, pto.i64, pto.i64, pto.f16, pto.f32, pto.i64), - (pto.f32, pto.i64, pto.i64, pto.bf16, pto.f32, pto.i64), - (pto.f32, pto.i64, pto.i64, pto.i8, pto.f32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.f16, pto.i32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.bf16, pto.i32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.i8, pto.i32, pto.i64), - ], - constraints=[_acc_to_mat_constraint], -) -def template_tinsert_acc_to_mat( - src: pto.Tile, - index_row: pto.i64, index_col: pto.i64, - dst: pto.Tile, - fp: pto.Tile, pre_quant_scalar: pto.i64, -): - src_ptr = src.as_ptr() - dst_ptr = dst.as_ptr() - - src_elem = src.element_type - dst_elem = dst.element_type - elem_bytes = pto.bytewidth(dst_elem) - c0_size = BLOCK_BYTE_SIZE // elem_bytes - - valid_rows = src.valid_shape[0] - valid_cols = src.valid_shape[1] - n_size = (valid_cols + c0_size - 1) // c0_size * c0_size - - dst_rows = dst.shape[0] - col_block = index_col // c0_size - col_mod = index_col - col_block * c0_size - dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod - dst_ptr = pto.addptr(dst_ptr, dst_offset) - - dst_stride = dst_rows * c0_size * elem_bytes - src_stride = src.shape[0] * pto.bytewidth(src_elem) - - relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") - has_fp = fp is not None - has_scalar = pre_quant_scalar is not None - - if pto.constexpr(has_fp): - if pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_l1( - src_ptr, dst_ptr, - valid_rows, n_size, src_stride, dst_stride, - pre_relu=("normal_relu", None, None), - pre_quant=(fp, _pre_quant_vec_mode(src_elem, dst_elem)), - ) - else: - pto.mte_l0c_l1( - src_ptr, dst_ptr, - valid_rows, n_size, src_stride, dst_stride, - pre_quant=(fp, _pre_quant_vec_mode(src_elem, dst_elem)), - ) - elif pto.constexpr(has_scalar): - if pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_l1( - src_ptr, dst_ptr, - valid_rows, n_size, src_stride, dst_stride, - pre_relu=("normal_relu", None, None), - pre_quant=(pre_quant_scalar, _pre_quant_scalar_mode(src_elem, dst_elem)), - ) - else: - pto.mte_l0c_l1( - src_ptr, dst_ptr, - valid_rows, n_size, src_stride, dst_stride, - pre_quant=(pre_quant_scalar, _pre_quant_scalar_mode(src_elem, dst_elem)), - ) - elif pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_l1( - src_ptr, dst_ptr, - valid_rows, n_size, src_stride, dst_stride, - pre_relu=("normal_relu", None, None), - ) - else: - pto.mte_l0c_l1( - src_ptr, dst_ptr, - valid_rows, n_size, src_stride, dst_stride, - ) - return None - - -@pto.ckernel( - target="a5", - op="pto.tinsert", - dtypes=[ - (pto.f32, pto.i32, pto.i32, pto.f16), - (pto.f32, pto.i32, pto.i32, pto.bf16), - (pto.f32, pto.i32, pto.i32, pto.f32), - (pto.i32, pto.i32, pto.i32, pto.i32), - ], - constraints=[_acc_to_mat_constraint], -) -def template_tinsert_acc_to_mat_basic( - src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile, -): - src_ptr = src.as_ptr() - dst_ptr = dst.as_ptr() - - src_elem = src.element_type - dst_elem = dst.element_type - elem_bytes = pto.bytewidth(dst_elem) - c0_size = BLOCK_BYTE_SIZE // elem_bytes - - valid_rows = src.valid_shape[0] - valid_cols = src.valid_shape[1] - n_size = (valid_cols + c0_size - 1) // c0_size * c0_size - - dst_rows = dst.shape[0] - col_block = index_col // c0_size - col_mod = index_col - col_block * c0_size - dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod - dst_ptr = pto.addptr(dst_ptr, dst_offset) - - dst_stride = dst_rows * c0_size * elem_bytes - src_stride = src.shape[0] * pto.bytewidth(src_elem) - - relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") - - if pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_l1( - src_ptr, dst_ptr, - valid_rows, n_size, src_stride, dst_stride, - pre_relu=("normal_relu", None, None), - ) - else: - pto.mte_l0c_l1( - src_ptr, dst_ptr, - valid_rows, n_size, src_stride, dst_stride, - ) - return None - - -# --------------------------------------------------------------------------- -# Acc -> Vec (ND, NONE_BOX) -# --------------------------------------------------------------------------- - - -@pto.ckernel( - target="a5", - op="pto.tinsert", - dtypes=[ - (pto.f32, pto.i64, pto.i64, pto.f32, pto.f32, pto.i64), - (pto.f32, pto.i64, pto.i64, pto.f16, pto.f32, pto.i64), - (pto.f32, pto.i64, pto.i64, pto.bf16, pto.f32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.i32, pto.i32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.f16, pto.i32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.bf16, pto.i32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.i8, pto.i32, pto.i64), - ], - priority=10, - constraints=[_acc_to_vec_nd_constraint], -) -def template_tinsert_acc_to_vec_nd( - src: pto.Tile, - index_row: pto.i64, index_col: pto.i64, - dst: pto.Tile, - fp: pto.Tile, pre_quant_scalar: pto.i64, -): - src_ptr = src.as_ptr() - dst_ptr = dst.as_ptr() - - src_elem = src.element_type - dst_elem = dst.element_type - elem_bytes = pto.bytewidth(dst_elem) - c0_size = 32 // elem_bytes - - valid_rows = src.valid_shape[0] - valid_cols_raw = src.valid_shape[1] - valid_cols = (valid_cols_raw + c0_size - 1) // c0_size * c0_size - - dst_cols = dst.shape[1] - dst_offset = index_row * dst_cols + index_col - dst_ptr = pto.addptr(dst_ptr, dst_offset) - - dst_stride = dst_cols - src_stride = (valid_rows + 15) // 16 * 16 - - relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") - acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") - has_fp = fp is not None - has_scalar = pre_quant_scalar is not None - - dst_mode = 0 - if pto.constexpr(acc_mode_name == "single_mode_vec1"): - dst_mode = 1 - elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): - dst_mode = "split_m" - elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): - dst_mode = "split_n" - - if pto.constexpr(has_fp): - if pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout="nz2nd", - pre_relu=("normal_relu", None, None), - pre_quant=(fp, _pre_quant_vec_mode(src_elem, dst_elem)), - ) - else: - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout="nz2nd", - pre_quant=(fp, _pre_quant_vec_mode(src_elem, dst_elem)), - ) - elif pto.constexpr(has_scalar): - if pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout="nz2nd", - pre_relu=("normal_relu", None, None), - pre_quant=(pre_quant_scalar, _pre_quant_scalar_mode(src_elem, dst_elem)), - ) - else: - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout="nz2nd", - pre_quant=(pre_quant_scalar, _pre_quant_scalar_mode(src_elem, dst_elem)), - ) - elif pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout="nz2nd", - pre_relu=("normal_relu", None, None), - ) - else: - if pto.constexpr(src_elem == pto.f32 and dst_elem == pto.f16): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout="nz2nd", - pre_quant=(pto.f16(1.0), "f32_f16"), - ) - elif pto.constexpr(src_elem == pto.f32 and dst_elem == pto.bf16): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout="nz2nd", - pre_quant=(pto.bf16(1.0), "f32_bf16"), - ) - else: - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout="nz2nd", - ) - return None - - -@pto.ckernel( - target="a5", - op="pto.tinsert", - dtypes=[ - (pto.f32, pto.i32, pto.i32, pto.f32), - (pto.f32, pto.i32, pto.i32, pto.f16), - (pto.f32, pto.i32, pto.i32, pto.bf16), - (pto.i32, pto.i32, pto.i32, pto.i32), - ], - priority=10, - constraints=[_acc_to_vec_nd_constraint], -) -def template_tinsert_acc_to_vec_nd_basic( - src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile, -): - src_ptr = src.as_ptr() - dst_ptr = dst.as_ptr() - - src_elem = src.element_type - elem_bytes = pto.bytewidth(dst.element_type) - c0_size = 32 // elem_bytes - - valid_rows = src.valid_shape[0] - valid_cols_raw = src.valid_shape[1] - valid_cols = (valid_cols_raw + c0_size - 1) // c0_size * c0_size - - dst_cols = dst.shape[1] - dst_offset = index_row * dst_cols + index_col - dst_ptr = pto.addptr(dst_ptr, dst_offset) - - dst_stride = dst_cols - src_stride = (valid_rows + 15) // 16 * 16 - - relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") - acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") - - dst_mode = 0 - if pto.constexpr(acc_mode_name == "single_mode_vec1"): - dst_mode = 1 - elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): - dst_mode = "split_m" - elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): - dst_mode = "split_n" - - if pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout="nz2nd", - pre_relu=("normal_relu", None, None), - ) - else: - if pto.constexpr(src_elem == pto.f32 and dst.element_type == pto.f16): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout="nz2nd", - pre_quant=(pto.f16(1.0), "f32_f16"), - ) - elif pto.constexpr(src_elem == pto.f32 and dst.element_type == pto.bf16): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout="nz2nd", - pre_quant=(pto.bf16(1.0), "f32_bf16"), - ) - else: - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout="nz2nd", - ) - return None - - -# --------------------------------------------------------------------------- -# Acc -> Vec (DN, NONE_BOX) -# --------------------------------------------------------------------------- - - -@pto.ckernel( - target="a5", - op="pto.tinsert", - dtypes=[ - (pto.f32, pto.i64, pto.i64, pto.f32, pto.f32, pto.i64), - (pto.f32, pto.i64, pto.i64, pto.f16, pto.f32, pto.i64), - (pto.f32, pto.i64, pto.i64, pto.bf16, pto.f32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.i32, pto.i32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.f16, pto.i32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.bf16, pto.i32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.i8, pto.i32, pto.i64), - ], - priority=10, - constraints=[_acc_to_vec_dn_constraint], -) -def template_tinsert_acc_to_vec_dn( - src: pto.Tile, - index_row: pto.i64, index_col: pto.i64, - dst: pto.Tile, - fp: pto.Tile, pre_quant_scalar: pto.i64, -): - src_ptr = src.as_ptr() - dst_ptr = dst.as_ptr() - - src_elem = src.element_type - dst_elem = dst.element_type - elem_bytes = pto.bytewidth(dst_elem) - c0_size = 32 // elem_bytes - - valid_rows_raw = src.valid_shape[0] - valid_rows = (valid_rows_raw + c0_size - 1) // c0_size * c0_size - valid_cols = src.valid_shape[1] - - dst_rows = dst.shape[0] - dst_offset = index_col * dst_rows + index_row - dst_ptr = pto.addptr(dst_ptr, dst_offset) - - dst_stride = dst_rows * elem_bytes - src_stride = (valid_rows + 15) // 16 * 16 * pto.bytewidth(src_elem) - - relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") - acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") - has_fp = fp is not None - has_scalar = pre_quant_scalar is not None - - dst_mode = 0 - if pto.constexpr(acc_mode_name == "single_mode_vec1"): - dst_mode = 1 - - if pto.constexpr(has_fp): - if pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2dn", pto.i64(1)), - pre_relu=("normal_relu", None, None), - pre_quant=(fp, _pre_quant_vec_mode(src_elem, dst_elem)), - ) - else: - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2dn", pto.i64(1)), - pre_quant=(fp, _pre_quant_vec_mode(src_elem, dst_elem)), - ) - elif pto.constexpr(has_scalar): - if pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2dn", pto.i64(1)), - pre_relu=("normal_relu", None, None), - pre_quant=(pre_quant_scalar, _pre_quant_scalar_mode(src_elem, dst_elem)), - ) - else: - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2dn", pto.i64(1)), - pre_quant=(pre_quant_scalar, _pre_quant_scalar_mode(src_elem, dst_elem)), - ) - elif pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2dn", pto.i64(1)), - pre_relu=("normal_relu", None, None), - ) - else: - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2dn", pto.i64(1)), - ) - return None - - -@pto.ckernel( - target="a5", - op="pto.tinsert", - dtypes=[ - (pto.f32, pto.i32, pto.i32, pto.f32), - (pto.f32, pto.i32, pto.i32, pto.f16), - (pto.f32, pto.i32, pto.i32, pto.bf16), - (pto.i32, pto.i32, pto.i32, pto.i32), - ], - priority=10, - constraints=[_acc_to_vec_dn_constraint], -) -def template_tinsert_acc_to_vec_dn_basic( - src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile, -): - src_ptr = src.as_ptr() - dst_ptr = dst.as_ptr() - - src_elem = src.element_type - elem_bytes = pto.bytewidth(dst.element_type) - c0_size = 32 // elem_bytes - - valid_rows_raw = src.valid_shape[0] - valid_rows = (valid_rows_raw + c0_size - 1) // c0_size * c0_size - valid_cols = src.valid_shape[1] - - dst_rows = dst.shape[0] - dst_offset = index_col * dst_rows + index_row - dst_ptr = pto.addptr(dst_ptr, dst_offset) - - dst_stride = dst_rows * elem_bytes - src_stride = (valid_rows + 15) // 16 * 16 * pto.bytewidth(src_elem) - - relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") - - dst_mode = 0 - acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") - if pto.constexpr(acc_mode_name == "single_mode_vec1"): - dst_mode = 1 - - if pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2dn", pto.i64(1)), - pre_relu=("normal_relu", None, None), - ) - else: - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2dn", pto.i64(1)), - ) - return None - - -# --------------------------------------------------------------------------- -# Acc -> Vec (NZ, ROW_MAJOR) -# --------------------------------------------------------------------------- - - -@pto.ckernel( - target="a5", - op="pto.tinsert", - dtypes=[ - (pto.f32, pto.i64, pto.i64, pto.f32, pto.f32, pto.i64), - (pto.f32, pto.i64, pto.i64, pto.f16, pto.f32, pto.i64), - (pto.f32, pto.i64, pto.i64, pto.bf16, pto.f32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.i32, pto.i32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.f16, pto.i32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.bf16, pto.i32, pto.i64), - (pto.i32, pto.i64, pto.i64, pto.i8, pto.i32, pto.i64), - ], - priority=10, - constraints=[_acc_to_vec_nz_constraint], -) -def template_tinsert_acc_to_vec_nz( - src: pto.Tile, - index_row: pto.i64, index_col: pto.i64, - dst: pto.Tile, - fp: pto.Tile, pre_quant_scalar: pto.i64, -): - src_ptr = src.as_ptr() - dst_ptr = dst.as_ptr() - - src_elem = src.element_type - dst_elem = dst.element_type - elem_bytes = pto.bytewidth(dst_elem) - c0_size = BLOCK_BYTE_SIZE // elem_bytes - - valid_rows = src.valid_shape[0] - valid_cols_raw = src.valid_shape[1] - - if pto.constexpr(dst_elem == pto.f32): - valid_cols_align = FRACTAL_NZ_ROW - else: - valid_cols_align = c0_size - valid_cols = (valid_cols_raw + valid_cols_align - 1) // valid_cols_align * valid_cols_align - - dst_rows = dst.shape[0] - col_block = index_col // c0_size - col_mod = index_col - col_block * c0_size - dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod - dst_ptr = pto.addptr(dst_ptr, dst_offset) - - dst_stride = dst_rows * c0_size * elem_bytes - src_stride = (valid_rows + 15) // 16 * 16 * pto.bytewidth(src_elem) - - relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") - acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") - has_fp = fp is not None - has_scalar = pre_quant_scalar is not None - - dst_mode = 0 - if pto.constexpr(acc_mode_name == "single_mode_vec1"): - dst_mode = 1 - elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): - dst_mode = "split_m" - elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): - dst_mode = "split_n" - - if pto.constexpr(has_fp): - if pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2nz", pto.i64(0)), - pre_relu=("normal_relu", None, None), - pre_quant=(fp, _pre_quant_vec_mode(src_elem, dst_elem)), - ) - else: - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2nz", pto.i64(0)), - pre_quant=(fp, _pre_quant_vec_mode(src_elem, dst_elem)), - ) - elif pto.constexpr(has_scalar): - if pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2nz", pto.i64(0)), - pre_relu=("normal_relu", None, None), - pre_quant=(pre_quant_scalar, _pre_quant_scalar_mode(src_elem, dst_elem)), - ) - else: - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2nz", pto.i64(0)), - pre_quant=(pre_quant_scalar, _pre_quant_scalar_mode(src_elem, dst_elem)), - ) - elif pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2nz", pto.i64(0)), - pre_relu=("normal_relu", None, None), - ) - else: - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2nz", pto.i64(0)), - ) - return None - - -@pto.ckernel( - target="a5", - op="pto.tinsert", - dtypes=[ - (pto.f32, pto.i32, pto.i32, pto.f32), - (pto.f32, pto.i32, pto.i32, pto.f16), - (pto.f32, pto.i32, pto.i32, pto.bf16), - (pto.i32, pto.i32, pto.i32, pto.i32), - ], - priority=10, - constraints=[_acc_to_vec_nz_constraint], -) -def template_tinsert_acc_to_vec_nz_basic( - src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile, -): - src_ptr = src.as_ptr() - dst_ptr = dst.as_ptr() - - src_elem = src.element_type - dst_elem = dst.element_type - elem_bytes = pto.bytewidth(dst_elem) - c0_size = BLOCK_BYTE_SIZE // elem_bytes - - valid_rows = src.valid_shape[0] - valid_cols_raw = src.valid_shape[1] - - if pto.constexpr(dst_elem == pto.f32): - valid_cols_align = FRACTAL_NZ_ROW - else: - valid_cols_align = c0_size - valid_cols = (valid_cols_raw + valid_cols_align - 1) // valid_cols_align * valid_cols_align - - dst_rows = dst.shape[0] - col_block = index_col // c0_size - col_mod = index_col - col_block * c0_size - dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod - dst_ptr = pto.addptr(dst_ptr, dst_offset) - - dst_stride = dst_rows * c0_size * elem_bytes - src_stride = (valid_rows + 15) // 16 * 16 * pto.bytewidth(src_elem) - - relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") - acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") - - dst_mode = 0 - if pto.constexpr(acc_mode_name == "single_mode_vec1"): - dst_mode = 1 - elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): - dst_mode = "split_m" - elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): - dst_mode = "split_n" - - if pto.constexpr(relu_mode_name == "normal_relu"): - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2nz", pto.i64(0)), - pre_relu=("normal_relu", None, None), - ) - else: - pto.mte_l0c_ub( - src_ptr, dst_ptr, - valid_rows, valid_cols, src_stride, dst_stride, dst_mode, - layout=("nz2nz", pto.i64(0)), - ) - return None - - -# --------------------------------------------------------------------------- -# Vec -> Vec (ND, ROW_MAJOR + NONE_BOX) - O1 -# --------------------------------------------------------------------------- - - -_VEC_TO_VEC_DTYPES = [ - (pto.f16, pto.i64, pto.i64, pto.f16), - (pto.bf16, pto.i64, pto.i64, pto.bf16), - (pto.f32, pto.i64, pto.i64, pto.f32), - (pto.i32, pto.i64, pto.i64, pto.i32), - (pto.i8, pto.i64, pto.i64, pto.i8), -] - -_VEC_TO_VEC_BASIC_DTYPES = [ - (pto.f16, pto.i32, pto.i32, pto.f16), - (pto.bf16, pto.i32, pto.i32, pto.bf16), - (pto.f32, pto.i32, pto.i32, pto.f32), - (pto.i32, pto.i32, pto.i32, pto.i32), - (pto.i8, pto.i32, pto.i32, pto.i8), -] - -_VEC_TO_MAT_SPLIT_DTYPES = [ - (pto.f16, pto.i32, pto.i32, pto.f16), - (pto.bf16, pto.i32, pto.i32, pto.bf16), - (pto.f32, pto.i32, pto.i32, pto.f32), - (pto.i32, pto.i32, pto.i32, pto.i32), - (pto.i8, pto.i32, pto.i32, pto.i8), -] - - -@pto.vkernel( - target="a5", - op="pto.tinsert", - dtypes=_VEC_TO_VEC_DTYPES, - constraints=[_vec_to_vec_nd_constraint], - advanced=True, -) -def template_tinsert_vec_to_vec_nd( - src: pto.Tile, - index_row: pto.i64, index_col: pto.i64, - dst: pto.Tile, -): - dtype = dst.element_type - elem_bytes = pto.bytewidth(dtype) - lanes = pto.get_lanes(dtype) - - valid_rows, valid_cols = src.valid_shape - src_stride = src.shape[1] - dst_stride = dst.shape[1] - - src_ptr = src.as_ptr() - - dst_offset = index_row * dst_stride + index_col - dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) - - src_stride_bytes = src_stride * elem_bytes - dst_stride_bytes = dst_stride * elem_bytes - strides_aligned = src_stride_bytes % BLOCK_BYTE_SIZE == 0 and dst_stride_bytes % BLOCK_BYTE_SIZE == 0 - - if pto.constexpr(strides_aligned): - if index_col * elem_bytes % BLOCK_BYTE_SIZE == 0: - if pto.constexpr(valid_cols * elem_bytes % BLOCK_BYTE_SIZE == 0): - row_bytes = valid_cols * elem_bytes - total_bytes = valid_rows * row_bytes - row_burst_len = row_bytes // BLOCK_BYTE_SIZE - if pto.constexpr(valid_cols == src_stride and valid_cols == dst_stride and total_bytes >= BLOCK_BYTE_SIZE): - burst_len = total_bytes // BLOCK_BYTE_SIZE - pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, 1, burst_len, 0, 0) - elif pto.constexpr(row_bytes >= BLOCK_BYTE_SIZE): - src_gap = (src_stride - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE - dst_gap = (dst_stride - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE - pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, valid_rows, row_burst_len, src_gap, dst_gap) - else: - burst_len = (total_bytes + BLOCK_BYTE_SIZE - 1) // BLOCK_BYTE_SIZE - pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, 1, burst_len, 0, 0) - else: - repeat_times = (valid_cols + lanes - 1) // lanes - for i in range(0, valid_rows, 1): - remained = valid_cols - for j in range(0, repeat_times, 1): - pred, remained = pto.make_mask(dtype, remained) - src_off = i * src_stride + j * lanes - dst_off = i * dst_stride + j * lanes - data = pto.vlds(pto.addptr(src_ptr, src_off), 0) - pto.vsts(data, pto.addptr(dst_ptr, dst_off), 0, pred) - else: - full_repeats = valid_cols // lanes - remainder = valid_cols % lanes - for i in range(0, valid_rows, 1): - ureg = pto.init_align() - src_row_off = i * src_stride - dst_row_ptr = pto.addptr(dst_ptr, i * dst_stride) - for j in range(0, full_repeats, 1): - data = pto.vlds(pto.addptr(src_ptr, src_row_off + j * lanes), 0) - ureg = pto.vstus(ureg, lanes, data, dst_row_ptr) - dst_row_ptr = pto.addptr(dst_row_ptr, lanes) - if pto.constexpr(remainder > 0): - data = pto.vlds(pto.addptr(src_ptr, src_row_off + full_repeats * lanes), 0) - ureg = pto.vstus(ureg, remainder, data, dst_row_ptr) - pto.vstas(ureg, dst_row_ptr, 0) - else: - full_repeats = valid_cols // lanes - remainder = valid_cols % lanes - for i in range(0, valid_rows, 1): - ureg = pto.init_align() - src_row_off = i * src_stride - dst_row_ptr = pto.addptr(dst_ptr, i * dst_stride) - for j in range(0, full_repeats, 1): - data = pto.vlds(pto.addptr(src_ptr, src_row_off + j * lanes), 0) - ureg = pto.vstus(ureg, lanes, data, dst_row_ptr) - dst_row_ptr = pto.addptr(dst_row_ptr, lanes) - if pto.constexpr(remainder > 0): - data = pto.vlds(pto.addptr(src_ptr, src_row_off + full_repeats * lanes), 0) - ureg = pto.vstus(ureg, remainder, data, dst_row_ptr) - pto.vstas(ureg, dst_row_ptr, 0) - return None - - -@pto.vkernel( - target="a5", - op="pto.tinsert", - dtypes=_VEC_TO_VEC_BASIC_DTYPES, - constraints=[_vec_to_vec_nd_constraint], - advanced=True, -) -def template_tinsert_vec_to_vec_nd_basic( - src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile, -): - dtype = dst.element_type - elem_bytes = pto.bytewidth(dtype) - lanes = pto.get_lanes(dtype) - - valid_rows, valid_cols = src.valid_shape - src_stride = src.shape[1] - dst_stride = dst.shape[1] - - src_ptr = src.as_ptr() - - dst_offset = index_row * dst_stride + index_col - dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) - - src_stride_bytes = src_stride * elem_bytes - dst_stride_bytes = dst_stride * elem_bytes - strides_aligned = src_stride_bytes % BLOCK_BYTE_SIZE == 0 and dst_stride_bytes % BLOCK_BYTE_SIZE == 0 - - if pto.constexpr(strides_aligned): - if index_col * elem_bytes % BLOCK_BYTE_SIZE == 0: - if pto.constexpr(valid_cols * elem_bytes % BLOCK_BYTE_SIZE == 0): - row_bytes = valid_cols * elem_bytes - total_bytes = valid_rows * row_bytes - row_burst_len = row_bytes // BLOCK_BYTE_SIZE - if pto.constexpr(valid_cols == src_stride and valid_cols == dst_stride and total_bytes >= BLOCK_BYTE_SIZE): - burst_len = total_bytes // BLOCK_BYTE_SIZE - pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, 1, burst_len, 0, 0) - elif pto.constexpr(row_bytes >= BLOCK_BYTE_SIZE): - src_gap = (src_stride - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE - dst_gap = (dst_stride - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE - pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, valid_rows, row_burst_len, src_gap, dst_gap) - else: - burst_len = (total_bytes + BLOCK_BYTE_SIZE - 1) // BLOCK_BYTE_SIZE - pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, 1, burst_len, 0, 0) - else: - repeat_times = (valid_cols + lanes - 1) // lanes - for i in range(0, valid_rows, 1): - remained = valid_cols - for j in range(0, repeat_times, 1): - pred, remained = pto.make_mask(dtype, remained) - src_off = i * src_stride + j * lanes - dst_off = i * dst_stride + j * lanes - data = pto.vlds(pto.addptr(src_ptr, src_off), 0) - pto.vsts(data, pto.addptr(dst_ptr, dst_off), 0, pred) - else: - full_repeats = valid_cols // lanes - remainder = valid_cols % lanes - for i in range(0, valid_rows, 1): - ureg = pto.init_align() - src_row_off = i * src_stride - dst_row_ptr = pto.addptr(dst_ptr, i * dst_stride) - for j in range(0, full_repeats, 1): - data = pto.vlds(pto.addptr(src_ptr, src_row_off + j * lanes), 0) - ureg = pto.vstus(ureg, lanes, data, dst_row_ptr) - dst_row_ptr = pto.addptr(dst_row_ptr, lanes) - if pto.constexpr(remainder > 0): - data = pto.vlds(pto.addptr(src_ptr, src_row_off + full_repeats * lanes), 0) - ureg = pto.vstus(ureg, remainder, data, dst_row_ptr) - pto.vstas(ureg, dst_row_ptr, 0) - else: - full_repeats = valid_cols // lanes - remainder = valid_cols % lanes - for i in range(0, valid_rows, 1): - ureg = pto.init_align() - src_row_off = i * src_stride - dst_row_ptr = pto.addptr(dst_ptr, i * dst_stride) - for j in range(0, full_repeats, 1): - data = pto.vlds(pto.addptr(src_ptr, src_row_off + j * lanes), 0) - ureg = pto.vstus(ureg, lanes, data, dst_row_ptr) - dst_row_ptr = pto.addptr(dst_row_ptr, lanes) - if pto.constexpr(remainder > 0): - data = pto.vlds(pto.addptr(src_ptr, src_row_off + full_repeats * lanes), 0) - ureg = pto.vstus(ureg, remainder, data, dst_row_ptr) - pto.vstas(ureg, dst_row_ptr, 0) - return None - - -@pto.vkernel( - target="a5", - op="pto.tinsert", - dtypes=_VEC_TO_VEC_DTYPES, - constraints=[_vec_to_vec_nd_scalar_constraint], - advanced=True, -) -def template_tinsert_vec_to_vec_nd_scalar( - src: pto.Tile, - index_row: pto.i64, index_col: pto.i64, - dst: pto.Tile, -): - dst_stride = dst.shape[1] - src_ptr = src.as_ptr() - dst_ptr = dst.as_ptr() - src_val = pto.load_scalar(src_ptr, 0) - dst_elem_offset = index_row * dst_stride + index_col - pto.store_scalar(src_val, dst_ptr, dst_elem_offset) - return None - - -@pto.vkernel( - target="a5", - op="pto.tinsert", - dtypes=_VEC_TO_VEC_BASIC_DTYPES, - constraints=[_vec_to_vec_nd_scalar_constraint], - advanced=True, -) -def template_tinsert_vec_to_vec_nd_scalar_basic( - src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile, -): - dst_stride = dst.shape[1] - src_ptr = src.as_ptr() - dst_ptr = dst.as_ptr() - src_val = pto.load_scalar(src_ptr, 0) - dst_elem_offset = index_row * dst_stride + index_col - pto.store_scalar(src_val, dst_ptr, dst_elem_offset) - return None - - -# --------------------------------------------------------------------------- -# Vec -> Vec (NZ, COL_MAJOR + ROW_MAJOR) - O2 -# --------------------------------------------------------------------------- - - -@pto.vkernel( - target="a5", - op="pto.tinsert", - dtypes=_VEC_TO_VEC_DTYPES, - constraints=[_vec_to_vec_nz_constraint], - advanced=True, -) -def template_tinsert_vec_to_vec_nz( - src: pto.Tile, - index_row: pto.i64, index_col: pto.i64, - dst: pto.Tile, -): - dtype = dst.element_type - elem_bytes = pto.bytewidth(dtype) - c0_size = BLOCK_BYTE_SIZE // elem_bytes - - valid_rows = src.valid_shape[0] - valid_cols = src.valid_shape[1] - dst_rows = dst.shape[0] - - src_ptr = src.as_ptr() - - col_block = index_col // c0_size - col_mod = index_col - col_block * c0_size - dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod - dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) - - burst_num = (valid_cols + c0_size - 1) // c0_size - burst_len = valid_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE - - compact = src.config.compact_mode - if pto.constexpr(compact == pto.CompactMode.NULL): - src_stride_rows = src.shape[0] - elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): - src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + 1 - else: - src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW - src_gap = src_stride_rows - valid_rows - dst_gap = dst_rows - valid_rows - - pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, burst_num, burst_len, src_gap, dst_gap) - return None - - -@pto.vkernel( - target="a5", - op="pto.tinsert", - dtypes=_VEC_TO_VEC_BASIC_DTYPES, - constraints=[_vec_to_vec_nz_constraint], - advanced=True, -) -def template_tinsert_vec_to_vec_nz_basic( - src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile, -): - dtype = dst.element_type - elem_bytes = pto.bytewidth(dtype) - c0_size = BLOCK_BYTE_SIZE // elem_bytes - - valid_rows = src.valid_shape[0] - valid_cols = src.valid_shape[1] - dst_rows = dst.shape[0] - - src_ptr = src.as_ptr() - - col_block = index_col // c0_size - col_mod = index_col - col_block * c0_size - dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod - dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) - - burst_num = (valid_cols + c0_size - 1) // c0_size - burst_len = valid_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE - - compact = src.config.compact_mode - if pto.constexpr(compact == pto.CompactMode.NULL): - src_stride_rows = src.shape[0] - elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): - src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + 1 - else: - src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW - src_gap = src_stride_rows - valid_rows - dst_gap = dst_rows - valid_rows - - pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, burst_num, burst_len, src_gap, dst_gap) - return None - - -# --------------------------------------------------------------------------- -# Vec -> Mat (NZ, COL_MAJOR + ROW_MAJOR) - O3 -# --------------------------------------------------------------------------- - - -_VEC_TO_MAT_DTYPES = [ - (pto.f16, pto.i64, pto.i64, pto.f16), - (pto.bf16, pto.i64, pto.i64, pto.bf16), - (pto.f32, pto.i64, pto.i64, pto.f32), - (pto.i32, pto.i64, pto.i64, pto.i32), - (pto.i8, pto.i64, pto.i64, pto.i8), -] - -_VEC_TO_MAT_BASIC_DTYPES = [ - (pto.f16, pto.i32, pto.i32, pto.f16), - (pto.bf16, pto.i32, pto.i32, pto.bf16), - (pto.f32, pto.i32, pto.i32, pto.f32), - (pto.i32, pto.i32, pto.i32, pto.i32), - (pto.i8, pto.i32, pto.i32, pto.i8), -] - - -@pto.vkernel( - target="a5", - op="pto.tinsert", - dtypes=_VEC_TO_MAT_DTYPES, - constraints=[_vec_to_mat_nz_constraint], - advanced=True, -) -def template_tinsert_vec_to_mat_nz( - src: pto.Tile, - index_row: pto.i64, index_col: pto.i64, - dst: pto.Tile, -): - dtype = dst.element_type - elem_bytes = pto.bytewidth(dtype) - c0_size = BLOCK_BYTE_SIZE // elem_bytes - - valid_rows = src.valid_shape[0] - valid_cols = src.valid_shape[1] - dst_rows = dst.shape[0] - - src_ptr = src.as_ptr() - - col_block = index_col // c0_size - col_mod = index_col - col_block * c0_size - dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod - dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) - - burst_num = (valid_cols + c0_size - 1) // c0_size - burst_len = valid_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE - - compact = src.config.compact_mode - if pto.constexpr(compact == pto.CompactMode.NULL): - src_stride_rows = src.shape[0] - elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): - src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + 1 - else: - src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW - src_gap = src_stride_rows - valid_rows - dst_gap = dst_rows - valid_rows - - pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(burst_num, src_gap, dst_gap)) - return None - - -@pto.vkernel( - target="a5", - op="pto.tinsert", - dtypes=_VEC_TO_MAT_BASIC_DTYPES, - constraints=[_vec_to_mat_nz_constraint], - advanced=True, -) -def template_tinsert_vec_to_mat_nz_basic( - src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile, -): - dtype = dst.element_type - elem_bytes = pto.bytewidth(dtype) - c0_size = BLOCK_BYTE_SIZE // elem_bytes - - valid_rows = src.valid_shape[0] - valid_cols = src.valid_shape[1] - dst_rows = dst.shape[0] - - src_ptr = src.as_ptr() - - col_block = index_col // c0_size - col_mod = index_col - col_block * c0_size - dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod - dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) - - burst_num = (valid_cols + c0_size - 1) // c0_size - burst_len = valid_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE - - compact = src.config.compact_mode - if pto.constexpr(compact == pto.CompactMode.NULL): - src_stride_rows = src.shape[0] - elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): - src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + 1 - else: - src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW - src_gap = src_stride_rows - valid_rows - dst_gap = dst_rows - valid_rows - - pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(burst_num, src_gap, dst_gap)) - return None - - -# --------------------------------------------------------------------------- -# Vec -> Mat (ND, ROW_MAJOR + NONE_BOX) - O4 -# --------------------------------------------------------------------------- - - -@pto.vkernel( - target="a5", - op="pto.tinsert", - dtypes=_VEC_TO_MAT_DTYPES, - constraints=[_vec_to_mat_nd_constraint], - advanced=True, -) -def template_tinsert_vec_to_mat_nd( - src: pto.Tile, - index_row: pto.i64, index_col: pto.i64, - dst: pto.Tile, -): - dtype = dst.element_type - elem_bytes = pto.bytewidth(dtype) - - valid_rows, valid_cols = src.valid_shape - src_cols = src.shape[1] - dst_cols = dst.shape[1] - - src_ptr = src.as_ptr() - - dst_offset = index_row * dst_cols + index_col - dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) - - row_bytes = valid_cols * elem_bytes - total_bytes = valid_rows * row_bytes - - if pto.constexpr(valid_cols == src_cols and valid_cols == dst_cols and total_bytes >= BLOCK_BYTE_SIZE): - burst_len = total_bytes // BLOCK_BYTE_SIZE - pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(1, 0, 0)) - elif pto.constexpr(row_bytes >= BLOCK_BYTE_SIZE): - row_burst_len = row_bytes // BLOCK_BYTE_SIZE - src_row_gap = (src_cols - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE - dst_row_gap = (dst_cols - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE - pto.mte_ub_l1(src_ptr, dst_ptr, row_burst_len, nburst=(valid_rows, src_row_gap, dst_row_gap)) - else: - burst_len = (total_bytes + BLOCK_BYTE_SIZE - 1) // BLOCK_BYTE_SIZE - pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(1, 0, 0)) - return None - - -@pto.vkernel( - target="a5", - op="pto.tinsert", - dtypes=_VEC_TO_MAT_BASIC_DTYPES, - constraints=[_vec_to_mat_nd_constraint], - advanced=True, -) -def template_tinsert_vec_to_mat_nd_basic( - src: pto.Tile, - index_row: pto.i32, index_col: pto.i32, - dst: pto.Tile, -): - dtype = dst.element_type - elem_bytes = pto.bytewidth(dtype) - - valid_rows, valid_cols = src.valid_shape - src_cols = src.shape[1] - dst_cols = dst.shape[1] - - src_ptr = src.as_ptr() - - dst_offset = index_row * dst_cols + index_col - dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) - - row_bytes = valid_cols * elem_bytes - total_bytes = valid_rows * row_bytes - - if pto.constexpr(valid_cols == src_cols and valid_cols == dst_cols and total_bytes >= BLOCK_BYTE_SIZE): - burst_len = total_bytes // BLOCK_BYTE_SIZE - pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(1, 0, 0)) - elif pto.constexpr(row_bytes >= BLOCK_BYTE_SIZE): - row_burst_len = row_bytes // BLOCK_BYTE_SIZE - src_row_gap = (src_cols - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE - dst_row_gap = (dst_cols - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE - pto.mte_ub_l1(src_ptr, dst_ptr, row_burst_len, nburst=(valid_rows, src_row_gap, dst_row_gap)) - else: - burst_len = (total_bytes + BLOCK_BYTE_SIZE - 1) // BLOCK_BYTE_SIZE - pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(1, 0, 0)) - return None - - -# --------------------------------------------------------------------------- -# Vec -> Mat (NZ, Split2/Split4) - O5 -# Split large-tile NZ DMA into 2/4 independent segments to reduce L1 bank -# conflicts (mirrors pto-isa TInsertMode::SPLIT2 / SPLIT4). -# --------------------------------------------------------------------------- - - -def _make_split_template(split_count): - @pto.vkernel( - target="a5", - op="pto.tinsert", - dtypes=_VEC_TO_MAT_SPLIT_DTYPES, - name=f"template_tinsert_vec_to_mat_nz_split{split_count}", - constraints=[_vec_to_mat_nz_constraint], - advanced=True, - ) - def _split_fn(src: pto.Tile, index_row: pto.i32, index_col: pto.i32, dst: pto.Tile): - dtype = dst.element_type - elem_bytes = pto.bytewidth(dtype) - c0_size = BLOCK_BYTE_SIZE // elem_bytes - - valid_rows = src.valid_shape[0] - valid_cols = src.valid_shape[1] - dst_rows = dst.shape[0] - aligned_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW - - src_ptr = src.as_ptr() - - col_block = index_col // c0_size - col_mod = index_col - col_block * c0_size - dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod - dst_base = pto.addptr(dst.as_ptr(), dst_offset) - - total_burst_num = (valid_cols + c0_size - 1) // c0_size - burst_len = aligned_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE - - compact = src.config.compact_mode - if pto.constexpr(compact == pto.CompactMode.NULL): - src_stride_rows = src.shape[0] - elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): - src_stride_rows = aligned_rows + 1 - else: - src_stride_rows = aligned_rows - src_gap = src_stride_rows - aligned_rows - dst_gap = dst_rows - aligned_rows - - part_num = total_burst_num // split_count - last_num = total_burst_num - part_num * (split_count - 1) - src_block_size = (burst_len + src_gap) * BLOCK_BYTE_SIZE // elem_bytes - dst_block_size = dst_rows * c0_size - - pto.mte_ub_l1(src_ptr, dst_base, burst_len, nburst=(part_num, src_gap, dst_gap)) - - src_ptr1 = pto.addptr(src_ptr, part_num * src_block_size) - dst_ptr1 = pto.addptr(dst_base, part_num * dst_block_size) - if pto.constexpr(split_count == 2): - pto.mte_ub_l1(src_ptr1, dst_ptr1, burst_len, nburst=(last_num, src_gap, dst_gap)) - else: - pto.mte_ub_l1(src_ptr1, dst_ptr1, burst_len, nburst=(part_num, src_gap, dst_gap)) - - if pto.constexpr(split_count == 4): - src_ptr2 = pto.addptr(src_ptr, 2 * part_num * src_block_size) - dst_ptr2 = pto.addptr(dst_base, 2 * part_num * dst_block_size) - pto.mte_ub_l1(src_ptr2, dst_ptr2, burst_len, nburst=(part_num, src_gap, dst_gap)) - - src_ptr3 = pto.addptr(src_ptr, 3 * part_num * src_block_size) - dst_ptr3 = pto.addptr(dst_base, 3 * part_num * dst_block_size) - pto.mte_ub_l1(src_ptr3, dst_ptr3, burst_len, nburst=(last_num, src_gap, dst_gap)) - return None - - return _split_fn - - -template_tinsert_vec_to_mat_nz_split2 = _make_split_template(2) -template_tinsert_vec_to_mat_nz_split4 = _make_split_template(4) diff --git a/lib/TileOps/tload_template.py b/lib/TileOps/tload_template.py deleted file mode 100644 index 310c8d243e..0000000000 --- a/lib/TileOps/tload_template.py +++ /dev/null @@ -1,528 +0,0 @@ -# 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. - -"""TileLang DSL template for `pto.tload`""" - -import tilelang_dsl as pto - - -def _constraint_scalar(value): - return value.value if hasattr(value, "value") else value - - -def _known_eq(lhs, rhs) -> bool: - lhs_value = _constraint_scalar(lhs) - rhs_value = _constraint_scalar(rhs) - if lhs_value is None or rhs_value is None: - return True - return lhs_value == rhs_value - - -def _known_le(lhs, rhs) -> bool: - lhs_value = _constraint_scalar(lhs) - rhs_value = _constraint_scalar(rhs) - if lhs_value is None or rhs_value is None: - return True - return lhs_value <= rhs_value - - -def _match_tile_layout(dst, *, row_major: bool, s_layout) -> bool: - b_layout_ok = ( - dst.config.b_layout == pto.BLayout.ROW_MAJOR - if row_major - else dst.config.b_layout != pto.BLayout.ROW_MAJOR - ) - return b_layout_ok and dst.config.s_layout == s_layout - - -def _check_load_bounds(src, dst, *, logical_rows, logical_cols=None, stride_axis=None) -> bool: - if src.rank != 5: - return False - if stride_axis is not None and not _known_eq(src.strides[stride_axis], 1): - return False - if not _known_le(dst.valid_shape[0], logical_rows): - return False - if not _known_le(logical_rows, dst.shape[0]): - return False - if not _known_le(dst.valid_shape[0], dst.shape[0]): - return False - if logical_cols is not None: - if not _known_le(dst.valid_shape[1], logical_cols): - return False - if not _known_le(logical_cols, dst.shape[1]): - return False - if not _known_le(dst.valid_shape[1], dst.shape[1]): - return False - return True - - -def _tload_preconditions_nd2nd(src, dst) -> bool: - logical_rows = src.shape[0] * src.shape[1] * src.shape[2] * src.shape[3] - logical_cols = src.shape[4] - return _match_tile_layout( - dst, row_major=True, s_layout=pto.SLayout.NONE_BOX - ) and _check_load_bounds( - src, dst, logical_rows=logical_rows, logical_cols=logical_cols, stride_axis=4 - ) - - -def _tload_preconditions_dn2dn(src, dst) -> bool: - logical_rows = src.shape[3] - logical_cols = src.shape[0] * src.shape[1] * src.shape[2] * src.shape[4] - return _match_tile_layout( - dst, row_major=False, s_layout=pto.SLayout.NONE_BOX - ) and _check_load_bounds( - src, dst, logical_rows=logical_rows, logical_cols=logical_cols, stride_axis=3 - ) - -def _tload_preconditions_nz2nz(src, dst) -> bool: - logical_rows = src.shape[2] - return _match_tile_layout( - dst, row_major=False, s_layout=pto.SLayout.ROW_MAJOR - ) and _check_load_bounds( - src, dst, logical_rows=logical_rows - ) - - -@pto.vkernel( - target="a5", - op="pto.tload", - advanced=True, - constraints=[_tload_preconditions_nd2nd], -) -def template_tload_nd2nd(src: pto.PartitionTensorView, dst: pto.Tile): - dtype = dst.element_type - elem_bytes = pto.bytewidth(dtype) - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - pto.set_mov_pad_val(dst.pad_value.eval()) - - g0, g1, g2, g3, g4 = src.shape - s0, s1, s2, s3, s4 = src.strides - - valid_rows, valid_cols = dst.valid_shape - ub_rows, ub_cols = dst.shape - - n_burst = g3 - len_burst = g4 * elem_bytes - gm_stride = s3 * elem_bytes - ub_stride = ub_cols * elem_bytes - - dst_stride2 = g3 * ub_cols - dst_stride1 = g2 * dst_stride2 - dst_stride0 = g1 * dst_stride1 - - loop1 = g2 - loop2 = g1 - loop1_src_stride = s2 * elem_bytes - loop1_dst_stride = dst_stride2 * elem_bytes - loop2_src_stride = s1 * elem_bytes - loop2_dst_stride = dst_stride1 * elem_bytes - - gm_ptr = src.as_ptr() - ub_ptr = dst.as_ptr() - - if loop1 != 1 or loop2 != 1: - pto.set_loop2_stride_outtoub( - src_stride=loop2_src_stride, dst_stride=loop2_dst_stride - ) - pto.set_loop1_stride_outtoub( - src_stride=loop1_src_stride, dst_stride=loop1_dst_stride - ) - pto.set_loop_size_outtoub(loop1=loop1, loop2=loop2) - - for i in range(0, g0, 1): - src_i = pto.addptr(gm_ptr, i * s0) - dst_i = pto.addptr(ub_ptr, i * dst_stride0) - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - pto.copy_gm_to_ubuf( - dst=dst_i, - src=src_i, - n_burst=n_burst, - len_burst=len_burst, - gm_stride=gm_stride, - ub_stride=ub_stride, - enable_ub_pad=True, - ) - else: - pto.copy_gm_to_ubuf( - dst=dst_i, - src=src_i, - n_burst=n_burst, - len_burst=len_burst, - gm_stride=gm_stride, - ub_stride=ub_stride, - enable_ub_pad=False, - ) - - if loop1 != 1 or loop2 != 1: - pto.set_loop_size_outtoub(loop1=1, loop2=1) - return - -@pto.vkernel( - target="a5", - op="pto.tload", - advanced=True, - constraints=[_tload_preconditions_dn2dn], -) -def template_tload_dn2dn(src: pto.PartitionTensorView, dst: pto.Tile): - dtype = dst.element_type - elem_bytes = pto.bytewidth(dtype) - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - pto.set_mov_pad_val(dst.pad_value.eval()) - - # rank-5 partition view metadata. - g0, g1, g2, g3, g4 = src.shape - s0, s1, s2, s3, s4 = src.strides - - tile_rows, tile_cols = dst.shape - valid_rows, valid_cols = dst.valid_shape - - n_burst = g4 - len_burst = valid_rows * elem_bytes - gm_stride = s4 * elem_bytes - ub_stride = tile_rows * elem_bytes - - # The UB destination tile has a compact col-major layout with column - # height `tile_rows`. From the innermost `g4 × tile_rows` block, three - # levels of stride are derived recursively. - dst_stride2 = g4 * tile_rows - dst_stride1 = g2 * dst_stride2 - dst_stride0 = g1 * dst_stride1 - - # loop1 <-> g2 (inner), loop2 <-> g1 (outer), software for <-> g0. - loop1 = g2 - loop2 = g1 - loop1_src_stride = s2 * elem_bytes - loop1_dst_stride = dst_stride2 * elem_bytes - loop2_src_stride = s1 * elem_bytes - loop2_dst_stride = dst_stride1 * elem_bytes - - gm_ptr = src.as_ptr() - ub_ptr = dst.as_ptr() - - if loop1 != 1 or loop2 != 1: - pto.set_loop2_stride_outtoub( - src_stride=loop2_src_stride, dst_stride=loop2_dst_stride - ) - pto.set_loop1_stride_outtoub( - src_stride=loop1_src_stride, dst_stride=loop1_dst_stride - ) - pto.set_loop_size_outtoub(loop1=loop1, loop2=loop2) - - for i in range(0, g0, 1): - src_i = pto.addptr(gm_ptr, i * s0) - dst_i = pto.addptr(ub_ptr, i * dst_stride0) - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - pto.copy_gm_to_ubuf( - dst=dst_i, - src=src_i, - n_burst=n_burst, - len_burst=len_burst, - gm_stride=gm_stride, - ub_stride=ub_stride, - enable_ub_pad=True, - ) - else: - pto.copy_gm_to_ubuf( - dst=dst_i, - src=src_i, - n_burst=n_burst, - len_burst=len_burst, - gm_stride=gm_stride, - ub_stride=ub_stride, - enable_ub_pad=False, - ) - - if loop1 != 1 or loop2 != 1: - pto.set_loop_size_outtoub(loop1=1, loop2=1) - return - -@pto.vkernel( - target="a5", - op="pto.tload", - advanced=True, - constraints=[_tload_preconditions_nz2nz], -) -def template_tload_nz2nz(src: pto.PartitionTensorView, dst: pto.Tile): - dtype = dst.element_type - elem_bytes = pto.bytewidth(dtype) - - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - pto.set_mov_pad_val(dst.pad_value.eval()) - - # rank-5 partition view metadata. NZ static tile constraints (g3/g4 vs - # dtype relationship) are enforced by higher-layer schema / static checks; - # here we only keep the runtime DMA formula. - g0, g1, g2, g3, g4 = src.shape - s0, s1, s2, s3, s4 = src.strides - - tile_rows, tile_cols = dst.shape - valid_rows, valid_cols = dst.valid_shape - - c0_size_bytes = 32 - n_burst = g1 - len_burst = valid_rows * c0_size_bytes - gm_stride = s1 * elem_bytes - ub_stride = tile_rows * c0_size_bytes - - # Each g0 block in UB contains `g1` NZ sub-blocks; each sub-block has `g4` columns. - tile_stride = g1 * tile_rows * g4 - - gm_ptr = src.as_ptr() - ub_ptr = dst.as_ptr() - - # NZ2NZ always uses normal mode; do not reuse loop1/loop2 registers. - pto.set_loop_size_outtoub(loop1=1, loop2=1) - for i in range(0, g0, 1): - src_i = pto.addptr(gm_ptr, i * s0) - dst_i = pto.addptr(ub_ptr, i * tile_stride) - if pto.constexpr(dst.pad_value != pto.PadValue.NULL): - pto.copy_gm_to_ubuf( - dst=dst_i, - src=src_i, - n_burst=n_burst, - len_burst=len_burst, - gm_stride=gm_stride, - ub_stride=ub_stride, - enable_ub_pad=True, - ) - else: - pto.copy_gm_to_ubuf( - dst=dst_i, - src=src_i, - n_burst=n_burst, - len_burst=len_burst, - gm_stride=gm_stride, - ub_stride=ub_stride, - enable_ub_pad=False, - ) - return - - -# ============================================================================ -# Cube Matrix Templates: TLOAD.MAT (GM → L1) -# ============================================================================ - -def _constraint_tload_mat_base(src, dst) -> bool: - """TLOAD.MAT base constraint check""" - # src must be GM - src_space = src.memory_space - src_space_value = "gm" if src_space is None else (src_space.value if hasattr(src_space, "value") else src_space) - if src_space_value not in {"gm", "GM"}: - return False - # dst must be MemorySpace.MAT - dst_space = dst.memory_space - if dst_space is None: - return False - dst_space_value = dst_space.value if hasattr(dst_space, "value") else dst_space - if dst_space_value not in {"mat", "MAT"}: - return False - # dst must be a 2D Tile - if dst.rank != 2: - return False - # dtype check - dst_dtype = dst.dtype - if dst_dtype is None: - return False - dtype_name = dst_dtype.name if hasattr(dst_dtype, "name") else str(dst_dtype) - supported_dtypes = {"f16", "bf16", "f32", "i8", "si8", "ui8", "i16", "si16", "ui16", "i32", "si32"} - if dtype_name not in supported_dtypes: - return False - return True - - -def _constraint_tload_mat_nd2nz(src, dst) -> bool: - """TLOAD.MAT ND2NZ fractal load constraint""" - if not _constraint_tload_mat_base(src, dst): - return False - # dst layout must be col_major + row_major slayout (NZ format) - config = dst.config - if config is None: - return False - b_layout = config.b_layout - s_layout = config.s_layout - if b_layout is None or s_layout is None: - return False - b_layout_value = b_layout.value if hasattr(b_layout, "value") else b_layout - s_layout_value = s_layout.value if hasattr(s_layout, "value") else s_layout - # COL_MAJOR + ROW_MAJOR corresponds to NZ format - if b_layout_value not in {"col_major", "COL_MAJOR"} or s_layout_value not in {"row_major", "ROW_MAJOR"}: - return False - # ND2NZ: source is in ND (row-major) format where the inner dimension (g4) - # corresponds to the tile column count. Disambiguates from DN format where - # g4 corresponds to the tile row count. Use the *static* dst.shape - # (rows/cols from the tile_buf type) rather than dst.valid_shape, which - # may be dynamic ([null, null]) and would otherwise let both ND2NZ and - # DN2NZ match when the shape comparison is skipped. - if not (hasattr(src, 'rank') and src.rank == 5): - return False - dst_cols = dst.shape[1] if hasattr(dst, 'shape') and dst.shape is not None else None - if dst_cols is not None and hasattr(src, 'shape') and src.shape is not None: - src_inner = src.shape[4] if len(src.shape) >= 5 else None - if src_inner is not None: - if not _known_eq(dst_cols, src_inner): - return False - return True - - -def _constraint_tload_mat_dn2nz(src, dst) -> bool: - """TLOAD.MAT DN2NZ fractal load constraint""" - if not _constraint_tload_mat_base(src, dst): - return False - config = dst.config - if config is None: - return False - b_layout = config.b_layout - s_layout = config.s_layout - if b_layout is None or s_layout is None: - return False - b_layout_value = b_layout.value if hasattr(b_layout, "value") else b_layout - s_layout_value = s_layout.value if hasattr(s_layout, "value") else s_layout - if b_layout_value not in {"col_major", "COL_MAJOR"} or s_layout_value not in {"row_major", "ROW_MAJOR"}: - return False - # DN2NZ: source is in DN (col-major) format where the inner dimension (g4) - # corresponds to the tile row count. Disambiguates from ND format where - # g4 corresponds to the tile column count. Use the *static* dst.shape - # (rows/cols from the tile_buf type) rather than dst.valid_shape, which - # may be dynamic and would otherwise let both templates match. - if not (hasattr(src, 'rank') and src.rank == 5): - return False - dst_rows = dst.shape[0] if hasattr(dst, 'shape') and dst.shape is not None else None - if dst_rows is not None and hasattr(src, 'shape') and src.shape is not None: - src_inner = src.shape[4] if len(src.shape) >= 5 else None - if src_inner is not None: - if not _known_eq(dst_rows, src_inner): - return False - return True - - -@pto.ckernel( - target="a5", - op="pto.tload", - priority=1, - dtypes=[ - (pto.f16, pto.f16), - (pto.bf16, pto.bf16), - (pto.f32, pto.f32), - ], - constraints=[_constraint_tload_mat_nd2nz], - name="tload_gm_to_mat_nd2nz", -) -def template_tload_gm_to_mat_nd2nz(src: pto.PartitionTensorView, dst: pto.Tile): - """GM -> MAT ND2NZ fractal load template - - Load Row-Major (ND) format data from GM into L1 MAT Buffer in NZ format. - - Args: - src: PartitionTensorView with GM memory_space (row-major source) - dst: Tile with MAT memory_space, shape=(M, K), col_major layout - - Uses: - pto.mte_gm_l1_frac with mode=FractalMode.ND2NZ - """ - m, k = dst.valid_shape - - gm_ptr = src.as_ptr() - mat_ptr = dst.as_ptr() - - # Element byte width (src and dst share dtype per the template dtypes table). - elem_bytes = pto.bytewidth(dst.element_type) - - # rank-5 partition view metadata (g3 = M rows, g4 = K cols for ND source). - g0, g1, g2, g3, g4 = src.shape - s0, s1, s2, s3, s4 = src.strides - - # ND2NZ parameter calculation - # n_value = M (row count), d_value = K (column count) - n_value = m - d_value = k - - # src_layout: inner stride in BYTES = src physical row stride (s3). - # Hardware loop1_src_stride is in bytes (pto-isa a5/TLoad.hpp: - # GetByteSize(gStride3)); a strided GM slice has s3 != K, so the - # element stride must be scaled by elem_bytes. - src_inner_stride = s3 * elem_bytes - - # dst_group: (group_count, loop2_stride, loop3_stride, loop4_stride) - # For simple single-block case: (1, 1, m, 0) - - # ctrl: (l2_cache_ctrl, smallc0_en) - # (0, False) → l2_cache_ctrl=0, smallc0_en=False - - pto.mte_gm_l1_frac( - gm_ptr, mat_ptr, pto.FractalMode.ND2NZ, - shape=(n_value, d_value), - src_layout=(src_inner_stride,), - dst_group=(1, 1, m, 0), - ctrl=(0, False) - ) - return - - -@pto.ckernel( - target="a5", - op="pto.tload", - priority=1, - dtypes=[ - (pto.f16, pto.f16), - (pto.bf16, pto.bf16), - (pto.f32, pto.f32), - ], - constraints=[_constraint_tload_mat_dn2nz], - name="tload_gm_to_mat_dn2nz", -) -def template_tload_gm_to_mat_dn2nz(src: pto.PartitionTensorView, dst: pto.Tile): - """GM -> MAT DN2NZ fractal load template - - Load Col-Major (DN) format data from GM into L1 MAT Buffer in NZ format. - The output is still logically N x D; only the memory layout changes to NZ - (fractal). No logical shape conversion is needed. - - Args: - src: PartitionTensorView with GM memory_space (col-major source) - dst: Tile with MAT memory_space, shape=(M, K), col_major layout - - Uses: - pto.mte_gm_l1_frac with mode=FractalMode.DN2NZ - """ - m, k = dst.valid_shape - - gm_ptr = src.as_ptr() - mat_ptr = dst.as_ptr() - - # Element byte width (src and dst share dtype per the template dtypes table). - elem_bytes = pto.bytewidth(dst.element_type) - - # rank-5 partition view metadata (g3 = K, g4 = M for DN col-major source). - g0, g1, g2, g3, g4 = src.shape - s0, s1, s2, s3, s4 = src.strides - - # DN2NZ parameter calculation - # For DN format, the original shape is (K, M) -- no logical conversion - # needed. dn2nz writes the same logical N x D result into NZ layout. - # n_value = K, d_value = M - n_value = k - d_value = m - - # src_layout: inner stride in BYTES = src physical stride along the M - # direction (s4). pto-isa TLoadCubeND2NZ uses gStride4 when layout == DN. - src_inner_stride = s4 * elem_bytes - - # dst_group: (group_count, loop2_stride, loop3_stride, loop4_stride) - # (1, 1, k, 0) - # ctrl: (l2_cache_ctrl, smallc0_en) - # (0, False) - - pto.mte_gm_l1_frac( - gm_ptr, mat_ptr, pto.FractalMode.DN2NZ, - shape=(n_value, d_value), - src_layout=(src_inner_stride,), - dst_group=(1, 1, k, 0), - ctrl=(0, False) - ) - return diff --git a/lib/TileOps/tlog_template.py b/lib/TileOps/tlog_template.py deleted file mode 100644 index e34b2faa54..0000000000 --- a/lib/TileOps/tlog_template.py +++ /dev/null @@ -1,66 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tlog""" - -import tilelang_dsl as pto - - -@pto.inline_proc -def _tlog_high_precision(src: pto.Tile, dst: pto.Tile, dtype, valid_rows, valid_cols): - if pto.constexpr(dtype == pto.f16): - subnormal_threshold = pto.f16("0x03FF") - mul_factor = pto.f16("0x6400") - compensation = pto.f16(-6.931471805599453094172) - elif pto.constexpr(dtype == pto.f32): - subnormal_threshold = pto.f32("0x007FFFFF") - mul_factor = pto.f32("0x4B000000") - compensation = pto.f32(-15.9423851528787421) - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vinput = pto.vlds(src[row, col:]) - cmp_mask = pto.vcmps(vinput, subnormal_threshold, mask, pto.CmpMode.LT) - scaled = pto.vmuls(vinput, mul_factor, mask) - selected_input = pto.vsel(scaled, vinput, cmp_mask) - log_result = pto.vln(selected_input, mask) - compensated = pto.vadds(log_result, compensation, mask) - result = pto.vsel(compensated, log_result, cmp_mask) - pto.vsts(result, dst[row, col:], mask) - return None - - -@pto.inline_proc -def _tlog_default(src: pto.Tile, dst: pto.Tile, dtype, valid_rows, valid_cols): - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vinput = pto.vlds(src[row, col:]) - result = pto.vln(vinput, mask) - pto.vsts(result, dst[row, col:], mask) - return None - - -@pto.vkernel( - target="a5", - op="pto.tlog", - advanced=True -) -def template_tlog(src: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - precision_type = pto.get_op_attr("precisionType", "default") - - if pto.constexpr(precision_type == "high_precision"): - _tlog_high_precision(src, dst, dtype, valid_rows, valid_cols) - else: - _tlog_default(src, dst, dtype, valid_rows, valid_cols) - return \ No newline at end of file diff --git a/lib/TileOps/tlrelu_template.py b/lib/TileOps/tlrelu_template.py deleted file mode 100644 index 33087ca513..0000000000 --- a/lib/TileOps/tlrelu_template.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tlrelu (Leaky ReLU with scalar slope)""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tlrelu", - advanced=True -) -def template_tlrelu(src: pto.Tile, slope: pto.f32, dst: pto.Tile): - """Leaky ReLU: dst = src if src > 0 else src * slope. - - Semantics: - For each element (i, j): - dst[i, j] = src[i, j] > 0 ? src[i, j] : slope * src[i, j] - - Supported data types: f16, f32 - """ - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - lanes = pto.get_lanes(dtype) - if pto.constexpr(dtype == pto.f16): - slope_scalar = pto.f16(slope) - else: - slope_scalar = slope - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - src_vec = pto.vlds(src[row, col:]) - result = pto.vlrelu(src_vec, slope_scalar, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tmatmul_acc_template.py b/lib/TileOps/tmatmul_acc_template.py deleted file mode 100644 index 5143ef2ac9..0000000000 --- a/lib/TileOps/tmatmul_acc_template.py +++ /dev/null @@ -1,27 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmatmul.acc.""" - -import tilelang_dsl as pto - - -@pto.ckernel( - target="a5", - op="pto.tmatmul.acc", - dtypes=[ - (pto.f32, pto.f16, pto.f16, pto.f32), - (pto.f32, pto.bf16, pto.bf16, pto.f32), - (pto.f32, pto.f32, pto.f32, pto.f32), - ], -) -def template_tmatmul_acc(acc_in: pto.Tile, lhs: pto.Tile, rhs: pto.Tile, dst: pto.Tile): - m, k = lhs.valid_shape - _, n = rhs.valid_shape - pto.mad_acc(lhs.as_ptr(), rhs.as_ptr(), dst.as_ptr(), m, n, k, disable_gemv=True) - return None \ No newline at end of file diff --git a/lib/TileOps/tmatmul_bias_template.py b/lib/TileOps/tmatmul_bias_template.py deleted file mode 100644 index 19d8530cae..0000000000 --- a/lib/TileOps/tmatmul_bias_template.py +++ /dev/null @@ -1,33 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmatmul.bias.""" - -import tilelang_dsl as pto - - -@pto.ckernel( - target="a5", - op="pto.tmatmul.bias", - dtypes=[ - (pto.f16, pto.f16, pto.f32, pto.f32), - (pto.bf16, pto.bf16, pto.f32, pto.f32), - (pto.f32, pto.f32, pto.f32, pto.f32), - (pto.i8, pto.i8, pto.i32, pto.i32), - (pto.f8e4m3, pto.f8e4m3, pto.f32, pto.f32), - (pto.f8e4m3, pto.f8e5m2, pto.f32, pto.f32), - (pto.f8e5m2, pto.f8e4m3, pto.f32, pto.f32), - (pto.f8e5m2, pto.f8e5m2, pto.f32, pto.f32), - (pto.hif8, pto.hif8, pto.f32, pto.f32), - ], -) -def template_tmatmul_bias(lhs: pto.Tile, rhs: pto.Tile, bias: pto.Tile, dst: pto.Tile): - m, k = lhs.valid_shape # (validM, validK) - _, n = rhs.valid_shape # (validK, validN) → n = validN - pto.mad_bias(lhs.as_ptr(), rhs.as_ptr(), dst.as_ptr(), bias.as_ptr(), m, n, k, disable_gemv=True) - return None \ No newline at end of file diff --git a/lib/TileOps/tmatmul_mx_template.py b/lib/TileOps/tmatmul_mx_template.py deleted file mode 100644 index cece695c58..0000000000 --- a/lib/TileOps/tmatmul_mx_template.py +++ /dev/null @@ -1,100 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmatmul.mx.""" - -import tilelang_dsl as pto - - -_LOW_PRECISION_PAIRS = ( - (pto.f8e4m3, pto.f8e4m3), - (pto.f8e4m3, pto.f8e5m2), - (pto.f8e5m2, pto.f8e4m3), - (pto.f8e5m2, pto.f8e5m2), - (pto.f4e1m2x2, pto.f4e1m2x2), - (pto.f4e1m2x2, pto.f4e2m1x2), - (pto.f4e2m1x2, pto.f4e1m2x2), - (pto.f4e2m1x2, pto.f4e2m1x2), -) - -_TMATMUL_MX_DTYPES = tuple( - (lhs_dtype, pto.f16, rhs_dtype, pto.f16, pto.f32) - for lhs_dtype, rhs_dtype in _LOW_PRECISION_PAIRS -) - -_TMATMUL_MX_ACC_DTYPES = tuple( - (pto.f32, lhs_dtype, pto.f16, rhs_dtype, pto.f16, pto.f32) - for lhs_dtype, rhs_dtype in _LOW_PRECISION_PAIRS -) - -_TMATMUL_MX_BIAS_DTYPES = tuple( - (lhs_dtype, pto.f16, rhs_dtype, pto.f16, pto.f32, pto.f32) - for lhs_dtype, rhs_dtype in _LOW_PRECISION_PAIRS -) - - -@pto.ckernel( - target="a5", - op="pto.tmatmul.mx", - dtypes=_TMATMUL_MX_DTYPES, -) -def template_tmatmul_mx( - lhs: pto.Tile, - lhs_scale: pto.Tile, - rhs: pto.Tile, - rhs_scale: pto.Tile, - dst: pto.Tile, -): - m, k = lhs.valid_shape - _, n = rhs.valid_shape - pto.mad_mx( - lhs.as_ptr(), rhs.as_ptr(), dst.as_ptr(), m, n, k, - disable_gemv=True, sat="sat") - return None - - -@pto.ckernel( - target="a5", - op="pto.tmatmul.mx.acc", - dtypes=_TMATMUL_MX_ACC_DTYPES, -) -def template_tmatmul_mx_acc( - acc_in: pto.Tile, - lhs: pto.Tile, - lhs_scale: pto.Tile, - rhs: pto.Tile, - rhs_scale: pto.Tile, - dst: pto.Tile, -): - m, k = lhs.valid_shape - _, n = rhs.valid_shape - pto.mad_mx_acc( - lhs.as_ptr(), rhs.as_ptr(), dst.as_ptr(), m, n, k, - disable_gemv=True, sat="sat") - return None - - -@pto.ckernel( - target="a5", - op="pto.tmatmul.mx.bias", - dtypes=_TMATMUL_MX_BIAS_DTYPES, -) -def template_tmatmul_mx_bias( - lhs: pto.Tile, - lhs_scale: pto.Tile, - rhs: pto.Tile, - rhs_scale: pto.Tile, - bias: pto.Tile, - dst: pto.Tile, -): - m, k = lhs.valid_shape - _, n = rhs.valid_shape - pto.mad_mx_bias( - lhs.as_ptr(), rhs.as_ptr(), dst.as_ptr(), bias.as_ptr(), m, n, k, - disable_gemv=True, sat="sat") - return None diff --git a/lib/TileOps/tmatmul_template.py b/lib/TileOps/tmatmul_template.py deleted file mode 100644 index 5fee19b8a4..0000000000 --- a/lib/TileOps/tmatmul_template.py +++ /dev/null @@ -1,33 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmatmul.""" - -import tilelang_dsl as pto - - -@pto.ckernel( - target="a5", - op="pto.tmatmul", - dtypes=[ - (pto.f16, pto.f16, pto.f32), - (pto.bf16, pto.bf16, pto.f32), - (pto.f32, pto.f32, pto.f32), - (pto.i8, pto.i8, pto.i32), - (pto.f8e4m3, pto.f8e4m3, pto.f32), - (pto.f8e4m3, pto.f8e5m2, pto.f32), - (pto.f8e5m2, pto.f8e4m3, pto.f32), - (pto.f8e5m2, pto.f8e5m2, pto.f32), - (pto.hif8, pto.hif8, pto.f32), - ], -) -def template_tmatmul(lhs: pto.Tile, rhs: pto.Tile, acc: pto.Tile): - m, k = lhs.valid_shape - _, n = rhs.valid_shape - pto.mad(lhs.as_ptr(), rhs.as_ptr(), acc.as_ptr(), m, n, k, disable_gemv=True) - return None diff --git a/lib/TileOps/tmax_template.py b/lib/TileOps/tmax_template.py deleted file mode 100644 index 8831d3eef9..0000000000 --- a/lib/TileOps/tmax_template.py +++ /dev/null @@ -1,32 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmax""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tmax" -) -def template_tmax(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - max_val = pto.vmax(lhs, rhs, mask) - pto.vsts(max_val, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tmaxs_template.py b/lib/TileOps/tmaxs_template.py deleted file mode 100644 index 5c9e409a3d..0000000000 --- a/lib/TileOps/tmaxs_template.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmaxs""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tmaxs", -) -def template_tmaxs(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - result = pto.vmaxs(vec, scalar, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tmin_template.py b/lib/TileOps/tmin_template.py deleted file mode 100644 index 61664b14d1..0000000000 --- a/lib/TileOps/tmin_template.py +++ /dev/null @@ -1,32 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmin""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tmin" -) -def template_tmin(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - min_val = pto.vmin(lhs, rhs, mask) - pto.vsts(min_val, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tmins_template.py b/lib/TileOps/tmins_template.py deleted file mode 100644 index bda0df5f92..0000000000 --- a/lib/TileOps/tmins_template.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmins""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tmins", -) -def template_tmins(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - result = pto.vmins(vec, scalar, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tmov2bias_template.py b/lib/TileOps/tmov2bias_template.py deleted file mode 100644 index 8cc5367d9c..0000000000 --- a/lib/TileOps/tmov2bias_template.py +++ /dev/null @@ -1,92 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmov - Mat to Bias buffer movement. - -This template implements the TMOV_M2B scenario for cube kernels: - - Source: L1 Mat buffer (memory_space="mat", 1xN row-major layout) - - Destination: L0 Bias Table buffer (memory_space="bias") - - Uses bias_load (mte_l1_bt) intrinsic operation - -The Bias Table is a special 4KB buffer in L0 used for bias addition -in matmul operations. Requirements: - - Row dimension must be 1 - - Column dimension * sizeof(dtype) must be aligned to 64 bits - - Total size must not exceed 4KB (4096 bytes) - -Constraint: This template is selected when dst.memory_space == BIAS. -""" - -import tilelang_dsl as pto - - -def _tmov_m2b_constraint(src: pto.Tile, dst: pto.Tile) -> bool: - """Constraint: Mat to Bias transfer scenario. - - Supported scenario: - - src.memory_space == MAT - - dst.memory_space == BIAS - """ - src_ms = src.memory_space - dst_ms = dst.memory_space - - # Check src is MAT - if isinstance(src_ms, str): - src_is_mat = src_ms == "mat" - elif isinstance(src_ms, pto.MemorySpace): - src_is_mat = src_ms == pto.MemorySpace.MAT - else: - src_is_mat = hasattr(src_ms, "value") and src_ms.value == "mat" - - # Check dst is BIAS - if isinstance(dst_ms, str): - dst_is_bias = dst_ms == "bias" - elif isinstance(dst_ms, pto.MemorySpace): - dst_is_bias = dst_ms == pto.MemorySpace.BIAS - else: - dst_is_bias = hasattr(dst_ms, "value") and dst_ms.value == "bias" - - return src_is_mat and dst_is_bias - - -@pto.ckernel( - target="a5", - op="pto.tmov", - constraints=[_tmov_m2b_constraint], - dtypes=[ - (pto.f32, pto.f32), - (pto.f16, pto.f32), - (pto.bf16, pto.f32), - (pto.i32, pto.i32), - ], -) -def template_tmov_m2b(src: pto.Tile, dst: pto.Tile): - """Move data from Mat buffer to Bias Table buffer. - - Args: - src: Source tile in L1 Mat location (1xN row-major) - dst: Destination tile in Bias Table location - - The bias data is moved using burst transfer to the Bias Table. - - mte_l1_bt semantics: - - len_burst = N (number of bias load units, where N = column count) - - Each load unit corresponds to one bias channel/column - - For 1x16 f32 bias: len_burst = 16 - """ - # Bias has shape 1xN, we derive N from the valid shape - _, n = dst.valid_shape - # len_burst = N (number of bias channels/columns) - # See test/tilelang_st/npu/a5/src/st/testcase/tmatmul_bias/tmatmul_bias.pto:79 - len_burst = n - # nburst = (n_burst, src_gap, dst_gap) - single burst with no gaps - n_burst = 1 - src_gap = 0 - dst_gap = 0 - pto.mte_l1_bt(src.as_ptr(), dst.as_ptr(), len_burst, nburst=(n_burst, src_gap, dst_gap)) - return \ No newline at end of file diff --git a/lib/TileOps/tmov2left_template.py b/lib/TileOps/tmov2left_template.py deleted file mode 100644 index 8f1863817e..0000000000 --- a/lib/TileOps/tmov2left_template.py +++ /dev/null @@ -1,78 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmov - Mat to Left buffer movement. - -This template implements the TMOV_M2L scenario for cube kernels: - - Source: L1 Mat buffer (memory_space="mat") - - Destination: L0A Left buffer (memory_space="left") - - Uses left_load (mte_l1_l0a) intrinsic operation - -The operation is part of the cube matmul data flow where: - 1. Data is loaded from GM to L1 Mat via TLOAD - 2. TMOV moves data from L1 Mat to L0A Left - 3. The Left buffer is then used as input for TMATMUL - -Constraint: This template is selected when dst.memory_space == LEFT. -""" - -import tilelang_dsl as pto - - -def _tmov_m2l_constraint(src: pto.Tile, dst: pto.Tile) -> bool: - """Constraint: Mat to Left transfer scenario. - - Supported scenario: - - src.memory_space == MAT - - dst.memory_space == LEFT - """ - src_ms = src.memory_space - dst_ms = dst.memory_space - - # Check src is MAT - if isinstance(src_ms, str): - src_is_mat = src_ms == "mat" - elif isinstance(src_ms, pto.MemorySpace): - src_is_mat = src_ms == pto.MemorySpace.MAT - else: - src_is_mat = hasattr(src_ms, "value") and src_ms.value == "mat" - - # Check dst is LEFT - if isinstance(dst_ms, str): - dst_is_left = dst_ms == "left" - elif isinstance(dst_ms, pto.MemorySpace): - dst_is_left = dst_ms == pto.MemorySpace.LEFT - else: - dst_is_left = hasattr(dst_ms, "value") and dst_ms.value == "left" - - return src_is_mat and dst_is_left - - -@pto.ckernel( - target="a5", - op="pto.tmov", - constraints=[_tmov_m2l_constraint], - dtypes=[ - (pto.f16, pto.f16), - (pto.bf16, pto.bf16), - (pto.f32, pto.f32), - (pto.i8, pto.i8), - ], -) -def template_tmov_m2l(src: pto.Tile, dst: pto.Tile): - """Move data from Mat buffer to Left buffer. - - Args: - src: Source tile in L1 Mat location - dst: Destination tile in L0A Left location - - The m, k dimensions are derived from the tile shapes. - """ - m, k = dst.valid_shape - pto.mte_l1_l0a(src.as_ptr(), dst.as_ptr(), m, k) - return \ No newline at end of file diff --git a/lib/TileOps/tmov2right_template.py b/lib/TileOps/tmov2right_template.py deleted file mode 100644 index e72cd5bd3f..0000000000 --- a/lib/TileOps/tmov2right_template.py +++ /dev/null @@ -1,79 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmov - Mat to Right buffer movement. - -This template implements the TMOV_M2R scenario for cube kernels: - - Source: L1 Mat buffer (memory_space="mat") - - Destination: L0B Right buffer (memory_space="right") - - Uses right_load (mte_l1_l0b) intrinsic operation - -The operation is part of the cube matmul data flow where: - 1. Data is loaded from GM to L1 Mat via TLOAD - 2. TMOV moves data from L1 Mat to L0B Right - 3. The Right buffer is then used as input for TMATMUL - -Constraint: This template is selected when dst.memory_space == RIGHT. -""" - -import tilelang_dsl as pto - - -def _tmov_m2r_constraint(src: pto.Tile, dst: pto.Tile) -> bool: - """Constraint: Mat to Right transfer scenario. - - Supported scenario: - - src.memory_space == MAT - - dst.memory_space == RIGHT - """ - src_ms = src.memory_space - dst_ms = dst.memory_space - - # Check src is MAT - if isinstance(src_ms, str): - src_is_mat = src_ms == "mat" - elif isinstance(src_ms, pto.MemorySpace): - src_is_mat = src_ms == pto.MemorySpace.MAT - else: - src_is_mat = hasattr(src_ms, "value") and src_ms.value == "mat" - - # Check dst is RIGHT - if isinstance(dst_ms, str): - dst_is_right = dst_ms == "right" - elif isinstance(dst_ms, pto.MemorySpace): - dst_is_right = dst_ms == pto.MemorySpace.RIGHT - else: - dst_is_right = hasattr(dst_ms, "value") and dst_ms.value == "right" - - return src_is_mat and dst_is_right - - -@pto.ckernel( - target="a5", - op="pto.tmov", - constraints=[_tmov_m2r_constraint], - dtypes=[ - (pto.f16, pto.f16), - (pto.bf16, pto.bf16), - (pto.f32, pto.f32), - (pto.i8, pto.i8), - ], -) -def template_tmov_m2r(src: pto.Tile, dst: pto.Tile): - """Move data from Mat buffer to Right buffer. - - Args: - src: Source tile in L1 Mat location - dst: Destination tile in L0B Right location - - The k, n dimensions are derived from the tile shapes. - Transpose is typically enabled for Right buffer layout. - """ - k, n = dst.valid_shape - pto.mte_l1_l0b(src.as_ptr(), dst.as_ptr(), k, n, transpose=True) - return \ No newline at end of file diff --git a/lib/TileOps/tmov2scale_template.py b/lib/TileOps/tmov2scale_template.py deleted file mode 100644 index f997a176fd..0000000000 --- a/lib/TileOps/tmov2scale_template.py +++ /dev/null @@ -1,96 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmov - Mat to Scaling/FB buffer movement. - -This template implements the TMOV_M2S scenario for cube kernels: - - Source: L1 Mat buffer (memory_space="mat", 1xN row-major layout) - - Destination: L0 Fixpipe Buffer (memory_space="scaling") - - Uses mte_l1_fb (fixpipe buffer load) intrinsic operation - -The Fixpipe Buffer (FB) is a 4KB buffer used for storing quantization -scale parameters in the fixpipe quantization flow. Requirements: - - Row dimension must be 1 - - Column dimension * sizeof(dtype) must be aligned to 128 bits (16 bytes) - - Total size must not exceed 4KB (4096 bytes) - -Constraint: This template is selected when dst.memory_space == SCALING. - -Data format: - - Each f32 scale value is stored as ui64 (f32 bits in lower 32 bits, upper 32 bits = 0) - - Hardware interprets each ui64's lower 32 bits as one f32 scale value - - Total bytes = N * 8 (N ui64 elements) - - len_burst = N (number of ui64 elements to transfer) -""" - -import tilelang_dsl as pto - - -def _tmov_m2s_constraint(src: pto.Tile, dst: pto.Tile) -> bool: - """Constraint: Mat to Scaling transfer scenario. - - Supported scenario: - - src.memory_space == MAT - - dst.memory_space == SCALING - """ - src_ms = src.memory_space - dst_ms = dst.memory_space - - # Check src is MAT - if isinstance(src_ms, str): - src_is_mat = src_ms == "mat" - elif isinstance(src_ms, pto.MemorySpace): - src_is_mat = src_ms == pto.MemorySpace.MAT - else: - src_is_mat = hasattr(src_ms, "value") and src_ms.value == "mat" - - # Check dst is SCALING - if isinstance(dst_ms, str): - dst_is_scaling = dst_ms == "scaling" - elif isinstance(dst_ms, pto.MemorySpace): - dst_is_scaling = dst_ms == pto.MemorySpace.SCALING - else: - dst_is_scaling = hasattr(dst_ms, "value") and dst_ms.value == "scaling" - - return src_is_mat and dst_is_scaling - - -@pto.ckernel( - target="a5", - op="pto.tmov", - constraints=[_tmov_m2s_constraint], - dtypes=[ - (pto.f32, pto.f32), - ], -) -def template_tmov_m2s(src: pto.Tile, dst: pto.Tile): - """Move data from Mat buffer to Fixpipe Buffer (Scaling). - - Args: - src: Source tile in L1 Mat location (1xN row-major) - dst: Destination tile in Scaling/FB location - - The scale parameters are loaded into FB for fixpipe quantization. - - mte_l1_fb semantics: - - len_burst = M (number of rows in the scaling tile) - - Each row contains one set of scale parameters for all columns - - For 16x16 f32 scaling: len_burst = 16 (see textract_fp.pto:128) - - For 1x16 f32 scaling: len_burst = 1 (single row of parameters) - - nburst = (n_burst, src_gap, dst_gap) - single burst with no gaps - """ - # Scale tile has shape MxN (typically 1xN for per-column scales) - m, _ = dst.valid_shape - # len_burst = M (number of rows/parameter sets) - len_burst = m - n_burst = 1 - src_gap = 0 - dst_gap = 0 - pto.mte_l1_fb(src.as_ptr(), dst.as_ptr(), len_burst, nburst=(n_burst, src_gap, dst_gap)) - return \ No newline at end of file diff --git a/lib/TileOps/tmov2vec_template.py b/lib/TileOps/tmov2vec_template.py deleted file mode 100644 index d5f4bde2bb..0000000000 --- a/lib/TileOps/tmov2vec_template.py +++ /dev/null @@ -1,85 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmov - Acc to Vec/UB buffer movement. - -This template implements the TMOV_A2V scenario for cube kernels: - - Source: L0C Accumulator buffer (memory_space="acc") - - Destination: UB (Unified Buffer) Vec location (memory_space="ub") - - Uses acc_store_ub (mte_l0c_ub) intrinsic operation - -This is part of the fixpipe path where accumulator results are -moved from the cube unit to the vector unit for further processing. - -Constraint: This template is selected when src.memory_space == ACC -and dst.memory_space == UB. -""" - -import tilelang_dsl as pto - - -def _tmov_a2v_constraint(src: pto.Tile, dst: pto.Tile) -> bool: - """Constraint: Acc to Vec/UB transfer scenario. - - Supported scenario: - - src.memory_space == ACC - - dst.memory_space == UB - """ - src_ms = src.memory_space - dst_ms = dst.memory_space - - # Check src is ACC - if isinstance(src_ms, str): - src_is_acc = src_ms == "acc" - elif isinstance(src_ms, pto.MemorySpace): - src_is_acc = src_ms == pto.MemorySpace.ACC - else: - src_is_acc = hasattr(src_ms, "value") and src_ms.value == "acc" - - # Check dst is UB (or "vec" which maps to UB) - if isinstance(dst_ms, str): - dst_is_ub = dst_ms == "ub" or dst_ms == "vec" - elif isinstance(dst_ms, pto.MemorySpace): - dst_is_ub = dst_ms == pto.MemorySpace.UB - else: - dst_is_ub = hasattr(dst_ms, "value") and (dst_ms.value == "ub" or dst_ms.value == "vec") - - return src_is_acc and dst_is_ub - - -@pto.ckernel( - target="a5", - op="pto.tmov", - constraints=[_tmov_a2v_constraint], - dtypes=[ - (pto.f32, pto.f32), - (pto.i32, pto.i32), - ], -) -def template_tmov_a2v(src: pto.Tile, dst: pto.Tile): - """Move data from Acc buffer to Vec/UB buffer. - - Args: - src: Source tile in Acc location - dst: Destination tile in Vec/UB location - - The m, n dimensions and strides are derived from the tile shapes. - This performs NZ2ND layout conversion. - """ - m, n = dst.valid_shape - src_stride = (m + 15) // 16 * 16 # Align to 16 blocks - dst_stride = n # Row-major stride - # mte_l0c_ub takes 7 positional args: src, dst, m, n, src_stride, dst_stride, dst_mode - # dst_mode: integer for sub_blockid mode (0 or 1) - # layout: keyword arg for layout conversion ("nz2nd" for NZ to ND) - pto.mte_l0c_ub( - src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, - 0, # dst_mode: sub_blockid value 0 - layout="nz2nd", - ) - return \ No newline at end of file diff --git a/lib/TileOps/tmov_fp_template.py b/lib/TileOps/tmov_fp_template.py deleted file mode 100644 index f7c3a6d8e3..0000000000 --- a/lib/TileOps/tmov_fp_template.py +++ /dev/null @@ -1,145 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmov - Acc to Mat with quantization (fixpipe). - -This template implements the TMOV_FP scenario for cube kernels: - - Source: L0C Accumulator buffer (memory_space="acc") - - Scaling: FB buffer with quantization parameters (memory_space="scaling") - - Destination: L1 Mat buffer (memory_space="mat", quantized output) - - Uses fixpipe intrinsic operation - -This is part of the fixpipe quantization path where: - 1. Matmul results are accumulated in L0C (int32/float) - 2. Scale parameters are loaded into FB buffer - 3. TMOV with fp parameter performs quantization: Acc * scale -> quantized output - -Constraint: This template is selected when src.memory_space == ACC, -scale.memory_space == SCALING, and dst.memory_space == MAT. - -Supported scenarios: - - float32 accumulator -> float16/bf16 output with scale -""" - -import tilelang_dsl as pto - - -def _tmov_fp_constraint(src: pto.Tile, dst: pto.Tile, fp: pto.Tile) -> bool: - """Constraint: Fixpipe quantization scenario. - - Supported scenario: - - src.memory_space == ACC - - dst.memory_space == MAT - - fp.memory_space == SCALING - """ - src_ms = src.memory_space - dst_ms = dst.memory_space - fp_ms = fp.memory_space - - # Check src is ACC - if isinstance(src_ms, str): - src_is_acc = src_ms == "acc" - elif isinstance(src_ms, pto.MemorySpace): - src_is_acc = src_ms == pto.MemorySpace.ACC - else: - src_is_acc = hasattr(src_ms, "value") and src_ms.value == "acc" - - # Check dst is MAT - if isinstance(dst_ms, str): - dst_is_mat = dst_ms == "mat" - elif isinstance(dst_ms, pto.MemorySpace): - dst_is_mat = dst_ms == pto.MemorySpace.MAT - else: - dst_is_mat = hasattr(dst_ms, "value") and dst_ms.value == "mat" - - # Check fp is SCALING - if isinstance(fp_ms, str): - fp_is_scaling = fp_ms == "scaling" - elif isinstance(fp_ms, pto.MemorySpace): - fp_is_scaling = fp_ms == pto.MemorySpace.SCALING - else: - fp_is_scaling = hasattr(fp_ms, "value") and fp_ms.value == "scaling" - - return src_is_acc and dst_is_mat and fp_is_scaling - - -def _make_fp_constraint(dst_dtype): - """Create a constraint that checks both memory spaces and dst dtype.""" - def _fp_constraint(src: pto.Tile, dst: pto.Tile, fp: pto.Tile) -> bool: - # Check memory spaces - if not _tmov_fp_constraint(src, dst, fp): - return False - # Check dst dtype matches (use .dtype in constraint, .element_type in template body) - return dst.dtype == dst_dtype - return _fp_constraint - - -@pto.ckernel( - target="a5", - op="pto.tmov", - constraints=[_make_fp_constraint(pto.f16)], - dtypes=[ - (pto.f32, pto.f16, pto.f32), # (src, dst, fp) - IR operand order - ], -) -def template_tmov_fp_f32_f16(src: pto.Tile, dst: pto.Tile, fp: pto.Tile): - """Move and quantize data from Acc to Mat with scaling parameters (f32 -> f16). - - Args: - src: Source tile in Acc location (accumulator, f32) - dst: Destination tile in Mat location (quantized output, f16) - fp: Scaling tile in FB location (quantization params, f32) - - The tmov with fp parameter performs fixpipe quantization using mte_l0c_l1 - with pre_quant keyword argument. - """ - # Get dimensions from destination tile - m, n = dst.valid_shape - # Strides: src is in Acc (fractal layout), dst is in Mat (row-major) - src_stride = (m + 15) // 16 * 16 # Align to 16 blocks for fractal - dst_stride = n # Row-major stride - - # Use mte_l0c_l1 with pre_quant for fixpipe quantization (f32 -> f16) - pto.mte_l0c_l1( - src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, - pre_quant=(fp.as_ptr(), "qf322f16_pre_vec"), - ) - return - - -@pto.ckernel( - target="a5", - op="pto.tmov", - constraints=[_make_fp_constraint(pto.bf16)], - dtypes=[ - (pto.f32, pto.bf16, pto.f32), # (src, dst, fp) - IR operand order - ], -) -def template_tmov_fp_f32_bf16(src: pto.Tile, dst: pto.Tile, fp: pto.Tile): - """Move and quantize data from Acc to Mat with scaling parameters (f32 -> bf16). - - Args: - src: Source tile in Acc location (accumulator, f32) - dst: Destination tile in Mat location (quantized output, bf16) - fp: Scaling tile in FB location (quantization params, f32) - - The tmov with fp parameter performs fixpipe quantization using mte_l0c_l1 - with pre_quant keyword argument. - """ - # Get dimensions from destination tile - m, n = dst.valid_shape - # Strides: src is in Acc (fractal layout), dst is in Mat (row-major) - src_stride = (m + 15) // 16 * 16 # Align to 16 blocks for fractal - dst_stride = n # Row-major stride - - # Use mte_l0c_l1 with pre_quant for fixpipe quantization (f32 -> bf16) - pto.mte_l0c_l1( - src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, - pre_quant=(fp.as_ptr(), "qf322bf16_pre_vec"), - ) - return \ No newline at end of file diff --git a/lib/TileOps/tmov_template.py b/lib/TileOps/tmov_template.py deleted file mode 100644 index fda9d7c26c..0000000000 --- a/lib/TileOps/tmov_template.py +++ /dev/null @@ -1,100 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmov - tile data movement - -This template implements UB2UB ND2ND tile data movement: - - UB2UB: Both src and dst must be in Unified Buffer (memory_space="ub") - - ND2ND: Both tiles must have N-dimensional layout (s_layout=NONE_BOX) - -For other transfer scenarios (GM2UB, UB2GM, or specialized layouts), -different templates/implementation paths should be used. -""" - -import tilelang_dsl as pto - - -def _tmov_ub2ub_nd2nd_constraint(src: pto.Tile, dst: pto.Tile) -> bool: - """Constraint: Both src and dst must be UB location with ND layout. - - Supported scenario: - - UB2UB: src and dst both have memory_space="ub" - - ND2ND: src and dst both have s_layout=NONE_BOX (N-dimensional format) - - Unsupported scenarios (require different implementation paths): - - GM2UB: src from Global Memory, dst to Unified Buffer - - UB2GM: src from Unified Buffer, dst to Global Memory - - Specialized layouts (e.g., cube formats with non-NONE_BOX s_layout) - """ - # Check memory_space for both tiles (UB2UB constraint) - src_ms = src.memory_space - dst_ms = dst.memory_space - if isinstance(src_ms, str): - src_is_ub = src_ms == "ub" - else: - src_is_ub = hasattr(src_ms, "value") and src_ms.value == "ub" - if isinstance(dst_ms, str): - dst_is_ub = dst_ms == "ub" - else: - dst_is_ub = hasattr(dst_ms, "value") and dst_ms.value == "ub" - - if not (src_is_ub and dst_is_ub): - return False - - # Check s_layout for both tiles (ND2ND constraint) - # ND layout uses NONE_BOX, specialized layouts (cube, etc.) use different values - src_config = src.config - dst_config = dst.config - if src_config is None or dst_config is None: - return False - if src_config.s_layout != pto.SLayout.NONE_BOX: - return False - if dst_config.s_layout != pto.SLayout.NONE_BOX: - return False - if pto.bytewidth(src.dtype) != pto.bytewidth(dst.dtype): - return False - - return True - - -@pto.vkernel( - target="a5", - op="pto.tmov", - constraints=[_tmov_ub2ub_nd2nd_constraint], - advanced=True, -) -def template_tmov_basic(src: pto.Tile, dst: pto.Tile): - """Basic tile-to-tile data movement using vlds/vsts. - - Based on TMovVecToVec in TMov.hpp (lines 378-405): - - Iterate over valid_row rows - - Each row processed in chunks of nRepeatElem elements - - Use predicate mask for partial chunks - - Args: - src: Source tile (Vec location) - dst: Destination tile (Vec location) - """ - src_dtype = src.element_type - dst_dtype = dst.element_type - lanes = pto.get_lanes(dst_dtype) - - # Use dst.valid_shape as the copy dimensions - # The dst tile defines how many elements to write - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dst_dtype, remained) - data = pto.vlds(src[row, col:]) - if pto.constexpr(src_dtype != dst_dtype): - data = pto.vbitcast(data, dst_dtype) - pto.vsts(data, dst[row, col:], mask) - - return None diff --git a/lib/TileOps/tmrgsort_template.py b/lib/TileOps/tmrgsort_template.py deleted file mode 100644 index ac0c8e1caa..0000000000 --- a/lib/TileOps/tmrgsort_template.py +++ /dev/null @@ -1,273 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmrgsort""" - -import tilelang_dsl as pto - -STRUCT_SIZE = 8 # bytes per structure (value + index) -STRUCT_SIZE_SHIFT = 3 # log2(8) -BLOCK_NUM = 4 - - -@pto.inline_proc -def tmrgsort_single_list_instr(dst: pto.Tile, src: pto.Tile, - num_structures, repeat_times): - dst_ptr = dst.as_ptr() - src_ptr = src.as_ptr() - - count = pto.i64(num_structures) - count = count | (pto.i64(num_structures) << pto.i64(16)) - count = count | (pto.i64(num_structures) << pto.i64(32)) - count = count | (pto.i64(num_structures) << pto.i64(48)) - - offset = num_structures * STRUCT_SIZE // pto.bytewidth(dst.element_type) - src0 = src_ptr - src1 = pto.addptr(src_ptr, offset) - src2 = pto.addptr(src_ptr, offset * 2) - src3 = pto.addptr(src_ptr, offset * 3) - - config = pto.i64(repeat_times) - config = config | (pto.i64(0b1111) << pto.i64(8)) - config = config | (pto.i64(0b0) << pto.i64(12)) - - pto.vmrgsort4(dst_ptr, src0, src1, src2, src3, pto.i64(count), pto.i64(config)) - return - - -@pto.inline_proc -def tmrgsort_multi_list2_instr(tmp: pto.Tile, src0: pto.Tile, src1: pto.Tile, - src0_structures: int, src1_structures: int): - tmp_ptr = tmp.as_ptr() - src0_ptr = src0.as_ptr() - src1_ptr = src1.as_ptr() - - count = pto.i64(src0_structures) - count = count | (pto.i64(src1_structures) << pto.i64(16)) - - repeat_time = 1 - list_mask = 0b0011 - exhausted_bit = 0 - - exhausted_str = pto.get_op_attr("exhausted", "0") - if pto.constexpr(exhausted_str == "1"): - exhausted_bit = 1 - - config = pto.i64(repeat_time) - config = config | (pto.i64(list_mask) << pto.i64(8)) - config = config | (pto.i64(exhausted_bit) << pto.i64(12)) - - pto.vmrgsort4(tmp_ptr, src0_ptr, src1_ptr, src0_ptr, src0_ptr, - count, config) - - return - - -@pto.inline_proc -def tmrgsort_multi_list3_instr(tmp: pto.Tile, src0: pto.Tile, src1: pto.Tile, src2: pto.Tile, - src0_structures: int, src1_structures: int, src2_structures: int): - tmp_ptr = tmp.as_ptr() - src0_ptr = src0.as_ptr() - src1_ptr = src1.as_ptr() - src2_ptr = src2.as_ptr() - - count = pto.i64(src0_structures) - count = count | (pto.i64(src1_structures) << pto.i64(16)) - count = count | (pto.i64(src2_structures) << pto.i64(32)) - - repeat_time = 1 - list_mask = 0b0111 - exhausted_bit = 0 - - exhausted_str = pto.get_op_attr("exhausted", "0") - if pto.constexpr(exhausted_str == "1"): - exhausted_bit = 1 - - config = pto.i64(repeat_time) - config = config | (pto.i64(list_mask) << pto.i64(8)) - config = config | (pto.i64(exhausted_bit) << pto.i64(12)) - - pto.vmrgsort4(tmp_ptr, src0_ptr, src1_ptr, src2_ptr, src0_ptr, - count, config) - - return - - -@pto.inline_proc -def tmrgsort_multi_list4_instr(tmp: pto.Tile, src0: pto.Tile, src1: pto.Tile, - src2: pto.Tile, src3: pto.Tile, - src0_structures: int, src1_structures: int, - src2_structures: int, src3_structures: int): - dtype = tmp.element_type - - tmp_ptr = tmp.as_ptr() - src0_ptr = src0.as_ptr() - src1_ptr = src1.as_ptr() - src2_ptr = src2.as_ptr() - src3_ptr = src3.as_ptr() - - count = pto.i64(src0_structures) - count = count | (pto.i64(src1_structures) << pto.i64(16)) - count = count | (pto.i64(src2_structures) << pto.i64(32)) - count = count | (pto.i64(src3_structures) << pto.i64(48)) - - repeat_time = 1 - list_mask = 0b1111 - exhausted_bit = 0 - - exhausted_str = pto.get_op_attr("exhausted", "0") - if pto.constexpr(exhausted_str == "1"): - exhausted_bit = 1 - - config = pto.i64(repeat_time) - config = config | (pto.i64(list_mask) << pto.i64(8)) - config = config | (pto.i64(exhausted_bit) << pto.i64(12)) - - pto.vmrgsort4(tmp_ptr, src0_ptr, src1_ptr, src2_ptr, src3_ptr, - count, config) - - return - - -@pto.vkernel( - target="a5", - op="pto.tmrgsort", - advanced=True, -) -def template_tmrgsort_single_list(src: pto.Tile, block_len: pto.AnyInt, dst: pto.Tile): - """Format1 template: single list internal block sorting. - - Standard Format1: single vmrgsort4 for block sorting. - TopK variant is handled by ST kernel via iterative tmrgsort + tmov calls. - """ - src_valid_col = src.valid_shape[1] - - # Block length in structures - num_structures = block_len * pto.bytewidth(src.element_type) >> STRUCT_SIZE_SHIFT - - # Repeat times: how many groups of 4 blocks need merging - repeat_times = src_valid_col // (block_len * BLOCK_NUM) - - # Standard Format1: single merge operation - tmrgsort_single_list_instr(dst, src, num_structures, repeat_times) - - return None - - -@pto.vkernel( - target="a5", - op="pto.tmrgsort", - advanced=True, -) -def template_tmrgsort_multi_list2(src0: pto.Tile, src1: pto.Tile, - tmp: pto.Tile, dst: pto.Tile, ex_vec: pto.AnyInt): - dtype = dst.element_type - bw = pto.bytewidth(dtype) - - src0_valid_col = src0.valid_shape[1] - src1_valid_col = src1.valid_shape[1] - dst_valid_col = dst.valid_shape[1] - - if pto.constexpr(bw == 4): - src0_structures = src0_valid_col // 2 - src1_structures = src1_valid_col // 2 - else: - src0_structures = src0_valid_col // 4 - src1_structures = src1_valid_col // 4 - - dst_elements = dst_valid_col - - tmrgsort_multi_list2_instr(tmp, src0, src1, src0_structures, src1_structures) - - lanes = pto.get_lanes(dtype) - for col in range(0, dst_elements, lanes): - remained = dst_elements - col - mask, remained = pto.make_mask(dtype, remained) - data = pto.vlds(tmp[0, col:]) - pto.vsts(data, dst[0, col:], mask) - - return None - - -@pto.vkernel( - target="a5", - op="pto.tmrgsort", - advanced=True, -) -def template_tmrgsort_multi_list3(src0: pto.Tile, src1: pto.Tile, src2: pto.Tile, - tmp: pto.Tile, dst: pto.Tile, ex_vec: pto.AnyInt): - dtype = dst.element_type - bw = pto.bytewidth(dtype) - - src0_valid_col = src0.valid_shape[1] - src1_valid_col = src1.valid_shape[1] - src2_valid_col = src2.valid_shape[1] - dst_valid_col = dst.valid_shape[1] - - if pto.constexpr(bw == 4): - src0_structures = src0_valid_col // 2 - src1_structures = src1_valid_col // 2 - src2_structures = src2_valid_col // 2 - else: - src0_structures = src0_valid_col // 4 - src1_structures = src1_valid_col // 4 - src2_structures = src2_valid_col // 4 - - dst_elements = dst_valid_col - - tmrgsort_multi_list3_instr(tmp, src0, src1, src2, src0_structures, src1_structures, src2_structures) - - lanes = pto.get_lanes(dtype) - for col in range(0, dst_elements, lanes): - remained = dst_elements - col - mask, remained = pto.make_mask(dtype, remained) - data = pto.vlds(tmp[0, col:]) - pto.vsts(data, dst[0, col:], mask) - - return None - - -@pto.vkernel( - target="a5", - op="pto.tmrgsort", - advanced=True, -) -def template_tmrgsort_multi_list4(src0: pto.Tile, src1: pto.Tile, src2: pto.Tile, src3: pto.Tile, - tmp: pto.Tile, dst: pto.Tile, ex_vec: pto.AnyInt): - dtype = dst.element_type - bw = pto.bytewidth(dtype) - - src0_valid_col = src0.valid_shape[1] - src1_valid_col = src1.valid_shape[1] - src2_valid_col = src2.valid_shape[1] - src3_valid_col = src3.valid_shape[1] - dst_valid_col = dst.valid_shape[1] - - if pto.constexpr(bw == 4): - src0_structures = src0_valid_col // 2 - src1_structures = src1_valid_col // 2 - src2_structures = src2_valid_col // 2 - src3_structures = src3_valid_col // 2 - else: - src0_structures = src0_valid_col // 4 - src1_structures = src1_valid_col // 4 - src2_structures = src2_valid_col // 4 - src3_structures = src3_valid_col // 4 - - dst_elements = dst_valid_col - - tmrgsort_multi_list4_instr(tmp, src0, src1, src2, src3, src0_structures, src1_structures, src2_structures, src3_structures) - - lanes = pto.get_lanes(dtype) - for col in range(0, dst_elements, lanes): - remained = dst_elements - col - mask, remained = pto.make_mask(dtype, remained) - data = pto.vlds(tmp[0, col:]) - pto.vsts(data, dst[0, col:], mask) - - return None \ No newline at end of file diff --git a/lib/TileOps/tmul_template.py b/lib/TileOps/tmul_template.py deleted file mode 100644 index ae7adf44ef..0000000000 --- a/lib/TileOps/tmul_template.py +++ /dev/null @@ -1,32 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmul""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tmul" -) -def template_tmul(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - multiplied = pto.vmul(lhs, rhs, mask) - pto.vsts(multiplied, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tmuls_template.py b/lib/TileOps/tmuls_template.py deleted file mode 100644 index 8d02ea8264..0000000000 --- a/lib/TileOps/tmuls_template.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tmuls""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tmuls", -) -def template_tmuls(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - result = pto.vmuls(vec, scalar, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tneg_template.py b/lib/TileOps/tneg_template.py deleted file mode 100644 index 8e10ce4cac..0000000000 --- a/lib/TileOps/tneg_template.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tneg""" - -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tneg" -) -def template_tneg(src: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vinput = pto.vlds(src[row, col:]) - result = pto.vneg(vinput, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tnot_template.py b/lib/TileOps/tnot_template.py deleted file mode 100644 index f4728e8530..0000000000 --- a/lib/TileOps/tnot_template.py +++ /dev/null @@ -1,30 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tnot""" - -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tnot", - dtypes=[(pto.AnyInt, pto.AnyInt)] -) -def template_tnot(src: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vinput = pto.vlds(src[row, col:]) - result = pto.vnot(vinput, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tor_template.py b/lib/TileOps/tor_template.py deleted file mode 100644 index e8be63d5e9..0000000000 --- a/lib/TileOps/tor_template.py +++ /dev/null @@ -1,32 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tor""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tor" -) -def template_tor(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - result = pto.vor(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tors_template.py b/lib/TileOps/tors_template.py deleted file mode 100644 index 542ee4044f..0000000000 --- a/lib/TileOps/tors_template.py +++ /dev/null @@ -1,40 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tors - -Note: A5 hardware implements tors as: - TEXPANDS_IMPL(dst, scalar); // broadcast scalar to dst - TOR_IMPL(dst, src, dst); // dst = src | dst - -This template uses vbr + vor to achieve element-wise bitwise OR. -Only supports tile, scalar order. -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tors", -) -def template_tors(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - scalar_vec = pto.vbr(scalar) - result = pto.vor(vec, scalar_vec, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tpartadd_template.py b/lib/TileOps/tpartadd_template.py deleted file mode 100644 index 8e9716389a..0000000000 --- a/lib/TileOps/tpartadd_template.py +++ /dev/null @@ -1,79 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tpartadd""" - -import tilelang_dsl as pto - - -@pto.inline_proc -def tpart_op_instr(dst: pto.Tile, src0: pto.Tile, src1: pto.Tile, valid_rows, valid_cols): - dtype = dst.element_type - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - summed = pto.vadd(lhs, rhs, mask) - pto.vsts(summed, dst[row, col:], mask) - return None - -@pto.inline_proc -def tpart_copy_instr(dst: pto.Tile, src: pto.Tile, valid_rows, valid_cols, start_row): - dtype = dst.element_type - for row in range(start_row, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - val = pto.vlds(src[row, col:]) - pto.vsts(val, dst[row, col:], mask) - return None - -@pto.inline_proc -def tpart_op(dst: pto.Tile, src0: pto.Tile, src1: pto.Tile, - dst_valid_rows, dst_valid_cols, - src1_valid_rows, src1_valid_cols): - - src1_eq_dst = (src1_valid_rows == dst_valid_rows and src1_valid_cols == dst_valid_cols) - src1_row_lt_dst = (src1_valid_rows < dst_valid_rows and src1_valid_cols == dst_valid_cols) - src1_col_lt_dst = (src1_valid_rows <= dst_valid_rows and src1_valid_cols < dst_valid_cols) - - if src1_eq_dst: - tpart_op_instr(dst, src0, src1, dst_valid_rows, dst_valid_cols) - elif src1_col_lt_dst: - tpart_copy_instr(dst, src0, dst_valid_rows, dst_valid_cols, 0) - if src1_valid_cols > 0: - tpart_op_instr(dst, src0, src1, src1_valid_rows, src1_valid_cols) - elif src1_row_lt_dst: - if src1_valid_cols > 0: - tpart_op_instr(dst, src0, src1, src1_valid_rows, src1_valid_cols) - tpart_copy_instr(dst, src0, dst_valid_rows, dst_valid_cols, src1_valid_rows) - - return - -@pto.vkernel( - target="a5", - op="pto.tpartadd" -) -def template_tpartadd(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dst_valid_rows, dst_valid_cols = dst.valid_shape - src0_valid_rows, src0_valid_cols = src0.valid_shape - src1_valid_rows, src1_valid_cols = src1.valid_shape - - src0_eq_dst = (src0_valid_rows == dst_valid_rows and src0_valid_cols == dst_valid_cols) - src1_eq_dst = (src1_valid_rows == dst_valid_rows and src1_valid_cols == dst_valid_cols) - - if src0_eq_dst or src1_eq_dst: - if src0_eq_dst: - tpart_op(dst, src0, src1, dst_valid_rows, dst_valid_cols, src1_valid_rows, src1_valid_cols) - elif src1_eq_dst: - tpart_op(dst, src1, src0, dst_valid_rows, dst_valid_cols, src0_valid_rows, src0_valid_cols) - # TODO: raise an error later - - return \ No newline at end of file diff --git a/lib/TileOps/tpartmax_template.py b/lib/TileOps/tpartmax_template.py deleted file mode 100644 index 87060cd704..0000000000 --- a/lib/TileOps/tpartmax_template.py +++ /dev/null @@ -1,54 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tpartmax""" - -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tpartmax", - advanced=True, -) -def template_tpartmax(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - src0_valid_rows, src0_valid_cols = src0.valid_shape - src1_valid_rows, src1_valid_cols = src1.valid_shape - lanes = pto.get_lanes(dtype) - - pad_scalar = pto.PadValue.MIN.eval(dtype) - pad_vec = pto.vbr(pad_scalar) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - pto.vsts(pad_vec, dst[row, col:], mask) - - pto.mem_bar(pto.BarrierType.VST_VLD) - - for row in range(0, src0_valid_rows, 1): - remained = src0_valid_cols - for col in range(0, src0_valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - vec0 = pto.vlds(src0[row, col:]) - pto.vsts(vec0, dst[row, col:], mask) - - pto.mem_bar(pto.BarrierType.VST_VLD) - - for row in range(0, src1_valid_rows, 1): - remained = src1_valid_cols - for col in range(0, src1_valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - vec_dst = pto.vlds(dst[row, col:]) - vec1 = pto.vlds(src1[row, col:]) - result = pto.vmax(vec_dst, vec1, mask) - pto.vsts(result, dst[row, col:], mask) - - return \ No newline at end of file diff --git a/lib/TileOps/tpartmin_template.py b/lib/TileOps/tpartmin_template.py deleted file mode 100644 index cc21c609cb..0000000000 --- a/lib/TileOps/tpartmin_template.py +++ /dev/null @@ -1,54 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tpartmin""" - -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tpartmin", - advanced=True, -) -def template_tpartmin(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - src0_valid_rows, src0_valid_cols = src0.valid_shape - src1_valid_rows, src1_valid_cols = src1.valid_shape - lanes = pto.get_lanes(dtype) - - pad_scalar = pto.PadValue.MAX.eval(dtype) - pad_vec = pto.vbr(pad_scalar) - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - pto.vsts(pad_vec, dst[row, col:], mask) - - pto.mem_bar(pto.BarrierType.VST_VLD) - - for row in range(0, src0_valid_rows, 1): - remained = src0_valid_cols - for col in range(0, src0_valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - vec0 = pto.vlds(src0[row, col:]) - pto.vsts(vec0, dst[row, col:], mask) - - pto.mem_bar(pto.BarrierType.VST_VLD) - - for row in range(0, src1_valid_rows, 1): - remained = src1_valid_cols - for col in range(0, src1_valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - vec_dst = pto.vlds(dst[row, col:]) - vec1 = pto.vlds(src1[row, col:]) - result = pto.vmin(vec_dst, vec1, mask) - pto.vsts(result, dst[row, col:], mask) - - return \ No newline at end of file diff --git a/lib/TileOps/tpartmul_template.py b/lib/TileOps/tpartmul_template.py deleted file mode 100644 index fe39597cad..0000000000 --- a/lib/TileOps/tpartmul_template.py +++ /dev/null @@ -1,79 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tpartmul""" - -import tilelang_dsl as pto - - -@pto.inline_proc -def tpart_op_instr(dst: pto.Tile, src0: pto.Tile, src1: pto.Tile, valid_rows, valid_cols): - dtype = dst.element_type - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - summed = pto.vmul(lhs, rhs, mask) - pto.vsts(summed, dst[row, col:], mask) - return None - -@pto.inline_proc -def tpart_copy_instr(dst: pto.Tile, src: pto.Tile, valid_rows, valid_cols, start_row): - dtype = dst.element_type - for row in range(start_row, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - val = pto.vlds(src[row, col:]) - pto.vsts(val, dst[row, col:], mask) - return None - -@pto.inline_proc -def tpart_op(dst: pto.Tile, src0: pto.Tile, src1: pto.Tile, - dst_valid_rows, dst_valid_cols, - src1_valid_rows, src1_valid_cols): - - src1_eq_dst = (src1_valid_rows == dst_valid_rows and src1_valid_cols == dst_valid_cols) - src1_row_lt_dst = (src1_valid_rows < dst_valid_rows and src1_valid_cols == dst_valid_cols) - src1_col_lt_dst = (src1_valid_rows <= dst_valid_rows and src1_valid_cols < dst_valid_cols) - - if src1_eq_dst: - tpart_op_instr(dst, src0, src1, dst_valid_rows, dst_valid_cols) - elif src1_col_lt_dst: - tpart_copy_instr(dst, src0, dst_valid_rows, dst_valid_cols, 0) - if src1_valid_cols > 0: - tpart_op_instr(dst, src0, src1, src1_valid_rows, src1_valid_cols) - elif src1_row_lt_dst: - if src1_valid_cols > 0: - tpart_op_instr(dst, src0, src1, src1_valid_rows, src1_valid_cols) - tpart_copy_instr(dst, src0, dst_valid_rows, dst_valid_cols, src1_valid_rows) - - return - -@pto.vkernel( - target="a5", - op="pto.tpartmul" -) -def template_tpartmul(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dst_valid_rows, dst_valid_cols = dst.valid_shape - src0_valid_rows, src0_valid_cols = src0.valid_shape - src1_valid_rows, src1_valid_cols = src1.valid_shape - - src0_eq_dst = (src0_valid_rows == dst_valid_rows and src0_valid_cols == dst_valid_cols) - src1_eq_dst = (src1_valid_rows == dst_valid_rows and src1_valid_cols == dst_valid_cols) - - if src0_eq_dst or src1_eq_dst: - if src0_eq_dst: - tpart_op(dst, src0, src1, dst_valid_rows, dst_valid_cols, src1_valid_rows, src1_valid_cols) - elif src1_eq_dst: - tpart_op(dst, src1, src0, dst_valid_rows, dst_valid_cols, src0_valid_rows, src0_valid_cols) - # TODO: raise an error later - - return \ No newline at end of file diff --git a/lib/TileOps/tprelu_template.py b/lib/TileOps/tprelu_template.py deleted file mode 100644 index 0f700259dc..0000000000 --- a/lib/TileOps/tprelu_template.py +++ /dev/null @@ -1,45 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tprelu (Parametric ReLU)""" - -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tprelu", - dtypes=[(pto.f16, pto.f16, pto.f16, pto.f16), (pto.f32, pto.f32, pto.f32, pto.f32), - (pto.f16, pto.f16, pto.i8, pto.f16), (pto.f32, pto.f32, pto.i8, pto.f32)], - advanced=True -) -def template_tprelu(src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - """Parametric ReLU: dst = src0 > 0 ? src0 : src0 * src1. - - Semantics: - For each element (i, j): - dst[i, j] = src0[i, j] > 0 ? src0[i, j] : src0[i, j] * src1[i, j] - - Supported data types: f16, f32 - A5 keeps the tmp operand in the ABI for cross-arch compatibility, but this - implementation does not read or write it. - """ - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - lanes = pto.get_lanes(dtype) - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - vec0 = pto.vlds(src0[row, col:]) - vec1 = pto.vlds(src1[row, col:]) - result = pto.vprelu(vec0, vec1, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/trandom_template.py b/lib/TileOps/trandom_template.py deleted file mode 100644 index af7bdd0c6c..0000000000 --- a/lib/TileOps/trandom_template.py +++ /dev/null @@ -1,285 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trandom - unified template using constexpr""" - -import tilelang_dsl as pto - -TRANDOM_ONCE_REPEAT = 4 -TRANDOM_CONST_0 = 0xD2511F53 -TRANDOM_CONST_1 = 0xCD9E8D57 -TRANDOM_CONST_KEY_ADD_0 = 0x9E3779B9 -TRANDOM_CONST_KEY_ADD_1 = 0xBB67AE85 - - -def _check_row_major(dst) -> bool: - return dst.config.b_layout == pto.BLayout.ROW_MAJOR - - -@pto.vkernel( - target="a5", - op="pto.trandom", - dtypes=[ - (pto.i32, pto.i32, pto.i32, pto.i32, pto.i32, pto.i32, pto.ui32), - ], - constraints=[_check_row_major], - advanced=True, -) -def template_trandom( - key0: pto.i32, - key1: pto.i32, - counter0: pto.i32, - counter1: pto.i32, - counter2: pto.i32, - counter3: pto.i32, - dst: pto.Tile, -): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - lanes = pto.get_lanes(dtype) - n_loop = (valid_cols + TRANDOM_ONCE_REPEAT * lanes - 1) // (TRANDOM_ONCE_REPEAT * lanes) - rounds_str = pto.get_op_attr("rounds", "10") - - pg = pto.pset_b32(pto.PAT.ALL) - - ctr0_init = pto.vbitcast(pto.vbr(counter0), pto.ui32) - ctr1_init = pto.vbitcast(pto.vbr(counter1), pto.ui32) - ctr2_init = pto.vbitcast(pto.vbr(counter2), pto.ui32) - ctr3_init = pto.vbitcast(pto.vbr(counter3), pto.ui32) - key0_v = pto.vbitcast(pto.vbr(key0), pto.ui32) - key1_v = pto.vbitcast(pto.vbr(key1), pto.ui32) - zeros = pto.vbr(pto.ui32(0)) - const0 = pto.vbr(pto.ui32(TRANDOM_CONST_0)) - const1 = pto.vbr(pto.ui32(TRANDOM_CONST_1)) - inc_idx = pto.vbitcast(pto.vci(pto.i32(0)), pto.ui32) - - ctr0, pd = pto.vaddc(ctr0_init, inc_idx, pg) - ctr1, pd = pto.vaddcs(ctr1_init, zeros, pd, pg) - ctr2, pd = pto.vaddcs(ctr2_init, zeros, pd, pg) - ctr3, pd = pto.vaddcs(ctr3_init, zeros, pd, pg) - - for i in range(0, valid_rows, 1): - s_reg = valid_cols - counter_add_val = lanes - for j in range(0, n_loop, 1): - tmp_ctr0 = ctr0 - tmp_ctr1 = ctr1 - tmp_ctr2 = ctr2 - tmp_ctr3 = ctr3 - tmp_key0 = key0_v - tmp_key1 = key1_v - - if pto.constexpr(rounds_str == "10"): - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - elif pto.constexpr(rounds_str == "7"): - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - tmpL0, tmpH0 = pto.vmull(tmp_ctr0, const0, pg) - tmpL1, tmpH1 = pto.vmull(tmp_ctr2, const1, pg) - tmpH1 = pto.vxor(tmpH1, tmp_ctr1, pg) - tmp_ctr0 = pto.vxor(tmpH1, tmp_key0, pg) - tmpH0 = pto.vxor(tmpH0, tmp_ctr3, pg) - tmp_ctr2 = pto.vxor(tmpH0, tmp_key1, pg) - tmp_key0 = pto.vadds(tmp_key0, pto.ui32(TRANDOM_CONST_KEY_ADD_0), pg) - tmp_key1 = pto.vadds(tmp_key1, pto.ui32(TRANDOM_CONST_KEY_ADD_1), pg) - tmp_ctr1 = tmpL1 - tmp_ctr3 = tmpL0 - - tmpL0, tmpH0 = pto.vintlv(tmp_ctr0, tmp_ctr2) - tmpL1, tmpH1 = pto.vintlv(tmp_ctr1, tmp_ctr3) - tmp_ctr0, tmp_ctr1 = pto.vintlv(tmpL0, tmpL1) - tmp_ctr2, tmp_ctr3 = pto.vintlv(tmpH0, tmpH1) - - remained = s_reg - mask0, remained = pto.make_mask(dtype, remained) - mask1, remained = pto.make_mask(dtype, remained) - mask2, remained = pto.make_mask(dtype, remained) - mask3, remained = pto.make_mask(dtype, remained) - - pto.vsts(tmp_ctr0, dst[i, TRANDOM_ONCE_REPEAT * j * lanes:], mask0) - pto.vsts(tmp_ctr1, dst[i, (TRANDOM_ONCE_REPEAT * j + 1) * lanes:], mask1) - pto.vsts(tmp_ctr2, dst[i, (TRANDOM_ONCE_REPEAT * j + 2) * lanes:], mask2) - pto.vsts(tmp_ctr3, dst[i, (TRANDOM_ONCE_REPEAT * j + 3) * lanes:], mask3) - - if s_reg >= TRANDOM_ONCE_REPEAT * lanes: - s_reg = s_reg - TRANDOM_ONCE_REPEAT * lanes - else: - s_reg = 0 - - if j != n_loop - 1: - counter_add_val = lanes - else: - counter_add_val = (valid_cols - 1) % lanes + 1 - v_ele_stride = pto.vbr(pto.ui32(counter_add_val)) - ctr0_next, pd_next = pto.vaddc(ctr0, v_ele_stride, pg) - ctr1_next, pd_next2 = pto.vaddcs(ctr1, zeros, pd_next, pg) - ctr2_next, pd_next3 = pto.vaddcs(ctr2, zeros, pd_next2, pg) - ctr3_next, pd_next4 = pto.vaddcs(ctr3, zeros, pd_next3, pg) - ctr0 = ctr0_next - ctr1 = ctr1_next - ctr2 = ctr2_next - ctr3 = ctr3_next - - return diff --git a/lib/TileOps/trecip_template.py b/lib/TileOps/trecip_template.py deleted file mode 100644 index 2c42a0b453..0000000000 --- a/lib/TileOps/trecip_template.py +++ /dev/null @@ -1,61 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trecip with IEEE 754 high-precision support - -Computes reciprocal: dst = 1 / src -High-precision mode uses IEEE 754 compliant division algorithms. -""" - -import tilelang_dsl as pto - -# Import shared high-precision division algorithms -from div_hp import _div_ieee754_f32_impl, _div_ieee754_f16_impl - - -@pto.vkernel( - target="a5", - op="pto.trecip", - dtypes=[(pto.f16, pto.f16), (pto.f32, pto.f32)] -) -def template_trecip(src: pto.Tile, dst: pto.Tile): - """Reciprocal with optional high-precision mode: dst = 1 / src""" - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - precision_type = pto.get_op_attr("precisionType", "default") - if pto.constexpr(precision_type == "high_precision"): - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vinput = pto.vlds(src[row, col:]) - if pto.constexpr(dtype == pto.f16): - one_scalar = pto.f16(1.0) - else: - one_scalar = pto.f32(1.0) - one = pto.vbr(one_scalar) - if pto.constexpr(dtype == pto.f32): - result = _div_ieee754_f32_impl(one, vinput, mask) - else: # dtype == pto.f16 (guaranteed by MLIR validation) - result = _div_ieee754_f16_impl(one, vinput, mask) - pto.vsts(result, dst[row, col:], mask) - else: - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vinput = pto.vlds(src[row, col:]) - if pto.constexpr(dtype == pto.f16): - one_scalar = pto.f16(1.0) - else: - one_scalar = pto.f32(1.0) - one = pto.vbr(one_scalar) - result = pto.vdiv(one, vinput, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/trelu_template.py b/lib/TileOps/trelu_template.py deleted file mode 100644 index 0cdd0e7eb3..0000000000 --- a/lib/TileOps/trelu_template.py +++ /dev/null @@ -1,43 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trelu (Elementwise ReLU)""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.trelu", - dtypes=[(pto.f16, pto.f16), (pto.f32, pto.f32), (pto.i32, pto.i32)], - advanced=True -) -def template_trelu(src: pto.Tile, dst: pto.Tile): - """Elementwise ReLU: dst = max(0, src). - - Semantics: - For each element (i, j): - dst[i, j] = max(0, src[i, j]) - - Supported data types: f16, f32, i32 - """ - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - lanes = pto.get_lanes(dtype) - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - src_vec = pto.vlds(src[row, col:]) - result = pto.vrelu(src_vec, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/trem_template.py b/lib/TileOps/trem_template.py deleted file mode 100644 index 258638c83b..0000000000 --- a/lib/TileOps/trem_template.py +++ /dev/null @@ -1,66 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trem""" - -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.trem", - dtypes=[ - (pto.f32, pto.f32, pto.f32, pto.f32), - (pto.f16, pto.f16, pto.f16, pto.f16), - (pto.i32, pto.i32, pto.i32, pto.i32), - ], - advanced=True, -) -def template_trem(src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - # A5 keeps tmp in the ABI for compatibility with A2/A3, but the template - # computes directly in registers and never touches it. - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - if pto.constexpr(dtype == pto.f32): - quotient = pto.vdiv(lhs, rhs, mask) - quotient = pto.vtrc(quotient, mask, rnd=pto.VcvtRoundMode.F) - floored_mul = pto.vmul(quotient, rhs, mask) - result = pto.vsub(lhs, floored_mul, mask) - sign_diff_mask = pto.vcmps(pto.vmul(rhs, result, mask), 0.0, mask, pto.CmpMode.LT) - corrected = pto.vadd(result, rhs, sign_diff_mask) - result = pto.vsel(corrected, result, sign_diff_mask) - elif pto.constexpr(dtype == pto.f16): - lhs_even = pto.vcvt(lhs, pto.f32, mask, part=pto.VcvtPartMode.EVEN) - rhs_even = pto.vcvt(rhs, pto.f32, mask, part=pto.VcvtPartMode.EVEN) - lhs_odd = pto.vcvt(lhs, pto.f32, mask, part=pto.VcvtPartMode.ODD) - rhs_odd = pto.vcvt(rhs, pto.f32, mask, part=pto.VcvtPartMode.ODD) - q_even = pto.vdiv(lhs_even, rhs_even, mask) - q_odd = pto.vdiv(lhs_odd, rhs_odd, mask) - q_even = pto.vtrc(q_even, mask, rnd=pto.VcvtRoundMode.F) - q_odd = pto.vtrc(q_odd, mask, rnd=pto.VcvtRoundMode.F) - fm_even = pto.vmul(q_even, rhs_even, mask) - fm_odd = pto.vmul(q_odd, rhs_odd, mask) - r_even = pto.vsub(lhs_even, fm_even, mask) - r_odd = pto.vsub(lhs_odd, fm_odd, mask) - dst_even = pto.vcvt(r_even, pto.f16, mask, rnd=pto.VcvtRoundMode.Z, sat=pto.VcvtSatMode.RS_ENABLE, part=pto.VcvtPartMode.EVEN) - dst_odd = pto.vcvt(r_odd, pto.f16, mask, rnd=pto.VcvtRoundMode.Z, sat=pto.VcvtSatMode.RS_ENABLE, part=pto.VcvtPartMode.ODD) - result = pto.vor(dst_even, dst_odd, mask) - sign_diff_mask = pto.vcmps(pto.vmul(rhs, result, mask), 0.0, mask, pto.CmpMode.LT) - corrected = pto.vadd(result, rhs, sign_diff_mask) - result = pto.vsel(corrected, result, sign_diff_mask) - elif pto.constexpr(dtype == pto.i32): - result = pto.vmod(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/trems_template.py b/lib/TileOps/trems_template.py deleted file mode 100644 index b508fd99d9..0000000000 --- a/lib/TileOps/trems_template.py +++ /dev/null @@ -1,96 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trems - -Note: A5 hardware implements trems as: - - float: dst = src - trunc(src/scalar) * scalar (f32 precision) - - half: dst = src - trunc(src/scalar) * scalar (computed in f32 precision, then converted back to f16) - - integer: dst = src % scalar (using vmod) - NOT YET SUPPORTED - -f16 path aligns: convert f16 to f32 (even/odd), compute in f32 with vtrc(ROUND_F), -convert back to f16 with ROUND_Z, merge with vor. -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.trems", - dtypes=[(pto.f32, pto.f32, pto.f32, pto.f32), (pto.f16, pto.f16, pto.f16, pto.f16)], - advanced=True, -) -def template_trems(src: pto.Tile, scalar: pto.AnyType, tmp: pto.Tile, dst: pto.Tile): - # A5 keeps tmp in the ABI for compatibility with A2/A3, but this template - # computes directly in registers and never touches it. - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - if pto.constexpr(dtype == pto.f32): - # f32 path: direct f32 computation - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f32)): - mask, remained = pto.make_mask(pto.f32, remained) - vec = pto.vlds(src[row, col:]) - scalar_vec = pto.vbr(scalar) - quotient = pto.vdiv(vec, scalar_vec, mask) - truncated = pto.vtrc(quotient, mask, rnd="Z") - product = pto.vmuls(truncated, scalar, mask) - result = pto.vsub(vec, product, mask) - pto.vsts(result, dst[row, col:], mask) - else: - # f16 path: compute in f32 precision, then convert back to f16 - # The predicate register is shared across f16 and f32 operations. - full_mask_b16 = pto.make_mask(pto.f16, pto.PAT.ALL) - full_mask_b32 = pto.make_mask(pto.f32, pto.PAT.ALL) - # Broadcast f16 scalar, convert to f32 vector for f32 arithmetic - scalar_vec_f16 = pto.vbr(scalar) - scalar_f32_vec = pto.vcvt(scalar_vec_f16, pto.f32, full_mask_b16, part=pto.VcvtPartMode.EVEN) - - for row in range(0, valid_rows, 1): - remained = valid_cols - remained_f32 = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f16)): - mask_f16, remained = pto.make_mask(pto.f16, remained) - mask_f32, remained_f32 = pto.make_mask(pto.f32, remained_f32) - vec = pto.vlds(src[row, col:]) - - # Convert f16 to f32 (even and odd parts) - vec_even = pto.vcvt(vec, pto.f32, full_mask_b16, part=pto.VcvtPartMode.EVEN) - vec_odd = pto.vcvt(vec, pto.f32, full_mask_b16, part=pto.VcvtPartMode.ODD) - - # Even part: f32 computation (use vmul with f32 scalar vector) - quotient_even = pto.vdiv(vec_even, scalar_f32_vec, mask_f32) - truncated_even = pto.vtrc(quotient_even, mask_f32, rnd="F") - product_even = pto.vmul(truncated_even, scalar_f32_vec, mask_f32) - result_even = pto.vsub(vec_even, product_even, mask_f32) - - # Odd part: f32 computation - quotient_odd = pto.vdiv(vec_odd, scalar_f32_vec, mask_f32) - truncated_odd = pto.vtrc(quotient_odd, mask_f32, rnd="F") - product_odd = pto.vmul(truncated_odd, scalar_f32_vec, mask_f32) - result_odd = pto.vsub(vec_odd, product_odd, mask_f32) - - # Convert f32 results back to f16 with ROUND_Z + saturation - result_f16_even = pto.vcvt(result_even, pto.f16, full_mask_b32, - rnd=pto.VcvtRoundMode.Z, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.EVEN) - result_f16_odd = pto.vcvt(result_odd, pto.f16, full_mask_b32, - rnd=pto.VcvtRoundMode.Z, - sat=pto.VcvtSatMode.SAT, - part=pto.VcvtPartMode.ODD) - - # Merge even and odd parts - result_f16 = pto.vor(result_f16_even, result_f16_odd, mask_f16) - pto.vsts(result_f16, dst[row, col:], mask_f16) - return diff --git a/lib/TileOps/trowargmax_template.py b/lib/TileOps/trowargmax_template.py deleted file mode 100644 index c0df597ed3..0000000000 --- a/lib/TileOps/trowargmax_template.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowargmax""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - -@pto.vkernel( - target="a5", - op="pto.trowargmax", - advanced=True, -) -def template_trowargmax(src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - src_dtype = src.element_type - idx_dtype = dst.element_type - lanes = pto.get_lanes(src_dtype) - valid_rows, valid_cols = src.valid_shape - elem_bytes = pto.bytewidth(idx_dtype) - - # Initialize with dtype-specific minimum value (aligned with pto-isa Padding::Min) - init_val = pto.PadValue.MIN.eval(src_dtype) - - # Select one-point store dist based on index dtype size - if pto.constexpr(elem_bytes == 4): - store_dist = pto.VStoreDist.ONE_POINT_B32 - elif pto.constexpr(elem_bytes == 2): - store_dist = pto.VStoreDist.ONE_POINT_B16 - else: - store_dist = pto.VStoreDist.ONE_POINT_B8 - - for row in range(0, valid_rows, 1): - remained = valid_cols - - v_val_acc = pto.vbr(init_val) - init_zero_idx = idx_dtype(0) - v_idx_acc = pto.vbr(init_zero_idx) - - # Masks: src_dtype for data ops and final store (matches pto-isa CreatePredicate) - # idx_dtype for index arithmetic operations - mask_1, _ = pto.make_mask(src_dtype, 1) - mask_1_idx, _ = pto.make_mask(idx_dtype, 1) - - # Process all column chunks - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(src_dtype, remained) - v_src = pto.vlds(src[row, col:]) - v_reduced = pto.vcmax(v_src, mask) - - v_val, v_idx = pto.vdintlv(v_reduced, pto.vbr(src_dtype(0))) - v_idx = pto.vbitcast(v_idx, idx_dtype) - - # Add absolute col offset to the chunk's local index - col_offset = idx_dtype(col) - v_idx = pto.vadds(v_idx, col_offset, mask_1_idx) - - # Compare current chunk max with global max so far - cmp_mask = pto.vcmp(v_val_acc, v_val, mask_1, "lt") - - # Update global max and global argmax - v_val_acc = pto.vsel(v_val, v_val_acc, cmp_mask) - # v_idx_acc is ui32, requires b32 mask; convert cmp_mask from src_dtype's mask to b32 - cmp_mask_b32 = pto.pbitcast(cmp_mask, pto.mask_b32) - v_idx_acc = pto.vsel(v_idx, v_idx_acc, cmp_mask_b32) - - # Store index accumulator to destination tile using one-point mode - pto.vsts(v_idx_acc, dst[row, 0:], mask_1_idx, dist=store_dist) - return \ No newline at end of file diff --git a/lib/TileOps/trowargmin_template.py b/lib/TileOps/trowargmin_template.py deleted file mode 100644 index f23ab7137f..0000000000 --- a/lib/TileOps/trowargmin_template.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowargmin""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - -@pto.vkernel( - target="a5", - op="pto.trowargmin", - advanced=True, -) -def template_trowargmin(src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - src_dtype = src.element_type - idx_dtype = dst.element_type - lanes = pto.get_lanes(src_dtype) - valid_rows, valid_cols = src.valid_shape - elem_bytes = pto.bytewidth(idx_dtype) - - # Initialize with dtype-specific maximum value (aligned with pto-isa Padding::Max) - init_val = pto.PadValue.MAX.eval(src_dtype) - - # Select one-point store dist based on index dtype size - if pto.constexpr(elem_bytes == 4): - store_dist = pto.VStoreDist.ONE_POINT_B32 - elif pto.constexpr(elem_bytes == 2): - store_dist = pto.VStoreDist.ONE_POINT_B16 - else: - store_dist = pto.VStoreDist.ONE_POINT_B8 - - for row in range(0, valid_rows, 1): - remained = valid_cols - - v_val_acc = pto.vbr(init_val) - init_zero_idx = idx_dtype(0) - v_idx_acc = pto.vbr(init_zero_idx) - - # Masks: src_dtype for data ops and final store (matches pto-isa CreatePredicate) - # idx_dtype for index arithmetic operations - mask_1, _ = pto.make_mask(src_dtype, 1) - mask_1_idx, _ = pto.make_mask(idx_dtype, 1) - - # Process all column chunks - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(src_dtype, remained) - v_src = pto.vlds(src[row, col:]) - v_reduced = pto.vcmin(v_src, mask) - - v_val, v_idx = pto.vdintlv(v_reduced, pto.vbr(src_dtype(0))) - v_idx = pto.vbitcast(v_idx, idx_dtype) - - # Add absolute col offset to the chunk's local index - col_offset = idx_dtype(col) - v_idx = pto.vadds(v_idx, col_offset, mask_1_idx) - - # Compare current chunk min with global min so far - cmp_mask = pto.vcmp(v_val_acc, v_val, mask_1, "gt") - - # Update global min and global argmin - v_val_acc = pto.vsel(v_val, v_val_acc, cmp_mask) - # v_idx_acc is ui32, requires b32 mask; cast cmp_mask from src_dtype's mask to b32 - cmp_mask_b32 = pto.pbitcast(cmp_mask, pto.mask_b32) - v_idx_acc = pto.vsel(v_idx, v_idx_acc, cmp_mask_b32) - - # Store index accumulator to destination tile using one-point mode - pto.vsts(v_idx_acc, dst[row, 0:], mask_1_idx, dist=store_dist) - return diff --git a/lib/TileOps/trowexpand_template.py b/lib/TileOps/trowexpand_template.py deleted file mode 100644 index 4dc6c6a791..0000000000 --- a/lib/TileOps/trowexpand_template.py +++ /dev/null @@ -1,48 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowexpand""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -def _constraint_trowexpand_row_major(src: pto.Tile, dst: pto.Tile) -> bool: - """Constraint for RowMajor layout trowexpand template.""" - # Both src and dst must be RowMajor layout - src_row_major = src.config.b_layout == pto.BLayout.ROW_MAJOR - dst_row_major = dst.config.b_layout == pto.BLayout.ROW_MAJOR - return src_row_major and dst_row_major - - -@pto.vkernel( - target="a5", - op="pto.trowexpand", - dtypes=[(pto.AnyFloat, pto.AnyFloat), (pto.AnyInt, pto.AnyInt)], - constraints=[_constraint_trowexpand_row_major], -) -def template_trowexpand(src: pto.Tile, dst: pto.Tile): - """Template for pto.trowexpand. - - Broadcast src[row, 0] to entire dst[row, :] for each row. - Semantics: dst[row, col] = src[row, 0] for all col. - """ - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - # Load the first element of each row (src has cols=1, so entire row is the scalar) - # vdup broadcasts the first element to the full vector width - scalar_vec = pto.vlds(src[row, :]) - broadcasted = pto.vdup(scalar_vec, mask) - pto.vsts(broadcasted, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/trowexpandadd_template.py b/lib/TileOps/trowexpandadd_template.py deleted file mode 100644 index 96019d5017..0000000000 --- a/lib/TileOps/trowexpandadd_template.py +++ /dev/null @@ -1,52 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowexpandadd""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -def _constraint_trowexpandadd_row_major(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile) -> bool: - """Constraint for RowMajor layout trowexpandadd template.""" - # All tiles must be RowMajor layout - src0_row_major = src0.config.b_layout == pto.BLayout.ROW_MAJOR - src1_row_major = src1.config.b_layout == pto.BLayout.ROW_MAJOR - dst_row_major = dst.config.b_layout == pto.BLayout.ROW_MAJOR - return src0_row_major and src1_row_major and dst_row_major - - -@pto.vkernel( - target="a5", - op="pto.trowexpandadd", - dtypes=[(pto.AnyFloat, pto.AnyFloat, pto.AnyFloat), (pto.AnyInt, pto.AnyInt, pto.AnyInt)], - constraints=[_constraint_trowexpandadd_row_major], -) -def template_trowexpandadd(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - """Template for pto.trowexpandadd. - - Add a per-row scalar from src1[row, 0] to each row of src0. - Semantics: dst[row, col] = src0[row, col] + src1[row, 0] - """ - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - # Load the scalar vector from src1[row, :] - # For row-major src1, valid_shape[1] is 32/sizeof(dtype) (e.g., 8 for f32) - # vdup broadcasts the first element to the full vector width - scalar_vec = pto.vlds(src1[row, :]) - broadcasted = pto.vdup(scalar_vec, mask) - lhs = pto.vlds(src0[row, col:]) - result = pto.vadd(lhs, broadcasted, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/trowexpanddiv_template.py b/lib/TileOps/trowexpanddiv_template.py deleted file mode 100644 index 888353014a..0000000000 --- a/lib/TileOps/trowexpanddiv_template.py +++ /dev/null @@ -1,131 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowexpanddiv with IEEE 754 high-precision support - -Divide each row of src0 by a per-row scalar from src1[row, 0]. -Semantics: dst[row, col] = src0[row, col] / src1[row, 0] -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - -# Import shared high-precision division algorithms -from div_hp import _div_ieee754_f32_impl, _div_ieee754_f16_impl - - -def _config_value(config, name): - if config is None: - return None - if isinstance(config, dict): - return config.get(name) - return getattr(config, name, None) - - -def _matches_layout(value, expected, expected_name): - if value is None: - return False - return value == expected or value == expected_name or str(value).lower().endswith(expected_name) - - -def _is_row_major(tile: pto.Tile) -> bool: - return _matches_layout(_config_value(tile.config, "b_layout"), pto.BLayout.ROW_MAJOR, "row_major") - - -def _is_col_major(tile: pto.Tile) -> bool: - return _matches_layout(_config_value(tile.config, "b_layout"), pto.BLayout.COL_MAJOR, "col_major") - - -def _is_col_major_row_scalar(tile: pto.Tile) -> bool: - shape = tuple(tile.shape) - return len(shape) == 2 and _is_col_major(tile) and shape[1] == 1 - - -def _constraint_trowexpanddiv_row_major(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile) -> bool: - """Constraint for RowMajor layout trowexpanddiv template.""" - src1_supported = _is_row_major(src1) or _is_col_major_row_scalar(src1) - return _is_row_major(src0) and src1_supported and _is_row_major(dst) - - -@pto.vkernel( - target="a5", - op="pto.trowexpanddiv", - dtypes=[(pto.f32, pto.f32, pto.f32)], - constraints=[_constraint_trowexpanddiv_row_major], -) -def template_trowexpanddiv_f32(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - """Template for pto.trowexpanddiv with f32 dtype and optional high-precision mode.""" - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - precision_type = pto.get_op_attr("precisionType", "default") - if pto.constexpr(precision_type == "high_precision"): - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - # Load the scalar vector from src1[row, :] - # For row-major src1, valid_shape[1] is 32/sizeof(dtype) (e.g., 8 for f32) - # vdup broadcasts the first element to the full vector width - scalar_vec = pto.vlds(src1[row, :]) - broadcasted = pto.vdup(scalar_vec, mask) - lhs = pto.vlds(src0[row, col:]) - result = _div_ieee754_f32_impl(lhs, broadcasted, mask) - pto.vsts(result, dst[row, col:], mask) - else: - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - # Load the scalar vector from src1[row, :] - scalar_vec = pto.vlds(src1[row, :]) - broadcasted = pto.vdup(scalar_vec, mask) - lhs = pto.vlds(src0[row, col:]) - result = pto.vdiv(lhs, broadcasted, mask) - pto.vsts(result, dst[row, col:], mask) - return - - -@pto.vkernel( - target="a5", - op="pto.trowexpanddiv", - dtypes=[(pto.f16, pto.f16, pto.f16)], - constraints=[_constraint_trowexpanddiv_row_major], -) -def template_trowexpanddiv_f16(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - """Template for pto.trowexpanddiv with f16 dtype and optional high-precision mode.""" - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - precision_type = pto.get_op_attr("precisionType", "default") - if pto.constexpr(precision_type == "high_precision"): - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - # Load the scalar vector from src1[row, :] - # For row-major src1, valid_shape[1] is 32/sizeof(dtype) (e.g., 16 for f16) - # vdup broadcasts the first element to the full vector width - scalar_vec = pto.vlds(src1[row, :]) - broadcasted = pto.vdup(scalar_vec, mask) - lhs = pto.vlds(src0[row, col:]) - result = _div_ieee754_f16_impl(lhs, broadcasted, mask) - pto.vsts(result, dst[row, col:], mask) - else: - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - # Load the scalar vector from src1[row, :] - scalar_vec = pto.vlds(src1[row, :]) - broadcasted = pto.vdup(scalar_vec, mask) - lhs = pto.vlds(src0[row, col:]) - result = pto.vdiv(lhs, broadcasted, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/trowexpandexpdif_template.py b/lib/TileOps/trowexpandexpdif_template.py deleted file mode 100644 index fed08feab5..0000000000 --- a/lib/TileOps/trowexpandexpdif_template.py +++ /dev/null @@ -1,78 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowexpandexpdif""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -def _constraint_trowexpandexpdif_row_major(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile) -> bool: - """Constraint for RowMajor layout trowexpandexpdif template.""" - # All tiles must be RowMajor layout - src0_row_major = src0.config.b_layout == pto.BLayout.ROW_MAJOR - src1_row_major = src1.config.b_layout == pto.BLayout.ROW_MAJOR - dst_row_major = dst.config.b_layout == pto.BLayout.ROW_MAJOR - return src0_row_major and src1_row_major and dst_row_major - - -@pto.vkernel( - target="a5", - op="pto.trowexpandexpdif", - dtypes=[(pto.f32, pto.f32, pto.f32)], - constraints=[_constraint_trowexpandexpdif_row_major], -) -def template_trowexpandexpdif_f32(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - """Template for pto.trowexpandexpdif with f32 dtype. - - Compute exp(src0 - scalar) for each row using per-row scalars from src1[row, 0]. - Semantics: dst[row, col] = exp(src0[row, col] - src1[row, 0]) - Used in numerically stable softmax computation. - """ - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(pto.f32)): - mask, remained = pto.make_mask(pto.f32, remained) - scalar_vec = pto.vlds(src1[row, :]) - broadcasted = pto.vdup(scalar_vec, mask) - lhs = pto.vlds(src0[row, col:]) - result = pto.vexpdif(lhs, broadcasted, mask, pto.VcvtPartMode.EVEN) - pto.vsts(result, dst[row, col:], mask) - return - - -@pto.vkernel( - target="a5", - op="pto.trowexpandexpdif", - dtypes=[(pto.f16, pto.f16, pto.f16)], - constraints=[_constraint_trowexpandexpdif_row_major], -) -def template_trowexpandexpdif_f16(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - """Template for pto.trowexpandexpdif with f16 dtype. - - Compute exp(src0 - scalar) for each row using per-row scalars from src1[row, 0]. - Semantics: dst[row, col] = exp(src0[row, col] - src1[row, 0]) - Used in numerically stable softmax computation. - """ - dtype = pto.f16 - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - scalar_vec = pto.vlds(src1[row, :]) - broadcasted = pto.vdup(scalar_vec, mask) - lhs = pto.vlds(src0[row, col:]) - diff = pto.vsub(lhs, broadcasted, mask) - result = pto.vexp(diff, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/trowexpandmax_template.py b/lib/TileOps/trowexpandmax_template.py deleted file mode 100644 index 2e5d53124e..0000000000 --- a/lib/TileOps/trowexpandmax_template.py +++ /dev/null @@ -1,52 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowexpandmax""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -def _constraint_trowexpandmax_row_major(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile) -> bool: - """Constraint for RowMajor layout trowexpandmax template.""" - # All tiles must be RowMajor layout - src0_row_major = src0.config.b_layout == pto.BLayout.ROW_MAJOR - src1_row_major = src1.config.b_layout == pto.BLayout.ROW_MAJOR - dst_row_major = dst.config.b_layout == pto.BLayout.ROW_MAJOR - return src0_row_major and src1_row_major and dst_row_major - - -@pto.vkernel( - target="a5", - op="pto.trowexpandmax", - dtypes=[(pto.AnyFloat, pto.AnyFloat, pto.AnyFloat), (pto.AnyInt, pto.AnyInt, pto.AnyInt)], - constraints=[_constraint_trowexpandmax_row_major], -) -def template_trowexpandmax(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - """Template for pto.trowexpandmax. - - Compute element-wise max of each row of src0 with a per-row scalar from src1[row, 0]. - Semantics: dst[row, col] = max(src0[row, col], src1[row, 0]) - """ - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - # Load the scalar vector from src1[row, :] - # For row-major src1, valid_shape[1] is 32/sizeof(dtype) (e.g., 8 for f32) - # vdup broadcasts the first element to the full vector width - scalar_vec = pto.vlds(src1[row, :]) - broadcasted = pto.vdup(scalar_vec, mask) - lhs = pto.vlds(src0[row, col:]) - result = pto.vmax(lhs, broadcasted, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/trowexpandmin_template.py b/lib/TileOps/trowexpandmin_template.py deleted file mode 100644 index eae99ff30b..0000000000 --- a/lib/TileOps/trowexpandmin_template.py +++ /dev/null @@ -1,52 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowexpandmin""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -def _constraint_trowexpandmin_row_major(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile) -> bool: - """Constraint for RowMajor layout trowexpandmin template.""" - # All tiles must be RowMajor layout - src0_row_major = src0.config.b_layout == pto.BLayout.ROW_MAJOR - src1_row_major = src1.config.b_layout == pto.BLayout.ROW_MAJOR - dst_row_major = dst.config.b_layout == pto.BLayout.ROW_MAJOR - return src0_row_major and src1_row_major and dst_row_major - - -@pto.vkernel( - target="a5", - op="pto.trowexpandmin", - dtypes=[(pto.AnyFloat, pto.AnyFloat, pto.AnyFloat), (pto.AnyInt, pto.AnyInt, pto.AnyInt)], - constraints=[_constraint_trowexpandmin_row_major], -) -def template_trowexpandmin(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - """Template for pto.trowexpandmin. - - Compute element-wise min of each row of src0 with a per-row scalar from src1[row, 0]. - Semantics: dst[row, col] = min(src0[row, col], src1[row, 0]) - """ - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - # Load the scalar vector from src1[row, :] - # For row-major src1, valid_shape[1] is 32/sizeof(dtype) (e.g., 8 for f32) - # vdup broadcasts the first element to the full vector width - scalar_vec = pto.vlds(src1[row, :]) - broadcasted = pto.vdup(scalar_vec, mask) - lhs = pto.vlds(src0[row, col:]) - result = pto.vmin(lhs, broadcasted, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/trowexpandmul_template.py b/lib/TileOps/trowexpandmul_template.py deleted file mode 100644 index 504e22392d..0000000000 --- a/lib/TileOps/trowexpandmul_template.py +++ /dev/null @@ -1,76 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowexpandmul""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -def _config_value(config, name): - if config is None: - return None - if isinstance(config, dict): - return config.get(name) - return getattr(config, name, None) - - -def _matches_layout(value, expected, expected_name): - if value is None: - return False - return value == expected or value == expected_name or str(value).lower().endswith(expected_name) - - -def _is_row_major(tile: pto.Tile) -> bool: - return _matches_layout(_config_value(tile.config, "b_layout"), pto.BLayout.ROW_MAJOR, "row_major") - - -def _is_col_major(tile: pto.Tile) -> bool: - return _matches_layout(_config_value(tile.config, "b_layout"), pto.BLayout.COL_MAJOR, "col_major") - - -def _is_col_major_row_scalar(tile: pto.Tile) -> bool: - shape = tuple(tile.shape) - return len(shape) == 2 and _is_col_major(tile) and shape[1] == 1 - - -def _constraint_trowexpandmul_row_major(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile) -> bool: - """Constraint for RowMajor layout trowexpandmul template.""" - src1_supported = _is_row_major(src1) or _is_col_major_row_scalar(src1) - return _is_row_major(src0) and src1_supported and _is_row_major(dst) - - -@pto.vkernel( - target="a5", - op="pto.trowexpandmul", - dtypes=[(pto.AnyFloat, pto.AnyFloat, pto.AnyFloat), (pto.AnyInt, pto.AnyInt, pto.AnyInt)], - constraints=[_constraint_trowexpandmul_row_major], -) -def template_trowexpandmul(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - """Template for pto.trowexpandmul. - - Multiply each row of src0 by a per-row scalar from src1[row, 0]. - Semantics: dst[row, col] = src0[row, col] * src1[row, 0] - """ - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - # Load the scalar vector from src1[row, :] - # For row-major src1, valid_shape[1] is 32/sizeof(dtype) (e.g., 8 for f32) - # vdup broadcasts the first element to the full vector width - scalar_vec = pto.vlds(src1[row, :]) - broadcasted = pto.vdup(scalar_vec, mask) - lhs = pto.vlds(src0[row, col:]) - result = pto.vmul(lhs, broadcasted, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/trowexpandsub_template.py b/lib/TileOps/trowexpandsub_template.py deleted file mode 100644 index 22cc94b6e4..0000000000 --- a/lib/TileOps/trowexpandsub_template.py +++ /dev/null @@ -1,76 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowexpandsub""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -def _config_value(config, name): - if config is None: - return None - if isinstance(config, dict): - return config.get(name) - return getattr(config, name, None) - - -def _matches_layout(value, expected, expected_name): - if value is None: - return False - return value == expected or value == expected_name or str(value).lower().endswith(expected_name) - - -def _is_row_major(tile: pto.Tile) -> bool: - return _matches_layout(_config_value(tile.config, "b_layout"), pto.BLayout.ROW_MAJOR, "row_major") - - -def _is_col_major(tile: pto.Tile) -> bool: - return _matches_layout(_config_value(tile.config, "b_layout"), pto.BLayout.COL_MAJOR, "col_major") - - -def _is_col_major_row_scalar(tile: pto.Tile) -> bool: - shape = tuple(tile.shape) - return len(shape) == 2 and _is_col_major(tile) and shape[1] == 1 - - -def _constraint_trowexpandsub_row_major(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile) -> bool: - """Constraint for RowMajor layout trowexpandsub template.""" - src1_supported = _is_row_major(src1) or _is_col_major_row_scalar(src1) - return _is_row_major(src0) and src1_supported and _is_row_major(dst) - - -@pto.vkernel( - target="a5", - op="pto.trowexpandsub", - dtypes=[(pto.AnyFloat, pto.AnyFloat, pto.AnyFloat), (pto.AnyInt, pto.AnyInt, pto.AnyInt)], - constraints=[_constraint_trowexpandsub_row_major], -) -def template_trowexpandsub(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - """Template for pto.trowexpandsub. - - Subtract a per-row scalar from src1[row, 0] from each row of src0. - Semantics: dst[row, col] = src0[row, col] - src1[row, 0] - """ - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - # Load the scalar vector from src1[row, :] - # For row-major src1, valid_shape[1] is 32/sizeof(dtype) (e.g., 8 for f32) - # vdup broadcasts the first element to the full vector width - scalar_vec = pto.vlds(src1[row, :]) - broadcasted = pto.vdup(scalar_vec, mask) - lhs = pto.vlds(src0[row, col:]) - result = pto.vsub(lhs, broadcasted, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/trowmax_template.py b/lib/TileOps/trowmax_template.py deleted file mode 100644 index 522ce699ed..0000000000 --- a/lib/TileOps/trowmax_template.py +++ /dev/null @@ -1,61 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowmax""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - -@pto.vkernel( - target="a5", - op="pto.trowmax", - advanced=True, -) -def template_trowmax(src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - lanes = pto.get_lanes(dtype) - valid_rows, valid_cols = src.valid_shape - elem_bytes = pto.bytewidth(dtype) - - # Initialize with dtype-specific minimum value (aligned with pto-isa Padding::Min) - init_val = pto.PadValue.MIN.eval(dtype) - - # Select one-point store dist based on dtype size - if pto.constexpr(elem_bytes == 4): - store_dist = pto.VStoreDist.ONE_POINT_B32 - elif pto.constexpr(elem_bytes == 2): - store_dist = pto.VStoreDist.ONE_POINT_B16 - else: - store_dist = pto.VStoreDist.ONE_POINT_B8 - - for row in range(0, valid_rows, 1): - remained = valid_cols - - mask_1, _ = pto.make_mask(dtype, 1) - - # Initialize the accumulator for ROWMAX - v_acc = pto.vbr(init_val) - - # Process column chunks - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - v_src = pto.vlds(src[row, col:]) - - # vcmax reduces src_dtype to acc_dtype - v_reduced = pto.vcmax(v_src, mask) - - # Clear masked lanes to init_val for float types so vmax doesn't see NaN - if pto.constexpr(dtype == pto.f32 or dtype == pto.f16): - v_reduced = pto.vsel(v_reduced, v_acc, mask) - - v_acc = pto.vmax(v_acc, v_reduced, mask_1) - - # Write final reduction to dest buffer once using one-point mode - pto.vsts(v_acc, dst[row, 0:], mask_1, dist=store_dist) - return diff --git a/lib/TileOps/trowmin_template.py b/lib/TileOps/trowmin_template.py deleted file mode 100644 index f74b798c4d..0000000000 --- a/lib/TileOps/trowmin_template.py +++ /dev/null @@ -1,63 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowmin""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.trowmin", - advanced=True, -) -def template_trowmin(src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - lanes = pto.get_lanes(dtype) - valid_rows, valid_cols = src.valid_shape - elem_bytes = pto.bytewidth(dtype) - - # Initialize with dtype-specific maximum value (aligned with pto-isa Padding::Max) - init_val = pto.PadValue.MAX.eval(dtype) - - # Select one-point store dist based on dtype size - if pto.constexpr(elem_bytes == 4): - store_dist = pto.VStoreDist.ONE_POINT_B32 - elif pto.constexpr(elem_bytes == 2): - store_dist = pto.VStoreDist.ONE_POINT_B16 - else: - store_dist = pto.VStoreDist.ONE_POINT_B8 - - mask_1, _ = pto.make_mask(dtype, 1) - - for row in range(0, valid_rows, 1): - remained = valid_cols - - # Initialize the accumulator for ROWMIN - v_acc = pto.vbr(init_val) - - # Process column chunks - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - v_src = pto.vlds(src[row, col:]) - - # vcmin reduces src_dtype to acc_dtype - v_reduced = pto.vcmin(v_src, mask) - - # Clear masked lanes to init_val for float types so vmin doesn't see NaN - if pto.constexpr(dtype == pto.f32 or dtype == pto.f16): - v_reduced = pto.vsel(v_reduced, v_acc, mask) - - # accumulate using the accumulator's mask logic - v_acc = pto.vmin(v_acc, v_reduced, mask_1) - - # Write final reduction to dest buffer once using one-point mode - pto.vsts(v_acc, dst[row, 0:], mask_1, dist=store_dist) - return diff --git a/lib/TileOps/trowprod_template.py b/lib/TileOps/trowprod_template.py deleted file mode 100644 index 55208d328b..0000000000 --- a/lib/TileOps/trowprod_template.py +++ /dev/null @@ -1,71 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowprod""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - -@pto.vkernel( - target="a5", - op="pto.trowprod", - advanced=True, -) -def template_trowprod(src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - lanes = pto.get_lanes(dtype) - valid_rows, valid_cols = src.valid_shape - elem_bytes = pto.bytewidth(dtype) - - # nLoop from C++ constants: TROW_PROD_LOOP_B16=7, TROW_PROD_LOOP_B32=6 - TROW_PROD_LOOP_B16 = 7 - TROW_PROD_LOOP_B32 = 6 - if pto.constexpr(dtype == pto.f16 or dtype == pto.i16): - n_loop = TROW_PROD_LOOP_B16 - else: - n_loop = TROW_PROD_LOOP_B32 - - # Select one-point store dist based on dtype size - if pto.constexpr(elem_bytes == 4): - store_dist = pto.VStoreDist.ONE_POINT_B32 - elif pto.constexpr(elem_bytes == 2): - store_dist = pto.VStoreDist.ONE_POINT_B16 - else: - store_dist = pto.VStoreDist.ONE_POINT_B8 - - mask_1, _ = pto.make_mask(dtype, 1) - - for row in range(0, valid_rows, 1): - remained = valid_cols - - one_val = dtype(1) - v_acc = pto.vbr(one_val) - v_one = pto.vbr(one_val) - - # Multiply across column chunks - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(dtype, remained) - v_src = pto.vlds(src[row, col:]) - - # Element-wise product - v_prod = pto.vmul(v_acc, v_src, mask) - - # Simulate MODE_MERGING with vsel (keep v_acc outside mask) - v_acc = pto.vsel(v_prod, v_acc, mask) - - # Log2 reduction phase across the vector - reduce_mask, _ = pto.make_mask(dtype, lanes) # all lanes active for inner reduction - - for k in range(0, n_loop, 1): - v_intlv1, v_intlv2 = pto.vintlv(v_acc, v_one) - v_acc = pto.vmul(v_intlv1, v_intlv2, reduce_mask) - - # Write final result at lane 0 using one-point mode - pto.vsts(v_acc, dst[row, 0:], mask_1, dist=store_dist) - return diff --git a/lib/TileOps/trowsum_template.py b/lib/TileOps/trowsum_template.py deleted file mode 100644 index cada565ffa..0000000000 --- a/lib/TileOps/trowsum_template.py +++ /dev/null @@ -1,75 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trowsum""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.trowsum", -) -def template_trowsum(src: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - src_dtype = src.element_type - dst_dtype = dst.element_type - - # vcadd widens i16 -> i32; floats/i32 unchanged - if pto.constexpr(src_dtype == pto.i16): - acc_dtype = pto.i32 - else: - acc_dtype = src_dtype - - lanes = pto.get_lanes(src_dtype) - valid_rows, valid_cols = src.valid_shape - - # Use type-appropriate zero for accumulator initialization - zero_val = acc_dtype(0) - - # Select one-point store dist based on dst dtype size - elem_bytes = pto.bytewidth(dst_dtype) - if pto.constexpr(elem_bytes == 4): - store_dist = pto.VStoreDist.ONE_POINT_B32 - elif pto.constexpr(elem_bytes == 2): - store_dist = pto.VStoreDist.ONE_POINT_B16 - else: - store_dist = pto.VStoreDist.ONE_POINT_B8 - - dst_mask_1, _ = pto.make_mask(dst_dtype, 1) - - for row in range(0, valid_rows, 1): - remained = valid_cols - - acc_mask_1, _ = pto.make_mask(acc_dtype, 1) - - # Initialize the accumulator with type-appropriate zero - v_acc = pto.vbr(zero_val) - - # Process column chunks - for col in range(0, valid_cols, lanes): - mask, remained = pto.make_mask(src_dtype, remained) - v_src = pto.vlds(src[row, col:]) - - # vcadd widens src_dtype to acc_dtype for integer types - v_reduced = pto.vcadd(v_src, mask) - - # accumulate using the accumulator's mask logic - v_acc = pto.vadd(v_acc, v_reduced, acc_mask_1) - - # Store the accumulated result safely once per row using one-point mode - if pto.constexpr(src_dtype == pto.i16): - # Truncate i32 accumulator back to i16 - # Non-saturation mode (wrap-around), matching pto-isa CTRL[59:60] behavior - acc_mask_for_cvt, _ = pto.make_mask(acc_dtype, 1) - v_acc_casted = pto.vcvt(v_acc, dst_dtype, acc_mask_for_cvt, sat=pto.VcvtSatMode.NOSAT, part=pto.VcvtPartMode.EVEN) - pto.vsts(v_acc_casted, dst[row, 0:], dst_mask_1, dist=store_dist) - else: - pto.vsts(v_acc, dst[row, 0:], dst_mask_1, dist=store_dist) - return diff --git a/lib/TileOps/trsqrt_template.py b/lib/TileOps/trsqrt_template.py deleted file mode 100644 index 87368adcaf..0000000000 --- a/lib/TileOps/trsqrt_template.py +++ /dev/null @@ -1,36 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.trsqrt""" - -import tilelang_dsl as pto - -# TODO: Add implementation for HIGH_PRECISION type -@pto.vkernel( - target="a5", - op="pto.trsqrt", - dtypes=[(pto.f16, pto.f16), (pto.f32, pto.f32)] -) -def template_trsqrt(src: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vinput = pto.vlds(src[row, col:]) - if pto.constexpr(dtype == pto.f16): - one_scalar = pto.f16(1.0) - else: - one_scalar = pto.f32(1.0) - one = pto.vbr(one_scalar) - sqrt_result = pto.vsqrt(vinput, mask) - result = pto.vdiv(one, sqrt_result, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tsel_template.py b/lib/TileOps/tsel_template.py deleted file mode 100644 index 92285716b4..0000000000 --- a/lib/TileOps/tsel_template.py +++ /dev/null @@ -1,99 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tsel - -NOTE: This template uses pto.plds for mask loading which directly -loads predicate mask from UB without vcmps comparison. -This approach matches the TSel.hpp implementation in pto-isa. - -Mask tile format: -- Packed predicate bytes in UB (`i8` tile data). -- Each row stores `ceil(valid_cols / 8)` valid bytes; tile row stride may be padded. - -REQUIRES: tilelang_dsl support for plds, astype(mask), pintlv_b16, castptr operations -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tsel", - dtypes=[ - (pto.i8, pto.f32, pto.f32, pto.f32, pto.f32), - (pto.i8, pto.f16, pto.f16, pto.f16, pto.f16), - (pto.i8, pto.i8, pto.i8, pto.i8, pto.i8), - ], - advanced=True -) -def template_tsel(mask: pto.Tile, src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - lanes = pto.get_lanes(dtype) - mask_row_stride = mask.shape[1] - mask_ptr = pto.castptr(mask.as_ptr(), pto.ptr(pto.ui8, pto.MemorySpace.UB)) - - if pto.constexpr(dtype == pto.f32): - full_mask_b16 = pto.pset_b16(pto.MaskPattern.ALL) - pair_width = lanes * 2 - paired_cols = (valid_cols // pair_width) * pair_width - for row in range(0, valid_rows, 1): - for col in range(0, paired_cols, pair_width): - mask_offset = row * mask_row_stride + col // 8 - select_mask_raw = pto.plds(mask_ptr, mask_offset, pto.PredicateDist.US) - select_mask = select_mask_raw.astype(pto.mask_b16) - pred0, _ = pto.make_mask(dtype, pair_width) - pred1, _ = pto.make_mask(dtype, lanes) - select_mask0, select_mask1 = pto.pintlv_b16(select_mask, full_mask_b16) - select_mask0 = select_mask0.astype(pto.mask_b32) - select_mask1 = select_mask1.astype(pto.mask_b32) - lhs0 = pto.vlds(src0[row, col:]) - rhs0 = pto.vlds(src1[row, col:]) - lhs1 = pto.vlds(src0[row, col + lanes:]) - rhs1 = pto.vlds(src1[row, col + lanes:]) - selected0 = pto.vsel(lhs0, rhs0, select_mask0) - selected1 = pto.vsel(lhs1, rhs1, select_mask1) - pto.vsts(selected0, dst[row, col:], pred0) - pto.vsts(selected1, dst[row, col + lanes:], pred1) - tail_cols = valid_cols - paired_cols - if tail_cols > 0: - col = paired_cols - mask_offset = row * mask_row_stride + col // 8 - select_mask_raw = pto.plds(mask_ptr, mask_offset, pto.PredicateDist.US) - select_mask = select_mask_raw.astype(pto.mask_b16) - select_mask0 = pto.punpack(select_mask, pto.PredicatePart.LOWER) - select_mask0 = select_mask0.astype(pto.mask_b32) - pred0, _ = pto.make_mask(dtype, tail_cols) - lhs0 = pto.vlds(src0[row, col:]) - rhs0 = pto.vlds(src1[row, col:]) - selected0 = pto.vsel(lhs0, rhs0, select_mask0) - pto.vsts(selected0, dst[row, col:], pred0) - else: - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - pred_mask, remained = pto.make_mask(dtype, remained) - mask_offset = row * mask_row_stride + col // 8 - if pto.constexpr(dtype == pto.f16): - select_mask_raw = pto.plds(mask_ptr, mask_offset, pto.PredicateDist.US) - select_mask = select_mask_raw.astype(pto.mask_b16) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - selected = pto.vsel(lhs, rhs, select_mask) - pto.vsts(selected, dst[row, col:], pred_mask) - else: - select_mask = pto.plds(mask_ptr, mask_offset, pto.PredicateDist.NORM) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - selected = pto.vsel(lhs, rhs, select_mask) - pto.vsts(selected, dst[row, col:], pred_mask) - return \ No newline at end of file diff --git a/lib/TileOps/tsels_template.py b/lib/TileOps/tsels_template.py deleted file mode 100644 index 905d0dd061..0000000000 --- a/lib/TileOps/tsels_template.py +++ /dev/null @@ -1,121 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tsels - -NOTE: This template uses pto.plds for mask loading which directly -loads predicate mask from UB without vcmps comparison. - -TSels: Select between source tile and scalar based on mask. -- mask=true: select from src -- mask=false: select scalar value - -Mask tile format: -- Packed predicate bytes in UB. -- Each row stores ceil(valid_cols / 8) valid bytes; tile row stride may be padded. -- mask_dtype determines the storage format (i8/i16/i32), but the actual - predicate bits are packed and accessed as bytes. - -IMPORTANT: mask_row_stride is always mask.shape[1] (element count), -because mask tile stride equals cols in element units regardless of mask_dtype. -Byte offset for plds is col // 8 (one byte covers 8 elements). -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tsels", - dtypes=[ - (pto.i8, pto.i8, pto.i8, pto.i8, pto.i8), - (pto.i16, pto.i8, pto.i8, pto.i8, pto.i8), - (pto.i32, pto.i8, pto.i8, pto.i8, pto.i8), - (pto.i8, pto.i16, pto.i16, pto.i16, pto.i16), - (pto.i16, pto.i16, pto.i16, pto.i16, pto.i16), - (pto.i32, pto.i16, pto.i16, pto.i16, pto.i16), - (pto.i8, pto.i32, pto.i32, pto.i32, pto.i32), - (pto.i16, pto.i32, pto.i32, pto.i32, pto.i32), - (pto.i32, pto.i32, pto.i32, pto.i32, pto.i32), - (pto.i8, pto.f16, pto.f16, pto.f16, pto.f16), - (pto.i16, pto.f16, pto.f16, pto.f16, pto.f16), - (pto.i32, pto.f16, pto.f16, pto.f16, pto.f16), - (pto.i8, pto.f32, pto.f32, pto.f32, pto.f32), - (pto.i16, pto.f32, pto.f32, pto.f32, pto.f32), - (pto.i32, pto.f32, pto.f32, pto.f32, pto.f32), - ], - advanced=True -) -def template_tsels(mask: pto.Tile, src: pto.Tile, tmp: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - mask_dtype = mask.element_type - - lanes = pto.get_lanes(dtype) - mask_row_stride = mask.shape[1] * pto.bytewidth(mask_dtype) - mask_ptr = pto.castptr(mask.as_ptr(), pto.ptr(pto.ui8, pto.MemorySpace.UB)) - - scalar_mask, _ = pto.make_mask(dtype, lanes) - vreg_scalar = pto.vdup(scalar, scalar_mask) - - if pto.constexpr(lanes == 64): - full_mask_b16 = pto.pset_b16(pto.MaskPattern.ALL) - pair_width = lanes * 2 - paired_cols = (valid_cols // pair_width) * pair_width - for row in range(0, valid_rows, 1): - for col in range(0, paired_cols, pair_width): - mask_offset = row * mask_row_stride + col // 8 - select_mask_raw = pto.plds(mask_ptr, mask_offset, pto.PredicateDist.US) - select_mask = select_mask_raw.astype(pto.mask_b16) - pred0, _ = pto.make_mask(dtype, pair_width) - pred1, _ = pto.make_mask(dtype, lanes) - select_mask0, select_mask1 = pto.pintlv_b16(select_mask, full_mask_b16) - select_mask0 = select_mask0.astype(pto.mask_b32) - select_mask1 = select_mask1.astype(pto.mask_b32) - src0 = pto.vlds(src[row, col:]) - src1 = pto.vlds(src[row, col + lanes:]) - selected0 = pto.vsel(src0, vreg_scalar, select_mask0) - selected1 = pto.vsel(src1, vreg_scalar, select_mask1) - pto.vsts(selected0, dst[row, col:], pred0) - pto.vsts(selected1, dst[row, col + lanes:], pred1) - tail_cols = valid_cols - paired_cols - if tail_cols > 0: - col = paired_cols - mask_offset = row * mask_row_stride + col // 8 - select_mask_raw = pto.plds(mask_ptr, mask_offset, pto.PredicateDist.US) - select_mask = select_mask_raw.astype(pto.mask_b16) - select_mask0 = pto.punpack(select_mask, pto.PredicatePart.LOWER) - select_mask0 = select_mask0.astype(pto.mask_b32) - pred0, _ = pto.make_mask(dtype, tail_cols) - src0 = pto.vlds(src[row, col:]) - selected0 = pto.vsel(src0, vreg_scalar, select_mask0) - pto.vsts(selected0, dst[row, col:], pred0) - elif pto.constexpr(lanes == 128): - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - pred_mask, remained = pto.make_mask(dtype, remained) - mask_offset = row * mask_row_stride + col // 8 - select_mask_raw = pto.plds(mask_ptr, mask_offset, pto.PredicateDist.US) - select_mask = select_mask_raw.astype(pto.mask_b16) - src_vec = pto.vlds(src[row, col:]) - selected = pto.vsel(src_vec, vreg_scalar, select_mask) - pto.vsts(selected, dst[row, col:], pred_mask) - else: - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, lanes): - pred_mask, remained = pto.make_mask(dtype, remained) - mask_offset = row * mask_row_stride + col // 8 - select_mask = pto.plds(mask_ptr, mask_offset, pto.PredicateDist.NORM) - src_vec = pto.vlds(src[row, col:]) - selected = pto.vsel(src_vec, vreg_scalar, select_mask) - pto.vsts(selected, dst[row, col:], pred_mask) - return \ No newline at end of file diff --git a/lib/TileOps/tshl_template.py b/lib/TileOps/tshl_template.py deleted file mode 100644 index d236c8940d..0000000000 --- a/lib/TileOps/tshl_template.py +++ /dev/null @@ -1,32 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tshl""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tshl" -) -def template_tshl(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - result = pto.vshl(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tshls_template.py b/lib/TileOps/tshls_template.py deleted file mode 100644 index def0b03537..0000000000 --- a/lib/TileOps/tshls_template.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tshls""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tshls", -) -def template_tshls(src: pto.Tile, scalar: pto.i16, dst: pto.Tile): - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - result = pto.vshls(vec, scalar, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tshr_template.py b/lib/TileOps/tshr_template.py deleted file mode 100644 index f16ba9abe2..0000000000 --- a/lib/TileOps/tshr_template.py +++ /dev/null @@ -1,32 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tshr""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tshr" -) -def template_tshr(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - result = pto.vshr(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tshrs_template.py b/lib/TileOps/tshrs_template.py deleted file mode 100644 index 8366a638bd..0000000000 --- a/lib/TileOps/tshrs_template.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tshrs""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tshrs", -) -def template_tshrs(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - result = pto.vshrs(vec, scalar, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tsort32_template.py b/lib/TileOps/tsort32_template.py deleted file mode 100644 index aaf5b9f065..0000000000 --- a/lib/TileOps/tsort32_template.py +++ /dev/null @@ -1,241 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tsort32""" - -import tilelang_dsl as pto - -BLOCK_SIZE = 32 -FLOAT_DST_STRIDE_COEF = 2 -HALF_DST_STRIDE_COEF = 4 -MAX_UB_TMP = 32 * 255 # 8160 bytes -REPEAT_MAX = 255 - - -def _constraint_aligned( - src_shape=(), - src_valid_shape=(), - idx_shape=(), - idx_valid_shape=(), - dst_shape=(), - dst_valid_shape=(), - src_config=None, - idx_config=None, - dst_config=None, -) -> bool: - """Constraint for Format1: valid_cols % 32 == 0 (aligned, no tmp needed).""" - if len(src_valid_shape) != 2: - return False - valid_cols = src_valid_shape[1] - return valid_cols % BLOCK_SIZE == 0 - - -def _constraint_unaligned( - src_shape=(), - src_valid_shape=(), - idx_shape=(), - idx_valid_shape=(), - tmp_shape=(), - tmp_valid_shape=(), - dst_shape=(), - dst_valid_shape=(), - src_config=None, - idx_config=None, - tmp_config=None, - dst_config=None, -) -> bool: - """Constraint for Format2: valid_cols % 32 != 0 (unaligned, tmp needed).""" - if len(src_valid_shape) != 2: - return False - valid_cols = src_valid_shape[1] - return valid_cols % BLOCK_SIZE != 0 - - -@pto.vkernel( - target="a5", - advanced=True, - op="pto.tsort32", - constraints=[_constraint_aligned] -) -def template_tsort32(src: pto.Tile, idx: pto.Tile, dst: pto.Tile): - """ - TSort32 Format1: Bitonic sort for aligned cols (valid_cols % 32 == 0). - - Semantics (matching pto-isa TSort32.hpp Format1): - - Sorts src values into dst, generating indices in idx - - Direct sort without tmp buffer when src.valid_cols % 32 == 0 - - No padding needed - """ - dtype = dst.element_type - valid_rows = dst.valid_shape[0] - valid_cols = src.valid_shape[1] - - dst_ptr = dst.as_ptr() - src_ptr = src.as_ptr() - idx_ptr = idx.as_ptr() - - elem_bytes = pto.bytewidth(dtype) - dst_stride = ((dst.shape[1] * elem_bytes + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE) // elem_bytes - src_stride = ((src.shape[1] * elem_bytes + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE) // elem_bytes - idx_stride = ((idx.shape[1] * 4 + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE) // 4 - if idx.valid_shape[0] == 1: - idx_stride = 0 - - type_coef = HALF_DST_STRIDE_COEF - if pto.constexpr(dtype == pto.f32): - type_coef = FLOAT_DST_STRIDE_COEF - - repeat_num_per_row = (valid_cols + BLOCK_SIZE - 1) // BLOCK_SIZE - - if repeat_num_per_row <= REPEAT_MAX: - for i in range(0, valid_rows, 1): - pto.vbitsort( - pto.addptr(dst_ptr, i * dst_stride), - pto.addptr(src_ptr, i * src_stride), - pto.addptr(idx_ptr, i * idx_stride), - repeat_num_per_row - ) - else: - loop_num = (repeat_num_per_row + REPEAT_MAX - 1) // REPEAT_MAX - tail_repeat_num = repeat_num_per_row % REPEAT_MAX - for i in range(0, valid_rows, 1): - for j in range(0, loop_num, 1): - repeat_num = REPEAT_MAX - if j == loop_num - 1: - repeat_num = tail_repeat_num - - pto.vbitsort( - pto.addptr(dst_ptr, i * dst_stride + j * REPEAT_MAX * BLOCK_SIZE * type_coef), - pto.addptr(src_ptr, i * src_stride + j * REPEAT_MAX * BLOCK_SIZE), - pto.addptr(idx_ptr, i * idx_stride + j * REPEAT_MAX * BLOCK_SIZE), - repeat_num - ) - return - - -@pto.vkernel( - target="a5", - advanced=True, - op="pto.tsort32", - constraints=[_constraint_unaligned] -) -def template_tsort32_with_tmp(src: pto.Tile, idx: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - """ - TSort32 Format2: Bitonic sort with tmp buffer for unaligned cols. - - Semantics (matching pto-isa TSort32.hpp Format2): - - Sorts src values into dst, generating indices in idx - - Uses tmp buffer when src.valid_cols % 32 != 0 (padding needed) - - Pads unaligned tail with NaN to ensure correct sorting - """ - dtype = dst.element_type - valid_rows = dst.valid_shape[0] - valid_cols = src.valid_shape[1] - - dst_ptr = dst.as_ptr() - src_ptr = src.as_ptr() - idx_ptr = idx.as_ptr() - tmp_ptr = tmp.as_ptr() - - elem_bytes = pto.bytewidth(dtype) - dst_stride = ((dst.shape[1] * elem_bytes + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE) // elem_bytes - src_stride = ((src.shape[1] * elem_bytes + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE) // elem_bytes - idx_stride = ((idx.shape[1] * 4 + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE) // 4 - if idx.valid_shape[0] == 1: - idx_stride = 0 - - type_coef = HALF_DST_STRIDE_COEF - if pto.constexpr(dtype == pto.f32): - type_coef = FLOAT_DST_STRIDE_COEF - - repeat_num_per_row = (valid_cols + BLOCK_SIZE - 1) // BLOCK_SIZE - src_tail_per_row = valid_cols % BLOCK_SIZE - src_tail_repeat_num = ((valid_cols + BLOCK_SIZE - 1) // BLOCK_SIZE) % REPEAT_MAX - - if pto.constexpr(dtype == pto.f16): - min_val = pto.f16(0xFC00) - elif pto.constexpr(dtype == pto.bf16): - min_val = pto.bf16(0xFF80) - else: - min_val = pto.f32(0xFF800000) - - src_shape_bytes_per_row = valid_cols * elem_bytes - - if src_shape_bytes_per_row <= MAX_UB_TMP: - len_burst = (src_shape_bytes_per_row + BLOCK_SIZE - 1) // BLOCK_SIZE - - for i in range(0, valid_rows, 1): - pto.copy_ubuf_to_ubuf( - pto.addptr(src_ptr, i * src_stride), - tmp_ptr, - 0, 1, len_burst, 0, 0 - ) - - tmp_last_offset = ((valid_cols + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE) - BLOCK_SIZE - - vec = pto.vlds(tmp[0, tmp_last_offset:]) - pad_mask, _ = pto.make_mask(dtype, BLOCK_SIZE - src_tail_per_row) - vec = pto.vdup(min_val, pad_mask) - pto.vsts(vec, tmp[0, tmp_last_offset:], pad_mask) - - pto.vbitsort( - pto.addptr(dst_ptr, i * dst_stride), - tmp_ptr, - pto.addptr(idx_ptr, i * idx_stride), - repeat_num_per_row - ) - else: - loop_num = (repeat_num_per_row + REPEAT_MAX - 1) // REPEAT_MAX - - for i in range(0, valid_rows, 1): - for j in range(0, loop_num, 1): - if j < loop_num - 1: - pto.vbitsort( - pto.addptr(dst_ptr, i * dst_stride + j * REPEAT_MAX * BLOCK_SIZE * type_coef), - pto.addptr(src_ptr, i * src_stride + j * REPEAT_MAX * BLOCK_SIZE), - pto.addptr(idx_ptr, i * idx_stride + j * REPEAT_MAX * BLOCK_SIZE), - REPEAT_MAX - ) - else: - if src_tail_repeat_num > 0: - sort_repeat_num = 0 - if src_tail_repeat_num > 1: - sort_repeat_num = src_tail_repeat_num - 1 - - pto.vbitsort( - pto.addptr(dst_ptr, i * dst_stride + j * REPEAT_MAX * BLOCK_SIZE * type_coef), - pto.addptr(src_ptr, i * src_stride + j * REPEAT_MAX * BLOCK_SIZE), - pto.addptr(idx_ptr, i * idx_stride + j * REPEAT_MAX * BLOCK_SIZE), - sort_repeat_num - ) - - tail_src_offset = (j * REPEAT_MAX + (src_tail_repeat_num - 1)) * BLOCK_SIZE - tail_dst_offset = (j * REPEAT_MAX + (src_tail_repeat_num - 1)) * BLOCK_SIZE * type_coef - len_burst = (src_tail_per_row * elem_bytes + BLOCK_SIZE - 1) // BLOCK_SIZE - - pto.copy_ubuf_to_ubuf( - pto.addptr(src_ptr, i * src_stride + tail_src_offset), - tmp_ptr, - 0, 1, len_burst, 0, 0 - ) - - tmp_last_offset = ((src_tail_per_row + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE) - BLOCK_SIZE - - vec = pto.vlds(tmp[0, tmp_last_offset:]) - pad_mask, _ = pto.make_mask(dtype, BLOCK_SIZE - src_tail_per_row) - vec = pto.vdup(min_val, pad_mask) - pto.vsts(vec, tmp[0, tmp_last_offset:], pad_mask) - - pto.vbitsort( - pto.addptr(dst_ptr, i * dst_stride + tail_dst_offset), - tmp_ptr, - pto.addptr(idx_ptr, i * idx_stride + tail_src_offset), - 1 - ) - - return \ No newline at end of file diff --git a/lib/TileOps/tsqrt_template.py b/lib/TileOps/tsqrt_template.py deleted file mode 100644 index 592e8c62c8..0000000000 --- a/lib/TileOps/tsqrt_template.py +++ /dev/null @@ -1,33 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tsqrt""" - -import tilelang_dsl as pto -from sqrt_hp import _tl_sqrt_precision - -@pto.vkernel( - target="a5", - op="pto.tsqrt" -) -def template_tsqrt(src: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - precision_type = pto.get_op_attr("precisionType", "default") - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vinput = pto.vlds(src[row, col:]) - if pto.constexpr(precision_type == "high_precision"): - result = _tl_sqrt_precision(vinput, mask, dtype) - else: - result = pto.vsqrt(vinput, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/tstore_template.py b/lib/TileOps/tstore_template.py deleted file mode 100644 index ef80400215..0000000000 --- a/lib/TileOps/tstore_template.py +++ /dev/null @@ -1,662 +0,0 @@ -# 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. - -"""TileLang DSL template for `pto.tstore`""" - -import tilelang_dsl as pto - - -def _constraint_scalar(value): - return value.value if hasattr(value, "value") else value - - -def _known_eq(lhs, rhs) -> bool: - lhs_value = _constraint_scalar(lhs) - rhs_value = _constraint_scalar(rhs) - if lhs_value is None or rhs_value is None: - return True - return lhs_value == rhs_value - - -def _known_le(lhs, rhs) -> bool: - lhs_value = _constraint_scalar(lhs) - rhs_value = _constraint_scalar(rhs) - if lhs_value is None or rhs_value is None: - return True - return lhs_value <= rhs_value - - -def _match_store_tile_layout(src, *, row_major: bool, s_layout) -> bool: - b_layout_ok = ( - src.config.b_layout == pto.BLayout.ROW_MAJOR - if row_major - else src.config.b_layout != pto.BLayout.ROW_MAJOR - ) - return b_layout_ok and src.config.s_layout == s_layout - - -def _check_store_bounds(src, dst, *, logical_rows, logical_cols, stride_axis=None) -> bool: - if dst.rank != 5: - return False - if stride_axis is not None and not _known_eq(dst.strides[stride_axis], 1): - return False - if not _known_eq(src.valid_shape[0], logical_rows): - return False - if not _known_eq(src.valid_shape[1], logical_cols): - return False - if not _known_le(src.valid_shape[0], src.shape[0]): - return False - if not _known_le(src.valid_shape[1], src.shape[1]): - return False - return True - - -def _tstore_preconditions_nd(src, dst) -> bool: - logical_rows = dst.shape[0] * dst.shape[1] * dst.shape[2] * dst.shape[3] - logical_cols = dst.shape[4] - return _match_store_tile_layout( - src, row_major=True, s_layout=pto.SLayout.NONE_BOX - ) and _check_store_bounds( - src, dst, logical_rows=logical_rows, logical_cols=logical_cols, stride_axis=4 - ) - -def _tstore_preconditions_dn(src, dst) -> bool: - logical_rows = dst.shape[3] - logical_cols = dst.shape[0] * dst.shape[1] * dst.shape[2] * dst.shape[4] - return _match_store_tile_layout( - src, row_major=False, s_layout=pto.SLayout.NONE_BOX - ) and _check_store_bounds( - src, dst, logical_rows=logical_rows, logical_cols=logical_cols, stride_axis=3 - ) - -def _tstore_preconditions_nz(src, dst) -> bool: - logical_rows = dst.shape[2] * dst.shape[3] - logical_cols = dst.shape[0] * dst.shape[1] * dst.shape[4] - return _match_store_tile_layout( - src, row_major=False, s_layout=pto.SLayout.ROW_MAJOR - ) and _check_store_bounds( - src, dst, logical_rows=logical_rows, logical_cols=logical_cols - ) - -@pto.vkernel( - target="a5", - op="pto.tstore", - advanced=True, - constraints=[_tstore_preconditions_nd], -) -def template_tstore_nd(src: pto.Tile, dst: pto.PartitionTensorView): - dtype = src.element_type - elem_bytes = pto.bytewidth(dtype) - - g0, g1, g2, g3, g4 = dst.shape - s0, s1, s2, s3, s4 = dst.strides - - valid_rows, valid_cols = src.valid_shape - ub_rows, ub_cols = src.shape - - # These preconditions are expressed through the descriptor-level constraint - # callable above, using direct `src.*` / `dst.*` metadata syntax. - - n_burst = g3 - len_burst = valid_cols * elem_bytes - ub_stride = ub_cols * elem_bytes - gm_stride = s3 * elem_bytes - - src_stride2 = g3 * ub_cols - src_stride1 = g2 * src_stride2 - src_stride0 = g1 * src_stride1 - - loop1 = g2 - loop2 = g1 - loop1_src_stride = src_stride2 * elem_bytes - loop1_dst_stride = s2 * elem_bytes - loop2_src_stride = src_stride1 * elem_bytes - loop2_dst_stride = s1 * elem_bytes - - ub_ptr = src.as_ptr() - gm_ptr = dst.as_ptr() - - if loop1 != 1 or loop2 != 1: - pto.set_loop2_stride_ubtoout( - src_stride=loop2_src_stride, dst_stride=loop2_dst_stride - ) - pto.set_loop1_stride_ubtoout( - src_stride=loop1_src_stride, dst_stride=loop1_dst_stride - ) - pto.set_loop_size_ubtoout(loop1=loop1, loop2=loop2) - - for i in range(0, g0, 1): - src_i = pto.addptr(ub_ptr, i * src_stride0) - dst_i = pto.addptr(gm_ptr, i * s0) - pto.copy_ubuf_to_gm( - dst=dst_i, - src=src_i, - n_burst=n_burst, - len_burst=len_burst, - gm_stride=gm_stride, - ub_stride=ub_stride, - ) - - if loop1 != 1 or loop2 != 1: - pto.set_loop_size_ubtoout(loop1=1, loop2=1) - return - -@pto.vkernel( - target="a5", - op="pto.tstore", - advanced=True, - constraints=[_tstore_preconditions_dn], -) -def template_tstore_dn(src: pto.Tile, dst: pto.PartitionTensorView): - dtype = src.element_type - elem_bytes = pto.bytewidth(dtype) - - g0, g1, g2, g3, g4 = dst.shape - s0, s1, s2, s3, s4 = dst.strides - - valid_rows, valid_cols = src.valid_shape - ub_rows, ub_cols = src.shape - - n_burst = g4 - len_burst = valid_rows * elem_bytes - gm_stride = s4 * elem_bytes - ub_stride = ub_rows * elem_bytes - - # The UB source tile has a compact col-major layout with column height - # `ub_rows`. Like `TStoreVecDN`, three levels of stride are derived from - # `g4` / `g2` / `g1`. - src_stride2 = ub_rows * g4 - src_stride1 = g2 * src_stride2 - src_stride0 = g1 * src_stride1 - - loop1 = g2 - loop2 = g1 - loop1_src_stride = src_stride2 * elem_bytes - loop1_dst_stride = s2 * elem_bytes - loop2_src_stride = src_stride1 * elem_bytes - loop2_dst_stride = s1 * elem_bytes - - ub_ptr = src.as_ptr() - gm_ptr = dst.as_ptr() - - if loop1 != 1 or loop2 != 1: - pto.set_loop2_stride_ubtoout( - src_stride=loop2_src_stride, dst_stride=loop2_dst_stride - ) - pto.set_loop1_stride_ubtoout( - src_stride=loop1_src_stride, dst_stride=loop1_dst_stride - ) - pto.set_loop_size_ubtoout(loop1=loop1, loop2=loop2) - - for i in range(0, g0, 1): - src_i = pto.addptr(ub_ptr, i * src_stride0) - dst_i = pto.addptr(gm_ptr, i * s0) - pto.copy_ubuf_to_gm( - dst=dst_i, - src=src_i, - n_burst=n_burst, - len_burst=len_burst, - gm_stride=gm_stride, - ub_stride=ub_stride, - ) - - if loop1 != 1 or loop2 != 1: - pto.set_loop_size_ubtoout(loop1=1, loop2=1) - return - -@pto.vkernel( - target="a5", - op="pto.tstore", - advanced=True, - constraints=[_tstore_preconditions_nz], -) -def template_tstore_nz(src: pto.Tile, dst: pto.PartitionTensorView): - dtype = src.element_type - elem_bytes = pto.bytewidth(dtype) - - g0, g1, g2, g3, g4 = dst.shape - s0, s1, s2, s3, s4 = dst.strides - - valid_rows, valid_cols = src.valid_shape - ub_rows, ub_cols = src.shape - - # Corresponds to C++ `C0_SIZE_BYTE`. Each NZ burst always writes one - # complete C0 block. - c0_size_bytes = 32 - n_burst = g1 - len_burst = valid_rows * c0_size_bytes - gm_stride = s1 * elem_bytes - ub_stride = ub_rows * c0_size_bytes - - # Each g0 block in UB is composed of `g1` NZ blocks concatenated together. - tile_stride = g1 * ub_rows * g4 - - ub_ptr = src.as_ptr() - gm_ptr = dst.as_ptr() - - # NZ path itself does not use loop1/loop2; explicitly reset to normal mode - # to avoid inheriting stale state. - pto.set_loop_size_ubtoout(loop1=1, loop2=1) - for i in range(0, g0, 1): - src_i = pto.addptr(ub_ptr, i * tile_stride) - dst_i = pto.addptr(gm_ptr, i * s0) - pto.copy_ubuf_to_gm( - dst=dst_i, - src=src_i, - n_burst=n_burst, - len_burst=len_burst, - gm_stride=gm_stride, - ub_stride=ub_stride, - ) - return - - -# ============================================================================ -# Cube Templates: TSTORE.ACC (ACC → GM) -# ============================================================================ - -def _constraint_tstore_acc_base(src, dst) -> bool: - """TSTORE.ACC base constraint check""" - # src must be MemorySpace.ACC - src_space = src.memory_space - if src_space is None: - return False - src_space_value = src_space.value if hasattr(src_space, "value") else src_space - if src_space_value not in {"acc", "ACC"}: - return False - # dst must be GM (via PartitionTensorView) - dst_space = dst.memory_space - if dst_space is None: - dst_space_value = "gm" # PartitionTensorView defaults to GM - else: - dst_space_value = dst_space.value if hasattr(dst_space, "value") else dst_space - if dst_space_value not in {"gm", "GM"}: - return False - # ACC dtype must be f32 or i32 - src_dtype = src.dtype - if src_dtype is None: - return False - dtype_name = src_dtype.name if hasattr(src_dtype, "name") else str(src_dtype) - if dtype_name not in {"f32", "i32"}: - return False - # dst dtype can be f32, f16, bf16, i32 - dst_dtype = dst.dtype - if dst_dtype is None: - return True # allow dst dtype to be unspecified - dst_dtype_name = dst_dtype.name if hasattr(dst_dtype, "name") else str(dst_dtype) - supported_dst_dtypes = {"f32", "f16", "bf16", "i32"} - if dst_dtype_name not in supported_dst_dtypes: - return False - return True - - -def _extract_b_layout_value(config) -> str | None: - """Extract the effective b_layout string from a config object. - - Handles both TileConfig (with b_layout/s_layout attributes) and - ViewConfig (with a single layout attribute like 'nd' or 'dn'). - Returns a canonical layout string: 'row_major' or 'col_major', or None. - """ - if config is None: - return None - # TileConfig has explicit b_layout attribute - if hasattr(config, "b_layout") and config.b_layout is not None: - bl = config.b_layout - return bl.value if hasattr(bl, "value") else bl - # ViewConfig uses a single 'layout' attribute: - # 'nd' -> row_major (ND/row-major format) - # 'dn' -> col_major (DN/col-major format) - # 'nz' -> col_major (NZ/fractal format, same block layout as col_major) - if hasattr(config, "layout") and config.layout is not None: - layout = config.layout - layout_str = layout.value if hasattr(layout, "value") else str(layout) - normalized = layout_str.strip().lower().replace("-", "_") - if normalized in {"nd", "row_major"}: - return "row_major" - if normalized in {"dn", "col_major"}: - return "col_major" - if normalized in {"nz"}: - return "col_major" # NZ fractal uses col-major block layout - return None - - -def _extract_s_layout_value(config) -> str | None: - """Extract the effective s_layout string from a config object. - - For TileConfig this is the s_layout attribute. - For ViewConfig there is no slayout distinction; returns None. - """ - if config is None: - return None - if hasattr(config, "s_layout") and config.s_layout is not None: - sl = config.s_layout - return sl.value if hasattr(sl, "value") else sl - return None - - -def _constraint_tstore_acc_nz2nd(src, dst) -> bool: - """TSTORE.ACC NZ2ND constraint""" - if not _constraint_tstore_acc_base(src, dst): - return False - # dst must be row-major layout (ND format) - b_layout_value = _extract_b_layout_value(dst.config) - if b_layout_value is None: - return True # default is row-major - # ROW_MAJOR corresponds to ND format - if b_layout_value not in {"row_major", "ROW_MAJOR"}: - return False - return True - - -def _constraint_tstore_acc_nz2dn(src, dst) -> bool: - """TSTORE.ACC NZ2DN constraint""" - if not _constraint_tstore_acc_base(src, dst): - return False - b_layout_value = _extract_b_layout_value(dst.config) - if b_layout_value is None: - return False - if b_layout_value not in {"col_major", "COL_MAJOR"}: - return False - return True - - -def _constraint_tstore_acc_nz2nz(src, dst) -> bool: - """TSTORE.ACC NZ2NZ constraint""" - if not _constraint_tstore_acc_base(src, dst): - return False - # dst must be NZ layout (fractal): b_layout=col_major + s_layout=row_major - b_layout_value = _extract_b_layout_value(dst.config) - s_layout_value = _extract_s_layout_value(dst.config) - if b_layout_value is None: - return False - if s_layout_value is None: - return False - if b_layout_value not in {"col_major", "COL_MAJOR"}: - return False - if s_layout_value not in {"row_major", "ROW_MAJOR"}: - return False - return True - - -@pto.ckernel( - target="a5", - op="pto.tstore", - priority=1, - dtypes=[ - (pto.f32, pto.f32), - (pto.f32, pto.f16), - (pto.f32, pto.bf16), - (pto.i32, pto.i32), - ], - constraints=[_constraint_tstore_acc_nz2nd], - name="tstore_acc_to_gm_nz2nd", -) -def template_tstore_acc_to_gm_nz2nd(src: pto.Tile, dst: pto.PartitionTensorView): - """ACC -> GM (NZ2ND mode) - - Write NZ-format data from L0C Accumulator Buffer back to GM in - Row-Major (ND) format. - - Args: - src: Tile with ACC memory_space, shape=(M, N), dtype=f32/i32 - dst: PartitionTensorView with GM memory_space, row-major (ND) format - - Uses: - pto.mte_l0c_gm with layout="nz2nd" - """ - m, n = src.valid_shape - - acc_ptr = src.as_ptr() - gm_ptr = dst.as_ptr() - - # rank-5 partition view metadata (ND: g3 = M rows, g4 = N cols). - g0, g1, g2, g3, g4 = dst.shape - s0, s1, s2, s3, s4 = dst.strides - - # Strides are in ELEMENT units (pto-isa a5/TStore.hpp: TStoreAccND passes - # srcStride/dstD to copy_matrix_cc_to_gm without GetByteSize; the hardware - # scales by dtype internally). The earlier `src_stride = dst_stride = n` - # filled both fields with the tile column count, which is wrong for any - # strided GM slice and produces sparse/shifted GM writes. - # - # src_stride = TileData::Rows = src.shape[0] (NZ row stride of L0C acc) - # dst_stride = gStride3 = s3 (GM physical row stride) - src_stride = src.shape[0] - dst_stride = s3 - - pto.mte_l0c_gm( - acc_ptr, gm_ptr, - m, n, src_stride, dst_stride, - 0, 0, - layout="nz2nd" - ) - return - - -@pto.ckernel( - target="a5", - op="pto.tstore", - priority=1, - dtypes=[ - (pto.f32, pto.f32), - (pto.f32, pto.f16), - (pto.f32, pto.bf16), - (pto.i32, pto.i32), - ], - constraints=[_constraint_tstore_acc_nz2dn], - name="tstore_acc_to_gm_nz2dn", -) -def template_tstore_acc_to_gm_nz2dn(src: pto.Tile, dst: pto.PartitionTensorView): - """ACC -> GM (NZ2DN mode) - - Write NZ-format data from L0C Accumulator Buffer back to GM in - Col-Major (DN) format. - - Args: - src: Tile with ACC memory_space, shape=(M, N) - dst: PartitionTensorView with GM memory_space, col-major (DN) format - - Uses: - pto.mte_l0c_gm with layout="nz2dn" - """ - m, n = src.valid_shape - - acc_ptr = src.as_ptr() - gm_ptr = dst.as_ptr() - - # rank-5 partition view metadata (DN: g3 = M, g4 = N, col-major). - g0, g1, g2, g3, g4 = dst.shape - s0, s1, s2, s3, s4 = dst.strides - - # Strides are in ELEMENT units (see template_tstore_acc_to_gm_nz2nd). - # DN (col-major): adjacent rows step along the M direction, whose physical - # stride is s4 (PTOAS 5D DN view strides collapse to [M, M, M, 1, M]). - # src_stride = TileData::Rows = src.shape[0] (NZ row stride of L0C acc) - # dst_stride = s4 (GM physical stride along M) - src_stride = src.shape[0] - dst_stride = s4 - loop0_src_stride = 1 # loop0_src_stride for NZ2DN (NZ source) - - pto.mte_l0c_gm( - acc_ptr, gm_ptr, - m, n, src_stride, dst_stride, - 0, 0, - layout=("nz2dn", loop0_src_stride) - ) - return - - -@pto.ckernel( - target="a5", - op="pto.tstore", - dtypes=[ - (pto.f32, pto.f32), - (pto.f32, pto.f16), - (pto.f32, pto.bf16), - (pto.i32, pto.i32), - ], - constraints=[_constraint_tstore_acc_nz2nz], - name="tstore_acc_to_gm_nz2nz", -) -def template_tstore_acc_to_gm_nz2nz(src: pto.Tile, dst: pto.PartitionTensorView): - """ACC -> GM (NZ2NZ mode) - - Write NZ-format data from L0C Accumulator Buffer back to GM in NZ - format (no layout conversion). - - Args: - src: Tile with ACC memory_space, shape=(M, N) - dst: PartitionTensorView with GM memory_space, NZ (fractal) format - - Uses: - pto.mte_l0c_gm with layout="nz2nz" - """ - m, n = src.valid_shape - - acc_ptr = src.as_ptr() - gm_ptr = dst.as_ptr() - - # src_stride is in ELEMENT units (pto-isa a5/TStore.hpp: TStoreAccNZ sets - # srcStride = TileData::Rows). The earlier `src_stride = n` filled the L0C - # source stride with the column count, which is wrong whenever N differs - # from the acc tile row count. - # src_stride = TileData::Rows = src.shape[0] (NZ row stride of L0C acc) - # NOTE: dst_stride for NZ2NZ is the NZ-fractal physical row stride - # (ceil(M, FRACTAL_NZ_ROW) * c0 in pto-isa TStoreAccNZ). The compact NZ - # views used today set this equal to n, so the legacy value is retained - # until a strided NZ destination case is exercised to validate the exact - # stride index. - src_stride = src.shape[0] - dst_stride = n - split = 1 # NZ2NZ requires a split parameter - - pto.mte_l0c_gm( - acc_ptr, gm_ptr, - m, n, src_stride, dst_stride, - 0, 0, - layout=("nz2nz", split) - ) - return - - -# ============================================================================ -# Cube Templates: TSTORE.MAT (MAT → GM) — NOT YET IMPLEMENTED -# ============================================================================ -# NOTE: There is no direct MAT → GM DMA path (mte_l1_gm does not exist). -# The correct implementation requires a two-step intermediate path: -# 1. Copy from MAT to UB via mte_l1_ub (allocating a temporary UB tile) -# 2. Copy from UB to GM via copy_ubuf_to_gm -# -# However, UB buffer allocation is not yet supported in ckernel mode. -# This template is intentionally NOT registered until a proper -# MAT → UB → GM implementation can be provided. -# Adding it with dtypes=[] would crash the module import (_freeze_dtypes -# requires at least one signature tuple), so we omit it entirely. - - -# ============================================================================ -# Cube Templates: TSTORE_FP (ACC + FP → GM) -# ============================================================================ - -def _constraint_tstore_fp(src, fp, dst) -> bool: - """TSTORE_FP constraint check""" - # src must be MemorySpace.ACC - src_space = src.memory_space - if src_space is None: - return False - src_space_value = src_space.value if hasattr(src_space, "value") else src_space - if src_space_value not in {"acc", "ACC"}: - return False - # fp must be SCALING memory space or specific buffer - fp_space = fp.memory_space - if fp_space is None: - return False - fp_space_value = fp_space.value if hasattr(fp_space, "value") else fp_space - if fp_space_value not in {"scaling", "SCALING", "ub", "UB"}: - return False - # dst must be GM - dst_space = dst.memory_space - if dst_space is None: - dst_space_value = "gm" - else: - dst_space_value = dst_space.value if hasattr(dst_space, "value") else dst_space - if dst_space_value not in {"gm", "GM"}: - return False - # src dtype must be f32 - src_dtype = src.dtype - if src_dtype is None: - return False - dtype_name = src_dtype.name if hasattr(src_dtype, "name") else str(src_dtype) - if dtype_name != "f32": - return False - return True - - -@pto.ckernel( - target="a5", - op="pto.tstore_fp", - dtypes=[ - (pto.f32, pto.f16, pto.f16), - (pto.f32, pto.bf16, pto.bf16), - ], - constraints=[_constraint_tstore_fp], - name="tstore_fp_acc_to_gm", -) -def template_tstore_fp_acc_to_gm(src: pto.Tile, fp: pto.Tile, dst: pto.PartitionTensorView): - """ACC + FP -> GM with floating-point conversion (TSTORE_FP) - - Write f32 data from L0C Accumulator Buffer, combined with FP (scaling) - parameters, back to GM in f16/bf16 format. - - Args: - src: Tile with ACC memory_space, dtype=f32 - fp: Tile with SCALING/UB memory_space, dtype=f16/bf16 - dst: PartitionTensorView with GM memory_space, dtype=f16/bf16 - - Note: - Uses qf322f16_pre_vec / qf322bf16_pre_vec mode for pre_quant, which - supports vector (per-channel) scaling rows per the hardware spec. - The f32_f16/f32_bf16 modes only accept scalar payload and cannot be - used for per-channel quantization via a scaling buffer. - """ - m, n = src.valid_shape - - acc_ptr = src.as_ptr() - fp_ptr = fp.as_ptr() - gm_ptr = dst.as_ptr() - - # Determine pre_quant mode based on fp (scaling) buffer dtype: - # Use qf322*_pre_vec modes which support vector scale rows (scaling buffer). - # f32_f16/f32_bf16 modes only accept scalar payload and cannot be used here. - # Note: dst.dtype is not supported by the TileLang DSL v1 frontend for - # PartitionTensorView, so we use fp.element_type instead. - fp_dtype = fp.element_type - - if pto.constexpr(fp_dtype == pto.bf16): - quant_mode = "qf322bf16_pre_vec" - else: - quant_mode = "qf322f16_pre_vec" - - # TODO: Replace with tstore_fp DSL surface once available. - # Currently using mte_l0c_gm + pre_quant as a temporary workaround. - # Strides are in ELEMENT units (see template_tstore_acc_to_gm_nz2nd). - # src_stride = TileData::Rows = src.shape[0] (NZ row stride of L0C acc) - # dst_stride = gStride3 = s3 (GM physical row stride, ND) - g0, g1, g2, g3, g4 = dst.shape - s0, s1, s2, s3, s4 = dst.strides - src_stride = src.shape[0] - dst_stride = s3 - - pto.mte_l0c_gm( - acc_ptr, gm_ptr, - m, n, src_stride, dst_stride, - 0, 0, - layout="nz2nd", - pre_quant=(fp_ptr, quant_mode) - ) - return diff --git a/lib/TileOps/tsub_template.py b/lib/TileOps/tsub_template.py deleted file mode 100644 index 81d1b13ddd..0000000000 --- a/lib/TileOps/tsub_template.py +++ /dev/null @@ -1,32 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tsub""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tsub" -) -def template_tsub(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - subtracted = pto.vsub(lhs, rhs, mask) - pto.vsts(subtracted, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/tsubs_template.py b/lib/TileOps/tsubs_template.py deleted file mode 100644 index 84dc8bfbd0..0000000000 --- a/lib/TileOps/tsubs_template.py +++ /dev/null @@ -1,38 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.tsubs - -Note: A5 hardware implements tsubs as vadds with negated scalar: - dst = src - scalar = src + (-scalar) -This template uses vbr + vsub to achieve element-wise subtraction. -TODO: Use vadds(vec, -scalar) when DSL supports unary negation on scalars. -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.tsubs", -) -def template_tsubs(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - scalar_vec = pto.vbr(scalar) - result = pto.vsub(vec, scalar_vec, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/lib/TileOps/txor_template.py b/lib/TileOps/txor_template.py deleted file mode 100644 index d2ca4f1f73..0000000000 --- a/lib/TileOps/txor_template.py +++ /dev/null @@ -1,32 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.txor""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.txor" -) -def template_txor(src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, dst: pto.Tile): - dtype = dst.element_type - valid_rows, valid_cols = dst.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - lhs = pto.vlds(src0[row, col:]) - rhs = pto.vlds(src1[row, col:]) - result = pto.vxor(lhs, rhs, mask) - pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file diff --git a/lib/TileOps/txors_template.py b/lib/TileOps/txors_template.py deleted file mode 100644 index c060209500..0000000000 --- a/lib/TileOps/txors_template.py +++ /dev/null @@ -1,39 +0,0 @@ -# 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. - -"""TileLang DSL template for pto.txors - -Note: A5 hardware implements txors as: - TEXPANDS_IMPL(dst, scalar); // broadcast scalar to dst - TXOR_IMPL(dst, src, dst, tmp); // dst = src ^ dst - -This template uses vbr + vxor to achieve element-wise bitwise XOR. -""" - -import sys -from pathlib import Path -import tilelang_dsl as pto - - -@pto.vkernel( - target="a5", - op="pto.txors", -) -def template_txors(src: pto.Tile, scalar: pto.AnyType, tmp: pto.Tile, dst: pto.Tile): - dtype = src.element_type - valid_rows, valid_cols = src.valid_shape - - for row in range(0, valid_rows, 1): - remained = valid_cols - for col in range(0, valid_cols, pto.get_lanes(dtype)): - mask, remained = pto.make_mask(dtype, remained) - vec = pto.vlds(src[row, col:]) - scalar_vec = pto.vbr(scalar) - result = pto.vxor(vec, scalar_vec, mask) - pto.vsts(result, dst[row, col:], mask) - return diff --git a/ptodsl/CMakeLists.txt b/ptodsl/CMakeLists.txt index 4123b99996..cb4f27b978 100644 --- a/ptodsl/CMakeLists.txt +++ b/ptodsl/CMakeLists.txt @@ -16,16 +16,21 @@ set(PTOAS_PY_PACKAGE_SRC_DIR "${CMAKE_CURRENT_SOURCE_DIR}/ptoas") set(PTOAS_WHEEL_BOOTSTRAP_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/ptoas_wheel_bootstrap.py") +set(TILEOPS_PACKAGE_SRC_DIR + "${CMAKE_SOURCE_DIR}/lib/TileOps") set(PTODSL_BUILD_ROOT "${CMAKE_BINARY_DIR}/python") set(PTODSL_BUILD_PACKAGE_DIR "${PTODSL_BUILD_ROOT}/ptodsl") set(PTOAS_PY_BUILD_PACKAGE_DIR "${PTODSL_BUILD_ROOT}/ptoas") +set(TILEOPS_BUILD_PACKAGE_DIR + "${PTODSL_BUILD_ROOT}/TileOps") add_custom_target(PTODSLPackage ALL COMMAND ${CMAKE_COMMAND} -E make_directory "${PTODSL_BUILD_ROOT}" COMMAND ${CMAKE_COMMAND} -E remove_directory "${PTODSL_BUILD_PACKAGE_DIR}" COMMAND ${CMAKE_COMMAND} -E remove_directory "${PTOAS_PY_BUILD_PACKAGE_DIR}" + COMMAND ${CMAKE_COMMAND} -E remove_directory "${TILEOPS_BUILD_PACKAGE_DIR}" COMMAND ${CMAKE_COMMAND} -E copy_directory "${PTODSL_PACKAGE_SRC_DIR}" "${PTODSL_BUILD_PACKAGE_DIR}" @@ -35,7 +40,10 @@ add_custom_target(PTODSLPackage ALL COMMAND ${CMAKE_COMMAND} -E copy "${PTOAS_WHEEL_BOOTSTRAP_SOURCE}" "${PTODSL_BUILD_ROOT}/ptoas_wheel_bootstrap.py" - COMMENT "Staging ptodsl package into build/python" + COMMAND ${CMAKE_COMMAND} -E copy_directory + "${TILEOPS_PACKAGE_SRC_DIR}" + "${TILEOPS_BUILD_PACKAGE_DIR}" + COMMENT "Staging ptodsl, ptoas, and TileOps packages into build/python" VERBATIM ) @@ -60,3 +68,11 @@ install( DESTINATION "." COMPONENT PTOAS_Runtime ) + +install( + DIRECTORY "${TILEOPS_PACKAGE_SRC_DIR}" + DESTINATION "." + COMPONENT PTOAS_Runtime + PATTERN "__pycache__" EXCLUDE + PATTERN "*.pyc" EXCLUDE +) diff --git a/ptodsl/README.md b/ptodsl/README.md index db05638333..d0d499e76e 100644 --- a/ptodsl/README.md +++ b/ptodsl/README.md @@ -93,9 +93,7 @@ The source-tree build bakes in `$PTOAS_REPO_ROOT/ptodsl` as the package root. The `PTODSL_PYTHON_ROOT` environment variable from `scripts/ptoas_env.sh` overrides that default. Use `--ptodsl-pkg-path=/path/to/package/root` for an explicit command-line override. PTODSL daemon failures are reported as errors -and never fall back to the TileLang implementation. Use -`--tile-lib-backend=tilelang` only when explicitly comparing against the legacy -TileLangDSL template library. +and never fall back to the TileLang implementation. `InsertTemplateAttributes` queries legal-candidate metadata before fusion and stores an ordered `candidates` array containing only `id`, `name`, diff --git a/ptodsl/docs/developer_guide/tilelib-template-authoring.md b/ptodsl/docs/developer_guide/tilelib-template-authoring.md index a415214e62..df05cefd6c 100644 --- a/ptodsl/docs/developer_guide/tilelib-template-authoring.md +++ b/ptodsl/docs/developer_guide/tilelib-template-authoring.md @@ -1,7 +1,7 @@ # PTODSL TileLib Template Authoring Guide This guide is for developers adding or changing PTODSL TileLib templates under -`ptodsl/ptodsl/tilelib/templates`. +`lib/TileOps`. For the compiler-side design, see `docs/designs/ptodsl-tilelib-template-selection-design.md`. diff --git a/ptodsl/ptodsl/tilelib/registry.py b/ptodsl/ptodsl/tilelib/registry.py index d1fbf569c6..c437a0a8ed 100644 --- a/ptodsl/ptodsl/tilelib/registry.py +++ b/ptodsl/ptodsl/tilelib/registry.py @@ -127,7 +127,7 @@ def _load_default_templates(op: str, target: str) -> None: # Import lazily to avoid a registry/templates import cycle during package # initialization. The loader is cached and registers descriptors as a # module-import side effect. - from .templates import load_template + from TileOps import load_template load_template(op, target) diff --git a/ptodsl/ptodsl/tilelib/render.py b/ptodsl/ptodsl/tilelib/render.py index 6a11ec2650..97354fefef 100644 --- a/ptodsl/ptodsl/tilelib/render.py +++ b/ptodsl/ptodsl/tilelib/render.py @@ -10,7 +10,7 @@ ``render_best`` is the seam ``ExpandTileOp`` will call: select the legal/best template for an op+target+operand specialization, then render it via ptodsl's engine. -CLI (standalone, parallels lib/TileOps/render_template_mlir.py): +CLI (standalone): python3 -m ptodsl.tilelib.render --op pto.tadd --target a5 \\ --tile dst=8x64@ub:f32 --tile src0=8x64@ub:f32 --tile src1=8x64@ub:f32 \\ diff --git a/ptodsl/ptodsl/tilelib/serving/daemon.py b/ptodsl/ptodsl/tilelib/serving/daemon.py index 3177075adc..f52c4c6dcf 100644 --- a/ptodsl/ptodsl/tilelib/serving/daemon.py +++ b/ptodsl/ptodsl/tilelib/serving/daemon.py @@ -28,7 +28,7 @@ from .. import constraints as _constraints from .. import registry as _registry from ..metadata import ScalarSpec, ScalarType, TileSpec, VectorSpec, ViewSpec -from ..templates import load_template +from TileOps import load_template from .wire import recv_message, send_message diff --git a/ptodsl/ptodsl/tilelib/templates/__init__.py b/ptodsl/ptodsl/tilelib/templates/__init__.py index eee2fa0ae0..6ed28847c2 100644 --- a/ptodsl/ptodsl/tilelib/templates/__init__.py +++ b/ptodsl/ptodsl/tilelib/templates/__init__.py @@ -5,146 +5,18 @@ # 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. -"""Lazy loader for ported, per-architecture TileLib templates.""" +"""Compatibility loader for PTODSL TileLib templates. -from functools import lru_cache -from importlib import import_module +The template sources live in the top-level ``TileOps`` package. +""" +from __future__ import annotations -_TEMPLATE_MODULES = { - ("a5", "pto.tabs"): ".a5.tabs", - ("a5", "pto.tci"): ".a5.tci", - ("a5", "pto.tadd"): ".a5.tadd", - ("a5", "pto.tadds"): ".a5.tadds", - ("a5", "pto.tand"): ".a5.tand", - ("a5", "pto.tands"): ".a5.tands", - ("a5", "pto.tcmp"): ".a5.tcmp", - ("a5", "pto.tcmps"): ".a5.tcmps", - ("a5", "pto.tcolargmax"): ".a5.tcolargmax", - ("a5", "pto.tcolargmin"): ".a5.tcolargmin", - ("a5", "pto.tcolmax"): ".a5.tcolmax", - ("a5", "pto.tconcat"): ".a5.tconcat", - ("a5", "pto.tcolexpand"): ".a5.tcolexpand", - ("a5", "pto.tcolexpandadd"): ".a5.tcolexpandadd", - ("a5", "pto.tcolexpanddiv"): ".a5.tcolexpanddiv", - ("a5", "pto.tcolexpandexpdif"): ".a5.tcolexpandexpdif", - ("a5", "pto.tcolexpandmax"): ".a5.tcolexpandmax", - ("a5", "pto.tcolexpandmin"): ".a5.tcolexpandmin", - ("a5", "pto.tcolexpandmul"): ".a5.tcolexpandmul", - ("a5", "pto.tcolexpandsub"): ".a5.tcolexpandsub", - ("a5", "pto.tcolmin"): ".a5.tcolmin", - ("a5", "pto.tcolprod"): ".a5.tcolprod", - ("a5", "pto.tcolsum"): ".a5.tcolsum", - ("a5", "pto.tdiv"): ".a5.tdiv", - ("a5", "pto.tdivs"): ".a5.tdivs", - ("a5", "pto.tcvt"): ".a5.tcvt", - ("a5", "pto.tdequant"): ".a5.tdequant", - ("a5", "pto.texp"): ".a5.texp", - ("a5", "pto.texpands"): ".a5.texpand", - ("a5", "pto.textract"): ".a5.textract", - ("a5", "pto.textract_fp"): ".a5.textract_fp", - ("a5", "pto.tfmod"): ".a5.tfmod", - ("a5", "pto.tfmods"): ".a5.tfmods", - ("a5", "pto.tfillpad"): ".a5.tfillpad", - ("a5", "pto.tfillpad_expand"): ".a5.tfillpad_expand", - ("a5", "pto.tfillpad_inplace"): ".a5.tfillpad_inplace", - ("a5", "pto.tgatherb"): ".a5.tgatherb", - ("a5", "pto.tgemv"): ".a5.tgemv", - ("a5", "pto.tgemv.acc"): ".a5.tgemv_acc", - ("a5", "pto.tgemv.bias"): ".a5.tgemv_bias", - ("a5", "pto.tgemv.mx"): ".a5.tgemv_mx", - ("a5", "pto.tgemv.mx.acc"): ".a5.tgemv_mx", - ("a5", "pto.tgemv.mx.bias"): ".a5.tgemv_mx", - ("a5", "pto.tinsert"): ".a5.tinsert", - ("a5", "pto.tlrelu"): ".a5.tlrelu", - ("a5", "pto.tlog"): ".a5.tlog", - ("a5", "pto.tload"): ".a5.tload", - ("a5", "pto.tmax"): ".a5.tmax", - ("a5", "pto.tmaxs"): ".a5.tmaxs", - ("a5", "pto.tmin"): ".a5.tmin", - ("a5", "pto.tmins"): ".a5.tmins", - ("a5", "pto.tmatmul"): ".a5.tmatmul", - ("a5", "pto.tmatmul.acc"): ".a5.tmatmul_acc", - ("a5", "pto.tmatmul.bias"): ".a5.tmatmul_bias", - ("a5", "pto.tmatmul.mx"): ".a5.tmatmul_mx", - ("a5", "pto.tmatmul.mx.acc"): ".a5.tmatmul_mx", - ("a5", "pto.tmatmul.mx.bias"): ".a5.tmatmul_mx", - ("a5", "pto.tmov"): ( - ".a5.tmov", - ".a5.tmov2bias", - ".a5.tmov2left", - ".a5.tmov2right", - ".a5.tmov2scale", - ".a5.tmov2vec", - ".a5.tmov_fp", - ), - ("a5", "pto.tmul"): ".a5.tmul", - ("a5", "pto.tmuls"): ".a5.tmuls", - ("a5", "pto.tneg"): ".a5.tneg", - ("a5", "pto.tnot"): ".a5.tnot", - ("a5", "pto.tor"): ".a5.tor", - ("a5", "pto.tors"): ".a5.tors", - ("a5", "pto.tpartadd"): ".a5.tpartadd", - ("a5", "pto.tpartmax"): ".a5.tpartmax", - ("a5", "pto.tpartmin"): ".a5.tpartmin", - ("a5", "pto.tpartmul"): ".a5.tpartmul", - ("a5", "pto.tprelu"): ".a5.tprelu", - ("a5", "pto.trelu"): ".a5.trelu", - ("a5", "pto.trecip"): ".a5.trecip", - ("a5", "pto.trem"): ".a5.trem", - ("a5", "pto.trems"): ".a5.trems", - ("a5", "pto.trandom"): ".a5.trandom", - ("a5", "pto.trsqrt"): ".a5.trsqrt", - ("a5", "pto.trowargmax"): ".a5.trowargmax", - ("a5", "pto.trowargmin"): ".a5.trowargmin", - ("a5", "pto.trowexpand"): ".a5.trowexpand", - ("a5", "pto.trowexpandadd"): ".a5.trowexpandadd", - ("a5", "pto.trowexpanddiv"): ".a5.trowexpanddiv", - ("a5", "pto.trowexpandexpdif"): ".a5.trowexpandexpdif", - ("a5", "pto.trowexpandmax"): ".a5.trowexpandmax", - ("a5", "pto.trowexpandmin"): ".a5.trowexpandmin", - ("a5", "pto.trowexpandmul"): ".a5.trowexpandmul", - ("a5", "pto.trowexpandsub"): ".a5.trowexpandsub", - ("a5", "pto.trowmax"): ".a5.trowmax", - ("a5", "pto.trowmin"): ".a5.trowmin", - ("a5", "pto.trowprod"): ".a5.trowprod", - ("a5", "pto.trowsum"): ".a5.trowsum", - ("a5", "pto.tsel"): ".a5.tsel", - ("a5", "pto.tsels"): ".a5.tsels", - ("a5", "pto.tshl"): ".a5.tshl", - ("a5", "pto.tshls"): ".a5.tshls", - ("a5", "pto.tshr"): ".a5.tshr", - ("a5", "pto.tshrs"): ".a5.tshrs", - ("a5", "pto.tmrgsort"): ".a5.tmrgsort", - ("a5", "pto.tsort32"): ".a5.tsort32", - ("a5", "pto.tstore"): ".a5.tstore", - ("a5", "pto.tstore_fp"): ".a5.tstore", - ("a5", "pto.ttrans"): ".a5.ttrans", - ("a5", "pto.tsub"): ".a5.tsub", - ("a5", "pto.tsubs"): ".a5.tsubs", - ("a5", "pto.tsqrt"): ".a5.tsqrt", - ("a5", "pto.txor"): ".a5.txor", - ("a5", "pto.txors"): ".a5.txors", - ("a5", "pto.ttri"): ".a5.ttri", - ("a5", "pto.thistogram"): ".a5.thistogram", -} +from pathlib import Path +from TileOps import load_template +import TileOps as _tileops -@lru_cache(maxsize=None) -def load_template(op: str, target: str) -> bool: - """Import and register only the template module for ``(target, op)``. - - Both this cache and Python's module cache make repeated requests no-ops. - Returns ``False`` when this TileLib has no module for the requested pair. - """ - module_names = _TEMPLATE_MODULES.get((target, op)) - if module_names is None: - return False - if isinstance(module_names, str): - module_names = (module_names,) - for module_name in module_names: - import_module(module_name, package=__name__) - return True - +__path__ = [str(Path(_tileops.__file__).resolve().parent)] __all__ = ["load_template"] diff --git a/ptodsl/tests/test_tilelib_render.py b/ptodsl/tests/test_tilelib_render.py index d1405ef0a6..fc20c5ee1b 100644 --- a/ptodsl/tests/test_tilelib_render.py +++ b/ptodsl/tests/test_tilelib_render.py @@ -16,7 +16,7 @@ from pathlib import Path from ptodsl.tilelib import TileSpec, f32 -from ptodsl.tilelib.templates.a5.tadd import template_tadd +from TileOps.a5.tadd import template_tadd FIXTURE = Path(__file__).parent / "fixtures" / "tadd_a5_8x64_f32.golden.mlir" diff --git a/scripts/ptoas_env.sh b/scripts/ptoas_env.sh index b5e814ccda..0b8cb9a9d3 100644 --- a/scripts/ptoas_env.sh +++ b/scripts/ptoas_env.sh @@ -57,6 +57,7 @@ export PTOAS_PYTHON_SITE export PTO_PYTHON_ROOT="${PTO_PYTHON_ROOT:-${PTO_INSTALL_DIR}}" export PTO_PYTHON_BUILD_ROOT="${PTO_PYTHON_BUILD_ROOT:-${PTO_SOURCE_DIR}/build/python}" export PTODSL_PYTHON_ROOT="${PTODSL_PYTHON_ROOT:-${PTO_SOURCE_DIR}/ptodsl}" +export PTO_TILEOPS_PYTHON_ROOT="${PTO_TILEOPS_PYTHON_ROOT:-${PTO_SOURCE_DIR}/lib}" export PYBIND11_CMAKE_DIR=$(python3 -m pybind11 --cmakedir) export PTOAS_FLAGS="${PTOAS_FLAGS:-}" export PTOAS_OUT_DIR="${PTOAS_OUT_DIR:-${PTO_SOURCE_DIR}/build/output}" @@ -119,6 +120,7 @@ _ptoas_prepend_path PYTHONPATH "${PTOAS_PYTHON_SITE}" _ptoas_prepend_path PYTHONPATH "${MLIR_PYTHON_ROOT}" _ptoas_prepend_path PYTHONPATH "${PTO_PYTHON_BUILD_ROOT}" _ptoas_prepend_path PYTHONPATH "${PTODSL_PYTHON_ROOT}" +_ptoas_prepend_path PYTHONPATH "${PTO_TILEOPS_PYTHON_ROOT}" _ptoas_prepend_path LD_LIBRARY_PATH "${LLVM_BUILD_DIR}/lib" _ptoas_prepend_path LD_LIBRARY_PATH "${PTO_INSTALL_DIR}/lib" diff --git a/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt index e84eb78ae8..e38cf8dcd9 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt +++ b/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt @@ -211,7 +211,6 @@ set(ALL_TESTCASES textract textract_fp textract_v2v - fp8_load_test tmov2left tmov2right tmov2bias diff --git a/tools/ptoas/CMakeLists.txt b/tools/ptoas/CMakeLists.txt index 1e2918bd5a..db14abf2ce 100644 --- a/tools/ptoas/CMakeLists.txt +++ b/tools/ptoas/CMakeLists.txt @@ -42,9 +42,8 @@ function(ptoas_configure_runtime_target target_name) # --tilelang-pkg-path / --ptodsl-pkg-path. Installed layouts are expected # to launch through the wrapper/launcher flow, which injects explicit # runtime paths instead of relying on executable-path probing. - PTOAS_DEFAULT_TILELANG_PATH="${CMAKE_SOURCE_DIR}/lib/TileOps" - PTOAS_DEFAULT_TILELANG_PKG_PATH="${CMAKE_SOURCE_DIR}/tilelang-dsl/python" PTOAS_DEFAULT_PTODSL_PKG_PATH="${CMAKE_SOURCE_DIR}/ptodsl" + PTOAS_DEFAULT_TILEOPS_PKG_PATH="${CMAKE_SOURCE_DIR}/lib" ${ARGN} ) target_link_libraries(${target_name} PRIVATE diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 3ac02d6309..9492b6c29d 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -49,6 +49,7 @@ #include "llvm/Support/SourceMgr.h" #include "llvm/Support/ToolOutputFile.h" #include "llvm/Support/FileSystem.h" // [Fix] Required for OF_None +#include "llvm/Support/Path.h" #include "ptobc/ptobc_decode.h" #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h" #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" @@ -175,6 +176,8 @@ static std::string resolveEffectiveTargetArch(ModuleOp module, } // namespace +int main(int argc, char **argv); + void mlir::pto::registerPTOASDialects(DialectRegistry ®istry) { func::registerInlinerExtension(registry); LLVM::registerInlinerInterface(registry); @@ -227,6 +230,38 @@ void mlir::pto::loadPTOASDialects(MLIRContext &context) { context.getOrLoadDialect(); } +static bool pathExists(llvm::StringRef path) { + return !path.empty() && llvm::sys::fs::exists(path); +} + +static std::string getParentDir(llvm::StringRef path) { + llvm::SmallString<256> parent(path); + llvm::sys::path::remove_filename(parent); + llvm::sys::path::remove_dots(parent, true); + return std::string(parent); +} + +static std::string joinPath(llvm::StringRef lhs, llvm::StringRef rhs) { + llvm::SmallString<256> joined(lhs); + llvm::sys::path::append(joined, rhs); + llvm::sys::path::remove_dots(joined, true); + return std::string(joined); +} + +static std::string detectInstalledPythonPkgRoot(const char *argv0, + llvm::StringRef packageName) { + std::string exePath = llvm::sys::fs::getMainExecutable(argv0, (void *)&main); + if (exePath.empty()) + return {}; + + const std::string exeDir = getParentDir(exePath); + const std::string prefixDir = getParentDir(exeDir); + const std::string installedPkg = joinPath(prefixDir, packageName); + if (pathExists(installedPkg)) + return prefixDir; + return {}; +} + static bool hasCLIOption(int argc, char **argv, llvm::StringRef option) { const std::string optionWithValue = (option + "=").str(); for (int i = 1; i < argc; ++i) { @@ -411,27 +446,12 @@ static llvm::cl::opt enableTileOpExpand( "--pto-backend=vpto."), llvm::cl::init(false)); -#ifndef PTOAS_DEFAULT_TILELANG_PATH -#define PTOAS_DEFAULT_TILELANG_PATH "" -#endif -#ifndef PTOAS_DEFAULT_TILELANG_PKG_PATH -#define PTOAS_DEFAULT_TILELANG_PKG_PATH "" -#endif #ifndef PTOAS_DEFAULT_PTODSL_PKG_PATH #define PTOAS_DEFAULT_PTODSL_PKG_PATH "" #endif - -static llvm::cl::opt tilelangPath( - "tilelang-path", - llvm::cl::desc("Path to directory of .py tilelang DSL template files " - "(default: /lib/TileOps, baked in at build time)"), - llvm::cl::init(PTOAS_DEFAULT_TILELANG_PATH)); - -static llvm::cl::opt tilelangPkgPath( - "tilelang-pkg-path", - llvm::cl::desc("PYTHONPATH for tilelang_dsl package " - "(default: /tilelang-dsl/python, baked in at build time)"), - llvm::cl::init(PTOAS_DEFAULT_TILELANG_PKG_PATH)); +#ifndef PTOAS_DEFAULT_TILEOPS_PKG_PATH +#define PTOAS_DEFAULT_TILEOPS_PKG_PATH "" +#endif static llvm::cl::opt ptodslPkgPath( "ptodsl-pkg-path", @@ -439,6 +459,12 @@ static llvm::cl::opt ptodslPkgPath( "(default: /ptodsl, baked in at build time)"), llvm::cl::init(PTOAS_DEFAULT_PTODSL_PKG_PATH)); +static llvm::cl::opt tileopsPkgPath( + "tileops-pkg-path", + llvm::cl::desc("PYTHONPATH for the TileOps PTODSL template package " + "(default: /lib, baked in at build time)"), + llvm::cl::init(PTOAS_DEFAULT_TILEOPS_PKG_PATH)); + static llvm::cl::opt daemonSocketPath( "daemon-socket-path", llvm::cl::desc("Path to Unix domain socket for daemon RPC " @@ -446,7 +472,6 @@ static llvm::cl::opt daemonSocketPath( llvm::cl::init("")); enum class TileLibBackend { - TileLang, PTODSL, }; @@ -454,8 +479,6 @@ static llvm::cl::opt tileLibBackend( "tile-lib-backend", llvm::cl::desc("TileLib backend used by ExpandTileOp"), llvm::cl::values( - clEnumValN(TileLibBackend::TileLang, "tilelang", - "Use the legacy TileLang DSL TileLib"), clEnumValN(TileLibBackend::PTODSL, "ptodsl", "Use the PTODSL TileLib daemon")), llvm::cl::init(TileLibBackend::PTODSL)); @@ -470,66 +493,66 @@ static std::string resolveTileLibPythonExe() { static pto::ExpandTileOpOptions resolveExpandTileOpOptions(int argc, char **argv) { pto::ExpandTileOpOptions expandOpts; - expandOpts.tilelangPath = tilelangPath; - expandOpts.tilelangPkgPath = tilelangPkgPath; expandOpts.pythonExe = resolveTileLibPythonExe(); - const bool usePTODSLTileLib = tileLibBackend == TileLibBackend::PTODSL; + expandOpts.tilelangPath.clear(); + expandOpts.tilelangPkgPath.clear(); std::string resolvedPtodslPkgPath = ptodslPkgPath; + std::string resolvedTileOpsPkgPath = tileopsPkgPath; if (!hasCLIOption(argc, argv, "--ptodsl-pkg-path")) { const char *envPtodslRoot = ::getenv("PTODSL_PYTHON_ROOT"); if (envPtodslRoot && envPtodslRoot[0] != '\0') resolvedPtodslPkgPath = envPtodslRoot; + else { + std::string installedPtodslPkgPath = + detectInstalledPythonPkgRoot(argv[0], "ptodsl"); + if (!installedPtodslPkgPath.empty()) + resolvedPtodslPkgPath = installedPtodslPkgPath; + } } - if (usePTODSLTileLib) { - // The PTODSL backend is package-based and must not depend on legacy - // TileLang template or package paths. - expandOpts.tilelangPath.clear(); - expandOpts.tilelangPkgPath.clear(); + if (!hasCLIOption(argc, argv, "--tileops-pkg-path")) { + const char *envTileOpsRoot = ::getenv("PTO_TILEOPS_PYTHON_ROOT"); + if (envTileOpsRoot && envTileOpsRoot[0] != '\0') + resolvedTileOpsPkgPath = envTileOpsRoot; + else { + std::string installedTileOpsPkgPath = + detectInstalledPythonPkgRoot(argv[0], "TileOps"); + if (!installedTileOpsPkgPath.empty()) + resolvedTileOpsPkgPath = installedTileOpsPkgPath; + } } - expandOpts.tileLibBackend = usePTODSLTileLib ? "ptodsl" : "tilelang"; - expandOpts.daemonHelperModule = - usePTODSLTileLib ? "ptodsl.tilelib.serving.helper" - : "tilelang_dsl.daemon_helper"; - expandOpts.tileLibPkgPath = - usePTODSLTileLib ? resolvedPtodslPkgPath - : std::string(expandOpts.tilelangPkgPath); + expandOpts.tileLibBackend = "ptodsl"; + expandOpts.daemonHelperModule = "ptodsl.tilelib.serving.helper"; + expandOpts.tileLibPkgPath = resolvedPtodslPkgPath; + if (!resolvedTileOpsPkgPath.empty()) { + if (!expandOpts.tileLibPkgPath.empty()) + expandOpts.tileLibPkgPath += ":"; + expandOpts.tileLibPkgPath += resolvedTileOpsPkgPath; + } // Daemon mode is default (no CLI option needed) // Automatically start daemon for instance caching - if (usePTODSLTileLib || !expandOpts.tilelangPath.empty()) { - std::string socket = daemonSocketPath; - if (socket.empty()) - socket = ptoas::DaemonManager::generateSocketPath(); - - // Register cleanup handler (daemon will be stopped on PTOAS exit) - ptoas::registerDaemonCleanup(); - - const std::string daemonModule = - usePTODSLTileLib ? "ptodsl.tilelib.serving.daemon" - : "tilelang_dsl.daemon"; - const std::string templateDir = - usePTODSLTileLib ? "" : std::string(expandOpts.tilelangPath); - - // Try to start daemon automatically - if (ptoas::DaemonManager::start(socket, daemonModule, expandOpts.pythonExe, - expandOpts.tileLibPkgPath, templateDir)) { - expandOpts.daemonSocketPath = socket; - llvm::errs() << "Info: " << expandOpts.tileLibBackend - << " TileLib daemon started successfully\n"; - } else { - expandOpts.daemonSocketPath = ""; - if (usePTODSLTileLib) { - llvm::errs() - << "Error: Failed to start the PTODSL TileLib daemon; no TileLang " - "fallback will be used\n"; - } else { - llvm::errs() << "Warning: Failed to start daemon, using legacy " - "TileLang subprocess mode\n"; - } - } + std::string socket = daemonSocketPath; + if (socket.empty()) + socket = ptoas::DaemonManager::generateSocketPath(); + + // Register cleanup handler (daemon will be stopped on PTOAS exit) + ptoas::registerDaemonCleanup(); + + // Try to start daemon automatically + if (ptoas::DaemonManager::start(socket, "ptodsl.tilelib.serving.daemon", + expandOpts.pythonExe, + expandOpts.tileLibPkgPath, "")) { + expandOpts.daemonSocketPath = socket; + llvm::errs() << "Info: " << expandOpts.tileLibBackend + << " TileLib daemon started successfully\n"; + } else { + expandOpts.daemonSocketPath = ""; + llvm::errs() + << "Error: Failed to start the PTODSL TileLib daemon; no TileLang " + "fallback will be used\n"; } return expandOpts; From 7e43f515358805ea0e93b8dcdc93cab3af2ce565 Mon Sep 17 00:00:00 2001 From: Alan Wang Date: Mon, 27 Jul 2026 15:10:47 +0000 Subject: [PATCH 2/4] Modify ci, docker, etc. --- .github/workflows/ci_sim.yml | 65 ++++++++---------------- docker/collect_ptoas_dist.sh | 16 ++---- docker/collect_ptoas_dist_mac.sh | 16 ++---- docker/create_wheel.sh | 5 +- lib/PTO/Transforms/PTOToEmitC.cpp | 4 +- ptodsl/CMakeLists.txt | 8 +++ ptodsl/ptoas/_runtime_entry.py | 6 +-- ptodsl/tests/test_ptoas_runtime_entry.py | 6 +-- tools/ptoas/CMakeLists.txt | 4 +- 9 files changed, 44 insertions(+), 86 deletions(-) diff --git a/.github/workflows/ci_sim.yml b/.github/workflows/ci_sim.yml index dcbccabf9d..5466fa738f 100644 --- a/.github/workflows/ci_sim.yml +++ b/.github/workflows/ci_sim.yml @@ -61,7 +61,6 @@ jobs: - 'test/tilelang_st/**' - 'test/dsl-st/**' - 'test/tilelib-st/**' - - 'tilelang-dsl/**' - 'ptodsl/**' - 'tools/ptoas/**' - 'lib/TileOps/**' @@ -115,8 +114,7 @@ jobs: PTO_WHEELHOUSE: ${{ github.workspace }}/.work/ptoas-wheelhouse PTOAS_VENV: ${{ github.workspace }}/.work/ptoas-sim-venv VPTO_SIM_WORKSPACE: ${{ github.workspace }}/.work/vpto-sim-ci - TILELANG_DSL_WORKSPACE: ${{ github.workspace }}/.work/tilelang-dsl-ci - TILELANG_DSL_UT_WORKSPACE: ${{ github.workspace }}/.work/tilelang-dsl-ut-ci + TILELIB_ST_WORKSPACE: ${{ github.workspace }}/.work/tilelib-st-ci PYPTO_REF: ef6ce7cd8bd33b4c93b58dc34830a20b145736ac PYPTO_WORKSPACE: ${{ github.workspace }}/.work/pypto-ci PYPTO_RUN_WORKSPACE: ${{ github.workspace }}/.work/pypto-run-ci @@ -253,7 +251,7 @@ jobs: rm -rf "${PTO_WHEELHOUSE}" rm -rf "${PTOAS_VENV}" rm -rf "${VPTO_SIM_WORKSPACE}" - rm -rf "${TILELANG_DSL_WORKSPACE}" + rm -rf "${TILELIB_ST_WORKSPACE}" rm -rf "${PYPTO_WORKSPACE}" rm -rf "${PYPTO_RUN_WORKSPACE}" rm -rf "${PTO_ISA_ROOT}" @@ -703,23 +701,23 @@ jobs: echo "PTO_DSL_ST_PYTHON_TAG=${PYTHON_TAG}" >> "${GITHUB_ENV}" echo "PTO_DSL_ST_LLVM_DIR=${PTO_DSL_ST_LLVM_DIR}" >> "${GITHUB_ENV}" - - name: Run TileLang DSL CI + - name: Run TileOp ST CI shell: bash run: | set -euo pipefail - mkdir -p "${TILELANG_DSL_WORKSPACE}" + mkdir -p "${TILELIB_ST_WORKSPACE}" export LLVM_BUILD_DIR="${LLVM_DIR}" export MLIR_PYTHON_ROOT="${MLIR_PYTHONPATH}" if [[ "${{ github.event_name }}" == "pull_request" ]]; then ASCEND_HOME_PATH="${ASCEND_HOME_PATH}" \ PTOAS_BIN="${PTOAS_BIN}" \ bash test/tilelang_st/script/run_ci.sh -r sim -v a5 --jobs 64 --smoke \ - 2>&1 | tee "${TILELANG_DSL_WORKSPACE}/run_ci.log" + 2>&1 | tee "${TILELIB_ST_WORKSPACE}/run_ci.log" else ASCEND_HOME_PATH="${ASCEND_HOME_PATH}" \ PTOAS_BIN="${PTOAS_BIN}" \ bash test/tilelang_st/script/run_ci.sh -r sim -v a5 --jobs 64 \ - 2>&1 | tee "${TILELANG_DSL_WORKSPACE}/run_ci.log" + 2>&1 | tee "${TILELIB_ST_WORKSPACE}/run_ci.log" fi - name: Restore PTODSL LLVM build cache @@ -822,7 +820,7 @@ jobs: shell: bash run: | set -euo pipefail - mkdir -p "${TILELANG_DSL_WORKSPACE}" + mkdir -p "${TILELIB_ST_WORKSPACE}" export TORCH_DEVICE_BACKEND_AUTOLOAD=0 source "${ASCEND_HOME_PATH}/bin/setenv.bash" @@ -850,7 +848,7 @@ jobs: PY } - PROBE_LOG="${TILELANG_DSL_WORKSPACE}/ptodsl-python-probe.log" + PROBE_LOG="${TILELIB_ST_WORKSPACE}/ptodsl-python-probe.log" : > "${PROBE_LOG}" echo "Probing PTODSL DSL ST Python: ${PTO_DSL_ST_PYTHON_BIN}" | tee -a "${PROBE_LOG}" echo "Probing PTODSL DSL ST ptoas: ${PTO_DSL_ST_PTOAS_BIN}" | tee -a "${PROBE_LOG}" @@ -895,26 +893,24 @@ jobs: shell: bash run: | set -euo pipefail - mkdir -p "${TILELANG_DSL_WORKSPACE}" + mkdir -p "${TILELIB_ST_WORKSPACE}" export PYTHON_BIN="${PTO_DSL_ST_PYTHON_BIN}" export PTO_PYTHON_BIN="${PTO_DSL_ST_PYTHON_BIN}" export PATH="${PTO_DSL_ST_BIN_DIR}:${PATH}" export TORCH_DEVICE_BACKEND_AUTOLOAD=0 source "${ASCEND_HOME_PATH}/bin/setenv.bash" - TILELIB_ST_OUTPUT_ROOT="${TILELANG_DSL_WORKSPACE}/tilelib-st-a5" \ + TILELIB_ST_OUTPUT_ROOT="${TILELIB_ST_WORKSPACE}/tilelib-st-a5" \ ASCEND_HOME_PATH="${ASCEND_HOME_PATH}" \ PTOAS_BIN="${PTO_DSL_ST_PTOAS_BIN}" \ - "${PTO_DSL_ST_PYTHON_BIN}" -m pytest \ - -n "${TILELIB_ST_JOBS:-16}" --dist=worksteal -v \ - test/tilelib-st/test_tilelib_st.py \ - 2>&1 | tee "${TILELANG_DSL_WORKSPACE}/tilelib-st-a5.log" + scripts/sim_dsl.sh test/tilelib-st/run_tilelib_st.py -- test/tilelib-st/a5 \ + 2>&1 | tee "${TILELIB_ST_WORKSPACE}/tilelib-st-a5.log" - name: Run PTODSL ST CI shell: bash run: | set -euo pipefail - mkdir -p "${TILELANG_DSL_WORKSPACE}" + mkdir -p "${TILELIB_ST_WORKSPACE}" export PYTHON_BIN="${PTO_DSL_ST_PYTHON_BIN}" export PTO_PYTHON_BIN="${PTO_DSL_ST_PYTHON_BIN}" export PATH="${PTO_DSL_ST_BIN_DIR}:${PATH}" @@ -924,39 +920,18 @@ jobs: ASCEND_HOME_PATH="${ASCEND_HOME_PATH}" \ PTOAS_BIN="${PTO_DSL_ST_PTOAS_BIN}" \ scripts/sim_dsl.sh test/dsl-st \ - 2>&1 | tee "${TILELANG_DSL_WORKSPACE}/ptodsl-dsl-st.log" + 2>&1 | tee "${TILELIB_ST_WORKSPACE}/ptodsl-dsl-st.log" - - name: Upload TileLang DSL logs + - name: Upload TileLib ST logs if: always() uses: actions/upload-artifact@v4 with: - name: tilelang-dsl-ci-${{ github.run_id }} + name: tilelib-st-ci-${{ github.run_id }} path: | - ${{ env.TILELANG_DSL_WORKSPACE }}/run_ci.log - ${{ env.TILELANG_DSL_WORKSPACE }}/ptodsl-dsl-st.log - ${{ env.TILELANG_DSL_WORKSPACE }}/tilelib-st-a5.log - ${{ env.TILELANG_DSL_WORKSPACE }}/tilelib-st-a5/** - ${{ env.TILELANG_DSL_WORKSPACE }}/ptodsl-python-probe.log - if-no-files-found: warn - - - name: Run TileLang DSL unit tests - shell: bash - run: | - set -euo pipefail - mkdir -p "${TILELANG_DSL_UT_WORKSPACE}" - export PATH="${PTOAS_VENV}/bin:${PATH}" - export LD_LIBRARY_PATH="${LLVM_DIR}/lib:${PTO_INSTALL_DIR}/lib:${GITHUB_WORKSPACE}/build/lib:${LD_LIBRARY_PATH:-}" - cd tilelang-dsl/python - python3 -m unittest discover -s ../tests -p 'test_*.py' \ - 2>&1 | tee "${TILELANG_DSL_UT_WORKSPACE}/unittest.log" - - - name: Upload TileLang DSL unit test logs - if: always() - uses: actions/upload-artifact@v4 - with: - name: tilelang-dsl-ut-ci-${{ github.run_id }} - path: | - ${{ env.TILELANG_DSL_UT_WORKSPACE }}/unittest.log + ${{ env.TILELIB_ST_WORKSPACE }}/run_ci.log + ${{ env.TILELIB_ST_WORKSPACE }}/ptodsl-dsl-st.log + ${{ env.TILELIB_ST_WORKSPACE }}/tilelib-st-a5.log + ${{ env.TILELIB_ST_WORKSPACE }}/ptodsl-python-probe.log if-no-files-found: warn - name: Upload VPTO SIM logs diff --git a/docker/collect_ptoas_dist.sh b/docker/collect_ptoas_dist.sh index 83f9f7299f..6e25c4980d 100755 --- a/docker/collect_ptoas_dist.sh +++ b/docker/collect_ptoas_dist.sh @@ -24,8 +24,7 @@ # ptoas_wheel_bootstrap.py - Compatibility bootstrap module # lib/ptoas.so - Shared runtime loaded by the launcher # lib/*.so* - Required shared library dependencies -# share/ptoas/TileOps - TileLang template library -# tilelang_dsl/ - TileLang DSL Python package +# share/ptoas/TileOps - PTODSL TileOps template package # # This compiler-oriented binary artifact intentionally does not bundle the # PTODSL Python package. PTODSL is provided only by the `ptoas` wheel. @@ -56,8 +55,6 @@ PTOAS_DEPS_DIR="${PTOAS_DIST_DIR}/lib" PTOAS_SHARED_MODULE_DIST_PATH="${PTOAS_DEPS_DIR}/ptoas.so" PTOAS_TILEOPS_SRC_DIR="${PTO_INSTALL_DIR}/share/ptoas/TileOps" PTOAS_TILEOPS_DIST_DIR="${PTOAS_DIST_DIR}/share/ptoas/TileOps" -PTOAS_TILELANG_DSL_SRC_DIR="${PTO_INSTALL_DIR}/tilelang_dsl" -PTOAS_TILELANG_DSL_DIST_DIR="${PTOAS_DIST_DIR}/tilelang_dsl" PTOAS_WRAPPER_PKG_SRC_DIR="${PTO_INSTALL_DIR}/ptoas" PTOAS_WRAPPER_PKG_DIST_DIR="${PTOAS_DIST_DIR}/ptoas" PTOAS_WHEEL_BOOTSTRAP_SRC="${PTO_INSTALL_DIR}/ptoas_wheel_bootstrap.py" @@ -227,18 +224,13 @@ while read -r packaged; do harden_elf "$packaged" done < <(find "${PTOAS_DEPS_DIR}" -type f | sort) -echo "Copying TileLang runtime resources..." +echo "Copying TileOps runtime resources..." if [ ! -d "${PTOAS_TILEOPS_SRC_DIR}" ]; then echo "Error: TileOps resource directory not found at ${PTOAS_TILEOPS_SRC_DIR}" >&2 exit 1 fi -if [ ! -d "${PTOAS_TILELANG_DSL_SRC_DIR}" ]; then - echo "Error: tilelang_dsl package directory not found at ${PTOAS_TILELANG_DSL_SRC_DIR}" >&2 - exit 1 -fi -rm -rf "${PTOAS_TILEOPS_DIST_DIR}" "${PTOAS_TILELANG_DSL_DIST_DIR}" +rm -rf "${PTOAS_TILEOPS_DIST_DIR}" cp -R "${PTOAS_TILEOPS_SRC_DIR}" "${PTOAS_TILEOPS_DIST_DIR}" -cp -R "${PTOAS_TILELANG_DSL_SRC_DIR}" "${PTOAS_TILELANG_DSL_DIST_DIR}" echo "Smoke testing packaged ptoas dist..." VERSION_OUTPUT="$(env -u PYTHONPATH -u DYLD_LIBRARY_PATH -u LD_LIBRARY_PATH \ @@ -256,7 +248,6 @@ else fi test -d "${PTOAS_TILEOPS_DIST_DIR}" -test -f "${PTOAS_TILELANG_DSL_DIST_DIR}/__init__.py" test -f "${PTOAS_SHARED_MODULE_DIST_PATH}" test -f "${PTOAS_WRAPPER_PKG_DIST_DIR}/_runtime_entry.py" test -f "${PTOAS_WHEEL_BOOTSTRAP_DIST_PATH}" @@ -271,7 +262,6 @@ echo "=== ptoas distribution contents ===" ls -la "${PTOAS_DIST_DIR}/" ls -la "${PTOAS_DIST_DIR}/bin/" ls -la "${PTOAS_DIST_DIR}/share/ptoas/" -ls -la "${PTOAS_TILELANG_DSL_DIST_DIR}" SO_COUNT=$(find "${PTOAS_DEPS_DIR}" -name "*.so*" 2>/dev/null | wc -l) echo "=== Collected .so dependencies (${SO_COUNT} files) ===" du -sh "${PTOAS_DEPS_DIR}/" diff --git a/docker/collect_ptoas_dist_mac.sh b/docker/collect_ptoas_dist_mac.sh index ee4a1f1f7e..0b30a5808b 100644 --- a/docker/collect_ptoas_dist_mac.sh +++ b/docker/collect_ptoas_dist_mac.sh @@ -24,8 +24,7 @@ # ptoas_wheel_bootstrap.py - Compatibility bootstrap module # lib/ptoas.so - Shared runtime loaded by the launcher # lib/*.dylib - Required shared library dependencies -# share/ptoas/TileOps - TileLang template library -# tilelang_dsl/ - TileLang DSL Python package +# share/ptoas/TileOps - PTODSL TileOps template package # # This compiler-oriented binary artifact intentionally does not bundle the # PTODSL Python package. PTODSL is provided only by the `ptoas` wheel. @@ -54,8 +53,6 @@ PTOAS_DEPS_DIR="${PTOAS_DIST_DIR}/lib" PTOAS_SHARED_MODULE_DIST_PATH="${PTOAS_DEPS_DIR}/ptoas.so" PTOAS_TILEOPS_SRC_DIR="${PTO_INSTALL_DIR}/share/ptoas/TileOps" PTOAS_TILEOPS_DIST_DIR="${PTOAS_DIST_DIR}/share/ptoas/TileOps" -PTOAS_TILELANG_DSL_SRC_DIR="${PTO_INSTALL_DIR}/tilelang_dsl" -PTOAS_TILELANG_DSL_DIST_DIR="${PTOAS_DIST_DIR}/tilelang_dsl" PTOAS_WRAPPER_PKG_SRC_DIR="${PTO_INSTALL_DIR}/ptoas" PTOAS_WRAPPER_PKG_DIST_DIR="${PTOAS_DIST_DIR}/ptoas" PTOAS_WHEEL_BOOTSTRAP_SRC="${PTO_INSTALL_DIR}/ptoas_wheel_bootstrap.py" @@ -274,18 +271,13 @@ PY echo "Collecting dylib dependencies..." collect_dylibs "${PTOAS_SHARED_MODULE_DIST_PATH}" -echo "Copying TileLang runtime resources..." +echo "Copying TileOps runtime resources..." if [[ ! -d "${PTOAS_TILEOPS_SRC_DIR}" ]]; then echo "Error: TileOps resource directory not found at ${PTOAS_TILEOPS_SRC_DIR}" >&2 exit 1 fi -if [[ ! -d "${PTOAS_TILELANG_DSL_SRC_DIR}" ]]; then - echo "Error: tilelang_dsl package directory not found at ${PTOAS_TILELANG_DSL_SRC_DIR}" >&2 - exit 1 -fi -rm -rf "${PTOAS_TILEOPS_DIST_DIR}" "${PTOAS_TILELANG_DSL_DIST_DIR}" +rm -rf "${PTOAS_TILEOPS_DIST_DIR}" cp -R "${PTOAS_TILEOPS_SRC_DIR}" "${PTOAS_TILEOPS_DIST_DIR}" -cp -R "${PTOAS_TILELANG_DSL_SRC_DIR}" "${PTOAS_TILELANG_DSL_DIST_DIR}" if [[ -e "${PTOAS_DIST_DIR}/ptodsl" ]]; then echo "Error: compiler-oriented ptoas dist must not bundle PTODSL" >&2 exit 1 @@ -376,7 +368,6 @@ else echo "$VERSION_OUTPUT" | grep -Eq '^ptoas [0-9]+\.[0-9]+$' fi test -d "${PTOAS_TILEOPS_DIST_DIR}" -test -f "${PTOAS_TILELANG_DSL_DIST_DIR}/__init__.py" test -f "${PTOAS_SHARED_MODULE_DIST_PATH}" test -f "${PTOAS_WRAPPER_PKG_DIST_DIR}/_runtime_entry.py" test -f "${PTOAS_WHEEL_BOOTSTRAP_DIST_PATH}" @@ -390,7 +381,6 @@ echo "=== ptoas distribution contents ===" ls -la "${PTOAS_DIST_DIR}/" ls -la "${PTOAS_DIST_DIR}/bin/" ls -la "${PTOAS_DIST_DIR}/share/ptoas/" -ls -la "${PTOAS_TILELANG_DSL_DIST_DIR}" DYLIB_COUNT=$(find "${PTOAS_DEPS_DIR}" -name "*.dylib" 2>/dev/null | wc -l) echo "=== Collected .dylib dependencies (${DYLIB_COUNT} files) ===" du -sh "${PTOAS_DEPS_DIR}/" diff --git a/docker/create_wheel.sh b/docker/create_wheel.sh index 36a9f8fa6d..a4748b1748 100755 --- a/docker/create_wheel.sh +++ b/docker/create_wheel.sh @@ -139,9 +139,8 @@ echo "Overlaying PTO native extension..." mkdir -p "${WHEEL_STAGING_DIR}/mlir/_mlir_libs" find "${PTO_INSTALL_DIR}/mlir/_mlir_libs" -maxdepth 1 -type f -name '_pto*' -exec cp {} "${WHEEL_STAGING_DIR}/mlir/_mlir_libs/" \; -echo "Copying TileLang resources..." -rm -rf "${WHEEL_STAGING_DIR}/tilelang_dsl" "${WHEEL_STAGING_DIR}/TileOps" -cp -R "${PTO_INSTALL_DIR}/tilelang_dsl" "${WHEEL_STAGING_DIR}/tilelang_dsl" +echo "Copying TileOps resources..." +rm -rf "${WHEEL_STAGING_DIR}/TileOps" cp -R "${PTO_INSTALL_DIR}/share/ptoas/TileOps" "${WHEEL_STAGING_DIR}/TileOps" echo "Copying ptodsl package..." diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 02b2d84f0a..d547a63334 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -6465,7 +6465,7 @@ struct PTOSyncSetToEmitC : public OpConversionPattern { return success(); } - [[maybe_unused]] PTOArch targetArch; + PTOArch targetArch; }; struct PTOSyncWaitToEmitC : public OpConversionPattern { @@ -6500,7 +6500,7 @@ struct PTOSyncWaitToEmitC : public OpConversionPattern { return success(); } - [[maybe_unused]] PTOArch targetArch; + PTOArch targetArch; }; // GetBlockIdxOp Lowering (pto.get_block_idx -> get_block_idx()) diff --git a/ptodsl/CMakeLists.txt b/ptodsl/CMakeLists.txt index cb4f27b978..7a2f79cf1a 100644 --- a/ptodsl/CMakeLists.txt +++ b/ptodsl/CMakeLists.txt @@ -76,3 +76,11 @@ install( PATTERN "__pycache__" EXCLUDE PATTERN "*.pyc" EXCLUDE ) + +install( + DIRECTORY "${TILEOPS_PACKAGE_SRC_DIR}" + DESTINATION "share/ptoas" + COMPONENT PTOAS_Runtime + PATTERN "__pycache__" EXCLUDE + PATTERN "*.pyc" EXCLUDE +) diff --git a/ptodsl/ptoas/_runtime_entry.py b/ptodsl/ptoas/_runtime_entry.py index 2fa032727d..1ddd9b4388 100644 --- a/ptodsl/ptoas/_runtime_entry.py +++ b/ptodsl/ptoas/_runtime_entry.py @@ -126,12 +126,10 @@ def launch(layout: PTOASRuntimeLayout, user_args: Sequence[str]) -> int: _prepend_env_path(env, "DYLD_LIBRARY_PATH", runtime_lib_dir) argv = [str(layout.wrapper)] - if not _has_cli_option(user_args, "--tilelang-path"): - argv.extend(["--tilelang-path", str(layout.tileops_dir)]) - if not _has_cli_option(user_args, "--tilelang-pkg-path"): - argv.extend(["--tilelang-pkg-path", str(layout.python_root)]) if not _has_cli_option(user_args, "--ptodsl-pkg-path"): argv.extend(["--ptodsl-pkg-path", str(layout.python_root)]) + if not _has_cli_option(user_args, "--tileops-pkg-path"): + argv.extend(["--tileops-pkg-path", str(layout.tileops_dir.parent)]) argv.extend(user_args) entrypoint = _load_shared_entrypoint( diff --git a/ptodsl/tests/test_ptoas_runtime_entry.py b/ptodsl/tests/test_ptoas_runtime_entry.py index 11def2f490..ccdf7905c9 100644 --- a/ptodsl/tests/test_ptoas_runtime_entry.py +++ b/ptodsl/tests/test_ptoas_runtime_entry.py @@ -69,12 +69,10 @@ def test_launch_injects_runtime_environment_and_calls_shared_entrypoint(self): [c_argv[i].decode("utf-8") for i in range(argc)], [ str(layout.wrapper), - "--tilelang-path", - str(layout.tileops_dir), - "--tilelang-pkg-path", - str(layout.python_root), "--ptodsl-pkg-path", str(layout.python_root), + "--tileops-pkg-path", + str(layout.tileops_dir.parent), "--version", ], ) diff --git a/tools/ptoas/CMakeLists.txt b/tools/ptoas/CMakeLists.txt index db14abf2ce..9aafaffdbf 100644 --- a/tools/ptoas/CMakeLists.txt +++ b/tools/ptoas/CMakeLists.txt @@ -38,8 +38,8 @@ function(ptoas_configure_runtime_target target_name) target_compile_definitions(${target_name} PRIVATE PTOAS_RELEASE_VERSION="${PTOAS_CLI_VERSION}" # Source-tree defaults for TileLib expansion. These let ptoas run directly - # from the build tree without passing --tilelang-path / - # --tilelang-pkg-path / --ptodsl-pkg-path. Installed layouts are expected + # from the build tree without passing --ptodsl-pkg-path / + # --tileops-pkg-path. Installed layouts are expected # to launch through the wrapper/launcher flow, which injects explicit # runtime paths instead of relying on executable-path probing. PTOAS_DEFAULT_PTODSL_PKG_PATH="${CMAKE_SOURCE_DIR}/ptodsl" From 221227c1dbe4cc2c564ba698ea872e61f943e512 Mon Sep 17 00:00:00 2001 From: Alan Wang Date: Tue, 28 Jul 2026 14:47:53 +0000 Subject: [PATCH 3/4] Fix CI --- lib/TileOps/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/lib/TileOps/__init__.py b/lib/TileOps/__init__.py index 77ad910701..241ed0a2d9 100644 --- a/lib/TileOps/__init__.py +++ b/lib/TileOps/__init__.py @@ -33,9 +33,11 @@ ("a5", "pto.tcolmin"): ".a5.tcolmin", ("a5", "pto.tcolprod"): ".a5.tcolprod", ("a5", "pto.tcolsum"): ".a5.tcolsum", + ("a5", "pto.tconcat"): ".a5.tconcat", ("a5", "pto.tdiv"): ".a5.tdiv", ("a5", "pto.tdivs"): ".a5.tdivs", ("a5", "pto.tcvt"): ".a5.tcvt", + ("a5", "pto.tdequant"): ".a5.tdequant", ("a5", "pto.texp"): ".a5.texp", ("a5", "pto.texpands"): ".a5.texpand", ("a5", "pto.textract"): ".a5.textract", From 1d7176faf318ae8853cd1a67b7cff2ab022de842 Mon Sep 17 00:00:00 2001 From: Alan Wang Date: Tue, 28 Jul 2026 16:10:09 +0000 Subject: [PATCH 4/4] Fix the CI --- lib/TileOps/__init__.py | 4 ++++ ptodsl/tests/test_tilelib_catalog.py | 25 +++++++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/lib/TileOps/__init__.py b/lib/TileOps/__init__.py index 241ed0a2d9..25b4293761 100644 --- a/lib/TileOps/__init__.py +++ b/lib/TileOps/__init__.py @@ -17,6 +17,7 @@ ("a5", "pto.tadds"): ".a5.tadds", ("a5", "pto.tand"): ".a5.tand", ("a5", "pto.tands"): ".a5.tands", + ("a5", "pto.tci"): ".a5.tci", ("a5", "pto.tcmp"): ".a5.tcmp", ("a5", "pto.tcmps"): ".a5.tcmps", ("a5", "pto.tcolargmax"): ".a5.tcolargmax", @@ -54,6 +55,7 @@ ("a5", "pto.tgemv.mx"): ".a5.tgemv_mx", ("a5", "pto.tgemv.mx.acc"): ".a5.tgemv_mx", ("a5", "pto.tgemv.mx.bias"): ".a5.tgemv_mx", + ("a5", "pto.thistogram"): ".a5.thistogram", ("a5", "pto.tinsert"): ".a5.tinsert", ("a5", "pto.tlrelu"): ".a5.tlrelu", ("a5", "pto.tlog"): ".a5.tlog", @@ -121,6 +123,8 @@ ("a5", "pto.tsub"): ".a5.tsub", ("a5", "pto.tsubs"): ".a5.tsubs", ("a5", "pto.tsqrt"): ".a5.tsqrt", + ("a5", "pto.ttrans"): ".a5.ttrans", + ("a5", "pto.ttri"): ".a5.ttri", ("a5", "pto.txor"): ".a5.txor", ("a5", "pto.txors"): ".a5.txors", } diff --git a/ptodsl/tests/test_tilelib_catalog.py b/ptodsl/tests/test_tilelib_catalog.py index 2217849fd9..d852355379 100644 --- a/ptodsl/tests/test_tilelib_catalog.py +++ b/ptodsl/tests/test_tilelib_catalog.py @@ -7,6 +7,8 @@ # See LICENSE in the root of the software repository for the full text of the License. """Table-driven selection and render coverage for the PTODSL TileLib catalog.""" +import ast +from pathlib import Path import unittest import ptodsl.tilelib as tilelib @@ -476,6 +478,29 @@ def test_tilelib_does_not_duplicate_the_public_ptodsl_surface(self): with self.subTest(name=name): self.assertFalse(hasattr(tilelib, name)) + def test_declared_a5_template_ops_are_loadable(self): + import TileOps + import TileOps.a5 as tileops_a5 + + declared_ops = set() + for path in Path(tileops_a5.__file__).parent.glob("*.py"): + if path.name.startswith("_") or path.name == "__init__.py": + continue + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if not ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "tile_template" + ): + continue + for keyword in node.keywords: + if keyword.arg == "op" and isinstance(keyword.value, ast.Constant): + declared_ops.add(keyword.value.value) + + missing = sorted(op for op in declared_ops if not TileOps.load_template(op, "a5")) + self.assertEqual(missing, []) + def test_each_catalog_entry_selects_and_renders(self): for op, entry in CATALOG.items(): with self.subTest(op=op):